#!/usr/bin/env python3 """浏览器 MuJoCo 平台的本地 mjlab 训练桥接服务（仅绑定 loopback）。""" from __future__ import annotations import argparse import hmac import json import os import re import secrets import shutil import signal import subprocess import sys import threading import uuid from collections import deque from contextlib import suppress from dataclasses import dataclass, field from datetime import UTC, datetime from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any from urllib.parse import parse_qs, unquote, urlsplit from pretrained_sources import PretrainedSources, SourceError from task_config import ( OBSTACLE_TASK, TaskConfigError, deployment_metadata, task_metadata, validate_task_config, ) from tuning.manager import TuningError, TuningManager from tuning.process import GpuLease, ResourceBusyError from tuning.schema import RewardConfigError, validate_configuration from tuning.scoring import EvaluationError VERSION = "0.4.0" MAX_REQUEST_BYTES = 128 * 1024 # Bounded full boxes-v1 payload (<=257 boxes). # Rough 可以训练，但其高度扫描 actor 不允许冒充浏览器 Flat 部署。 DEFAULT_TASKS = ("Unitree-Go2-Flat", "Unitree-Go2-Rough", OBSTACLE_TASK) ACTIVE_STATES = {"queued", "running"} MAX_JOBS = 20 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+)?$") LOCAL_HOST = re.compile(r"^(?:localhost|127\.0\.0\.1)(?::\d+)?$") def now_iso() -> str: return datetime.now(UTC).isoformat() def termination_signal_handler(_signum: int, _frame: Any) -> None: """将 SIGTERM 转成受控退出，使 main 的 finally 能清理训练子进程。""" raise KeyboardInterrupt 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 reward_config: dict[str, Any] | None = None terrain_preset: str | None = None terrain_params: dict[str, Any] = field(default_factory=dict) sensor_cfg: dict[str, Any] | None = None task_config: dict[str, Any] | None = None deployment: dict[str, Any] = field(default_factory=dict) pretrained: dict[str, Any] | None = None @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, "deployment": self.config.deployment, "pretrained": self.config.pretrained, "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, lease: GpuLease | None = None, sources: PretrainedSources | None = None, ): 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 self.lease = lease or GpuLease() self.preset_resolver: Any = None self.sources = sources 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','wandb') " "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), "pretrainedSources": self.sources.catalog() if self.sources else [], "pretrainedUpload": { "enabled": self.sources is not None, "templateId": "go2-legacy47-v1", "formats": {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}, "endpoint": "/api/training/pretrained-sources/upload", }, "taskMetadata": task_metadata(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 对象") allowed = { "taskId", "numEnvs", "maxIterations", "seed", "runName", "device", "gpuIds", "wandbMode", "rewardPresetId", "pretrainedSourceId", "terrainPreset", "terrainParams", "sensorCfg", "sensorType", "customTerrainBoxes", } if payload.keys() - allowed: raise ApiError(HTTPStatus.BAD_REQUEST, "请求包含未知字段（不接受配置路径/MJCF）") 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") preset_id = payload.get("rewardPresetId") reward_config = None if preset_id is not None: if task_id != "Unitree-Go2-Flat": raise ApiError(HTTPStatus.BAD_REQUEST, "自动调参奖励 preset 仅支持 Flat 任务") if not isinstance(preset_id, str) or not re.fullmatch(r"[0-9a-f]{32}", preset_id): raise ApiError(HTTPStatus.BAD_REQUEST, "rewardPresetId 格式无效") if self.preset_resolver is None: raise ApiError(HTTPStatus.BAD_REQUEST, "奖励 preset 服务未就绪") try: reward_config = validate_configuration( self.preset_resolver(preset_id, task_id), task_id ) except KeyError as error: raise ApiError(HTTPStatus.BAD_REQUEST, "奖励 preset 不存在") from error except RewardConfigError as error: raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error seed = integer("seed", 0, 2_147_483_647) try: custom = validate_task_config(task_id, payload, seed) except TaskConfigError as error: raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error pretrained = None if "pretrainedSourceId" in payload: if self.sources is None: raise ApiError( HTTPStatus.BAD_REQUEST, "服务尚未注册基础策略，请配置--pretrained-sources" ) try: pretrained = self.sources.bind(payload["pretrainedSourceId"], task_id, custom) except SourceError as error: raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error return TrainingConfig( pretrained=pretrained, task_id=task_id, terrain_preset=custom["terrainPreset"] if custom else None, terrain_params=custom["terrainParams"] if custom else {}, sensor_cfg=custom["sensorCfg"] if custom else None, task_config=custom, deployment=deployment_metadata(task_id, custom, seed), num_envs=integer("numEnvs", 1, 16384), max_iterations=integer("maxIterations", 1, 1_000_000), seed=seed, run_name=run_name, device=device, gpu_ids=raw_gpu_ids, wandb_mode=wandb_mode, reward_config=reward_config, ) 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, "已有训练任务正在运行，请等待完成或先停止任务") while len(self.jobs) >= MAX_JOBS: completed = next( (job_id for job_id, job in self.jobs.items() if job.state not in ACTIVE_STATES), None, ) if completed is None: raise ApiError(HTTPStatus.CONFLICT, "训练任务历史已满，请稍后重试") del self.jobs[completed] job = TrainingJob(id=uuid.uuid4().hex, config=config) owner = f"training:{job.id}" try: self.lease.acquire(owner) except ResourceBusyError as error: raise ApiError(HTTPStatus.CONFLICT, str(error)) from error self.jobs[job.id] = job try: threading.Thread( target=self._run, args=(job,), name=f"training-{job.id[:8]}", daemon=True ).start() except Exception: self.jobs.pop(job.id, None) self.lease.release(owner) raise 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: with suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGTERM) 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: with suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) def shutdown(self) -> None: """同步停止活动训练，避免服务退出后遗留子进程。""" active = self.active_job_id() if not active: return self.cancel(active) with self.lock: process = self.jobs[active].process if process and process.poll() is None: try: process.wait(timeout=6) except subprocess.TimeoutExpired: with suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) process.wait(timeout=2) def command_for( self, config: TrainingConfig, task_config_path: Path | None = None ) -> 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 task_config_path is not None: command.extend(("--task-config", str(task_config_path))) if config.pretrained is not None: if self.sources is None: raise SourceError("基础策略快照服务未配置，拒绝随机初始化") command.extend(self.sources.arguments(config.pretrained)) if config.reward_config is not None: command.extend( ( "--reward-config-json", json.dumps(config.reward_config, ensure_ascii=False, separators=(",", ":")), ) ) 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() environment = os.environ.copy() # 默认离线记录，保留本地 W&B 指标但不要求 API Key；只有前端明确选择 # online 时才允许 wandb 发起登录/联网。 environment["WANDB_MODE"] = job.config.wandb_mode environment["WANDB_SILENT"] = "true" try: config_path = None if job.config.task_config is not None: job_dir = self.trainer_root / "logs" / "rsl_rl" / "web_jobs" / job.id job_dir.mkdir(parents=True, exist_ok=True) config_path = job_dir / "training_config.json" config_path.write_text(json.dumps(job.config.task_config), encoding="utf-8") command = self.command_for(job.config, config_path) if config_path is not None: command.extend(("--output-dir", str(config_path.parent))) # Popen 与 process 登记必须和取消检查处于同一个临界区：cancel() 要么在 # 创建前标记取消，要么在创建后取得进程并终止，不能落入二者之间。 with self.lock: if job.cancel_requested: job.state, job.ended_at, job.message = "cancelled", now_iso(), "训练已取消" return 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, ) job.process = process job.state, job.started_at, job.message = ( "running", now_iso(), "本地训练进程已启动", ) 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) finally: self.lease.release(f"training:{job.id}") class TrainingRequestHandler(BaseHTTPRequestHandler): manager: TrainingManager tuning_manager: TuningManager allowed_origins: tuple[str, ...] = () access_token = "" server_version = "MuJoCoLocalTraining/0.4" 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 _host_allowed(self) -> bool: host = self.headers.get("Host", "") return bool(LOCAL_HOST.fullmatch(host)) def _authorized(self) -> bool: authorization = self.headers.get("Authorization", "") prefix = "Bearer " return authorization.startswith(prefix) and hmac.compare_digest( authorization[len(prefix) :], self.access_token ) 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)}) elif isinstance(error, KeyError): self._json(HTTPStatus.NOT_FOUND, {"error": "调参 session、trial 或 proposal 不存在"}) elif isinstance(error, ResourceBusyError): self._json(HTTPStatus.CONFLICT, {"error": str(error)}) elif isinstance(error, (TuningError, RewardConfigError, EvaluationError, SourceError)): self._json(HTTPStatus.BAD_REQUEST, {"error": str(error)}) else: self._json( HTTPStatus.INTERNAL_SERVER_ERROR, {"error": f"本地训练服务内部错误：{error}"} ) def _ensure_origin(self) -> None: if not self._host_allowed(): raise ApiError(HTTPStatus.FORBIDDEN, "不允许的 Host") if not self._origin_allowed(): raise ApiError(HTTPStatus.FORBIDDEN, "不允许的浏览器来源") def _ensure_request(self) -> None: self._ensure_origin() if not self._authorized(): raise ApiError(HTTPStatus.UNAUTHORIZED, "训练服务访问令牌无效") 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 > MAX_REQUEST_BYTES: raise ApiError( HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "训练请求体不能为空且不能超过 128 KiB" ) try: return json.loads(self.rfile.read(length)) except (UnicodeDecodeError, json.JSONDecodeError) as error: raise ApiError(HTTPStatus.BAD_REQUEST, "训练请求不是有效 JSON") from error def _upload(self): # This dedicated binary route must never pass through the 128KiB JSON reader. self.close_connection = True # no unread extra bytes may become a second request if self.manager.sources is None: raise ApiError(HTTPStatus.SERVICE_UNAVAILABLE, "基础策略上传服务不可用") if self.headers.get("Transfer-Encoding") or self.headers.get("Content-Encoding"): raise ApiError(HTTPStatus.BAD_REQUEST, "上传不接受chunked/压缩编码") lengths = self.headers.get_all("Content-Length", []) if len(lengths) != 1 or not re.fullmatch(r"[0-9]{1,10}", lengths[0]): raise ApiError(HTTPStatus.BAD_REQUEST, "上传必须提供唯一Content-Length") length = int(lengths[0]) if self.headers.get("Content-Type") != "application/octet-stream": raise ApiError(HTTPStatus.BAD_REQUEST, "上传Content-Type必须是application/octet-stream") query = parse_qs(urlsplit(self.path).query, keep_blank_values=True) if set(query) != {"format", "template", "name"} or any(len(v) != 1 for v in query.values()): raise ApiError(HTTPStatus.BAD_REQUEST, "上传必须显式提供format/template/name") fmt, template, name = (query[key][0] for key in ("format", "template", "name")) limit = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}.get(fmt) if limit is not None and (length == 0 or length > limit): raise ApiError(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "上传为空或超过格式大小上限") if len(name) > 255: raise ApiError(HTTPStatus.BAD_REQUEST, "上传显示名称过长") try: result = self.manager.sources.receive_upload( self.rfile, length, fmt, template, name, set_timeout=self.connection.settimeout, ) except OSError as error: raise ApiError( HTTPStatus.SERVICE_UNAVAILABLE, "上传存储写入失败，请检查本机磁盘空间/权限" ) from error self._json(HTTPStatus.CREATED, result) @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 _send_file(self, file_path: Path, filename: str) -> None: 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", f'attachment; filename="{filename}"') 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) 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, PUT, DELETE, OPTIONS") self.send_header("Access-Control-Allow-Headers", "Authorization, 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_request() parsed = urlsplit(self.path) path = parsed.path if path == "/api/training/health": health = self.manager.health() health["tuning"] = self.tuning_manager.capability() health["resourceOwner"] = self.manager.lease.public() self._json(HTTPStatus.OK, health) return if path == "/api/tuning/capabilities": self._json(HTTPStatus.OK, self.tuning_manager.capability()) return if path == "/api/tuning/sessions": self._json(HTTPStatus.OK, {"sessions": self.tuning_manager.list()}) return if path == "/api/tuning/presets": self._json(HTTPStatus.OK, {"presets": self.tuning_manager.storage.list_presets()}) return match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})", path) if match: self._json(HTTPStatus.OK, self.tuning_manager.detail(match.group(1))) return match = re.fullmatch( r"/api/tuning/sessions/([0-9a-f]{32})/trials/([0-9a-f]{32})/metrics", path ) if match: query = parse_qs(parsed.query) tags = [tag for value in query.get("tags", []) for tag in value.split(",") if tag] try: max_points = int(query.get("maxPoints", ["1000"])[0]) after_raw = query.get("afterStep", [None])[0] after_step = int(after_raw) if after_raw is not None else None except ValueError as error: raise TuningError("maxPoints/afterStep 必须是整数") from error if not 10 <= max_points <= 5000: raise TuningError("maxPoints 必须在 10–5000 之间") if after_step is not None and after_step < -1: raise TuningError("afterStep 不能小于 -1") self._json( HTTPStatus.OK, self.tuning_manager.metrics( match.group(1), match.group(2), tags or None, max_points, after_step ), ) return match = re.fullmatch( r"/api/tuning/sessions/([0-9a-f]{32})/artifacts/best/policy\.onnx", path ) if match: self._send_file(self.tuning_manager.best_artifact(match.group(1)), "policy.onnx") return job_id, artifact = self._route(path) if not job_id: raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在") if artifact: self._send_file(self.manager.artifact(job_id), "policy.onnx") 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_request() path = urlsplit(self.path).path if path == "/api/training/pretrained-sources/upload": self._upload() return if path == "/api/training/jobs": self._json(HTTPStatus.ACCEPTED, self.manager.start(self._payload())) return if path == "/api/tuning/agent/test": self._json(HTTPStatus.OK, self.tuning_manager.test_agent()) return if path == "/api/tuning/sessions": self._json(HTTPStatus.ACCEPTED, self.tuning_manager.create(self._payload())) return match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})/(pause|resume)", path) if match: action = ( self.tuning_manager.pause if match.group(2) == "pause" else self.tuning_manager.resume ) self._json(HTTPStatus.ACCEPTED, action(match.group(1))) return match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})/(step|rollback)", path) if match: action = ( self.tuning_manager.step if match.group(2) == "step" else self.tuning_manager.rollback ) self._json(HTTPStatus.ACCEPTED, action(match.group(1), self._payload())) return match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})/mode", path) if match: self._json( HTTPStatus.ACCEPTED, self.tuning_manager.set_mode(match.group(1), self._payload()), ) return match = re.fullmatch( r"/api/tuning/sessions/([0-9a-f]{32})/proposals/([0-9a-f]{32})/(approve|reject)", path, ) if match: action = ( self.tuning_manager.approve if match.group(3) == "approve" else self.tuning_manager.reject ) self._json( HTTPStatus.ACCEPTED, action(match.group(1), match.group(2), self._payload()) ) return raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在") except Exception as error: self._error(error) def do_PUT(self) -> None: try: self._ensure_request() path = urlsplit(self.path).path match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})/constraints", path) if match: self._json( HTTPStatus.ACCEPTED, self.tuning_manager.set_constraints(match.group(1), self._payload()), ) return raise ApiError(HTTPStatus.NOT_FOUND, "接口不存在") except Exception as error: self._error(error) def do_DELETE(self) -> None: try: self._ensure_request() path = urlsplit(self.path).path match = re.fullmatch(r"/api/tuning/sessions/([0-9a-f]{32})", path) if match: self._json(HTTPStatus.ACCEPTED, self.tuning_manager.cancel(match.group(1))) return job_id, artifact = self._route(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) return Path(__file__).resolve().parent / "rl" 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="训练工程目录；默认使用仓库内置的 Go2 训练器", ) parser.add_argument( "--trainer-python", default=sys.executable, help="已安装 mjlab/torch 的 Python 解释器" ) parser.add_argument( "--tuning-data-root", type=Path, default=None, help="调参 SQLite 与 trial 产物目录；默认位于训练工程 logs/auto_tuning", ) parser.add_argument( "--pretrained-sources", type=Path, default=None, help="可选旧管理员基础策略注册JSON；不配置也可直接上传单个.pt/.onnx", ) parser.add_argument( "--task", action="append", dest="tasks", help="允许前端启动的任务 ID；可重复" ) parser.add_argument( "--allow-origin", action="append", default=[], help="额外允许的前端 Origin；可重复" ) parser.add_argument( "--token", default=os.environ.get("MUJOCO_TRAINING_TOKEN"), help="访问令牌；默认随机生成，也可通过 MUJOCO_TRAINING_TOKEN 设置", ) return parser.parse_args() def main() -> None: args = parse_args() token = args.token or secrets.token_urlsafe(24) if len(token) < 16: raise SystemExit("训练服务访问令牌至少需要 16 个字符") lease = GpuLease() tuning_root = args.tuning_data_root or (Path(args.trainer_root) / "logs" / "auto_tuning") sources = PretrainedSources( args.pretrained_sources, tuning_root / "pretrained_sources", args.trainer_python, args.trainer_root, ) manager = TrainingManager( args.trainer_root, args.trainer_python, tuple(args.tasks or DEFAULT_TASKS), lease=lease, sources=sources, ) tuning_manager = TuningManager( args.trainer_root, args.trainer_python, tuning_root, lease, sources=sources ) manager.preset_resolver = tuning_manager.preset_config TrainingRequestHandler.manager = manager TrainingRequestHandler.tuning_manager = tuning_manager TrainingRequestHandler.allowed_origins = tuple(args.allow_origin) TrainingRequestHandler.access_token = token server = ThreadingHTTPServer((args.host, args.port), TrainingRequestHandler) print(f"本地训练服务：http://{args.host}:{args.port}") print(f"访问令牌：{token}") print(f"训练工程：{manager.trainer_root}") print(f"Python：{manager.python}") if manager.readiness_error(): print(f"警告：{manager.readiness_error()}", file=sys.stderr) signal.signal(signal.SIGTERM, termination_signal_handler) try: server.serve_forever() except KeyboardInterrupt: print("\n正在停止本地训练服务…") finally: tuning_manager.shutdown() manager.shutdown() server.server_close() if __name__ == "__main__": main()