Mujoco_WASM/training_server/server.py

#!/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 mobile_manipulator.config import MOBILE_TASKS, validate_mobile_params
from mobile_manipulator.config import TASK as MOBILE_CONTRACT
from mobile_manipulator.packages import MAX_UPLOAD, MobilePackages
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.6.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, *MOBILE_TASKS)
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
    mobile_package_id: str | None = None
    mobile_params: dict[str, Any] | None = None
    mobile_checkpoint: str | 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,
        mobile_python: str | None = None,
    ):
        self.trainer_root = trainer_root.expanduser().resolve()
        self.python = str(Path(python).expanduser()) if os.sep in python else python
        self.mobile_python = mobile_python or self.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
        self.mobile_packages = MobilePackages(self.trainer_root / "logs" / "mobile_packages")
        self._mobile_environment_error: str | None | bool = False

    def readiness_error(self, task_id: str | None = None) -> str | None:
        if task_id in MOBILE_TASKS:
            if not self.check_environment:
                return None
            if self._mobile_environment_error is False:
                try:
                    result = subprocess.run(
                        [
                            self.mobile_python,
                            "-c",
                            (
                                "import mujoco,gymnasium,torch,stable_baselines3,onnx,onnxruntime; "
                                "assert mujoco.__version__ == '3.11.0', "
                                "'移动操作需要 MuJoCo 3.11.0，请配置 --mobile-python；"
                                "不要升级 Go2 环境'"
                            ),
                        ],
                        capture_output=True,
                        text=True,
                        timeout=30,
                        check=False,
                    )
                    self._mobile_environment_error = (
                        "移动操作 Python 依赖不可用：" + result.stderr[-1500:]
                        if result.returncode
                        else None
                    )
                except (OSError, subprocess.TimeoutExpired) as error:
                    self._mobile_environment_error = f"无法检查移动操作环境：{error}"
            return self._mobile_environment_error or 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]:
        errors = {task: self.readiness_error(task) for task in self.tasks}
        ready = any(error is None for error in errors.values())
        error = None if ready else next(iter(errors.values()), "没有可用任务")
        metadata = task_metadata(self.tasks)
        for item in metadata:
            item.update(ready=errors[item["id"]] is None, error=errors[item["id"]])
        return {
            "version": VERSION,
            "ready": ready,
            "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": metadata,
            "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",
            "mobilePackageId",
            "mobileParams",
        }
        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)
        if task_id in MOBILE_TASKS:
            if any(
                k in payload
                for k in (
                    "terrainPreset",
                    "terrainParams",
                    "sensorCfg",
                    "sensorType",
                    "customTerrainBoxes",
                    "pretrainedSourceId",
                )
            ):
                raise ApiError(
                    HTTPStatus.BAD_REQUEST, "移动操作任务不接受 Go2 地形/传感器/预训练参数"
                )
            try:
                package_id = payload.get("mobilePackageId")
                package = self.mobile_packages.describe(package_id)
                if package["robotId"] != MOBILE_TASKS[task_id]:
                    raise ValueError("训练任务与场景机器人变体不匹配")
                params = validate_mobile_params(payload.get("mobileParams", {}))
                checkpoint = self.mobile_resume(params, task_id, package, package_id)
            except (ValueError, OSError) as error:
                raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error
            if device == "gpu" and len(raw_gpu_ids) != 1:
                raise ApiError(HTTPStatus.BAD_REQUEST, "移动操作 PPO 仅支持单个 GPU")
            return TrainingConfig(
                task_id=task_id,
                num_envs=integer("numEnvs", 1, 64),
                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,
                mobile_package_id=package_id,
                mobile_params=params,
                mobile_checkpoint=checkpoint,
                deployment={
                    "trainingStage": params["stage"],
                    "actionSemantics": MOBILE_CONTRACT["actionSemantics"],
                    "version": 1,
                    "browserCompatible": True,
                    "trainingTaskId": task_id,
                    "taskId": MOBILE_CONTRACT["id"],
                    **package,
                },
            )
        if "mobilePackageId" in payload or "mobileParams" in payload:
            raise ApiError(HTTPStatus.BAD_REQUEST, "Go2 任务不接受移动操作参数")
        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 mobile_resume(self, params, task_id, package, package_id) -> str | None:
        source_id = params.get("sourceJobId")
        if source_id is None:
            if params["stage"] != "navigate":
                raise ValueError("请先完成底盘接近训练，再选择通过评估的前一阶段作业")
            return None
        with self.lock:
            source = self.jobs.get(source_id)
            if (
                not source
                or source.state != "succeeded"
                or source.config.task_id != task_id
                or not source.artifact
            ):
                raise ValueError("接续训练需要同一机器人已完成的服务内作业")
            if source.config.mobile_package_id != package_id:
                raise ValueError("接续场景/资产快照不匹配；资产变化后请重新训练")
            deployment = source.config.deployment
            if (
                any(
                    deployment.get(key) != package.get(key)
                    for key in ("robotId", "sceneSha256", "robotConfigSha256")
                )
                or deployment.get("actionSemantics") != MOBILE_CONTRACT["actionSemantics"]
                or deployment.get("taskId") != MOBILE_CONTRACT["id"]
            ):
                raise ValueError("接续作业与当前场景/安全控制契约不匹配")
            stages = ("navigate", "reach", "pick-place")
            previous = deployment.get("trainingStage")
            if (
                previous not in stages
                or not stages.index(previous)
                <= stages.index(params["stage"])
                <= stages.index(previous) + 1
            ):
                raise ValueError("只支持同阶段续训或依次推进：底盘接近 → 末端接近 → 抓取放置")
            if previous != params["stage"]:
                if any(
                    params[key] != (source.config.mobile_params or {}).get(key)
                    for key in ("objectPosition", "goalPosition", "positionJitter")
                ):
                    raise ValueError(
                        "升级阶段必须保留已评估的初态分布；改变坐标/随机范围请先同阶段续训"
                    )
                evaluation = deployment.get("evaluation", {})
                if (
                    evaluation.get("episodes", 0) < 10
                    or evaluation.get("successRate", 0) < MOBILE_CONTRACT["navigationSuccessRate"]
                    or evaluation.get("safetyStops", 1) != 0
                ):
                    raise ValueError(
                        "上一阶段尚未达标：至少 10 回合独立评估、成功率 ≥80%、"
                        "无安全终止；请先同阶段续训"
                    )
            checkpoint = source.artifact.with_suffix(".ppo.zip")
            if not checkpoint.is_file():
                raise ValueError("接续作业缺少 PPO checkpoint")
            return str(checkpoint)

    def start(self, payload: Any) -> dict[str, Any]:
        config = self.parse_config(payload)
        error = self.readiness_error(config.task_id)
        if error:
            raise ApiError(HTTPStatus.SERVICE_UNAVAILABLE, error)
        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 deployment_artifact(self, job_id: str) -> Path:
        artifact = self.artifact(job_id).with_name("deployment.json")
        if not artifact.is_file():
            raise ApiError(HTTPStatus.NOT_FOUND, "该任务没有独立部署元数据")
        return 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]:
        if config.task_id in MOBILE_TASKS:
            return [
                self.mobile_python,
                "-u",
                "-m",
                "training_server.mobile_manipulator.train",
                "--package",
                str(self.mobile_packages.path(config.mobile_package_id)),
                "--iterations",
                str(config.max_iterations),
                "--num-envs",
                str(config.num_envs),
                "--seed",
                str(config.seed),
                "--device",
                "cpu" if config.device == "cpu" else f"cuda:{config.gpu_ids[0]}",
                "--params",
                json.dumps(config.mobile_params),
                "--task-id",
                config.task_id,
            ] + (["--resume", config.mobile_checkpoint] if config.mobile_checkpoint else [])
        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)
            mobile = job.config.task_id in MOBILE_TASKS
            if mobile:
                job_dir = self.trainer_root / "logs" / "rsl_rl" / "web_jobs" / job.id
                job_dir.mkdir(parents=True, exist_ok=True)
                command.extend(("--output", str(job_dir / "policy.onnx")))
                (job_dir / "training_config.json").write_text(
                    json.dumps(
                        {
                            "taskId": job.config.task_id,
                            "mobilePackageId": job.config.mobile_package_id,
                            "mobileParams": job.config.mobile_params,
                            "numEnvs": job.config.num_envs,
                            "maxIterations": job.config.max_iterations,
                            "seed": job.config.seed,
                            "device": job.config.device,
                            "gpuIds": job.config.gpu_ids,
                        }
                    ),
                    encoding="utf-8",
                )
            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=Path(__file__).resolve().parent.parent if mobile else 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 = (job_dir / "policy.onnx") if mobile else self._find_artifact(before)
            if artifact is not None and not artifact.is_file():
                artifact = None
            if mobile and return_code == 0 and not job.cancel_requested and artifact:
                deployment = json.loads((job_dir / "deployment.json").read_text())
                for key in (
                    "trainingTaskId",
                    "robotId",
                    "sceneSha256",
                    "robotConfigSha256",
                    "trainingStage",
                    "actionSemantics",
                ):
                    if deployment.get(key) != job.config.deployment.get(key):
                        raise ValueError(f"导出部署元数据不匹配：{key}")
                job.config.deployment = deployment
            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.6"

    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
            deployment_match = re.fullmatch(
                r"/api/training/jobs/([0-9a-f]{32})/artifacts/deployment\.json", path
            )
            if deployment_match:
                self._send_file(
                    self.manager.deployment_artifact(deployment_match.group(1)), "deployment.json"
                )
                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/mobile-packages":
                self.close_connection = True
                lengths = self.headers.get_all("Content-Length", [])
                if (
                    self.headers.get("Transfer-Encoding")
                    or self.headers.get("Content-Encoding")
                    or self.headers.get("Content-Type") != "application/zip"
                    or len(lengths) != 1
                    or not re.fullmatch(r"[0-9]{1,10}", lengths[0])
                ):
                    raise ApiError(
                        HTTPStatus.BAD_REQUEST, "场景上传需要 application/zip 和唯一 Content-Length"
                    )
                length = int(lengths[0])
                if not 0 < length <= MAX_UPLOAD:
                    raise ApiError(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "场景上传上限 128 MiB")
                self.connection.settimeout(60)
                try:
                    result = self.manager.mobile_packages.receive(self.rfile, length)
                except ValueError as error:
                    raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error
                self._json(HTTPStatus.CREATED, result)
                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(
        "--mobile-python",
        default=None,
        help="可选独立移动操作 Python（MuJoCo 3.11.0 + SB3），避免更改 Go2 环境",
    )
    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),
        mobile_python=args.mobile_python,
        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()