Update tools/hf_backup.py

tools/hf_backup.py

@@ -1,85 +1,86 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-import os, sys, time, tarfile, tempfile, shutil
-from huggingface_hub import HfApi
-
-def log(*a): print(time.strftime("[%Y-%m-%d %H:%M:%S]"), *a, flush=True)
-
-token = os.environ.get("HF_TOKEN")
-repo_id = os.environ.get("DATASET_ID")
-max_backups = int(os.environ.get("MAX_BACKUPS", "10"))
-backup_paths = [p.strip() for p in os.environ.get("BACKUP_PATHS", "/app/data,/app/logs").split(",") if p.strip()]
-prefix = os.environ.get("BACKUP_PREFIX", "crs_backup_")
 
-if not token or not repo_id:
-    log("HF_TOKEN/DATASET_ID 未配置，退出")
-    sys.exit(0)
-
-api = HfApi(token=token)
+import argparse
+import os
+import sys
+import tempfile
+from huggingface_hub import HfApi
 
-def list_backups():
+def list_backups(api: HfApi, repo_id: str, prefix: str):
     files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
     backs = [f for f in files if f.startswith(prefix) and f.endswith(".tar.gz")]
     backs.sort()
     return backs
 
-def create_archive():
-    ts = time.strftime("%Y%m%d_%H%M%S")
-    name = f"{prefix}{ts}.tar.gz"
-    tmpdir = tempfile.mkdtemp()
-    path = os.path.join(tmpdir, name)
+def upload(args):
+    api = HfApi(token=args.token)
+    # 可选：确认数据集存在
+    try:
+        api.dataset_info(repo_id=args.repo)
+    except Exception:
+        pass
 
-    with tarfile.open(path, "w:gz") as tar:
-        for d in backup_paths:
-            if not os.path.exists(d):
-                log("skip missing:", d)
-                continue
-            # 归档名尽量短，/app 前缀裁掉，方便恢复
-            arc = os.path.relpath(d, "/app") if d.startswith("/app/") else os.path.basename(d.rstrip("/"))
-            log("add to archive:", d, "as", arc)
-            tar.add(d, arcname=arc, recursive=True)
-    return path, name, tmpdir
+    # 上传
+    api.upload_file(
+        path_or_fileobj=args.file,
+        path_in_repo=os.path.basename(args.file),
+        repo_id=args.repo,
+        repo_type="dataset",
+    )
 
-def upload_and_prune(path, name):
-    log("Uploading", name)
-    api.upload_file(path_or_fileobj=path, path_in_repo=name, repo_id=repo_id, repo_type="dataset")
-    backs = list_backups()
-    excess = len(backs) - max_backups
-    for f in backs[:excess]:
-        try:
-            log("Pruning old backup:", f)
-            api.delete_file(path_in_repo=f, repo_id=repo_id, repo_type="dataset")
-        except Exception as e:
-            log("Delete failed:", f, e)
+    # 保留最新 N 份
+    if args.max and args.max > 0:
+      backs = list_backups(api, args.repo, args.prefix)
+      if len(backs) > args.max:
+          to_del = backs[: len(backs) - args.max]
+          for f in to_del:
+              try:
+                  api.delete_file(path_in_repo=f, repo_id=args.repo, repo_type="dataset")
+              except Exception:
+                  # 删除失败不致命
+                  pass
 
-def restore_latest():
-    backs = list_backups()
+def restore(args):
+    api = HfApi(token=args.token)
+    backs = list_backups(api, args.repo, args.prefix)
     if not backs:
-        log("No backups in dataset")
-        return
+        return  # 没有备份则安静退出，由 shell 脚本处理
+
     latest = backs[-1]
-    log("Restoring", latest)
-    tmpdir = tempfile.mkdtemp()
-    local = api.hf_hub_download(repo_id=repo_id, filename=latest, repo_type="dataset", local_dir=tmpdir)
-    with tarfile.open(local, "r:gz") as tar:
-        tar.extractall("/app")
-    log("Restored", latest)
+    with tempfile.TemporaryDirectory() as d:
+        path = api.hf_hub_download(
+            repo_id=args.repo,
+            filename=latest,
+            repo_type="dataset",
+            local_dir=d
+        )
+        # 将本地文件路径直接输出给 shell
+        print(path)
 
 def main():
-    if len(sys.argv) < 2:
-        log("Usage: hf_backup.py [restore|backup]")
-        return
-    cmd = sys.argv[1]
-    if cmd == "restore":
-        restore_latest()
-    elif cmd == "backup":
-        path, name, tmp = create_archive()
-        try:
-            upload_and_prune(path, name)
-        finally:
-            shutil.rmtree(os.path.dirname(path), ignore_errors=True)
-    else:
-        log("Unknown cmd:", cmd)
+    p = argparse.ArgumentParser()
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    up = sub.add_parser("upload")
+    up.add_argument("--token", required=True)
+    up.add_argument("--repo", required=True)
+    up.add_argument("--file", required=True)
+    up.add_argument("--prefix", required=True)
+    up.add_argument("--max", type=int, default=10)
+    up.set_defaults(func=upload)
+
+    rs = sub.add_parser("restore")
+    rs.add_argument("--token", required=True)
+    rs.add_argument("--repo", required=True)
+    rs.add_argument("--prefix", required=True)
+    rs.set_defaults(func=restore)
+
+    args = p.parse_args()
+    try:
+        args.func(args)
+    except KeyboardInterrupt:
+        sys.exit(130)
 
 if __name__ == "__main__":
     main()