v1.0.0 prev; add serve, client, search_seeds
This commit is contained in:
wty-yy
122
robogauge/scripts/client.py
Normal file
122
robogauge/scripts/client.py
Normal file
@@ -0,0 +1,122 @@
|
||||
import requests
|
||||
import time
|
||||
import json
|
||||
import sys
|
||||
from typing import Dict, Any, Optional
|
||||
from robogauge.scripts.server import ResponseStatus
|
||||
|
||||
class RoboGaugeClient:
|
||||
def __init__(self, base_url: str = "http://127.0.0.1:9973"):
|
||||
self.base_url = base_url
|
||||
self.processing_ids = []
|
||||
self.task_id2info = {}
|
||||
self.response_data = {}
|
||||
|
||||
def submit_task(self,
|
||||
model_path: str,
|
||||
step: int,
|
||||
task_name: str,
|
||||
experiment_name: str,
|
||||
wait_for_server: bool = True,
|
||||
retry_interval: int = 2
|
||||
) -> Optional[str]:
|
||||
""" Submit stress pipeline evaluation
|
||||
Args:
|
||||
model_path (str): Torch script model path.
|
||||
step (int): Model training step.
|
||||
task_name (str): Registered task name.
|
||||
experiment_name (str): Experiment name for logging.
|
||||
wait_for_server (bool): If True, will keep retrying until the server is available.
|
||||
retry_interval (int): Seconds to wait before retrying connection.
|
||||
Returns:
|
||||
Optional[str]: Task ID if submission is successful, else None.
|
||||
"""
|
||||
payload = {
|
||||
"model_path": model_path,
|
||||
"step": step,
|
||||
"task_name": task_name,
|
||||
"experiment_name": experiment_name
|
||||
}
|
||||
|
||||
print(f"[RoboGaugeClient]📤 Preparing to submit task: {task_name}")
|
||||
|
||||
while True:
|
||||
try:
|
||||
response = requests.post(f"{self.base_url}/submit_eval", json=payload)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
task_id = data["task_id"]
|
||||
print(f"[RoboGaugeClient]✅ Submission successful! Task ID: {task_id}")
|
||||
self.processing_ids.append(task_id)
|
||||
self.task_id2info[task_id] = f"{payload['task_name']}_step{payload['step']}_{payload['experiment_name']}_{task_id}"
|
||||
return task_id
|
||||
else:
|
||||
print(f"[RoboGaugeClient]❌ Server returned error: {response.text}")
|
||||
return None
|
||||
|
||||
except requests.exceptions.ConnectionError:
|
||||
if not wait_for_server:
|
||||
print("[RoboGaugeClient]❌ Unable to connect to server (Connection Refused).")
|
||||
return None
|
||||
|
||||
print(f"[RoboGaugeClient]⏳ Server not responding, retrying in {retry_interval} seconds... (Make sure server.py is running)")
|
||||
time.sleep(retry_interval)
|
||||
except Exception as e:
|
||||
print(f"[RoboGaugeClient]❌ Unknown error: {e}")
|
||||
return None
|
||||
|
||||
def monitor_tasks(self):
|
||||
print("[RoboGaugeClient]⏱️ Monitoring submitted tasks...")
|
||||
""" Monitor all submitted tasks until completion. """
|
||||
for task_id in reversed(self.processing_ids):
|
||||
respone = requests.get(f"{self.base_url}/get_result/{task_id}")
|
||||
if respone.status_code != 200:
|
||||
continue
|
||||
resp_data = respone.json()
|
||||
status = resp_data['status']
|
||||
|
||||
if status in [ResponseStatus.PENDING, ResponseStatus.PROCESSING]:
|
||||
print(f"[RoboGaugeClient]⏳ Task {self.task_id2info[task_id]} is still {status}.")
|
||||
continue
|
||||
|
||||
if status == ResponseStatus.FINISHED:
|
||||
self.response_data[task_id] = resp_data
|
||||
print(f"[RoboGaugeClient]🎉 Task {self.task_id2info[task_id]} finished successfully!")
|
||||
elif status == ResponseStatus.ERROR:
|
||||
print(f"[RoboGaugeClient]❌ Task {self.task_id2info[task_id]} encountered an error: {resp_data.get('error_msg')}")
|
||||
elif status == ResponseStatus.NOT_FOUND:
|
||||
print(f"[RoboGaugeClient]❓ Task {self.task_id2info[task_id]} not found on server.")
|
||||
self.processing_ids.remove(task_id)
|
||||
|
||||
if __name__ == "__main__":
|
||||
SERVER_URL = "http://127.0.0.1:9973"
|
||||
client = RoboGaugeClient(base_url=SERVER_URL)
|
||||
|
||||
test_payload = {
|
||||
"model_path": "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt",
|
||||
"step": 124000,
|
||||
"task_name": "go2_moe",
|
||||
"experiment_name": "client_debug_001"
|
||||
}
|
||||
|
||||
print("="*40)
|
||||
print(" RoboGauge Client Demo")
|
||||
print("="*40)
|
||||
|
||||
task_id = client.submit_task(
|
||||
model_path=test_payload["model_path"],
|
||||
step=test_payload["step"],
|
||||
task_name=test_payload["task_name"],
|
||||
experiment_name=test_payload["experiment_name"],
|
||||
wait_for_server=True
|
||||
)
|
||||
|
||||
while True:
|
||||
client.monitor_tasks()
|
||||
for task_id, resp in client.response_data.items():
|
||||
scores = resp['results']['scores']
|
||||
print("[RoboGaugeClient]📊 Scores:")
|
||||
print(json.dumps(scores, indent=2, ensure_ascii=False))
|
||||
client.response_data.clear()
|
||||
time.sleep(5)
|
||||
149
robogauge/scripts/server.py
Normal file
149
robogauge/scripts/server.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user