feat: udapte timeout feature

web_server.py

@@ -151,32 +151,6 @@ 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
@@ -399,6 +373,7 @@ class EnvHandler:
 class EnvHandlerManager:
     def __init__(self):
         self._env_handlers = {}
+        self._token_info_map = {}
         self._lock = threading.Lock()
         threading.Thread(target=self._clean_expired_env_handlers, daemon=True).start()
 
@@ -452,7 +427,7 @@ class EnvHandlerManager:
         with self._lock:
             return env_id in self._env_handlers
 
-    def get_env_handler(self, env_id: str) -> EnvHandler:
+    def get_env_handler(self, env_id: str, token_info: Dict[str, Any]) -> EnvHandler:
         """
         Get the environment handler for the given environment ID.
         Args:
@@ -463,6 +438,7 @@ class EnvHandlerManager:
         with self._lock:
             if env_id not in self._env_handlers:
                 self._env_handlers[env_id] = self._generate_env_handler(env_id)
+                self._token_info_map[env_id] = token_info
             return self._env_handlers[env_id]
 
     def close_env_handler(self, env_id: str):
@@ -493,8 +469,11 @@ class EnvHandlerManager:
                     ]
                 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"})
+                    token_info = self._token_info_map.pop(env_id, None)
+                    if token_info:
+                        update_submission_data(token_info["team_id"], token_info["submission_id"], {"status": SubmissionStatus.FAILED.value, "error_message": "SPACE_TIMEOUT"})
+                        delete_client_space(token_info["client_space_id"])
+
             except Exception as e:
                 print(f"Error in cleaning expired environment handlers: {e}")
             time.sleep(15)
@@ -520,11 +499,10 @@ def _get_env_handler(
     
     submission_id = token_info["submission_id"]
     team_id = token_info["team_id"]
-    env_id = get_env_id(team_id, submission_id)
-    if not env_manager.exists_env_handler(env_id):
+    if not env_manager.exists_env_handler(submission_id):
         update_submission_data(team_id, submission_id, {"status": SubmissionStatus.PROCESSING.value})
 
-    env_handler = env_manager.get_env_handler(env_id)
+    env_handler = env_manager.get_env_handler(submission_id, token_info)
     if env_handler is None:
         raise HTTPException(status_code=404, detail="Environment handler already closed.")
     return env_handler
@@ -590,7 +568,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(get_env_id(token_info["team_id"], token_info["submission_id"]))
+        env_manager.close_env_handler(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})