Update tools/hf_backup.py

tools/hf_backup.py

@@ -4,9 +4,11 @@
 import argparse
 import os
 import sys
-import tempfile
+from pathlib import Path
 from huggingface_hub import HfApi
 
+RESTORE_DIR = os.environ.get("HF_RESTORE_DIR", "/tmp/crs_backup")  # 与 shell 的 TMP_DIR 对齐
+
 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")]
@@ -15,16 +17,16 @@ def list_backups(api: HfApi, repo_id: str, prefix: str):
 
 def ensure_dataset(api: HfApi, repo_id: str):
     try:
-        api.dataset_info(repo_id=repo_id, use_auth_token=api.token)  # 兼容旧版
+        api.dataset_info(repo_id=repo_id)
     except Exception:
-        # 不存在则创建（可按需改为 public）
+        # 若不存在则创建为私有数据集；如需公开可改 private=False
         api.create_repo(repo_id=repo_id, repo_type="dataset", private=True, exist_ok=True)
 
 def upload(args):
     api = HfApi(token=args.token)
     ensure_dataset(api, args.repo)
 
-    # 上传
+    # 上传当前归档
     api.upload_file(
         path_or_fileobj=args.file,
         path_in_repo=os.path.basename(args.file),
@@ -32,7 +34,7 @@ def upload(args):
         repo_type="dataset",
     )
 
-    # 保留最新 N 份
+    # 仅保留最新 N 份
     if args.max and args.max > 0:
         backs = list_backups(api, args.repo, args.prefix)
         if len(backs) > args.max:
@@ -41,25 +43,25 @@ def upload(args):
                 try:
                     api.delete_file(path_in_repo=f, repo_id=args.repo, repo_type="dataset")
                 except Exception:
-                    # 删除失败不致命
-                    pass
+                    pass  # 删除失败不致命
 
 def restore(args):
     api = HfApi(token=args.token)
     backs = list_backups(api, args.repo, args.prefix)
     if not backs:
-        return  # 没有备份则安静退出，由 shell 脚本处理
+        return  # 无备份，安静退出，由 shell 打日志
 
     latest = backs[-1]
-    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)
+    # 用持久目录而不是临时目录（否则 Python 退出就删了）
+    Path(RESTORE_DIR).mkdir(parents=True, exist_ok=True)
+    path = api.hf_hub_download(
+        repo_id=args.repo,
+        filename=latest,
+        repo_type="dataset",
+        local_dir=RESTORE_DIR,
+        local_dir_use_symlinks=False,  # 更稳妥，避免符号链接指向被清理
+    )
+    print(path)  # 仅输出路径，供 shell 脚本接收
 
 def main():
     p = argparse.ArgumentParser()