agentscope / examples /redis_space /backup_manager.py
ghostdrive1's picture
Upload folder using huggingface_hub
86ba74e verified
Raw
History Blame Contribute Delete
2.16 kB
import os
import time
import schedule
from huggingface_hub import HfApi, hf_hub_download
import subprocess
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = os.getenv("DATASET_REPO", "augment17/agentscope-redis-backup")
DUMP_FILE = "/data/dump.rdb"
LOCAL_DUMP = "dump.rdb"
api = HfApi(token=HF_TOKEN)
def ensure_repo():
try:
api.repo_info(repo_id=DATASET_REPO, repo_type="dataset")
except Exception:
logger.info(f"Creating dataset repo {DATASET_REPO}...")
api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", private=True)
def restore_backup():
ensure_repo()
try:
logger.info("Checking for existing Redis backup...")
downloaded_file = hf_hub_download(
repo_id=DATASET_REPO,
repo_type="dataset",
filename=LOCAL_DUMP,
token=HF_TOKEN,
local_dir="/data"
)
logger.info(f"Restored backup from {downloaded_file}")
except Exception as e:
logger.info(f"No existing backup found or failed to restore: {e}")
def perform_backup():
logger.info("Performing Redis backup...")
try:
# Trigger Redis bgsave
subprocess.run(["redis-cli", "save"], check=True)
if os.path.exists(DUMP_FILE):
api.upload_file(
path_or_fileobj=DUMP_FILE,
path_in_repo=LOCAL_DUMP,
repo_id=DATASET_REPO,
repo_type="dataset",
token=HF_TOKEN
)
logger.info("Backup uploaded to Hugging Face Datasets successfully.")
else:
logger.error("dump.rdb not found after save!")
except Exception as e:
logger.error(f"Backup failed: {e}")
if __name__ == "__main__":
mode = os.getenv("MODE", "restore")
if mode == "restore":
restore_backup()
elif mode == "cron":
logger.info("Starting backup cron schedule (every 30 minutes)...")
schedule.every(30).minutes.do(perform_backup)
while True:
schedule.run_pending()
time.sleep(60)