Update sync_run.py

sync_run.py

@@ -1,86 +1,142 @@
+from __future__ import annotations
+
 import os
-import time
 import subprocess
 import threading
+import time
+from pathlib import Path
+
 from huggingface_hub import HfApi, snapshot_download
 
-# 配置
-TOKEN = os.getenv("HF_TOKEN")
-DATASET_ID = os.getenv("DATASET_ID")
 
-# 定义需要监控的目录和文件类型
-SYNC_FOLDERS = ["data", "output", "config"]
-SYNC_EXTENSIONS = [".json", ".db", ".yaml", ".yml", ".txt", ".env"]
+BASE_DIR = Path(__file__).resolve().parent
+PROJECT_DIR = BASE_DIR / "grok2api"
+DATA_DIR = Path(os.getenv("DATA_DIR", str(PROJECT_DIR / "data"))).expanduser()
+LOG_DIR = Path(os.getenv("LOG_DIR", str(PROJECT_DIR / "logs"))).expanduser()
+
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+DATASET_ID = os.getenv("DATASET_ID", "")
+SYNC_INTERVAL = max(int(os.getenv("HF_SYNC_INTERVAL", "1800")), 60)
+
+SERVER_HOST = os.getenv("SERVER_HOST", "0.0.0.0")
+SERVER_PORT = os.getenv("SERVER_PORT") or os.getenv("PORT") or "8000"
+SERVER_WORKERS = os.getenv("SERVER_WORKERS", "1")
+
+SYNC_ALLOW_PATTERNS = ["data/**"]
+SYNC_IGNORE_PATTERNS = [
+    "data/.locks/**",
+    "data/tmp/**",
+    "logs/**",
+    "**/__pycache__/**",
+]
+
+
+def log(message: str) -> None:
+    print(f"[HF-Space] {message}", flush=True)
 
-api = HfApi(token=TOKEN)
 
-def download_data():
-    """启动前：从 Dataset 下载所有备份"""
+def ensure_local_dirs() -> None:
+    DATA_DIR.mkdir(parents=True, exist_ok=True)
+    LOG_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def download_data() -> None:
     if not DATASET_ID:
-        print("[System] 警告: 未配置 DATASET_ID")
+        log("未配置 DATASET_ID，跳过启动数据同步。")
         return
+
     try:
-        print(f"[System] 正在从 Dataset ({DATASET_ID}) 拉取持久化数据...")
-        # 下载到当前目录，允许覆盖，但排除代码文件防止版本冲突
+        log(f"开始从 Dataset 拉取数据: {DATASET_ID}")
         snapshot_download(
             repo_id=DATASET_ID,
             repo_type="dataset",
-            local_dir=".", 
-            token=TOKEN,
-            ignore_patterns=["*.py", "Dockerfile", "requirements.txt", "logs/*", ".git/*"]
+            local_dir=str(PROJECT_DIR),
+            token=HF_TOKEN or None,
+            allow_patterns=SYNC_ALLOW_PATTERNS,
+            ignore_patterns=SYNC_IGNORE_PATTERNS,
         )
-        print("[System] 数据拉取完成。")
-    except Exception as e:
-        print(f"[System] 拉取失败 (首次运行或网络问题): {e}")
+        log("数据拉取完成。")
+    except Exception as exc:
+        log(f"数据拉取失败，继续本地启动: {exc}")
+
 
-def upload_data():
-    """运行中：定时将改动上传至 Dataset"""
+def upload_data(run_as_future: bool) -> None:
+    if not DATASET_ID:
+        return
+
+    if not HF_TOKEN:
+        log("已配置 DATASET_ID，但未配置 HF_TOKEN，跳过数据上传。")
+        return
+
+    try:
+        api = HfApi(token=HF_TOKEN)
+        api.upload_folder(
+            folder_path=str(PROJECT_DIR),
+            repo_id=DATASET_ID,
+            repo_type="dataset",
+            commit_message="chore: sync Grok2API data from Space",
+            allow_patterns=SYNC_ALLOW_PATTERNS,
+            ignore_patterns=SYNC_IGNORE_PATTERNS,
+            run_as_future=run_as_future,
+        )
+        if run_as_future:
+            log("已提交后台数据同步任务。")
+        else:
+            log("退出前数据同步完成。")
+    except Exception as exc:
+        if "No files have been modified" not in str(exc):
+            log(f"数据上传失败: {exc}")
+
+
+def upload_loop() -> None:
     while True:
-        # 每 2 分钟备份一次（缩短时间，防止改完密码还没备份就休眠了）
-        time.sleep(1800) 
-        if not DATASET_ID:
-            continue
-            
-        try:
-            # 扫描当前目录下所有符合条件的文件进行上传
-            # 我们直接上传整个根目录，但通过 allow_patterns 过滤出我们要的文件
-            api.upload_folder(
-                folder_path=".",
-                repo_id=DATASET_ID,
-                repo_type="dataset",
-                commit_message="Auto-backup data and configs",
-                allow_patterns=[
-                    "data/**", 
-                    "output/**", 
-                    "config/**",
-                    "*.json", 
-                    "*.db", 
-                    "*.yaml", 
-                    "*.yml", 
-                    ".env"
-                ],
-                ignore_patterns=["logs/**", "__pycache__/**"],
-                run_as_future=True
-            )
-            print(f"[Backup] {time.strftime('%H:%M:%S')} 备份同步任务已提交")
-        except Exception as e:
-            if "No files have been modified" not in str(e):
-                print(f"[Backup] 备份出错: {e}")
+        time.sleep(SYNC_INTERVAL)
+        upload_data(run_as_future=True)
 
-if __name__ == "__main__":
-    # 确保基础目录存在
-    for folder in SYNC_FOLDERS:
-        os.makedirs(folder, exist_ok=True)
 
-    # 1. 启动前同步云端数据到本地
+def init_storage() -> None:
+    subprocess.run(
+        ["sh", "scripts/init_storage.sh"],
+        cwd=PROJECT_DIR,
+        check=True,
+        env=os.environ.copy(),
+    )
+
+
+def run_server() -> None:
+    env = os.environ.copy()
+    env.setdefault("DATA_DIR", str(DATA_DIR))
+    env.setdefault("LOG_DIR", str(LOG_DIR))
+    env.setdefault("LOG_FILE_ENABLED", "false")
+
+    command = [
+        "granian",
+        "--interface",
+        "asgi",
+        "--host",
+        SERVER_HOST,
+        "--port",
+        SERVER_PORT,
+        "--workers",
+        SERVER_WORKERS,
+        "main:app",
+    ]
+
+    log(
+        f"启动 Grok2API: host={SERVER_HOST} port={SERVER_PORT} workers={SERVER_WORKERS}"
+    )
+    subprocess.run(command, cwd=PROJECT_DIR, check=True, env=env)
+
+
+if __name__ == "__main__":
+    ensure_local_dirs()
     download_data()
+    init_storage()
 
-    # 2. 启动后台备份线程
-    backup_thread = threading.Thread(target=upload_data, daemon=True)
+    backup_thread = threading.Thread(target=upload_loop, daemon=True)
     backup_thread.start()
 
-    # 3. 启动主程序
-    # 注意：根据你的日志，程序默认端口是 8000，但 HF 要求 7860
-    # 我们强制它运行在 7860
-    print("[System] 正在启动 OpenAI/Codex CLI...")
-    subprocess.run(["python", "webui.py", "--host", "0.0.0.0", "--port", "7860"])
+    try:
+        run_server()
+    finally:
+        upload_data(run_as_future=False)