db_syncer.py

#!/usr/bin/env python3
import hashlib
import logging
import os
import time
from huggingface_hub import hf_hub_download, upload_file

logging.basicConfig(
    format="%(asctime)s [db_syncer] %(levelname)s %(message)s",
    level=logging.INFO,
)
log = logging.getLogger("db_syncer")

HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "z90486091/LITELLM_DB")
DB_PATH = os.environ.get("DB_PATH", "/app/data/litellm.db")
REPO_DB_FILENAME = "litellm.db"
SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL", "300"))


def file_hash(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def download_db() -> bool:
    try:
        log.info(f"Downloading {REPO_DB_FILENAME} from {HF_DATASET_REPO}...")
        hf_hub_download(
            repo_id=HF_DATASET_REPO,
            filename=REPO_DB_FILENAME,
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=os.path.dirname(DB_PATH),
        )
        log.info(f"Downloaded DB to {DB_PATH}")
        return True
    except Exception as e:
        log.warning(f"Download failed (will use existing or fresh DB): {e}")
        return False


def upload_db() -> bool:
    try:
        upload_file(
            path_or_fileobj=DB_PATH,
            path_in_repo=REPO_DB_FILENAME,
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN,
            commit_message="db_syncer: auto-sync litellm.db",
        )
        log.info("Synced litellm.db to HF Dataset repo")
        return True
    except Exception as e:
        log.error(f"Upload failed: {e}")
        return False


def main():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    download_db()
    log.info(f"Starting db_syncer — watching {DB_PATH} every {SYNC_INTERVAL}s")
    last_hash = None

    while True:
        try:
            if not os.path.exists(DB_PATH):
                log.warning(f"{DB_PATH} not found — waiting...")
                time.sleep(30)
                continue
            current_hash = file_hash(DB_PATH)
            if current_hash != last_hash:
                log.info(f"Change detected (hash: {current_hash[:8]}...) — syncing")
                if upload_db():
                    last_hash = current_hash
            else:
                log.debug("No change — skipping")
        except Exception as e:
            log.error(f"Unexpected error: {e}")
        time.sleep(SYNC_INTERVAL)


if __name__ == "__main__":
    main()