update code

Dockerfile

@@ -15,7 +15,6 @@ ENV PATH /app/miniconda/bin:$PATH
 
 SHELL ["conda", "run","--no-capture-output", "-p","/app/env", "/bin/bash", "-c"]
 
-COPY --chown=1000:1000 ./test_server.py /app/test_server.py
 COPY --chown=1000:1000 ./web_server.py /app/web_server.py
 COPY --chown=1000:1000 ./docker/web_server_config/scene-0383-medium-00.yaml  /app/docker/web_server_config/scene-0383-medium-00.yaml
 COPY --chown=1000:1000 ./download_pre_datas.py /app/download_pre_datas.py
@@ -27,6 +26,10 @@ ENV TORCH_CUDA_ARCH_LIST "7.5"
 ENV http_proxy=
 ENV https_proxy=
 
+RUN ./.pixi/envs/default/bin/python3 -m ensurepip
+
+RUN ./.pixi/envs/default/bin/python3 -m pip install psutil
+
 RUN ./.pixi/envs/default/bin/python /app/download_pre_datas.py
 
 CMD ["./.pixi/envs/default/bin/python", "web_server.py"]

web_server.py

@@ -7,23 +7,26 @@ import time
 import io
 import enum
 import hugsim_env
+import subprocess as sp
 from collections import deque, OrderedDict
 from datetime import datetime
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 sys.path.append(os.getcwd())
 
-from fastapi import FastAPI, Body, Header, Depends, HTTPException
+from fastapi import FastAPI, Body, Header, Depends, HTTPException, Query
 from fastapi.responses import HTMLResponse, Response
 from omegaconf import OmegaConf
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi
 import open3d as o3d
 import numpy as np
 import gymnasium
 import uvicorn
+import psutil
 
 from sim.utils.sim_utils import traj2control, traj_transform_to_global
 from sim.utils.score_calculator import hugsim_evaluate
 
+ADMIN_TOKEN = os.getenv('ADMIN_TOKEN', None)
 HF_TOKEN = os.getenv('HF_TOKEN', None)
 COMPETITION_ID = os.getenv('COMPETITION_ID', None)
 
@@ -38,6 +41,38 @@ class SubmissionStatus(enum.Enum):
     FAILED = 4
 
 
+def get_gpu_memory():
+    output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
+    COMMAND = "nvidia-smi --query-gpu=memory.used --format=csv"
+    try:
+        memory_use_info = output_to_list(sp.check_output(COMMAND.split(),stderr=sp.STDOUT))[1:]
+    except sp.CalledProcessError as e:
+        raise RuntimeError("command '{}' return with error (code {}): {}".format(e.cmd, e.returncode, e.output))
+    memory_use_values = [int(x.split()[0]) for x in memory_use_info]
+    return memory_use_values
+
+
+def get_system_status():
+    cpu_percent = psutil.cpu_percent(interval=1)
+    cpu_count = psutil.cpu_count(logical=True)
+
+    virtual_mem = psutil.virtual_memory()
+    total_mem = virtual_mem.total / (1024 ** 3)
+    used_mem = virtual_mem.used / (1024 ** 3)
+    mem_percent = virtual_mem.percent
+
+    system_info = {
+        "cpu_percent": cpu_percent,
+        "cpu_count": cpu_count,
+        "total_memory_gb": round(total_mem, 2),
+        "used_memory_gb": round(used_mem, 2),
+        "memory_percent": mem_percent,
+        "gpus": get_gpu_memory(),
+    }
+
+    return system_info
+
+
 def get_token_info(token: str) -> Dict[str, Any]:
     token_info_path = hf_api.hf_hub_download(
         repo_id=COMPETITION_ID,
@@ -276,9 +311,15 @@ _result_dict= FifoDict(max_size=100)
 env_manager = EnvHandlerManager()
 
 
-def _get_env_handler(auth_token: str = Header(...)) -> EnvHandler:
+def _get_env_handler(
+    auth_token: Optional[str] = Header(None),
+    query_token: Optional[str] = Query(None)
+) -> EnvHandler:
+    token = auth_token or query_token
+    if not token:
+        raise HTTPException(status_code=401, detail="Authorization token is required.")
     try:
-        token_info = get_token_info(auth_token)
+        token_info = get_token_info(token)
     except Exception:
         raise HTTPException(status_code=401)
     return env_manager.get_env_handler(token_info["submission_id"])
@@ -376,13 +417,30 @@ def main_page_endpoint(env_handler: EnvHandler = Depends(_get_env_handler)):
 
 
 @app.get("/")
-def main_page_endpoint():
+def main_page_endpoint(
+    admin_token: Optional[str] = Query(None),
+):
     """
     Main page endpoint to display logs.
     """
+    if admin_token != ADMIN_TOKEN:
+        html_content = f"""
+            <html>
+            running
+            </html>
+        """
+        return HTMLResponse(content=html_content)
+    
+    system_info = get_system_status()
     html_content = f"""
         <html>
-        running
+        <head>
+            <title>System Status</title>
+        </head>
+        <body>
+            <h1>System Status</h1>
+            <pre>{json.dumps(system_info, indent=4)}</pre>
+        </body>
         </html>
     """
     return HTMLResponse(content=html_content)