feat: udapte web_server

web_server.py

@@ -255,10 +255,20 @@ class WebServer:
         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)
 
+    def register_env_handler(self, env_handler: EnvHandler):
+        """
+        Register an environment handler to the web server.
+        Args:
+            env_handler (EnvHandler): The environment handler to register.
+        """
+        self.env_handler = env_handler
+        self._ready = True
+
     def _reset_endpoint(self):
         self.env_handler.reset_env()
         return {"success": True}
@@ -312,6 +322,12 @@ class WebServer:
             </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:
@@ -323,14 +339,17 @@ class WebServer:
         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"])
 
 
-# TODO: add code to update submission info
-if __name__ == "__main__":
+def _register_env_handler_to_server(web_server: WebServer):
+    """
+    Register the environment handler to the web server.
+    Args:
+        web_server (WebServer): The web server instance.
+    """
     # Using fixed paths for web server
-    ad = "uniad"
     base_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_base.yaml')
-    # unknown config
     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')
@@ -354,10 +373,15 @@ if __name__ == "__main__":
     
     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 = WebServer(env_handler, auth_token=os.getenv('HUGSIM_AUTH_TOKEN'))
+    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()