Create hf_sync.py

hf_sync.py

@@ -0,0 +1,200 @@
+# hf_sync.py
+from huggingface_hub import HfApi
+import sys
+import os
+import datetime
+
+def manage_backups(api, repo_id, max_files=50):
+    """管理备份文件，保留最新的max_files个文件"""
+    try:
+        print(f"开始管理 Hugging Face 仓库 '{repo_id}' 的备份...")
+        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
+        # 筛选出符合备份文件命名规则的文件
+        backup_files = [f for f in files if f.startswith('webui_backup_') and f.endswith('.db')]
+        print(f"在仓库中找到 {len(backup_files)} 个备份文件。")
+
+        # 提取日期并排序
+        dated_files = []
+        for f in backup_files:
+            try:
+                # 假设文件名格式为 webui_backup_YYYYMMDD.db
+                date_str = f.split('_')[-1].split('.')[0]
+                # 严格匹配 YYYYMMDD 格式
+                if len(date_str) == 8 and date_str.isdigit():
+                    date_obj = datetime.datetime.strptime(date_str, '%Y%m%d').date()
+                    dated_files.append((date_obj, f))
+                else:
+                     print(f"警告: 文件名格式不符，跳过: {f}")
+            except (IndexError, ValueError) as e:
+                print(f"警告: 无法从文件名解析日期，跳过: {f} (错误: {e})")
+                continue # 跳过格式不正确的文件
+
+        # 按日期降序排序（最新的在前）
+        dated_files.sort(key=lambda item: item[0], reverse=True)
+
+        if len(dated_files) > max_files:
+            files_to_delete = [item[1] for item in dated_files[max_files:]]
+            print(f"找到 {len(files_to_delete)} 个需要删除的旧备份。")
+            for file_to_delete in files_to_delete:
+                try:
+                    print(f"正在删除旧备份: {file_to_delete}")
+                    api.delete_file(path_in_repo=file_to_delete, repo_id=repo_id, repo_type="dataset")
+                    print(f"已删除旧备份: {file_to_delete}")
+                except Exception as e_del:
+                    print(f"删除文件 {file_to_delete} 时出错: {str(e_del)}")
+        else:
+            print(f"备份文件数量 ({len(dated_files)}) 未超过限制 ({max_files})，无需删除。")
+
+    except Exception as e_list:
+        print(f"列出或管理备份文件时出错: {str(e_list)}")
+
+
+def upload_backup(file_path, file_name, token, repo_id):
+    """上传备份文件到Hugging Face"""
+    if not os.path.exists(file_path):
+        print(f"错误: 本地文件未找到，无法上传: {file_path}")
+        return
+
+    print(f"尝试将 {file_path} 中的文件上传为 {file_name} 到仓库 {repo_id}")
+    api = HfApi(token=token)
+    try:
+        # 上传新文件 (upload_file 会覆盖同名文件)
+        print(f"正在上传 {file_name}...")
+        api.upload_file(
+            path_or_fileobj=file_path,
+            path_in_repo=file_name,
+            repo_id=repo_id,
+            repo_type="dataset"
+        )
+        print(f"成功上传: {file_name}")
+        # 上传成功后管理备份
+        print("上传成功，开始管理备份...")
+        manage_backups(api, repo_id)
+    except Exception as e:
+        print(f"上传文件 {file_name} 失败: {str(e)}")
+
+
+def download_latest_backup(token, repo_id, target_dir):
+    """从Hugging Face下载最新备份"""
+    print(f"尝试从仓库 {repo_id} 下载最新备份到目录 {target_dir}")
+    api = HfApi(token=token)
+    try:
+        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
+        backup_files = [f for f in files if f.startswith('webui_backup_') and f.endswith('.db')]
+
+        if not backup_files:
+            print("在 Hugging Face 数据集中未找到备份文件。")
+            return False
+
+        # 根据文件名中的日期找到最新的文件
+        latest_file = None
+        latest_date = None
+        for f in backup_files:
+             try:
+                # 假设文件名格式为 webui_backup_YYYYMMDD.db
+                date_str = f.split('_')[-1].split('.')[0]
+                if len(date_str) == 8 and date_str.isdigit():
+                    date_obj = datetime.datetime.strptime(date_str, '%Y%m%d').date()
+                    if latest_date is None or date_obj > latest_date:
+                        latest_date = date_obj
+                        latest_file = f
+                else:
+                     continue # 跳过格式不正确的文件
+             except (IndexError, ValueError):
+                 continue # 跳过无法解析的文件
+
+        if latest_file is None:
+            print("无法确定最新的备份文件（可能所有文件名格式都不正确）。")
+            return False
+
+        print(f"找到最新的备份文件: {latest_file}")
+        print(f"正在下载 {latest_file}...")
+        # 先下载到临时路径
+        downloaded_path = api.hf_hub_download(
+            repo_id=repo_id,
+            filename=latest_file,
+            repo_type="dataset",
+            local_dir_use_symlinks=False, # 避免在容器中出现符号链接问题
+            cache_dir=os.path.join(target_dir, '.hf_cache') # 指定缓存目录，避免权限问题
+        )
+
+        if downloaded_path and os.path.exists(downloaded_path):
+            # 创建目标目录（如果不存在）
+            os.makedirs(target_dir, exist_ok=True)
+            # 目标文件路径
+            target_path = os.path.join(target_dir, "webui.db")
+            print(f"下载完成，路径: {downloaded_path}，准备移动到: {target_path}")
+            # 使用 os.replace 进行原子性移动/覆盖，如果失败则回退到复制
+            try:
+                 os.replace(downloaded_path, target_path)
+                 print(f"成功将 {latest_file} 移动到 {target_path}")
+            except OSError as e_mv:
+                 print(f"移动文件失败 (错误: {e_mv})，尝试复制...")
+                 import shutil
+                 try:
+                    shutil.copy2(downloaded_path, target_path)
+                    print(f"成功将 {latest_file} 复制到 {target_path}")
+                    # 尝试删除临时下载文件
+                    try:
+                        os.remove(downloaded_path)
+                        # 清理可能的缓存目录结构
+                        hf_cache_dir = os.path.join(target_dir, '.hf_cache')
+                        if os.path.exists(hf_cache_dir):
+                            shutil.rmtree(hf_cache_dir)
+                    except OSError as e_rm:
+                        print(f"警告: 删除临时下载文件 {downloaded_path} 失败: {e_rm}")
+                 except Exception as e_cp:
+                    print(f"复制文件也失败: {e_cp}")
+                    return False
+
+
+            print(f"成功从 Hugging Face 恢复: {latest_file}")
+            # 清理空的缓存目录（如果存在且为空）
+            try:
+                cache_root = os.path.dirname(downloaded_path)
+                if os.path.exists(cache_root) and not os.listdir(cache_root):
+                    os.rmdir(cache_root)
+            except OSError:
+                pass # 忽略清理错误
+            return True
+        else:
+            print(f"下载失败或文件路径无效: {downloaded_path}")
+            return False
+    except Exception as e:
+        print(f"下载或恢复过程中发生错误: {str(e)}")
+        # 打印更详细的堆栈跟踪信息以便调试
+        import traceback
+        traceback.print_exc()
+        return False
+
+if __name__ == "__main__":
+    # 基本的参数数量检查
+    if len(sys.argv) < 4:
+        print("使用方法:")
+        print("  上传: python3 hf_sync.py upload <HF_TOKEN> <DATASET_ID> <FILE_PATH> <FILE_NAME_IN_REPO>")
+        print("  下载: python3 hf_sync.py download <HF_TOKEN> <DATASET_ID> <TARGET_DIR>")
+        sys.exit(1)
+
+    action = sys.argv[1]
+    token = sys.argv[2]
+    repo_id = sys.argv[3]
+
+    if action == "upload":
+        if len(sys.argv) != 6:
+            print("上传命令参数错误。需要: upload <HF_TOKEN> <DATASET_ID> <FILE_PATH> <FILE_NAME_IN_REPO>")
+            sys.exit(1)
+        file_path = sys.argv[4]
+        file_name = sys.argv[5]
+        upload_backup(file_path, file_name, token, repo_id)
+    elif action == "download":
+        if len(sys.argv) != 5:
+            print("下载命令参数错误。需要: download <HF_TOKEN> <DATASET_ID> <TARGET_DIR>")
+            sys.exit(1)
+        target_dir = sys.argv[4] # 下载需要目标目录参数
+        # 调用下载函数，并根据返回值退出
+        if not download_latest_backup(token, repo_id, target_dir):
+             sys.exit(1) # 下载失败则脚本退出，表示恢复失败
+    else:
+        print(f"未知的操作: {action}")
+        sys.exit(1)
+