Upload sync_data.py

sync_data.py

@@ -1,7 +1,8 @@
 import os
+import sys
 import shutil
 import time
-import threading
+import signal
 from pathlib import Path
 from datetime import datetime
 from huggingface_hub import HfApi, hf_hub_download, list_repo_files
@@ -12,6 +13,15 @@ SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL", 300))
 ARCHIVE_NAME = "openclaw_backup.tar.gz"
 
 api = HfApi()
+running = True
+
+
+def signal_handler(signum, frame):
+    global running
+    print(f"[{datetime.now()}] Received signal {signum}, uploading before exit...")
+    upload_to_dataset()
+    running = False
+    sys.exit(0)
 
 
 def download_from_dataset():
@@ -59,25 +69,28 @@ def upload_to_dataset():
         print(f"[{datetime.now()}] Upload failed: {e}")
 
 
-def sync_loop():
-    while True:
+def upload_loop():
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+    while running:
         time.sleep(SYNC_INTERVAL)
         upload_to_dataset()
 
 
-def main():
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python sync_data.py [download|upload_loop]")
+        sys.exit(1)
+
+    cmd = sys.argv[1]
     print(f"Dataset: {DATASET_REPO_ID}")
     print(f"Data dir: {DATA_DIR}")
-    print(f"Sync interval: {SYNC_INTERVAL}s")
-
-    download_from_dataset()
 
-    sync_thread = threading.Thread(target=sync_loop, daemon=True)
-    sync_thread.start()
-
-    while True:
-        time.sleep(3600)
-
-
-if __name__ == "__main__":
-    main()
+    if cmd == "download":
+        download_from_dataset()
+    elif cmd == "upload_loop":
+        print(f"Sync interval: {SYNC_INTERVAL}s")
+        upload_loop()
+    else:
+        print(f"Unknown command: {cmd}")
+        sys.exit(1)