493 lines
19 KiB
Python
493 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
"""浏览器 MuJoCo 平台的本地 mjlab 训练桥接服务（仅绑定 loopback）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import signal
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from collections import deque
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime, timezone
|
||
from http import HTTPStatus
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from urllib.parse import unquote, urlsplit
|
||
|
||
VERSION = "0.1.0"
|
||
# 浏览器当前 ONNX 运行时只实现 Go2 的 47→12 部署契约；其他任务须由服务启动参数显式放行。
|
||
DEFAULT_TASKS = ("Unitree-Go2-Flat",)
|
||
ACTIVE_STATES = {"queued", "running"}
|
||
ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
|
||
ITERATION_PATTERNS = (
|
||
re.compile(r"(?:learning\s+)?iteration\D+(\d+)\s*/\s*(\d+)", re.I),
|
||
re.compile(r"(?:learning\s+)?iteration\D+(\d+)", re.I),
|
||
)
|
||
RUN_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")
|
||
LOCAL_ORIGIN = re.compile(r"^https?://(?:localhost|127\.0\.0\.1)(?::\d+)?$")
|
||
|
||
|
||
def now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
|
||
class ApiError(Exception):
|
||
def __init__(self, status: int, message: str):
|
||
super().__init__(message)
|
||
self.status = status
|
||
|
||
|
||
@dataclass
|
||
class TrainingConfig:
|
||
task_id: str
|
||
num_envs: int
|
||
max_iterations: int
|
||
seed: int
|
||
run_name: str
|
||
device: str
|
||
gpu_ids: list[int]
|
||
wandb_mode: str
|
||
|
||
|
||
@dataclass
|
||
class TrainingJob:
|
||
id: str
|
||
config: TrainingConfig
|
||
state: str = "queued"
|
||
created_at: str = field(default_factory=now_iso)
|
||
started_at: str | None = None
|
||
ended_at: str | None = None
|
||
iteration: int = 0
|
||
message: str = "等待本地训练进程启动"
|
||
logs: deque[str] = field(default_factory=lambda: deque(maxlen=200))
|
||
artifact: Path | None = None
|
||
process: subprocess.Popen[str] | None = None
|
||
cancel_requested: bool = False
|
||
|
||
def public(self) -> dict[str, Any]:
|
||
progress = min(1.0, max(0.0, self.iteration / self.config.max_iterations))
|
||
if self.state == "succeeded":
|
||
progress = 1.0
|
||
return {
|
||
"id": self.id,
|
||
"state": self.state,
|
||
"taskId": self.config.task_id,
|
||
"createdAt": self.created_at,
|
||
"startedAt": self.started_at,
|
||
"endedAt": self.ended_at,
|
||
"iteration": self.iteration,
|
||
"maxIterations": self.config.max_iterations,
|
||
"progress": progress,
|
||
"message": self.message,
|
||
"logs": list(self.logs),
|
||
"artifactReady": self.artifact is not None and self.artifact.is_file(),
|
||
"artifactName": self.artifact.name if self.artifact else None,
|
||
}
|
||
|
||
|
||
class TrainingManager:
|
||
def __init__(self, trainer_root: Path, python: str, tasks: tuple[str, ...], check_environment: bool = True):
|
||
self.trainer_root = trainer_root.expanduser().resolve()
|
||
self.python = str(Path(python).expanduser()) if os.sep in python else python
|
||
self.tasks = tasks
|
||
self.jobs: dict[str, TrainingJob] = {}
|
||
self.lock = threading.RLock()
|
||
self.check_environment = check_environment
|
||
self._environment_error: str | None | bool = False
|
||
|
||
def readiness_error(self) -> str | None:
|
||
if not self.trainer_root.is_dir():
|
||
return f"训练工程目录不存在：{self.trainer_root}"
|
||
if not (self.trainer_root / "scripts" / "train.py").is_file():
|
||
return f"训练入口不存在：{self.trainer_root / 'scripts/train.py'}"
|
||
executable = Path(self.python)
|
||
if not executable.is_file() and shutil.which(self.python) is None:
|
||
return f"Python 解释器不存在：{self.python}"
|
||
if self.check_environment and self._environment_error is False:
|
||
probe = "import importlib.util,sys; missing=[m for m in ('mjlab','torch','tyro') if importlib.util.find_spec(m) is None]; print(','.join(missing)); sys.exit(bool(missing))"
|
||
try:
|
||
result = subprocess.run([self.python, "-c", probe], cwd=self.trainer_root, capture_output=True, text=True, timeout=15, check=False)
|
||
missing = result.stdout.strip()
|
||
self._environment_error = f"训练 Python 缺少依赖：{missing}" if result.returncode else None
|
||
except (OSError, subprocess.TimeoutExpired) as error:
|
||
self._environment_error = f"无法检查训练 Python 环境：{error}"
|
||
return self._environment_error if isinstance(self._environment_error, str) else None
|
||
|
||
def active_job_id(self) -> str | None:
|
||
with self.lock:
|
||
return next((job.id for job in self.jobs.values() if job.state in ACTIVE_STATES), None)
|
||
|
||
def health(self) -> dict[str, Any]:
|
||
error = self.readiness_error()
|
||
return {
|
||
"version": VERSION,
|
||
"ready": error is None,
|
||
"trainerRoot": str(self.trainer_root),
|
||
"python": self.python,
|
||
"tasks": list(self.tasks),
|
||
"activeJobId": self.active_job_id(),
|
||
"error": error,
|
||
}
|
||
|
||
def parse_config(self, payload: Any) -> TrainingConfig:
|
||
if not isinstance(payload, dict):
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "请求体必须是 JSON 对象")
|
||
task_id = payload.get("taskId")
|
||
if task_id not in self.tasks:
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, f"不允许的训练任务：{task_id}")
|
||
|
||
def integer(name: str, minimum: int, maximum: int) -> int:
|
||
value = payload.get(name)
|
||
if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, f"{name} 必须在 {minimum}–{maximum} 之间")
|
||
return value
|
||
|
||
run_name = payload.get("runName", "web")
|
||
if not isinstance(run_name, str) or not RUN_NAME.fullmatch(run_name):
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "runName 只能包含字母、数字、点、下划线和连字符，最长 64 字符")
|
||
device = payload.get("device")
|
||
if device not in ("cpu", "gpu"):
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "device 必须是 cpu 或 gpu")
|
||
raw_gpu_ids = payload.get("gpuIds", [])
|
||
if not isinstance(raw_gpu_ids, list) or any(isinstance(value, bool) or not isinstance(value, int) or value < 0 or value > 255 for value in raw_gpu_ids):
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "gpuIds 必须是非负整数数组")
|
||
if device == "gpu" and not raw_gpu_ids:
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "GPU 训练至少需要一个 GPU 编号")
|
||
wandb_mode = payload.get("wandbMode", "offline")
|
||
if wandb_mode not in ("offline", "disabled", "online"):
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "wandbMode 必须是 offline、disabled 或 online")
|
||
return TrainingConfig(
|
||
task_id=task_id,
|
||
num_envs=integer("numEnvs", 1, 16384),
|
||
max_iterations=integer("maxIterations", 1, 1_000_000),
|
||
seed=integer("seed", 0, 2_147_483_647),
|
||
run_name=run_name,
|
||
device=device,
|
||
gpu_ids=raw_gpu_ids,
|
||
wandb_mode=wandb_mode,
|
||
)
|
||
|
||
def start(self, payload: Any) -> dict[str, Any]:
|
||
error = self.readiness_error()
|
||
if error:
|
||
raise ApiError(HTTPStatus.SERVICE_UNAVAILABLE, error)
|
||
config = self.parse_config(payload)
|
||
with self.lock:
|
||
if self.active_job_id():
|
||
raise ApiError(HTTPStatus.CONFLICT, "已有训练任务正在运行，请等待完成或先停止任务")
|
||
job = TrainingJob(id=uuid.uuid4().hex, config=config)
|
||
self.jobs[job.id] = job
|
||
threading.Thread(target=self._run, args=(job,), name=f"training-{job.id[:8]}", daemon=True).start()
|
||
return job.public()
|
||
|
||
def get(self, job_id: str) -> dict[str, Any]:
|
||
with self.lock:
|
||
job = self.jobs.get(job_id)
|
||
if not job:
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
|
||
return job.public()
|
||
|
||
def artifact(self, job_id: str) -> Path:
|
||
with self.lock:
|
||
job = self.jobs.get(job_id)
|
||
if not job:
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
|
||
if not job.artifact or not job.artifact.is_file():
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "该训练任务尚未生成 policy.onnx")
|
||
return job.artifact
|
||
|
||
def cancel(self, job_id: str) -> dict[str, Any]:
|
||
with self.lock:
|
||
job = self.jobs.get(job_id)
|
||
if not job:
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "训练任务不存在或服务已重启")
|
||
if job.state not in ACTIVE_STATES:
|
||
return job.public()
|
||
job.cancel_requested = True
|
||
job.message = "正在停止训练进程"
|
||
process = job.process
|
||
if process and process.poll() is None:
|
||
try:
|
||
os.killpg(process.pid, signal.SIGTERM)
|
||
except ProcessLookupError:
|
||
pass
|
||
threading.Thread(target=self._kill_later, args=(process,), daemon=True).start()
|
||
return self.get(job_id)
|
||
|
||
@staticmethod
|
||
def _kill_later(process: subprocess.Popen[str]) -> None:
|
||
try:
|
||
process.wait(timeout=5)
|
||
except subprocess.TimeoutExpired:
|
||
try:
|
||
os.killpg(process.pid, signal.SIGKILL)
|
||
except ProcessLookupError:
|
||
pass
|
||
|
||
def command_for(self, config: TrainingConfig) -> list[str]:
|
||
command = [
|
||
self.python, "-u", "scripts/train.py", config.task_id,
|
||
f"--env.scene.num-envs={config.num_envs}",
|
||
f"--agent.max-iterations={config.max_iterations}",
|
||
f"--agent.seed={config.seed}",
|
||
f"--agent.run-name={config.run_name}",
|
||
]
|
||
if config.device == "cpu":
|
||
command.extend(("--gpu-ids", "None"))
|
||
else:
|
||
# mjlab.TYRO_FLAGS 对 Union[list[int], Literal["all"], None] 使用 JSON 风格
|
||
# list token；传成多个独立参数会被解析为错误的 Union 分支。
|
||
command.extend(("--gpu-ids", json.dumps(config.gpu_ids, separators=(",", ":"))))
|
||
return command
|
||
|
||
def _update_from_log(self, job: TrainingJob, raw_line: str) -> None:
|
||
line = ANSI_ESCAPE.sub("", raw_line).strip()
|
||
if not line:
|
||
return
|
||
with self.lock:
|
||
job.logs.append(line[-4000:])
|
||
for pattern in ITERATION_PATTERNS:
|
||
match = pattern.search(line)
|
||
if match:
|
||
job.iteration = min(job.config.max_iterations, max(job.iteration, int(match.group(1))))
|
||
break
|
||
job.message = line[-240:]
|
||
|
||
def _artifact_snapshot(self) -> dict[Path, int]:
|
||
root = self.trainer_root / "logs" / "rsl_rl"
|
||
if not root.is_dir():
|
||
return {}
|
||
return {path: path.stat().st_mtime_ns for path in root.glob("**/policy.onnx") if path.is_file()}
|
||
|
||
def _find_artifact(self, before: dict[Path, int]) -> Path | None:
|
||
root = self.trainer_root / "logs" / "rsl_rl"
|
||
if not root.is_dir():
|
||
return None
|
||
changed = [path for path in root.glob("**/policy.onnx") if path.is_file() and before.get(path) != path.stat().st_mtime_ns]
|
||
return max(changed, key=lambda path: path.stat().st_mtime_ns) if changed else None
|
||
|
||
def _run(self, job: TrainingJob) -> None:
|
||
before = self._artifact_snapshot()
|
||
command = self.command_for(job.config)
|
||
with self.lock:
|
||
if job.cancel_requested:
|
||
job.state, job.ended_at, job.message = "cancelled", now_iso(), "训练已取消"
|
||
return
|
||
job.state, job.started_at, job.message = "running", now_iso(), "本地训练进程已启动"
|
||
try:
|
||
environment = os.environ.copy()
|
||
# 默认离线记录，保留本地 W&B 指标但不要求 API Key；只有前端明确选择
|
||
# online 时才允许 wandb 发起登录/联网。
|
||
environment["WANDB_MODE"] = job.config.wandb_mode
|
||
environment["WANDB_SILENT"] = "true"
|
||
process = subprocess.Popen(
|
||
command,
|
||
cwd=self.trainer_root,
|
||
env=environment,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
encoding="utf-8",
|
||
errors="replace",
|
||
bufsize=1,
|
||
start_new_session=True,
|
||
)
|
||
with self.lock:
|
||
job.process = process
|
||
assert process.stdout is not None
|
||
try:
|
||
for line in process.stdout:
|
||
self._update_from_log(job, line)
|
||
finally:
|
||
process.stdout.close()
|
||
return_code = process.wait()
|
||
artifact = self._find_artifact(before)
|
||
with self.lock:
|
||
job.process = None
|
||
job.ended_at = now_iso()
|
||
if job.cancel_requested:
|
||
job.state, job.message = "cancelled", "训练已由用户取消"
|
||
elif return_code != 0:
|
||
job.state, job.message = "failed", f"训练进程退出，返回码 {return_code}"
|
||
elif artifact is None:
|
||
job.state, job.message = "failed", "训练结束，但没有找到本次生成的 policy.onnx"
|
||
else:
|
||
job.state, job.artifact = "succeeded", artifact
|
||
job.iteration = job.config.max_iterations
|
||
job.message = f"训练完成：{artifact.relative_to(self.trainer_root)}"
|
||
except Exception as error: # 服务必须保留错误供前端诊断。
|
||
with self.lock:
|
||
job.process = None
|
||
job.ended_at = now_iso()
|
||
job.state = "cancelled" if job.cancel_requested else "failed"
|
||
job.message = f"启动训练失败：{error}"
|
||
job.logs.append(job.message)
|
||
|
||
|
||
class TrainingRequestHandler(BaseHTTPRequestHandler):
|
||
manager: TrainingManager
|
||
allowed_origins: tuple[str, ...] = ()
|
||
server_version = "MuJoCoLocalTraining/0.1"
|
||
|
||
def log_message(self, format: str, *args: Any) -> None:
|
||
sys.stderr.write(f"[{self.log_date_time_string()}] {format % args}\n")
|
||
|
||
def _origin_allowed(self) -> bool:
|
||
origin = self.headers.get("Origin")
|
||
return origin is None or bool(LOCAL_ORIGIN.fullmatch(origin)) or origin in self.allowed_origins
|
||
|
||
def _cors(self) -> None:
|
||
origin = self.headers.get("Origin")
|
||
if origin and self._origin_allowed():
|
||
self.send_header("Access-Control-Allow-Origin", origin)
|
||
self.send_header("Vary", "Origin")
|
||
|
||
def _json(self, status: int, payload: Any) -> None:
|
||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
self.send_response(status)
|
||
self._cors()
|
||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def _error(self, error: Exception) -> None:
|
||
if isinstance(error, ApiError):
|
||
self._json(error.status, {"error": str(error)})
|
||
else:
|
||
self._json(HTTPStatus.INTERNAL_SERVER_ERROR, {"error": f"本地训练服务内部错误：{error}"})
|
||
|
||
def _ensure_origin(self) -> None:
|
||
if not self._origin_allowed():
|
||
raise ApiError(HTTPStatus.FORBIDDEN, "不允许的浏览器来源")
|
||
|
||
def _payload(self) -> Any:
|
||
try:
|
||
length = int(self.headers.get("Content-Length", "0"))
|
||
except ValueError as error:
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "Content-Length 无效") from error
|
||
if length <= 0 or length > 32 * 1024:
|
||
raise ApiError(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "训练请求体不能为空且不能超过 32 KiB")
|
||
try:
|
||
return json.loads(self.rfile.read(length))
|
||
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
||
raise ApiError(HTTPStatus.BAD_REQUEST, "训练请求不是有效 JSON") from error
|
||
|
||
@staticmethod
|
||
def _route(path: str) -> tuple[str | None, bool]:
|
||
match = re.fullmatch(r"/api/training/jobs/([0-9a-f]{32})(/artifacts/policy\.onnx)?", path)
|
||
return (unquote(match.group(1)), bool(match.group(2))) if match else (None, False)
|
||
|
||
def do_OPTIONS(self) -> None:
|
||
try:
|
||
self._ensure_origin()
|
||
self.send_response(HTTPStatus.NO_CONTENT)
|
||
self._cors()
|
||
self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
|
||
self.send_header("Access-Control-Allow-Headers", "Content-Type")
|
||
self.send_header("Access-Control-Max-Age", "600")
|
||
self.end_headers()
|
||
except Exception as error:
|
||
self._error(error)
|
||
|
||
def do_GET(self) -> None:
|
||
try:
|
||
self._ensure_origin()
|
||
path = urlsplit(self.path).path
|
||
if path == "/api/training/health":
|
||
self._json(HTTPStatus.OK, self.manager.health())
|
||
return
|
||
job_id, artifact = self._route(path)
|
||
if not job_id:
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
|
||
if artifact:
|
||
file_path = self.manager.artifact(job_id)
|
||
size = file_path.stat().st_size
|
||
self.send_response(HTTPStatus.OK)
|
||
self._cors()
|
||
self.send_header("Content-Type", "application/octet-stream")
|
||
self.send_header("Content-Disposition", 'attachment; filename="policy.onnx"')
|
||
self.send_header("Content-Length", str(size))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.end_headers()
|
||
with file_path.open("rb") as source:
|
||
shutil.copyfileobj(source, self.wfile)
|
||
else:
|
||
self._json(HTTPStatus.OK, self.manager.get(job_id))
|
||
except Exception as error:
|
||
self._error(error)
|
||
|
||
def do_POST(self) -> None:
|
||
try:
|
||
self._ensure_origin()
|
||
if urlsplit(self.path).path != "/api/training/jobs":
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
|
||
self._json(HTTPStatus.ACCEPTED, self.manager.start(self._payload()))
|
||
except Exception as error:
|
||
self._error(error)
|
||
|
||
def do_DELETE(self) -> None:
|
||
try:
|
||
self._ensure_origin()
|
||
job_id, artifact = self._route(urlsplit(self.path).path)
|
||
if not job_id or artifact:
|
||
raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在")
|
||
self._json(HTTPStatus.ACCEPTED, self.manager.cancel(job_id))
|
||
except Exception as error:
|
||
self._error(error)
|
||
|
||
|
||
def default_trainer_root() -> Path:
|
||
configured = os.environ.get("UNITREE_RL_MJLAB_ROOT")
|
||
if configured:
|
||
return Path(configured)
|
||
repository = Path(__file__).resolve().parents[2]
|
||
return repository.parent.parent / "unitree_rl_mjlab"
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="MuJoCo Web 平台本地强化学习训练服务")
|
||
parser.add_argument("--host", default="127.0.0.1", choices=("127.0.0.1", "localhost"), help="仅允许绑定本机回环地址")
|
||
parser.add_argument("--port", type=int, default=8765)
|
||
parser.add_argument("--trainer-root", type=Path, default=default_trainer_root(), help="unitree_rl_mjlab 工程目录")
|
||
parser.add_argument("--trainer-python", default=sys.executable, help="已安装 mjlab/torch 的 Python 解释器")
|
||
parser.add_argument("--task", action="append", dest="tasks", help="允许前端启动的任务 ID；可重复")
|
||
parser.add_argument("--allow-origin", action="append", default=[], help="额外允许的前端 Origin；可重复")
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
manager = TrainingManager(args.trainer_root, args.trainer_python, tuple(args.tasks or DEFAULT_TASKS))
|
||
TrainingRequestHandler.manager = manager
|
||
TrainingRequestHandler.allowed_origins = tuple(args.allow_origin)
|
||
server = ThreadingHTTPServer((args.host, args.port), TrainingRequestHandler)
|
||
print(f"本地训练服务：http://{args.host}:{args.port}")
|
||
print(f"训练工程：{manager.trainer_root}")
|
||
print(f"Python：{manager.python}")
|
||
if manager.readiness_error():
|
||
print(f"警告：{manager.readiness_error()}", file=sys.stderr)
|
||
try:
|
||
server.serve_forever()
|
||
except KeyboardInterrupt:
|
||
print("\n正在停止本地训练服务…")
|
||
finally:
|
||
active = manager.active_job_id()
|
||
if active:
|
||
manager.cancel(active)
|
||
server.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|