feat: auto pause hugsim server

competitions/runner.py

@@ -5,17 +5,21 @@ import os
 import time
 import uuid
 import shutil
+import os
 from dataclasses import dataclass
 from typing import List, Dict, Any
 from collections import defaultdict
 
+import tenacity
 import pandas as pd
 from huggingface_hub import HfApi, hf_hub_download, snapshot_download
 from loguru import logger
 
 from competitions.enums import SubmissionStatus, ErrorMessage
 from competitions.info import CompetitionInfo
-from competitions.utils import user_token_api, space_cleaner, dockerfile_modifier
+from competitions.utils import user_token_api, space_cleaner, dockerfile_modifier, server_manager
+
+IS_AUTO_PAUSE_HUGSIM_SERVER = os.getenv("IS_AUTO_PAUSE_HUGSIM_SERVER", "true").lower() == "true"
 
 
 @dataclass
@@ -219,6 +223,7 @@ class JobRunner:
             shutil.rmtree(client_code_local_dir, ignore_errors=True)
         self._queue_submission(team_id, submission_id)
 
+    @tenacity.retry(stop=tenacity.stop_never, wait=tenacity.wait_fixed(15))
     def run(self):
         cur = 0
         while True:
@@ -231,6 +236,7 @@ class JobRunner:
             # Clean up spaces every 100 iterations
             if cur % 100 == 1:
                 logger.info("Cleaning up spaces...")
+                ready_submissions_count = 0
                 for space in all_submissions:
                     if space["status"] in {SubmissionStatus.QUEUED.value, SubmissionStatus.PROCESSING.value}:
                         logger.info(f"Cleaning up space {space['space_id']} for submission {space['submission_id']}")
@@ -239,6 +245,11 @@ class JobRunner:
                             space["team_id"],
                             space["submission_id"],
                         )
+                    if space["status"] in {SubmissionStatus.QUEUED.value, SubmissionStatus.PROCESSING.value, SubmissionStatus.PROCESSING.value}:
+                        ready_submissions_count += 1
+                if ready_submissions_count == 0:
+                    if IS_AUTO_PAUSE_HUGSIM_SERVER:
+                        server_manager.pause_all_servers()
 
             pending_submissions = self._get_pending_subs(all_submissions)
             if pending_submissions is None:
@@ -249,6 +260,8 @@ class JobRunner:
             if server_active_count[first_pending_sub["server_url"]] >= 1:
                 continue
             try:
+                if IS_AUTO_PAUSE_HUGSIM_SERVER:
+                    server_manager.start_all_servers()
                 self.create_space(first_pending_sub["team_id"], first_pending_sub["submission_id"], first_pending_sub["submission_repo"], first_pending_sub["space_id"], first_pending_sub["server_url"], first_pending_sub["hardware"])
             except Exception as e:
                 logger.error(

competitions/utils.py

@@ -8,10 +8,12 @@ import threading
 import uuid
 import base64
 import glob
+import time
 from typing import Optional, Dict, Any, List, Literal, Tuple
 from collections import defaultdict
 from datetime import datetime, timezone, timedelta
 
+import tenacity
 import requests
 import pandas as pd
 import jwt
@@ -485,8 +487,11 @@ user_token_api = UserTokenApi(
 
 
 class ServerManager:
-    def __init__(self, server_url_list: List[str]):
+    def __init__(self, hf_token: str, server_url_list: List[str], space_name_list: List[str]):
+        self.hf_token = hf_token
+        self.api = HfApi(token=hf_token)
         self.server_url_list = server_url_list
+        self.space_name_list = space_name_list
         self._cur_index = 0
         self._lock = threading.Lock()
 
@@ -496,8 +501,37 @@ class ServerManager:
             self._cur_index = (self._cur_index + 1) % len(self.server_url_list)
         return server_url
 
+    def pause_all_servers(self):
+        """Pause all servers."""
+        for space_name in self.space_name_list:
+            try:
+                self.api.pause_space(repo_id=space_name)
+            except Exception as e:
+                logger.error(f"Failed to pause space {space_name} - {e}")
+    
+    @tenacity.retry(stop=tenacity.stop_after_attempt(5), wait=tenacity.wait_fixed(15))
+    def start_all_servers(self):
+        """Start all servers."""
+        for space_name in self.space_name_list:
+            stage = self.api.space_info(space_name).runtime.stage.lower()
+            if stage == "running":
+                continue
+            self.api.restart_space(repo_id=space_name)
+            while True:
+                time.sleep(10)
+                stage = self.api.space_info(space_name).runtime.stage.lower()
+                if "error" in stage:
+                    self.api.restart_space(repo_id=space_name)
+                    continue
+                if stage == "running":
+                    break
+    
 
-server_manager = ServerManager(["https://xdimlab-hugsim-web-server-0.hf.space"])
+server_manager = ServerManager(
+    os.environ.get("HF_TOKEN", None),
+    ["https://xdimlab-hugsim-web-server-0.hf.space"],
+    ["XDimLab/hugsim_web_server_0"]
+)
 
 
 class SubmissionApi:
@@ -562,8 +596,7 @@ class SubmissionApi:
 
     def count_by_status(self, team_id: str, status_list: List[SubmissionStatus]) -> int:
         user_submission_info = self.download_submission_info(team_id)
-        count = sum(1 for submission in user_submission_info["submissions"] if
-                    SubmissionStatus(submission["status"]) in status_list)
+        count = sum(1 for submission in user_submission_info["submissions"] if SubmissionStatus(submission["status"]) in status_list)
         return count
 
 
@@ -684,8 +717,7 @@ class SpaceCleaner:
                 submission_api.update_submission_data(
                     team_id=team_id,
                     submission_id=submission_id,
-                    data={"status": SubmissionStatus.FAILED.value,
-                          "error_message": ErrorMessage.BUILD_SPACE_FAILED.value},
+                    data={"status": SubmissionStatus.FAILED.value, "error_message": ErrorMessage.BUILD_SPACE_FAILED.value},
                 )
             else:
                 self.api.restart_space(repo_id=space_id)
@@ -923,8 +955,7 @@ class DockerfileModifier:
         if last_user is None:
             modified_lines.insert(-1, f"COPY {self.source_so_path}" + f" {self.tatget_so_path}")
         else:
-            modified_lines.insert(-1,
-                                  f"COPY --chown={last_user}:{last_user} {self.source_so_path} {self.tatget_so_path}")
+            modified_lines.insert(-1, f"COPY --chown={last_user}:{last_user} {self.source_so_path} {self.tatget_so_path}")
             modified_lines.insert(-1, f"RUN chown -R {last_user}:{last_user} {self.tatget_so_dir}")
 
         return '\n'.join(modified_lines), changes

requirements.txt

@@ -18,3 +18,4 @@ hf-transfer>=0.1.9
 filelock>=3.13.3
 cachetools==6.0.0
 PyJWT>=2.10.1
+tenacity>=9.1.2