Spaces:
Paused
Paused
File size: 2,617 Bytes
ccec00a 4736107 ccec00a 4736107 ccec00a 4736107 ccec00a 581b000 ccec00a 4736107 581b000 4736107 581b000 4736107 581b000 4736107 ccec00a 4736107 ccec00a 4736107 ccec00a 4736107 ccec00a 581b000 ccec00a 581b000 ccec00a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/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() |