936 lines
39 KiB
Python
936 lines
39 KiB
Python
#!/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()
|