feat: update code

competitions/evaluate.py

@@ -1,135 +0,0 @@
-import argparse
-import json
-import os
-import shlex
-import shutil
-import subprocess
-
-from huggingface_hub import HfApi, hf_hub_download, snapshot_download
-from huggingface_hub.utils._errors import EntryNotFoundError
-from loguru import logger
-
-from competitions import utils
-from competitions.compute_metrics import compute_metrics
-from competitions.enums import SubmissionStatus
-from competitions.params import EvalParams
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--config", type=str, required=True)
-    return parser.parse_args()
-
-
-def upload_submission_file(params, file_path):
-    logger.info("Uploading submission file")
-    pass
-
-
-def generate_submission_file(params):
-    logger.info("Downloading submission dataset")
-    submission_dir = snapshot_download(
-        repo_id=params.submission_repo,
-        local_dir=params.output_path,
-        token=os.environ.get("USER_TOKEN"),
-        repo_type="model",
-    )
-    # submission_dir has a script.py file
-    # start a subprocess to run the script.py
-    # the script.py will generate a submission.csv file in the submission_dir
-    # push the submission.csv file to the repo using upload_submission_file
-    logger.info("Generating submission file")
-
-    # invalidate USER_TOKEN env var
-    os.environ["USER_TOKEN"] = ""
-
-    # Copy sandbox to submission_dir
-    shutil.copyfile("sandbox", f"{submission_dir}/sandbox")
-    sandbox_path = f"{submission_dir}/sandbox"
-    os.chmod(sandbox_path, 0o755)
-    os.chown(sandbox_path, os.getuid(), os.getgid())
-
-    # Define your command
-    cmd = f"{sandbox_path} python script.py"
-    cmd = shlex.split(cmd)
-
-    # Copy the current environment and modify it
-    env = os.environ.copy()
-
-    # Start the subprocess
-    process = subprocess.Popen(cmd, cwd=submission_dir, env=env)
-
-    # Wait for the process to complete or timeout
-    try:
-        process.wait(timeout=params.time_limit)
-    except subprocess.TimeoutExpired:
-        logger.info(f"Process exceeded {params.time_limit} seconds time limit. Terminating...")
-        process.kill()
-        process.wait()
-
-    # Check if process terminated due to timeout
-    if process.returncode and process.returncode != 0:
-        logger.error("Subprocess didn't terminate successfully")
-    else:
-        logger.info("Subprocess terminated successfully")
-
-    logger.info("contents of submission_dir")
-    logger.info(os.listdir(submission_dir))
-
-    api = HfApi(token=params.token)
-    for sub_file in params.submission_filenames:
-        logger.info(f"Uploading {sub_file} to the repository")
-        sub_file_ext = sub_file.split(".")[-1]
-        api.upload_file(
-            path_or_fileobj=f"{submission_dir}/{sub_file}",
-            path_in_repo=f"submissions/{params.team_id}-{params.submission_id}.{sub_file_ext}",
-            repo_id=params.competition_id,
-            repo_type="dataset",
-        )
-
-
-@utils.monitor
-def run(params):
-    logger.info(params)
-    if isinstance(params, dict):
-        params = EvalParams(**params)
-
-    utils.update_submission_status(params, SubmissionStatus.PROCESSING.value)
-
-    if params.competition_type == "script":
-        try:
-            requirements_fname = hf_hub_download(
-                repo_id=params.competition_id,
-                filename="requirements.txt",
-                token=params.token,
-                repo_type="dataset",
-            )
-        except EntryNotFoundError:
-            requirements_fname = None
-
-        if requirements_fname:
-            logger.info("Installing requirements")
-            utils.uninstall_requirements(requirements_fname)
-            utils.install_requirements(requirements_fname)
-        if len(str(params.dataset).strip()) > 0:
-            # _ = Repository(local_dir="/tmp/data", clone_from=params.dataset, token=params.token)
-            _ = snapshot_download(
-                repo_id=params.dataset,
-                local_dir="/tmp/data",
-                token=params.token,
-                repo_type="dataset",
-            )
-        generate_submission_file(params)
-
-    evaluation = compute_metrics(params)
-
-    utils.update_submission_score(params, evaluation["public_score"], evaluation["private_score"])
-    utils.update_submission_status(params, SubmissionStatus.SUCCESS.value)
-    utils.delete_space(params)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    _params = json.load(open(args.config, encoding="utf-8"))
-    _params = EvalParams(**_params)
-    run(_params)

competitions/runner.py

@@ -4,7 +4,10 @@ import json
 import os
 import time
 import uuid
+import shutil
 from dataclasses import dataclass
+from typing import List, Dict, Any
+from collections import defaultdict
 
 import pandas as pd
 from huggingface_hub import HfApi, hf_hub_download, snapshot_download
@@ -12,17 +15,7 @@ from loguru import logger
 
 from competitions.enums import SubmissionStatus
 from competitions.info import CompetitionInfo
-from competitions.utils import run_evaluation, user_token_api
-
-
-_DOCKERFILE = """
-FROM huggingface/competitions:latest
-
-CMD uvicorn competitions.api:api --port 7860 --host 0.0.0.0
-"""
-
-# format _DOCKERFILE
-_DOCKERFILE = _DOCKERFILE.replace("\n", " ").replace("  ", "\n").strip()
+from competitions.utils import user_token_api, space_cleaner
 
 
 @dataclass
@@ -43,7 +36,7 @@ class JobRunner:
         self.dataset = self.competition_info.dataset
         self.submission_filenames = self.competition_info.submission_filenames
 
-    def get_pending_subs(self):
+    def _get_all_submissions(self) -> List[Dict[str, Any]]:
         submission_jsons = snapshot_download(
             repo_id=self.competition_id,
             allow_patterns="submission_info/*.json",
@@ -51,21 +44,32 @@ class JobRunner:
             repo_type="dataset",
         )
         submission_jsons = glob.glob(os.path.join(submission_jsons, "submission_info/*.json"))
-        pending_submissions = []
-        for _json in submission_jsons:
-            _json = json.load(open(_json, "r", encoding="utf-8"))
+        all_submissions = []
+        for _json_path in submission_jsons:
+            with open(_json_path, "r", encoding="utf-8") as f:
+                _json = json.load(f)
             team_id = _json["id"]
             for sub in _json["submissions"]:
-                if sub["status"] == SubmissionStatus.PENDING.value:
-                    pending_submissions.append(
-                        {
-                            "team_id": team_id,
-                            "submission_id": sub["submission_id"],
-                            "datetime": sub["datetime"],
-                            "submission_repo": sub["submission_repo"],
-                            "space_id": sub["space_id"],
-                        }
-                    )
+                all_submissions.append(
+                    {
+                        "team_id": team_id,
+                        "submission_id": sub["submission_id"],
+                        "datetime": sub["datetime"],
+                        "status": sub["status"],
+                        "submission_repo": sub["submission_repo"],
+                        "space_id": sub["space_id"],
+                        "server_url": sub["server_url"],
+                        "hardware": sub["hardware"],
+                    }
+                )
+        return all_submissions
+
+
+    def _get_pending_subs(self, submissions: List[Dict[str, Any]]) -> pd.DataFrame:
+        pending_submissions = []
+        for sub in submissions:
+            if sub["status"] == SubmissionStatus.PENDING.value:
+                pending_submissions.append(sub)
         if len(pending_submissions) == 0:
             return None
         logger.info(f"Found {len(pending_submissions)} pending submissions.")
@@ -75,6 +79,13 @@ class JobRunner:
         pending_submissions = pending_submissions.reset_index(drop=True)
         return pending_submissions
 
+    def _get_server_active_count(self, submissions: List[Dict[str, Any]]) -> Dict[str, int]:
+        server_active_count = defaultdict(int)
+        for sub in submissions:
+            if sub["status"] in {SubmissionStatus.PROCESSING.value, SubmissionStatus.QUEUED.value}:
+                server_active_count[sub["server_url"]] += 1
+        return server_active_count
+
     def _queue_submission(self, team_id, submission_id):
         team_fname = hf_hub_download(
             repo_id=self.competition_id,
@@ -127,29 +138,7 @@ class JobRunner:
             repo_type="dataset",
         )
 
-    def run_local(self, team_id, submission_id, submission_repo):
-        self._queue_submission(team_id, submission_id)
-        eval_params = {
-            "competition_id": self.competition_id,
-            "competition_type": self.competition_type,
-            "metric": self.metric,
-            "token": self.token,
-            "team_id": team_id,
-            "submission_id": submission_id,
-            "submission_id_col": self.submission_id_col,
-            "submission_cols": self.submission_cols,
-            "submission_rows": self.submission_rows,
-            "output_path": self.output_path,
-            "submission_repo": submission_repo,
-            "time_limit": self.time_limit,
-            "dataset": self.dataset,
-            "submission_filenames": self.submission_filenames,
-        }
-        eval_params = json.dumps(eval_params)
-        eval_pid = run_evaluation(eval_params, local=True, wait=True)
-        logger.info(f"New evaluation process started with pid {eval_pid}.")
-
-    def _create_readme(self, project_name):
+    def _create_readme(self, project_name: str) -> str:
         _readme = "---\n"
         _readme += f"title: {project_name}\n"
         _readme += "emoji: 🚀\n"
@@ -158,92 +147,107 @@ class JobRunner:
         _readme += "sdk: docker\n"
         _readme += "pinned: false\n"
         _readme += "---\n"
-        _readme = io.BytesIO(_readme.encode())
         return _readme
 
-    def create_space(self, team_id, submission_id, submission_repo, space_id):
-        server_space_id = space_id + "-server"
-        client_space_id = space_id + "-client"
-        space_auth_token = uuid.uuid4().hex
+    def create_space(self, team_id, submission_id, submission_repo, space_id, server_url, hardware):
         user_token = user_token_api.get(team_id)
 
         api = HfApi(token=self.token)
         params = {
+            "space_id": space_id,
+            "client_space_id": space_id,
             "competition_id": self.competition_id,
-            "competition_type": self.competition_type,
-            "metric": self.metric,
-            "token": self.token,
             "team_id": team_id,
             "submission_id": submission_id,
-            "submission_id_col": self.submission_id_col,
-            "submission_cols": self.submission_cols,
-            "submission_rows": self.submission_rows,
             "output_path": self.output_path,
             "submission_repo": submission_repo,
             "time_limit": self.time_limit,
             "dataset": self.dataset,
             "submission_filenames": self.submission_filenames,
         }
+        token_info_json = json.dumps(params, indent=4)
+        token_info_json_bytes = token_info_json.encode("utf-8")
+        token_info_json_buffer = io.BytesIO(token_info_json_bytes)
+
+        api = HfApi(token=self.token)
+        client_token = uuid.uuid4().hex + uuid.uuid4().hex
 
-        api.add_space_secret(repo_id=server_space_id, key="PARAMS", value=json.dumps(params))
-        api.add_space_secret(repo_id=server_space_id, key="HUGSIM_AUTH_TOKEN", value=space_auth_token)
-        api.add_space_secret(repo_id=server_space_id, key="HF_TOKEN", value=self.token)
-        readme = self._create_readme(space_id.split("/")[-1])
         api.upload_file(
-            path_or_fileobj=readme,
-            path_in_repo="README.md",
-            repo_id=server_space_id,
-            repo_type="space",
+            path_or_fileobj=token_info_json_buffer,
+            path_in_repo=f"token_data_info/{client_token}.json",
+            repo_id=self.competition_id,
+            repo_type="dataset",
         )
-        api.upload_file(
-            path_or_fileobj=io.BytesIO(_DOCKERFILE.encode()),
-            path_in_repo="Dockerfile",
+        api.create_repo(
             repo_id=space_id,
             repo_type="space",
+            space_sdk="docker",
+            space_hardware=hardware,
+            private=True,
         )
 
-        api.add_space_secret(repo_id=client_space_id, key="HUGSIM_AUTH_TOKEN", value=space_auth_token)
+        api.add_space_secret(repo_id=space_id, key="HUGSIM_API_TOKEN", value=client_token)
+        api.add_space_secret(repo_id=space_id, key="HUGSIM_SERVER_HOST", value=server_url)
+        client_code_local_dir = f"/tmp/data/client_repo/{space_id}"
+        client_commits = api.list_repo_commits(submission_repo, repo_type="model")
         api.snapshot_download(
             repo_id=submission_repo,
             repo_type="model",
-            revision="main",
+            revision=client_commits[0].commit_id,
             token=user_token,
-            local_dir="/tmp/data/user_repo",
+            local_dir=client_code_local_dir,
             allow_patterns=["*"],
         )
-        api.upload_folder(
-            repo_id=client_space_id,
-            repo_type="space",
-            folder_path="/tmp/data/user_repo",
-        )
+        with open(f"{client_code_local_dir}/README.md", "w", encoding="utf-8") as f:
+            f.write(self._create_readme(space_id))
+        try:
+            api.upload_folder(
+                repo_id=space_id,
+                repo_type="space",
+                folder_path=client_code_local_dir,
+            )
+        finally:
+            shutil.rmtree(client_code_local_dir, ignore_errors=True)
         self._queue_submission(team_id, submission_id)
 
     def run(self):
+        cur = 0
         while True:
-            pending_submissions = self.get_pending_subs()
+            time.sleep(5)
+            if cur == 10000:
+                cur = 0
+            cur += 1
+            all_submissions = self._get_all_submissions()
+
+            # Clean up spaces every 100 iterations
+            if cur % 100 == 0:
+                for space in all_submissions:
+                    if space["status"] == SubmissionStatus.QUEUED.value:
+                        space_cleaner.clean_space(
+                            space["space_id"],
+                            space["team_id"],
+                            space["submission_id"],
+                        )
+
+            pending_submissions = self._get_pending_subs(all_submissions)
             if pending_submissions is None:
-                time.sleep(5)
                 continue
-            if self.competition_type == "generic":
-                for _, row in pending_submissions.iterrows():
-                    team_id = row["team_id"]
-                    submission_id = row["submission_id"]
-                    submission_repo = row["submission_repo"]
-                    self.run_local(team_id, submission_id, submission_repo)
-            elif self.competition_type == "script":
-                for _, row in pending_submissions.iterrows():
-                    team_id = row["team_id"]
-                    submission_id = row["submission_id"]
-                    submission_repo = row["submission_repo"]
-                    space_id = row["space_id"]
-                    try:
-                        self.create_space(team_id, submission_id, submission_repo, space_id)
-                    except Exception as e:
-                        logger.error(
-                            f"Failed to create space for {team_id} {submission_id} {submission_repo} {space_id}: {e}"
-                        )
-                        # mark submission as failed
-                        self.mark_submission_failed(team_id, submission_id)
-                        logger.error(f"Marked submission {submission_id} as failed.")
-                        continue
-            time.sleep(5)
+            first_pending_sub = pending_submissions.iloc[0]
+            server_active_count = self._get_server_active_count(all_submissions)
+
+            if server_active_count[first_pending_sub["server_url"]] >= 1:
+                continue
+            try:
+                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(
+                    f"Failed to create space for {first_pending_sub['submission_id']}: {e}"
+                )
+                # mark submission as failed
+                self.mark_submission_failed(first_pending_sub['team_id'], first_pending_sub['submission_id'])
+                try:
+                    space_cleaner.delete_space(first_pending_sub["space_id"])
+                except Exception as e:
+                    logger.error(f"Failed to delete space {first_pending_sub['space_id']}: {e}")
+                logger.error(f"Marked submission {first_pending_sub['submission_id']} as failed.")
+                continue

competitions/submissions.py

@@ -5,12 +5,12 @@ from dataclasses import dataclass
 from datetime import datetime
 
 import pandas as pd
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, hf_hub_download, SpaceHardware
 from huggingface_hub.utils._errors import EntryNotFoundError
 
 from competitions.enums import SubmissionStatus
 from competitions.errors import AuthenticationError, PastDeadlineError, SubmissionError, SubmissionLimitError
-from competitions.utils import token_information, team_file_api, user_token_api
+from competitions.utils import token_information, team_file_api, user_token_api, server_manager
 
 
 @dataclass
@@ -86,6 +86,8 @@ class Submissions:
                 "selected": False,
                 "public_score": {},
                 "private_score": {},
+                "server_url": server_manager.get_next_server(),
+                "hardware": self.hardware,
             }
         )
         # count the number of times user has submitted today
@@ -179,6 +181,15 @@ class Submissions:
         return self._get_team_subs(team_id, private=private)
 
     def _create_submission(self, team_id: str):
+        api = HfApi(token=self.token)
+
+        if api.file_exists(
+            repo_id=self.competition_id,
+            filename=f"submission_info/{team_id}.json",
+            repo_type="dataset",
+        ):
+            return
+
         team_submission_info = {}
         team_submission_info["id"] = team_id
         team_submission_info["submissions"] = []
@@ -186,7 +197,6 @@ class Submissions:
         team_submission_info_json_bytes = team_submission_info_json.encode("utf-8")
         team_submission_info_json_buffer = io.BytesIO(team_submission_info_json_bytes)
 
-        api = HfApi(token=self.token)
         api.upload_file(
             path_or_fileobj=team_submission_info_json_buffer,
             path_in_repo=f"submission_info/{team_id}.json",
@@ -231,31 +241,11 @@ class Submissions:
             return self.submission_limit - submissions_made
 
         user_api = HfApi(token=user_token)
-        # submission_id is the sha of the submitted model repo + "__" + submission_id
         submission_id = user_api.model_info(repo_id=uploaded_file).sha + "__" + submission_id
         competition_organizer = self.competition_id.split("/")[0]
         space_id = f"{competition_organizer}/comp-{submission_id}"
-        server_space_id = space_id + "-server"
-        client_space_id = space_id + "-client"
 
-        api = HfApi(token=self.token)
-        api.create_repo(
-            repo_id=server_space_id,
-            repo_type="space",
-            space_sdk="docker",
-            space_hardware=self.hardware,
-            private=False,
-        )
-        api.create_repo(
-            repo_id=client_space_id,
-            repo_type="space",
-            space_sdk="docker",
-            space_hardware=self.hardware,
-            private=True,
-        )
         user_token_api.put(team_id, user_token)
-        
-        # api.add_space_secret(repo_id=space_id, key="USER_TOKEN", value=user_token)
         submissions_made = self._increment_submissions(
             team_id=team_id,
             user_id=user_id,

competitions/utils.py

@@ -7,13 +7,15 @@ import traceback
 import threading
 import uuid
 import base64
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
+from collections import defaultdict
 
 import requests
 from fastapi import Request
 from huggingface_hub import HfApi, hf_hub_download
 from loguru import logger
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from huggingface_hub import SpaceStage
 
 from competitions.enums import SubmissionStatus
 from competitions.params import EvalParams
@@ -521,3 +523,109 @@ user_token_api = UserTokenApi(
     os.environ.get("USER_TOKEN_KEY_BASE64"),
     os.environ.get("COMPETITION_ID")
 )
+
+
+class ServerManager:
+    def __init__(self, server_url_list: List[str]):
+        self.server_url_list = server_url_list
+        self._cur_index = 0
+        self._lock = threading.Lock()
+    
+    def get_next_server(self) -> str:
+        with self._lock:
+            server_url = self.server_url_list[self._cur_index]
+            self._cur_index = (self._cur_index + 1) % len(self.server_url_list)
+        return server_url
+    
+
+server_manager = ServerManager(["https://xdimlab-hugsim-web-server-0.hf.space"])
+
+
+class SubmissionApi:
+    def __init__(self, hf_token: str, competition_id: str):
+        self.hf_token = hf_token
+        self.competition_id = competition_id
+        self.api = HfApi(token=hf_token)
+    
+    def download_submission_info(self, team_id: str) -> Dict[str, Any]:
+        """
+        Download the submission info from Hugging Face Hub.
+        Args:
+            team_id (str): The team ID.
+        Returns:
+            Dict[str, Any]: The submission info.
+        """
+        submission_info_path = self.api.hf_hub_download(
+            repo_id=self.competition_id,
+            filename=f"submission_info/{team_id}.json",
+            repo_type="dataset",
+        )
+        with open(submission_info_path, 'r') as f:
+            submission_info = json.load(f)
+        
+        return submission_info
+
+    def upload_submission_info(self, team_id: str, user_submission_info: Dict[str, Any]):
+        user_submission_info_json = json.dumps(user_submission_info, indent=4)
+        user_submission_info_json_bytes = user_submission_info_json.encode("utf-8")
+        user_submission_info_json_buffer = io.BytesIO(user_submission_info_json_bytes)
+        self.api.upload_file(
+            path_or_fileobj=user_submission_info_json_buffer,
+            path_in_repo=f"submission_info/{team_id}.json",
+            repo_id=self.competition_id,
+            repo_type="dataset",
+        )
+
+    def update_submission_status(self, team_id: str, submission_id: str, status: int):
+        user_submission_info = self.download_submission_info(team_id)
+        for submission in user_submission_info["submissions"]:
+            if submission["submission_id"] == submission_id:
+                submission["status"] = status
+                break
+        self.upload_submission_info(team_id, user_submission_info)
+
+
+submission_api = SubmissionApi(
+    hf_token=os.environ.get("HF_TOKEN", None),
+    competition_id=os.environ.get("COMPETITION_ID")
+)
+
+
+class SpaceCleaner:
+    def __init__(self, hf_token: str):
+        self.hf_token = hf_token
+        self.api = HfApi(token=hf_token)
+        self.space_build_error_count = defaultdict(int)
+    
+    def delete_space(self, space_id: str):
+        """Delete a space by its ID."""
+        self.api.delete_repo(repo_id=space_id, repo_type="space")
+
+    def clean_space(self, space_id: str, team_id: str, submission_id: str):
+        space_info = self.api.space_info(repo_id=space_id)
+        if space_info.runtime.stage == SpaceStage.BUILD_ERROR:
+            self.space_build_error_count[space_id] += 1
+            if self.space_build_error_count[space_id] >= 3:
+                self.delete_space(space_id)
+                submission_api.update_submission_status(
+                    team_id=team_id,
+                    submission_id=submission_id,
+                    status=SubmissionStatus.FAILED.value
+                )
+            else:
+                self.api.restart_space(repo_id=space_id)
+            return
+        
+        if space_info.runtime.stage == SpaceStage.RUNTIME_ERROR:
+            self.delete_space(space_id)
+            submission_api.update_submission_status(
+                team_id=team_id,
+                submission_id=submission_id,
+                status=SubmissionStatus.FAILED.value
+            )
+            return
+
+
+space_cleaner = SpaceCleaner(
+    os.environ.get("HF_TOKEN", None),
+)