v1.0.0 prev; add serve, client, search_seeds

2025-12-29 21:29:03 +08:00
parent feb9ddf907
commit c28bb64b0a
11 changed files with 319 additions and 27 deletions
--- a/robogauge/scripts/server.py
+++ b/robogauge/scripts/server.py
@@ -0,0 +1,149 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : server.py
+@Time    : 2025/12/29 10:58:05
+@Author  : wty-yy, Gemini3 Pro
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : Asynchronous stress pipeline evaluation server
+'''
+import os
+os.environ['MUJOCO_GL'] = 'glfw'  # avoid mujoco.Renderer EGL context error
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+
+import multiprocessing
+import uvicorn
+import queue
+import time
+import uuid
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Dict, Optional
+
+from dataclasses import dataclass
+from robogauge.utils.helpers import parse_args, class_to_dict
+from robogauge.tasks.pipeline.stress_pipeline import StressPipeline
+from pprint import pprint
+
+default_args_list = [
+    '--stress-benchmark',
+    # '--stress-terrain-names', 'flat', 'wave', 'slope', 'stairs_up', 'stairs_down', 'obstacle',
+    '--stress-terrain-names', 'flat', 'wave',
+    '--num-processes', '50',
+    '--seeds', '0', '1', '2',
+    '--search-seeds', '0', '1', '2', '3', '4',
+    '--frictions', '0.5', '0.75', '1.0', '1.25', '1.5', '1.75', '2.0', '2.25', '2.5',
+    '--compress-logs',
+    '--headless',
+]
+
+@dataclass
+class EvalTaskData:
+    model_path: str
+    step: int
+    task_name: str
+    experiment_name: str
+
+class EvalRequest(BaseModel):
+    model_path: str
+    step: int
+    task_name: str
+    experiment_name: str
+
+class ResponseStatus:
+    PENDING = "pending"
+    PROCESSING = "processing"
+    FINISHED = "finished"
+    ERROR = "error"
+    NOT_FOUND = "not_found"
+
+def run_api_server(input_queue, result_dict, port=9973):
+    """
+    Running in a separate subprocess.
+    I/O Process: submit requests -> put into queue -> return ID.
+    """
+    app = FastAPI()
+
+    @app.post("/submit_eval")
+    def submit_eval(req: EvalRequest):
+        task_id = str(uuid.uuid4())
+        task_data = EvalTaskData(
+            model_path=req.model_path,
+            step=req.step,
+            task_name=req.task_name,
+            experiment_name=req.experiment_name
+        )
+        input_queue.put((task_id, task_data))
+
+        result_dict[task_id] = {"status": ResponseStatus.PENDING}
+        return {"task_id": task_id, "message": "Queued"}
+
+    @app.get("/get_result/{task_id}")
+    def get_result(task_id: str):
+        if task_id not in result_dict:
+            return {"status": ResponseStatus.NOT_FOUND}
+        return result_dict[task_id]
+
+    print(f"📡 API Server listening on port {port}...")
+    uvicorn.run(app, host="127.0.0.1", port=port, log_level="error")
+
+def main():
+    print("🤖 RoboGauge Evaluation Server Starting...")
+    ctx = multiprocessing.get_context('spawn')
+    manager = ctx.Manager()
+    task_queue = manager.Queue()
+    results_store = manager.dict()
+
+    api_p = ctx.Process(
+        target=run_api_server, 
+        args=(task_queue, results_store),
+        daemon=True
+    )
+    api_p.start()
+
+    print("🚀 Main Process started. Waiting for tasks...")
+    print("   (StressPipeline will run directly in this Main Process)")
+
+    try:
+        while True:
+            try:
+                task_data: EvalTaskData
+                task_id, task_data = task_queue.get(timeout=1.0)
+                
+                print(f"\n🔄 [Main] Processing Task {task_id} (Step {task_data.step})...")
+                results_store[task_id] = {"status": ResponseStatus.PROCESSING}
+                
+                args_list = default_args_list.copy()
+                args_list += ['--model-path', task_data.model_path, '--task-name', task_data.task_name, '--experiment-name', task_data.experiment_name]
+                args = parse_args(args_list)
+
+                print(f"📋 Running with args:")
+                pprint(class_to_dict(args))
+                
+                pipeline = StressPipeline(args)
+                stress_results = pipeline.run() 
+
+                results_store[task_id] = {
+                    "status": ResponseStatus.FINISHED,
+                    "step": task_data.step,
+                    "results": stress_results
+                }
+                print(f"✅ [Main] Task {task_id} Finished.")
+
+            except queue.Empty:
+                continue
+            except Exception as e:
+                print(f"❌ [Main] Error: {e}")
+                import traceback
+                traceback.print_exc()
+                if 'task_id' in locals():
+                    results_store[task_id] = {"status": ResponseStatus.ERROR, "error": str(e), "error_msg": traceback.format_exc()}
+
+    except KeyboardInterrupt:
+        print("\n🛑 Shutting down...")
+        api_p.terminate()
+        api_p.join()
+
+if __name__ == "__main__":
+    main()