use new server

Dockerfile

@@ -31,4 +31,4 @@ RUN ./.pixi/envs/default/bin/python /app/download_pre_datas.py
 
 EXPOSE 7860
 
-CMD ["./.pixi/envs/default/bin/python", "test_server.py"]
+CMD ["./.pixi/envs/default/bin/python", "web_server.py"]

web_server.py

@@ -8,11 +8,11 @@ import io
 import enum
 import hugsim_env
 from collections import deque, OrderedDict
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import Any, Dict
 sys.path.append(os.getcwd())
 
-from fastapi import FastAPI, Body, Header, HTTPException, Depends
+from fastapi import FastAPI, Body, Header, Depends, HTTPException
 from fastapi.responses import HTMLResponse, Response
 from omegaconf import OmegaConf
 from huggingface_hub import HfApi, hf_hub_download
@@ -24,18 +24,10 @@ import uvicorn
 from sim.utils.sim_utils import traj2control, traj_transform_to_global
 from sim.utils.score_calculator import hugsim_evaluate
 
-IN_HUGGINGFACE_SPACE = os.getenv('IN_HUGGINGFACE_SPACE', 'false') == 'true'
-STOP_SPACE_TIMEOUT = int(os.getenv('STOP_SPACE_TIMEOUT', '7200'))
 HF_TOKEN = os.getenv('HF_TOKEN', None)
-SPACE_PARAMS = json.loads(os.getenv('PARAMS', '{}'))
-OUTPUT_DIR = "/app/app_datas/env_output"
+COMPETITION_ID = os.getenv('COMPETITION_ID', None)
 
-print("IN_HUGGINGFACE_SPACE:", IN_HUGGINGFACE_SPACE)
-print("STOP_SPACE_TIMEOUT:", STOP_SPACE_TIMEOUT)
-print("SPACE_PARAMS:", SPACE_PARAMS)
-
-class GlobalState:
-    done = False
+hf_api = HfApi(token=HF_TOKEN)
 
 
 class SubmissionStatus(enum.Enum):
@@ -46,7 +38,21 @@ class SubmissionStatus(enum.Enum):
     FAILED = 4
 
 
-def download_submission_info() -> Dict[str, Any]:
+def get_token_info(token: str) -> Dict[str, Any]:
+    token_info_path = hf_hub_download(
+        repo_id=COMPETITION_ID,
+        filename=f"token_data_info/{token}.json",
+        repo_type="dataset",
+        token=token
+    )
+    
+    with open(token_info_path, 'r') as f:
+        token_info = json.load(f)
+    
+    return token_info
+
+
+def download_submission_info(team_id: str) -> Dict[str, Any]:
     """
     Download the submission info from Hugging Face Hub.
     Args:
@@ -55,8 +61,8 @@ def download_submission_info() -> Dict[str, Any]:
         Dict[str, Any]: The submission info.
     """
     submission_info_path = hf_hub_download(
-        repo_id=SPACE_PARAMS["competition_id"],
-        filename=f"submission_info/{SPACE_PARAMS['team_id']}.json",
+        repo_id=COMPETITION_ID,
+        filename=f"submission_info/{team_id}.json",
         repo_type="dataset",
         token=HF_TOKEN
     )
@@ -66,69 +72,33 @@ def download_submission_info() -> Dict[str, Any]:
     return submission_info
 
 
-def upload_submission_info(user_submission_info: Dict[str, Any]):
+def upload_submission_info(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)
-    api = HfApi(token=HF_TOKEN)
-    api.upload_file(
+    hf_api.upload_file(
         path_or_fileobj=user_submission_info_json_buffer,
-        path_in_repo=f"submission_info/{SPACE_PARAMS['team_id']}.json",
-        repo_id=SPACE_PARAMS["competition_id"],
+        path_in_repo=f"submission_info/{team_id}.json",
+        repo_id=COMPETITION_ID,
         repo_type="dataset",
     )
 
 
-def update_submission_status(status):
-    user_submission_info = download_submission_info()
+def update_submission_status(team_id: str, submission_id: str, status: int):
+    user_submission_info = download_submission_info(team_id)
     for submission in user_submission_info["submissions"]:
-        if submission["submission_id"] == SPACE_PARAMS["submission_id"]:
+        if submission["submission_id"] == submission_id:
             submission["status"] = status
             break
-    upload_submission_info(user_submission_info)
+    upload_submission_info(team_id, user_submission_info)
 
 
-def auto_stop():
-    """
-    Automatically stop the server after a certain timeout.
-    """
-    stop_deadline = datetime.now() + timedelta(seconds=STOP_SPACE_TIMEOUT)
-    while 1:
-        if datetime.now() > stop_deadline:
-            update_submission_status(SubmissionStatus.FAILED.value)
-            break
-        if GlobalState.done:
-            update_submission_status(SubmissionStatus.SUCCESS.value)
-            break
-        time.sleep(60)
-
-    server_space_id = SPACE_PARAMS["server_space_id"]
-    client_space_id = SPACE_PARAMS["client_space_id"]
-    api = HfApi(token=HF_TOKEN)
-
-    if GlobalState.done:
-        api.upload_folder(
-            repo_id=SPACE_PARAMS["competition_id"],
-            folder_path=os.path.join(OUTPUT_DIR, "hugsim_env"),
-            repo_type="dataset",
-            path_in_repo=f"eval_results/{SPACE_PARAMS['submission_id']}",
-        )
-
-    api.delete_repo(
-        repo_id=server_space_id,
-        repo_type="space"
-    )
-    api.delete_repo(
+def delete_client_space(client_space_id: str):
+    hf_api.delete_repo(
         repo_id=client_space_id,
         repo_type="space"
     )
 
-if IN_HUGGINGFACE_SPACE:
-    # Start a thread to automatically stop the server after a timeout
-    auto_stop_thread = threading.Thread(target=auto_stop, daemon=True)
-    auto_stop_thread.start()
-    update_submission_status(SubmissionStatus.PROCESSING.value)
-
 
 class FifoDict:
     def __init__(self, max_size: int):
@@ -157,6 +127,13 @@ class EnvHandler:
         self._lock = threading.Lock()
         self.reset_env()
 
+    def close(self):
+        """
+        Close the environment and release resources.
+        """
+        self.env.close()
+        self._log("Environment closed.")
+
     def reset_env(self):
         """
         Reset the environment and initialize variables.
@@ -249,139 +226,149 @@ class EnvHandler:
             self._log_list.append(log_message)
 
 
-class WebServer:
-    def __init__(self, env_handler: EnvHandler, auth_token: str):
-        self.env_handler = env_handler
-        self.auth_token = auth_token
-        self._init_app()
-        self._result_dict= FifoDict(max_size=30)
-        self._ready = self.env_handler is not None
-    
-    def run(self):
-        uvicorn.run(self._app, host="0.0.0.0", port=7860, workers=1)
+class EnvHandlerManager:
+    def __init__(self):
+        self._env_handlers = {}
+        self._lock = threading.Lock()
+
+    def _generate_env_handler(self, env_id: str):
+        base_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_base.yaml')
+        scenario_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'scene-0383-medium-00.yaml')
+        camera_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_camera.yaml')
+        kinematic_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'kinematic.yaml')
+
+        scenario_config = OmegaConf.load(scenario_path)
+        base_config = OmegaConf.load(base_path)
+        camera_config = OmegaConf.load(camera_path)
+        kinematic_config = OmegaConf.load(kinematic_path)
+        cfg = OmegaConf.merge(
+            {"scenario": scenario_config},
+            {"base": base_config},
+            {"camera": camera_config},
+            {"kinematic": kinematic_config}
+        )
 
-    def register_env_handler(self, env_handler: EnvHandler):
+        model_path = os.path.join(cfg.base.model_base, cfg.scenario.scene_name)
+        model_config = OmegaConf.load(os.path.join(model_path, 'cfg.yaml'))
+        model_config.update({"model_path": "/app/app_datas/PAMI2024/release/ss/scenes/nuscenes/scene-0383"})
+        cfg.update(model_config)
+        cfg.base.output_dir = "/app/app_datas/env_output"
+        
+        output = os.path.join(cfg.base.output_dir, f"{env_id}_hugsim_env")
+        os.makedirs(output, exist_ok=True)
+        return EnvHandler(cfg, output)
+
+    def get_env_handler(self, env_id: str) -> EnvHandler:
         """
-        Register an environment handler to the web server.
+        Get the environment handler for the given environment ID.
         Args:
-            env_handler (EnvHandler): The environment handler to register.
+            env_id (str): The environment ID.
+        Returns:
+            EnvHandler: The environment handler instance.
         """
-        self.env_handler = env_handler
-        self._ready = True
+        with self._lock:
+            if env_id not in self._env_handlers:
+                self._env_handlers[env_id] = self._generate_env_handler(env_id)
+            return self._env_handlers[env_id]
 
-    def _reset_endpoint(self):
-        self.env_handler.reset_env()
-        return {"success": True}
 
-    def _get_current_state_endpoint(self):
-        state = self.env_handler.get_current_state()
-        return Response(content=pickle.dumps({"done": self.env_handler.has_done, "state": state}), media_type="application/octet-stream")
+app = FastAPI()
 
-    def _load_numpy_ndarray_json_str(self, json_str: str) -> np.ndarray:
-        """
-        Load a numpy ndarray from a JSON string.
-        """
-        data = json.loads(json_str)
-        return np.array(data["data"], dtype=data["dtype"]).reshape(data["shape"])
-
-    def _execute_action_endpoint(
-        self,
-        plan_traj: str = Body(..., embed=True),
-        transaction_id: str = Body(..., embed=True),
-    ):
-        cache_result = self._result_dict.get(transaction_id)
-        if cache_result is not None:
-            return Response(content=cache_result, media_type="application/octet-stream")
-
-        if self.env_handler.has_done:
-            result = pickle.dumps({"done": done, "state": None})
-            self._result_dict.push(transaction_id, result)
-            return Response(content=result, media_type="application/octet-stream")
-
-        plan_traj = self._load_numpy_ndarray_json_str(plan_traj)
-        done = self.env_handler.execute_action(plan_traj)
-        GlobalState.done = done
-        if done:
-            result = pickle.dumps({"done": done, "state": None})
-            self._result_dict.push(transaction_id, result)
-            return Response(content=result, media_type="application/octet-stream")
-        
-        state = self.env_handler.get_current_state()
-        result = pickle.dumps({"done": done, "state": state})
-        self._result_dict.push(transaction_id, result)
-        return Response(content=result, media_type="application/octet-stream")
+_result_dict= FifoDict(max_size=100)
+env_manager = EnvHandlerManager()
 
-    def _main_page_endpoint(self):
-        log_str = "\n".join(self.env_handler.log_list)
-        html_content = f"""
-            <html><body><pre>{log_str}</pre></body></html>
-            <script>
-                setTimeout(function() {{
-                    window.location.reload();
-                }}, 5000);
-            </script>
-        """
-        return HTMLResponse(content=html_content)
-    
-    def _ready_endpoint(self):
-        if self._ready:
-            return {"ready": True}
-        else:
-            raise HTTPException(status_code=503, detail="Server is not ready yet.")
-
-    def _verify_token(self, auth_token: str = Header(...)):
-        if self.auth_token and self.auth_token != auth_token:
-            raise HTTPException(status_code=401)
-
-    def _init_app(self):
-        self._app = FastAPI()
-        self._app.add_api_route("/reset", self._reset_endpoint, methods=["POST"], dependencies=[Depends(self._verify_token)])
-        self._app.add_api_route("/get_current_state", self._get_current_state_endpoint, methods=["GET"], dependencies=[Depends(self._verify_token)])
-        self._app.add_api_route("/execute_action", self._execute_action_endpoint, methods=["POST"], dependencies=[Depends(self._verify_token)])
-        self._app.add_api_route("/", self._main_page_endpoint, methods=["GET"])
-        self._app.add_api_route("/ready", self._ready_endpoint, methods=["GET"])
-
-
-def _register_env_handler_to_server(web_server: WebServer):
+
+def _get_env_handler(auth_token: str = Header(...)) -> EnvHandler:
+    try:
+        token_info = get_token_info(auth_token)
+    except Exception:
+        raise HTTPException(status_code=401)
+    return env_manager.get_env_handler(token_info["submission_id"])
+
+
+def _load_numpy_ndarray_json_str(json_str: str) -> np.ndarray:
+    """
+    Load a numpy ndarray from a JSON string.
+    """
+    data = json.loads(json_str)
+    return np.array(data["data"], dtype=data["dtype"]).reshape(data["shape"])
+
+
+@app.post("/reset")
+def reset_endpoint(env_handler: EnvHandler = Depends(_get_env_handler)):
+    """
+    Reset the environment.
+    """
+    env_handler.reset_env()
+    return {"success": True}
+
+
+@app.get("/get_current_state")
+def get_current_state_endpoint(env_handler: EnvHandler = Depends(_get_env_handler)):
+    """
+    Get the current state of the environment.
     """
-    Register the environment handler to the web server.
+    state = env_handler.get_current_state()
+    return Response(content=pickle.dumps(state), media_type="application/octet-stream")
+
+
+@app.post("/execute_action")
+def execute_action_endpoint(
+    plan_traj: str = Body(..., embed=True),
+    transaction_id: str = Body(..., embed=True),
+    auth_token: str = Header(...),
+    env_handler: EnvHandler = Depends(_get_env_handler)
+):
+    """
+    Execute the action based on the planned trajectory.
     Args:
-        web_server (WebServer): The web server instance.
+        plan_traj (str): The planned trajectory in JSON format.
+        transaction_id (str): The unique transaction ID for caching results.
+        env_handler (EnvHandler): The environment handler instance.
+    Returns:
+        Response: The response containing the execution result.
     """
-    # Using fixed paths for web server
-    base_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_base.yaml')
-    scenario_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'scene-0383-medium-00.yaml')
-    camera_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_camera.yaml')
-    kinematic_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'kinematic.yaml')
-
-    scenario_config = OmegaConf.load(scenario_path)
-    base_config = OmegaConf.load(base_path)
-    camera_config = OmegaConf.load(camera_path)
-    kinematic_config = OmegaConf.load(kinematic_path)
-    cfg = OmegaConf.merge(
-        {"scenario": scenario_config},
-        {"base": base_config},
-        {"camera": camera_config},
-        {"kinematic": kinematic_config}
-    )
+    cache_result = _result_dict.get(transaction_id)
+    if cache_result is not None:
+        return Response(content=cache_result, media_type="application/octet-stream")
+
+    if env_handler.has_done:
+        result = pickle.dumps({"done": done, "state": None})
+        _result_dict.push(transaction_id, result)
+        return Response(content=result, media_type="application/octet-stream")
 
-    model_path = os.path.join(cfg.base.model_base, cfg.scenario.scene_name)
-    model_config = OmegaConf.load(os.path.join(model_path, 'cfg.yaml'))
-    model_config.update({"model_path": "/app/app_datas/PAMI2024/release/ss/scenes/nuscenes/scene-0383"})
-    cfg.update(model_config)
-    cfg.base.output_dir = OUTPUT_DIR
+    plan_traj = _load_numpy_ndarray_json_str(plan_traj)
+    done = env_handler.execute_action(plan_traj)
+    if done:
+        token_info = get_token_info(auth_token)
+        env_manager.get_env_handler(token_info["submission_id"]).close()
+        delete_client_space(token_info["client_space_id"])
+        update_submission_status(token_info["team_id"], token_info["submission_id"], SubmissionStatus.SUCCESS.value)
+        result = pickle.dumps({"done": done, "state": None})
+        _result_dict.push(transaction_id, result)
+        return Response(content=result, media_type="application/octet-stream")
     
-    output = os.path.join(OUTPUT_DIR, "hugsim_env")
-    os.makedirs(output, exist_ok=True)
-    print("Output directory:", output)
-    env_handler = EnvHandler(cfg, output)
-    print("Environment handler initialized.")
-    web_server.register_env_handler(env_handler)
-
-
-if __name__ == "__main__":
-    # due to the limitation of huggingface space, we need to use a thread to register the environment handler.
-    web_server = WebServer(None, auth_token=os.getenv('HUGSIM_AUTH_TOKEN'))
-    print("Web server initialized.")
-    threading.Thread(target=_register_env_handler_to_server, args=(web_server,), daemon=True).start()
-    web_server.run()
+    state = env_handler.get_current_state()
+    result = pickle.dumps({"done": done, "state": state})
+    _result_dict.push(transaction_id, result)
+    return Response(content=result, media_type="application/octet-stream")
+
+
+@app.get("/")
+def main_page_endpoint(env_handler: EnvHandler = Depends(_get_env_handler)):
+    """
+    Main page endpoint to display logs.
+    """
+    log_str = "\n".join(env_handler.log_list)
+    html_content = f"""
+        <html><body><pre>{log_str}</pre></body></html>
+        <script>
+            setTimeout(function() {{
+                window.location.reload();
+            }}, 5000);
+        </script>
+    """
+    return HTMLResponse(content=html_content)
+
+
+uvicorn.run(app, port=7860, workers=1)