feat: udapte timeout feature

web_server.py

@@ -8,9 +8,10 @@ import enum
 import hugsim_env
 import subprocess as sp
 import shutil
+import time
 from collections import deque, OrderedDict
-from datetime import datetime
-from typing import Any, Dict, Optional, List
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional, List, Tuple
 from dataclasses import dataclass
 sys.path.append(os.getcwd())
 
@@ -150,6 +151,32 @@ def delete_client_space(client_space_id: str):
         print(f"Failed to delete space {client_space_id}. It may not exist or already deleted.")
 
 
+def get_env_id(team_id: str, submission_id: str) -> str:
+    """
+    Generate a unique environment ID based on team ID and submission ID.
+    Args:
+        team_id (str): The team ID.
+        submission_id (str): The submission ID.
+    Returns:
+        str: The unique environment ID.
+    """
+    return f"{team_id}____{submission_id}"
+
+
+def parse_env_id(env_id: str) -> Tuple[str, str]:
+    """
+    Parse the environment ID to extract team ID and submission ID.
+    Args:
+        env_id (str): The environment ID.
+    Returns:
+        Dict[str, str]: A dictionary containing team ID and submission ID.
+    """
+    parts = env_id.split('____')
+    if len(parts) != 2:
+        raise ValueError("Invalid environment ID format.")
+    return parts[0], parts[1]
+
+
 class FifoDict:
     def __init__(self, max_size: int):
         self.max_size = max_size
@@ -186,6 +213,8 @@ class EnvHandler:
     This can include multiple scene and configurations.
     """
     def __init__(self, scene_list: List[SceneConfig], base_output: str):
+        self._created_time = datetime.now(timezone.utc)
+        self._last_active_time = datetime.now(timezone.utc)
         self._lock = threading.Lock()
         self.scene_list = scene_list
         self.base_output = base_output
@@ -228,6 +257,7 @@ class EnvHandler:
         """
         Reset the environment and initialize variables.
         """
+        self._last_active_time = datetime.now(timezone.utc)
         self._log_list = deque(maxlen=100)
         self._done = False
         self._score_list = []
@@ -238,11 +268,30 @@ class EnvHandler:
         """
         Get the current state of the environment.
         """
+        self._last_active_time = datetime.now(timezone.utc)
         return {
             "obs": self._obs,
             "info": self._info,
         }
 
+    @property
+    def created_time(self) -> datetime:
+        """
+        Get the creation time of the environment handler.
+        Returns:
+            datetime: The creation time.
+        """
+        return self._created_time
+    
+    @property
+    def last_active_time(self) -> datetime:
+        """
+        Get the last active time of the environment handler.
+        Returns:
+            datetime: The last active time.
+        """
+        return self._last_active_time
+
     @property
     def has_done(self) -> bool:
         """
@@ -278,6 +327,7 @@ class EnvHandler:
         Returns:
             bool: True if the episode is done, False otherwise.
         """
+        self._last_active_time = datetime.now(timezone.utc)
         acc, steer_rate = traj2control(plan_traj, self._info)
         action = {'acc': acc, 'steer_rate': steer_rate}
         self._log("Executing action:", action)
@@ -350,8 +400,9 @@ class EnvHandlerManager:
     def __init__(self):
         self._env_handlers = {}
         self._lock = threading.Lock()
+        threading.Thread(target=self._clean_expired_env_handlers, daemon=True).start()
 
-    def _get_scene_list(self, env_id: str, base_output: str) -> List[SceneConfig]:
+    def _get_scene_list(self, base_output: str) -> List[SceneConfig]:
         """
         Load the scene configurations from the YAML files.
         Returns:
@@ -385,7 +436,7 @@ class EnvHandlerManager:
 
     def _generate_env_handler(self, env_id: str):
         base_output = "/app/app_datas/env_output"
-        scene_list = self._get_scene_list(env_id, base_output)
+        scene_list = self._get_scene_list(base_output)
         output = os.path.join(base_output, f"{env_id}_hugsim_env")
         os.makedirs(output, exist_ok=True)
         return EnvHandler(scene_list, base_output=output)
@@ -426,6 +477,27 @@ class EnvHandlerManager:
         if env is not None:
             env.close()
         torch.cuda.empty_cache()
+    
+    def _clean_expired_env_handlers(self):
+        """
+        Clean up expired environment handlers based on the last active time.
+        """
+        while 1:
+            try:
+                current_time = datetime.now(timezone.utc)
+                with self._lock:
+                    expired_env_ids = [
+                        env_id
+                        for env_id, handler in self._env_handlers.items()
+                        if handler and ((current_time - handler.created_time).total_seconds() > 3600 * 1.5 or (current_time - handler.last_active_time).total_seconds() > 180)
+                    ]
+                for env_id in expired_env_ids:
+                    self.close_env_handler(env_id)
+                    team_id, submission_id = parse_env_id(env_id)
+                    update_submission_data(team_id, submission_id, {"status": SubmissionStatus.FAILED.value, "error_message": "SPACE_TIMEOUT"})
+            except Exception as e:
+                print(f"Error in cleaning expired environment handlers: {e}")
+            time.sleep(15)
 
 
 app = FastAPI()
@@ -447,10 +519,12 @@ def _get_env_handler(
         raise HTTPException(status_code=401)
     
     submission_id = token_info["submission_id"]
-    if not env_manager.exists_env_handler(submission_id):
-        update_submission_data(token_info["team_id"], submission_id, {"status": SubmissionStatus.PROCESSING.value})
+    team_id = token_info["team_id"]
+    env_id = get_env_id(team_id, submission_id)
+    if not env_manager.exists_env_handler(env_id):
+        update_submission_data(team_id, submission_id, {"status": SubmissionStatus.PROCESSING.value})
 
-    env_handler = env_manager.get_env_handler(submission_id)
+    env_handler = env_manager.get_env_handler(env_id)
     if env_handler is None:
         raise HTTPException(status_code=404, detail="Environment handler already closed.")
     return env_handler
@@ -516,7 +590,7 @@ def execute_action_endpoint(
     execute_result = env_handler.execute_action(plan_traj)
     if execute_result.done:
         token_info = get_token_info(auth_token)
-        env_manager.close_env_handler(token_info["submission_id"])
+        env_manager.close_env_handler(get_env_id(token_info["team_id"], token_info["submission_id"]))
         delete_client_space(token_info["client_space_id"])
         final_score = env_handler.calculate_score()
         update_submission_data(token_info["team_id"], token_info["submission_id"], {"status": SubmissionStatus.SUCCESS.value, "score": final_score})