1234 lines
53 KiB
Python
1234 lines
53 KiB
Python
"""Persistent reward tuning session orchestrator."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import threading
|
||
from contextlib import suppress
|
||
from copy import deepcopy
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from pretrained_sources import PretrainedSources, SourceError
|
||
from task_config import TaskConfigError, validate_task_config
|
||
|
||
from . import obstacle_scoring
|
||
from .advisor import AdvisorUnavailable, DeepSeekAdvisor
|
||
from .process import GpuLease, ResourceBusyError, terminate_process
|
||
from .schema import (
|
||
FLAT_TASK,
|
||
OBSTACLE_TASK,
|
||
RewardConfigError,
|
||
base_configuration,
|
||
merge_proposal,
|
||
task_specs,
|
||
validate_configuration,
|
||
validate_configuration_constraints,
|
||
validate_constraints,
|
||
validate_proposal,
|
||
)
|
||
from .scoring import DEFAULT_OBJECTIVE_WEIGHTS, score_evaluation, validate_objective_weights
|
||
from .storage import StorageConflict, TuningStorage, now_iso
|
||
from .study import OptunaStudies
|
||
from .tensorboard import ingest_scalars
|
||
|
||
ACTIVE_SESSION_STATES = {
|
||
"queued",
|
||
"running",
|
||
"evaluating",
|
||
"awaiting_approval",
|
||
"paused",
|
||
"interrupted",
|
||
}
|
||
RUNNING_STATES = {"queued", "running", "evaluating"}
|
||
RUN_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")
|
||
|
||
|
||
class TuningError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class TuningManager:
|
||
def __init__(
|
||
self,
|
||
trainer_root: Path,
|
||
python: str,
|
||
data_root: Path,
|
||
lease: GpuLease,
|
||
advisor: Any | None = None,
|
||
sources: PretrainedSources | None = None,
|
||
):
|
||
self.trainer_root = trainer_root.expanduser().resolve()
|
||
self.python = python
|
||
self.data_root = data_root.expanduser().resolve()
|
||
self.data_root.mkdir(parents=True, exist_ok=True)
|
||
self.storage = TuningStorage(self.data_root / "tuning.sqlite3")
|
||
self.storage.recover_interrupted()
|
||
self.studies = OptunaStudies(self.data_root)
|
||
self.lease = lease
|
||
self.advisor = advisor or DeepSeekAdvisor()
|
||
self.lock = threading.RLock()
|
||
self.condition = threading.Condition(self.lock)
|
||
self.workers: dict[str, threading.Thread] = {}
|
||
self.processes: dict[str, subprocess.Popen[str]] = {}
|
||
self.cancel_events: dict[str, threading.Event] = {}
|
||
self.sources = sources
|
||
|
||
def capability(self) -> dict[str, Any]:
|
||
capability = self.advisor.capability()
|
||
capability.update(
|
||
{
|
||
"ready": (self.trainer_root / "scripts" / "evaluate.py").is_file(),
|
||
"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",
|
||
},
|
||
}
|
||
)
|
||
return capability
|
||
|
||
@staticmethod
|
||
def _integer(payload: dict, name: str, default: int, minimum: int, maximum: int) -> int:
|
||
value = payload.get(name, default)
|
||
if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
|
||
raise TuningError(f"{name} 必须在 {minimum}–{maximum} 之间")
|
||
return value
|
||
|
||
def parse_create(self, payload: Any) -> tuple[str, dict, dict, bool]:
|
||
if not isinstance(payload, dict):
|
||
raise TuningError("请求体必须是 JSON 对象")
|
||
mode = payload.get("mode", "automatic")
|
||
if mode not in ("automatic", "approval"):
|
||
raise TuningError("mode 必须是 automatic 或 approval")
|
||
task_id = payload.get("taskId", "Unitree-Go2-Flat")
|
||
task_specs(task_id)
|
||
allowed = {
|
||
"taskId",
|
||
"mode",
|
||
"runName",
|
||
"gpuIds",
|
||
"trialCount",
|
||
"initialIterations",
|
||
"middleIterations",
|
||
"finalIterations",
|
||
"numEnvs",
|
||
"seed",
|
||
"evalNumEnvs",
|
||
"evalSteps",
|
||
"earlyStopPatience",
|
||
"objectiveWeights",
|
||
"fallbackEnabled",
|
||
"terrainPreset",
|
||
"terrainParams",
|
||
"sensorType",
|
||
"sensorCfg",
|
||
"customTerrainBoxes",
|
||
"taskConfig",
|
||
"pretrainedSourceId",
|
||
}
|
||
if payload.keys() - allowed:
|
||
raise TuningError("未知session字段")
|
||
run_name = payload.get("runName", "auto-tune")
|
||
if not isinstance(run_name, str) or not RUN_NAME.fullmatch(run_name):
|
||
raise TuningError("runName 格式无效")
|
||
gpu_ids = payload.get("gpuIds", [0])
|
||
if (
|
||
not isinstance(gpu_ids, list)
|
||
or not gpu_ids
|
||
or any(
|
||
isinstance(value, bool) or not isinstance(value, int) or not 0 <= value <= 255
|
||
for value in gpu_ids
|
||
)
|
||
):
|
||
raise TuningError("gpuIds 必须是非空非负整数数组")
|
||
trial_count = self._integer(payload, "trialCount", 12, 1, 100)
|
||
rung0 = self._integer(payload, "initialIterations", 300, 1, 1000000)
|
||
rung1 = self._integer(payload, "middleIterations", 900, rung0, 1000000)
|
||
rung2 = self._integer(payload, "finalIterations", 2000, rung1, 1000000)
|
||
config = {
|
||
"taskId": task_id,
|
||
"numEnvs": self._integer(payload, "numEnvs", 4096, 1, 16384),
|
||
"seed": self._integer(payload, "seed", 42, 0, 2147483647),
|
||
"runName": run_name,
|
||
"gpuIds": gpu_ids,
|
||
"trialCount": trial_count,
|
||
"rungs": [rung0, rung1, rung2],
|
||
"promote": [trial_count, min(4, trial_count), min(2, trial_count)],
|
||
"evalNumEnvs": self._integer(payload, "evalNumEnvs", 256, 1, 4096),
|
||
"evalSteps": self._integer(payload, "evalSteps", 1000, 10, 100000),
|
||
"earlyStopPatience": self._integer(payload, "earlyStopPatience", 4, 1, 20),
|
||
}
|
||
if task_id == OBSTACLE_TASK:
|
||
if config["evalSteps"] != obstacle_scoring.STEPS:
|
||
raise TuningError("避障评估固定1000步，不能调整")
|
||
objective = deepcopy(obstacle_scoring.WEIGHTS)
|
||
if "objectiveWeights" in payload and payload["objectiveWeights"] != objective:
|
||
raise TuningError("避障评估权重固定")
|
||
task_payload = payload.get("taskConfig", payload)
|
||
if "taskConfig" in payload:
|
||
if not isinstance(task_payload, dict) or task_payload.keys() - {
|
||
"terrainPreset",
|
||
"terrainParams",
|
||
"sensorCfg",
|
||
"sensorType",
|
||
"seed",
|
||
"customTerrainBoxes",
|
||
}:
|
||
raise TuningError("taskConfig字段无效")
|
||
if any(
|
||
k in payload
|
||
for k in (
|
||
"terrainPreset",
|
||
"terrainParams",
|
||
"sensorCfg",
|
||
"sensorType",
|
||
"customTerrainBoxes",
|
||
)
|
||
):
|
||
raise TuningError("taskConfig不能与顶层场景配置混用")
|
||
task_seed = task_payload.get("seed", config["seed"])
|
||
if (
|
||
isinstance(task_seed, bool)
|
||
or not isinstance(task_seed, int)
|
||
or task_seed != config["seed"]
|
||
):
|
||
raise TuningError("taskConfig seed必须与session一致")
|
||
try:
|
||
config["taskConfig"] = validate_task_config(task_id, task_payload, config["seed"])
|
||
except TaskConfigError as error:
|
||
raise TuningError(str(error)) from error
|
||
base = base_configuration(task_id)
|
||
base["weights"]["avoidance_weight"] = config["taskConfig"]["sensorCfg"][
|
||
"avoidanceWeight"
|
||
]
|
||
validate_configuration_constraints(base, {}, task_id)
|
||
else:
|
||
if any(
|
||
k in payload
|
||
for k in (
|
||
"terrainPreset",
|
||
"terrainParams",
|
||
"sensorCfg",
|
||
"sensorType",
|
||
"customTerrainBoxes",
|
||
"taskConfig",
|
||
)
|
||
):
|
||
raise TuningError("Flat调参不接受避障场景字段")
|
||
objective = validate_objective_weights(
|
||
payload.get("objectiveWeights", DEFAULT_OBJECTIVE_WEIGHTS)
|
||
)
|
||
fallback = payload.get("fallbackEnabled", False)
|
||
if not isinstance(fallback, bool):
|
||
raise TuningError("fallbackEnabled 必须是布尔值")
|
||
if "pretrainedSourceId" in payload:
|
||
if self.sources is None:
|
||
raise TuningError("服务尚未注册基础策略，请配置--pretrained-sources")
|
||
try:
|
||
config["pretrained"] = self.sources.bind(
|
||
payload["pretrainedSourceId"], task_id, config.get("taskConfig")
|
||
)
|
||
except SourceError as error:
|
||
raise TuningError(str(error)) from error
|
||
return mode, config, objective, fallback
|
||
|
||
def create(self, payload: Any) -> dict:
|
||
mode, config, objective, fallback = self.parse_create(payload)
|
||
capability = self.capability()
|
||
if not capability["ready"]:
|
||
raise TuningError("评估入口未就绪")
|
||
if not capability["configured"] and not fallback:
|
||
raise TuningError("DeepSeek Agent 未配置；设置 DEEPSEEK_API_KEY 或显式启用 fallback")
|
||
with self.lock:
|
||
active = [
|
||
session
|
||
for session in self.storage.list_sessions()
|
||
if session["state"] in ACTIVE_SESSION_STATES
|
||
]
|
||
if active:
|
||
raise ResourceBusyError("已有调参 session 未结束")
|
||
session = self.storage.create_session(mode, config, objective, fallback)
|
||
self._start_worker(session["id"], resume=False)
|
||
return self.detail(session["id"])
|
||
|
||
def _start_worker(self, session_id: str, resume: bool) -> None:
|
||
cancel = threading.Event()
|
||
self.cancel_events[session_id] = cancel
|
||
worker = threading.Thread(
|
||
target=self._run_session,
|
||
args=(session_id, resume, cancel),
|
||
name=f"tuning-{session_id[:8]}",
|
||
daemon=True,
|
||
)
|
||
self.workers[session_id] = worker
|
||
worker.start()
|
||
|
||
def detail(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
session["trials"] = self.storage.list_trials(session_id)
|
||
session["proposals"] = self.storage.list_proposals(session_id)
|
||
session["audit"] = self.storage.audit_events(session_id)
|
||
control = self.storage.get_control(session_id)
|
||
current = next(
|
||
(trial for trial in session["trials"] if trial["id"] == session["currentTrialId"]),
|
||
None,
|
||
)
|
||
control["effectiveAfterCurrent"] = bool(
|
||
current and current["state"] in {"training", "evaluating"}
|
||
)
|
||
session["control"] = control
|
||
return session
|
||
|
||
def list(self) -> list[dict]:
|
||
return self.storage.list_sessions()
|
||
|
||
def _session_root(self, session_id: str) -> Path:
|
||
root = (self.data_root / "sessions" / session_id).resolve()
|
||
if not root.is_relative_to(self.data_root):
|
||
raise TuningError("非法 session 路径")
|
||
root.mkdir(parents=True, exist_ok=True)
|
||
return root
|
||
|
||
def _run_command(
|
||
self,
|
||
session_id: str,
|
||
command: list[str],
|
||
cwd: Path,
|
||
environment: dict[str, str],
|
||
log_path: Path,
|
||
) -> int:
|
||
owner = f"tuning:{session_id}"
|
||
self.lease.acquire(owner)
|
||
try:
|
||
process = subprocess.Popen(
|
||
command,
|
||
cwd=cwd,
|
||
env=environment,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
encoding="utf-8",
|
||
errors="replace",
|
||
bufsize=1,
|
||
start_new_session=True,
|
||
)
|
||
with self.lock:
|
||
self.processes[session_id] = process
|
||
assert process.stdout is not None
|
||
with log_path.open("a", encoding="utf-8") as log:
|
||
for line in process.stdout:
|
||
log.write(line)
|
||
log.flush()
|
||
if self.cancel_events[session_id].is_set():
|
||
terminate_process(process)
|
||
break
|
||
return process.wait()
|
||
finally:
|
||
with self.lock:
|
||
self.processes.pop(session_id, None)
|
||
self.lease.release(owner)
|
||
|
||
@staticmethod
|
||
def _latest_checkpoint(run_dir: Path) -> Path | None:
|
||
values = []
|
||
for path in run_dir.glob("model_*.pt"):
|
||
match = re.fullmatch(r"model_(\d+)\.pt", path.name)
|
||
if match:
|
||
values.append((int(match.group(1)), path))
|
||
return max(values, default=(0, None), key=lambda value: value[0])[1]
|
||
|
||
def _execute_trial(
|
||
self, session: dict, trial: dict, resume_checkpoint: Path | None = None
|
||
) -> dict:
|
||
session_id, trial_id = session["id"], trial["id"]
|
||
root = self._session_root(session_id)
|
||
run_dir = (root / trial["runDir"]).resolve()
|
||
if not run_dir.is_relative_to(root):
|
||
raise TuningError("trial 目录越界")
|
||
run_dir.mkdir(parents=True, exist_ok=True)
|
||
reward_path = run_dir / "reward_config.json"
|
||
reward_path.write_text(
|
||
json.dumps(trial["rewardConfig"], ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
||
)
|
||
config = session["config"]
|
||
command = [
|
||
self.python,
|
||
"-u",
|
||
"scripts/train.py",
|
||
config["taskId"],
|
||
f"--env.scene.num-envs={config['numEnvs']}",
|
||
f"--agent.max-iterations={trial['targetIterations']}",
|
||
f"--agent.seed={config['seed']}",
|
||
f"--agent.run-name={config['runName']}-t{trial['number']}-r{trial['rung']}",
|
||
"--agent.logger=tensorboard",
|
||
"--agent.upload-model=False",
|
||
"--gpu-ids",
|
||
json.dumps(config["gpuIds"], separators=(",", ":")),
|
||
"--output-dir",
|
||
str(run_dir),
|
||
"--reward-config",
|
||
str(reward_path),
|
||
]
|
||
task_path = None
|
||
if config["taskId"] == OBSTACLE_TASK:
|
||
task_path = run_dir / "task_config.json"
|
||
task_path.write_text(
|
||
json.dumps(config["taskConfig"], allow_nan=False), encoding="utf-8"
|
||
)
|
||
command.extend(("--task-config", str(task_path)))
|
||
pretrained = config.get("pretrained")
|
||
if pretrained is not None:
|
||
if self.sources is None:
|
||
raise TuningError("基础策略快照服务未配置，拒绝随机初始化")
|
||
self.sources.verify(pretrained)
|
||
if resume_checkpoint is not None:
|
||
if pretrained is not None:
|
||
origin_path = resume_checkpoint.parent / "initialization.json"
|
||
if (
|
||
not origin_path.is_file()
|
||
or origin_path.is_symlink()
|
||
or origin_path.stat().st_size > 64 * 1024
|
||
):
|
||
raise TuningError("续训checkpoint缺少有效基础策略来源记录，拒绝丢失来源")
|
||
origin = json.loads(origin_path.read_text(encoding="utf-8"))
|
||
if (
|
||
origin.get("source_id") != pretrained["sourceId"]
|
||
or origin.get("artifacts") != pretrained["manifest"]["artifacts"]
|
||
):
|
||
raise TuningError("续训checkpoint基础来源身份不匹配")
|
||
command.extend(("--resume-checkpoint", str(resume_checkpoint)))
|
||
elif pretrained is not None:
|
||
command.extend(self.sources.arguments(pretrained))
|
||
environment = os.environ.copy()
|
||
environment["WANDB_MODE"] = "disabled"
|
||
environment["WANDB_SILENT"] = "true"
|
||
self.storage.update_trial(
|
||
trial_id, state="training", started_at=now_iso(), message="正在训练"
|
||
)
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="running",
|
||
message=f"正在训练 trial {trial['number']} / rung {trial['rung']}",
|
||
)
|
||
return_code = self._run_command(
|
||
session_id, command, self.trainer_root, environment, run_dir / "train.log"
|
||
)
|
||
ingest_scalars(self.storage, trial_id, run_dir)
|
||
if self.cancel_events[session_id].is_set():
|
||
raise TuningError("session 已取消")
|
||
checkpoint = self._latest_checkpoint(run_dir)
|
||
policy = run_dir / "policy.onnx"
|
||
if return_code != 0 or checkpoint is None or not policy.is_file():
|
||
raise TuningError(f"训练失败（返回码 {return_code}）或缺少 checkpoint/policy.onnx")
|
||
# 暂停只关闭后续 Trial 调度门；已经开始的 Trial 必须连同固定评估一起
|
||
# 完成，避免把一次单步令牌错误地消耗在半个 Trial 上。
|
||
eval_output = run_dir / "evaluation.json"
|
||
eval_command = [
|
||
self.python,
|
||
"-u",
|
||
"scripts/evaluate.py",
|
||
config["taskId"],
|
||
"--checkpoint",
|
||
str(checkpoint),
|
||
"--output",
|
||
str(eval_output),
|
||
"--reward-config",
|
||
str(reward_path),
|
||
f"--num-envs={config['evalNumEnvs']}",
|
||
f"--steps-per-seed={config['evalSteps']}",
|
||
"--gpu-ids",
|
||
json.dumps(config["gpuIds"], separators=(",", ":")),
|
||
]
|
||
if task_path is not None:
|
||
eval_command.extend(("--task-config", str(task_path)))
|
||
self.storage.update_trial(trial_id, state="evaluating", message="正在固定协议评估")
|
||
self.storage.update_session(
|
||
session_id, state="evaluating", message=f"正在评估 trial {trial['number']}"
|
||
)
|
||
return_code = self._run_command(
|
||
session_id, eval_command, self.trainer_root, environment, run_dir / "evaluate.log"
|
||
)
|
||
ingest_scalars(self.storage, trial_id, run_dir / "evaluation-events")
|
||
if return_code != 0 or not eval_output.is_file():
|
||
raise TuningError(f"评估失败（返回码 {return_code}）")
|
||
evaluation = json.loads(eval_output.read_text(encoding="utf-8"))
|
||
baseline_trial = self.storage.list_trials(session_id)[0]
|
||
if config["taskId"] == OBSTACLE_TASK:
|
||
expected = obstacle_scoring.protocol(config["taskConfig"], config["evalNumEnvs"])
|
||
metrics = obstacle_scoring.validate_evaluation(evaluation, expected)
|
||
baseline = baseline_trial["evaluation"]
|
||
baseline_metrics = (
|
||
obstacle_scoring.validate_evaluation(baseline, expected) if baseline else None
|
||
)
|
||
scored = obstacle_scoring.score_evaluation(metrics, baseline_metrics)
|
||
elif baseline_trial["evaluation"] is None:
|
||
scored = {
|
||
"eligible": True,
|
||
"score": 0.0,
|
||
"components": {},
|
||
"metrics": evaluation["metrics"],
|
||
}
|
||
else:
|
||
scored = score_evaluation(
|
||
baseline_trial["evaluation"]["metrics"],
|
||
evaluation["metrics"],
|
||
session["objectiveWeights"],
|
||
)
|
||
evaluation["score"] = scored
|
||
rel_checkpoint = str(checkpoint.relative_to(root))
|
||
rel_policy = str(policy.relative_to(root))
|
||
self.storage.update_trial(
|
||
trial_id,
|
||
state="completed",
|
||
ended_at=now_iso(),
|
||
message="训练与评估完成",
|
||
checkpoint_path=rel_checkpoint,
|
||
policy_path=rel_policy,
|
||
evaluation=evaluation,
|
||
score=scored["score"],
|
||
eligible=scored["eligible"],
|
||
)
|
||
best = self._best_highest_rung(session_id)
|
||
if best is not None:
|
||
self.storage.update_session(session_id, best_trial_id=best["id"])
|
||
try:
|
||
optuna_number = self.studies.record(
|
||
session_id,
|
||
trial["rewardConfig"],
|
||
scored["score"],
|
||
scored["eligible"],
|
||
trial["rung"],
|
||
)
|
||
self.storage.audit(
|
||
session_id, "optuna_trial_recorded", {"trialId": trial_id, "number": optuna_number}
|
||
)
|
||
except Exception as error:
|
||
self.storage.audit(
|
||
session_id, "optuna_record_failed", {"trialId": trial_id, "error": str(error)[:500]}
|
||
)
|
||
return self.storage.get_trial(trial_id)
|
||
|
||
def _best(self, session_id: str, rung: int | None = None) -> dict | None:
|
||
trials = [
|
||
trial
|
||
for trial in self.storage.list_trials(session_id)
|
||
if trial["state"] == "completed" and trial["eligible"]
|
||
]
|
||
if rung is not None:
|
||
trials = [trial for trial in trials if trial["rung"] == rung]
|
||
return max(
|
||
trials,
|
||
key=lambda trial: trial["score"] if trial["score"] is not None else -999,
|
||
default=None,
|
||
)
|
||
|
||
def _best_highest_rung(self, session_id: str) -> dict | None:
|
||
for rung in (2, 1, 0):
|
||
best = self._best(session_id, rung=rung)
|
||
if best is not None:
|
||
return best
|
||
return None
|
||
|
||
def _proposal_context(self, session: dict) -> dict:
|
||
trials = self.storage.list_trials(session["id"])[-12:]
|
||
rejected_feedback = [
|
||
proposal["feedback"]
|
||
for proposal in self.storage.list_proposals(session["id"])
|
||
if proposal["state"] == "rejected" and proposal["feedback"]
|
||
][-4:]
|
||
return {
|
||
"task": session["config"]["taskId"],
|
||
"objectiveWeights": session["objectiveWeights"],
|
||
"allowlist": {
|
||
section: {
|
||
key: {"min": spec.minimum, "max": spec.maximum} for key, spec in specs.items()
|
||
}
|
||
for section, specs in zip(
|
||
("weights", "params"), task_specs(session["config"]["taskId"]), strict=True
|
||
)
|
||
},
|
||
"taskContext": (
|
||
(
|
||
"97维目标导航：48条前视ray，3x16层pitch=[0,-20,-45]deg；仍有侧后/层间/坑盲区。"
|
||
if session["config"]["taskConfig"]["sensorCfg"].get("sensorMode")
|
||
== "multi_ring_raycast"
|
||
else "81维目标导航：32条前视ray仅单高度切片，存在侧后方/矮障碍/跌落盲区。"
|
||
)
|
||
+ "避障权重对应近障平方惩罚，collision_penalty为非足端>10N接触惩罚；"
|
||
"target_velocity是真实导航command速度而非reward权重。关注擦碰、绕行、目标到达和动作抖动。"
|
||
"评估固定seed/1000步/客观权重，不允许修改地图、传感器、起终点、协议或以训练奖励代替指标。"
|
||
)
|
||
if session["config"]["taskId"] == OBSTACLE_TASK
|
||
else "Flat速度跟踪与步态稳定；保留原六目标评估。",
|
||
"taskConfig": session["config"].get("taskConfig"),
|
||
"parameterConstraints": self.storage.get_control(session["id"])["constraints"],
|
||
"rejectedFeedback": rejected_feedback,
|
||
"trials": [
|
||
{
|
||
"number": t["number"],
|
||
"rung": t["rung"],
|
||
"score": t["score"],
|
||
"eligible": t["eligible"],
|
||
"rewardConfig": t["rewardConfig"],
|
||
"evaluation": t["evaluation"] and t["evaluation"].get("metrics"),
|
||
}
|
||
for t in trials
|
||
],
|
||
}
|
||
|
||
def _fallback_patch(self, previous: dict, index: int, task_id="Unitree-Go2-Flat") -> dict:
|
||
names = ("track_linear_velocity", "action_rate_l2", "body_orientation_l2", "foot_slip")
|
||
if task_id == OBSTACLE_TASK:
|
||
names = ("avoidance_weight", "collision_penalty")
|
||
name = names[index % len(names)]
|
||
old = previous["weights"][name]
|
||
factor = 1.1 if index % 2 == 0 else 0.9
|
||
return validate_proposal({"weights": {name: old * factor}}, previous, task_id=task_id)
|
||
|
||
def _request_proposal(
|
||
self, session: dict, previous: dict, base_trial_id: str, index: int
|
||
) -> dict:
|
||
try:
|
||
result = self.advisor.propose(self._proposal_context(session), previous)
|
||
source = "agent"
|
||
except Exception as error:
|
||
if not session["fallbackEnabled"]:
|
||
raise AdvisorUnavailable(str(error)) from error
|
||
result = {
|
||
"patch": self._fallback_patch(previous, index, session["config"]["taskId"]),
|
||
"rationale": f"Agent 不可用，显式 fallback：{error}",
|
||
"expectedImpact": {},
|
||
"confidence": 0.2,
|
||
"promptHash": None,
|
||
"usage": {},
|
||
"model": "optuna-fallback",
|
||
}
|
||
source = "fallback"
|
||
constraints = self.storage.get_control(session["id"])["constraints"]
|
||
result["patch"] = validate_proposal(
|
||
result["patch"], previous, constraints, session["config"]["taskId"]
|
||
)
|
||
proposal = self.storage.create_proposal(
|
||
session["id"],
|
||
base_trial_id,
|
||
result["patch"],
|
||
result["rationale"],
|
||
result.get("expectedImpact", {}),
|
||
result["confidence"],
|
||
source,
|
||
)
|
||
self.storage.audit(
|
||
session["id"],
|
||
"proposal_created",
|
||
{
|
||
"proposalId": proposal["id"],
|
||
"source": source,
|
||
"promptHash": result.get("promptHash"),
|
||
"usage": result.get("usage", {}),
|
||
"model": result.get("model"),
|
||
},
|
||
)
|
||
return proposal
|
||
|
||
def _wait_for_approval(
|
||
self, session_id: str, proposal_id: str, cancel: threading.Event
|
||
) -> dict:
|
||
with self.condition:
|
||
while not cancel.is_set():
|
||
proposal = self.storage.get_proposal(proposal_id)
|
||
if proposal["state"] != "pending":
|
||
return proposal
|
||
self.condition.wait(timeout=1.0)
|
||
raise TuningError("session 已取消")
|
||
|
||
@staticmethod
|
||
def _base_configuration(session):
|
||
base = base_configuration(session["config"]["taskId"])
|
||
if session["config"]["taskId"] == OBSTACLE_TASK:
|
||
base["weights"]["avoidance_weight"] = session["config"]["taskConfig"]["sensorCfg"][
|
||
"avoidanceWeight"
|
||
]
|
||
return base
|
||
|
||
def _run_session(self, session_id: str, resume: bool, cancel: threading.Event) -> None:
|
||
try:
|
||
session = self.storage.get_session(session_id)
|
||
trials = self.storage.list_trials(session_id)
|
||
if resume:
|
||
for proposal in self.storage.list_proposals(session_id):
|
||
if proposal["state"] == "pending" and self.storage.decide_proposal(
|
||
proposal["id"],
|
||
"rejected",
|
||
"service_restart/recovery_invalidated：服务重启，需重新提案并审批（非用户拒绝）",
|
||
):
|
||
self.storage.audit(
|
||
session_id,
|
||
"recovery_proposal_invalidated",
|
||
{
|
||
"proposalId": proposal["id"],
|
||
"reason": "service_restart/recovery_invalidated",
|
||
},
|
||
)
|
||
root = self._session_root(session_id)
|
||
for interrupted in [trial for trial in trials if trial["state"] == "interrupted"]:
|
||
run_dir = (root / interrupted["runDir"]).resolve()
|
||
if run_dir.is_relative_to(root):
|
||
shutil.rmtree(run_dir, ignore_errors=True)
|
||
self.storage.delete_trial(interrupted["id"])
|
||
self.storage.audit(
|
||
session_id,
|
||
"session_resumed",
|
||
{
|
||
"discardedInterruptedTrials": [
|
||
t["id"] for t in trials if t["state"] == "interrupted"
|
||
]
|
||
},
|
||
)
|
||
trials = self.storage.list_trials(session_id)
|
||
if not trials:
|
||
baseline_dir = "trial-000-rung-0"
|
||
trial = self.storage.create_trial(
|
||
session_id,
|
||
0,
|
||
0,
|
||
session["config"]["rungs"][0],
|
||
self._base_configuration(session),
|
||
None,
|
||
baseline_dir,
|
||
)
|
||
self._execute_trial(session, trial)
|
||
session = self.storage.get_session(session_id)
|
||
completed_rung0 = [
|
||
t
|
||
for t in self.storage.list_trials(session_id)
|
||
if t["rung"] == 0 and t["state"] == "completed"
|
||
]
|
||
next_number = len({t["number"] for t in completed_rung0})
|
||
best_score = max(
|
||
(t["score"] if t["score"] is not None else float("-inf") for t in completed_rung0),
|
||
default=0.0,
|
||
)
|
||
no_improve = 0
|
||
while (
|
||
next_number < session["config"]["trialCount"]
|
||
and no_improve < session["config"]["earlyStopPatience"]
|
||
):
|
||
if cancel.is_set():
|
||
raise TuningError("session 已取消")
|
||
self._wait_for_dispatch(session_id, cancel)
|
||
control = self.storage.get_control(session_id)
|
||
base = (
|
||
self.storage.get_trial(control["activeBaseTrialId"])
|
||
if control["activeBaseTrialId"]
|
||
else self._best(session_id, rung=0) or completed_rung0[0]
|
||
)
|
||
proposal = self._request_proposal(
|
||
session, base["rewardConfig"], base["id"], next_number
|
||
)
|
||
# 模式允许在 session 运行期间切换，因此每次决策都读取最新持久化值。
|
||
current_mode = self.storage.get_session(session_id)["mode"]
|
||
if current_mode == "approval":
|
||
self.storage.update_session(
|
||
session_id, state="awaiting_approval", message="等待批准 Agent 建议"
|
||
)
|
||
proposal = self._wait_for_approval(session_id, proposal["id"], cancel)
|
||
if proposal["state"] == "rejected":
|
||
self.storage.audit(
|
||
session_id,
|
||
"proposal_rejected",
|
||
{"proposalId": proposal["id"], "feedback": proposal["feedback"]},
|
||
)
|
||
continue
|
||
else:
|
||
self.storage.decide_proposal(proposal["id"], "approved", "自动模式")
|
||
proposal = self.storage.get_proposal(proposal["id"])
|
||
self._claim_dispatch(session_id, cancel)
|
||
constraints = self.storage.get_control(session_id)["constraints"]
|
||
reward_config = merge_proposal(
|
||
base["rewardConfig"],
|
||
proposal["patch"],
|
||
constraints,
|
||
session["config"]["taskId"],
|
||
)
|
||
trial = self.storage.create_trial(
|
||
session_id,
|
||
next_number,
|
||
0,
|
||
session["config"]["rungs"][0],
|
||
reward_config,
|
||
proposal["id"],
|
||
f"trial-{next_number:03d}-rung-0",
|
||
)
|
||
result = self._execute_trial(session, trial)
|
||
if control["activeBaseTrialId"]:
|
||
self.storage.set_active_base(session_id, None)
|
||
self._pause_after_step(session_id)
|
||
result_score = result["score"]
|
||
if (
|
||
result["eligible"]
|
||
and result_score is not None
|
||
and result_score > best_score + 0.01
|
||
):
|
||
best_score = result_score
|
||
no_improve = 0
|
||
else:
|
||
no_improve += 1
|
||
self.storage.update_session(session_id, consecutive_no_improve=no_improve)
|
||
next_number += 1
|
||
|
||
# Promote top configurations; each new rung resumes its own previous checkpoint.
|
||
for rung in (1, 2):
|
||
previous = [
|
||
t
|
||
for t in self.storage.list_trials(session_id)
|
||
if t["rung"] == rung - 1 and t["state"] == "completed" and t["eligible"]
|
||
]
|
||
previous.sort(
|
||
key=lambda trial: (
|
||
trial["score"] if trial["score"] is not None else float("-inf")
|
||
),
|
||
reverse=True,
|
||
)
|
||
promoted_numbers = {
|
||
trial["number"]
|
||
for trial in self.storage.list_trials(session_id)
|
||
if trial["rung"] == rung and trial["state"] == "completed"
|
||
}
|
||
for parent in previous[: session["config"]["promote"][rung]]:
|
||
if parent["number"] in promoted_numbers:
|
||
continue
|
||
if cancel.is_set():
|
||
raise TuningError("session 已取消")
|
||
self._claim_dispatch(session_id, cancel)
|
||
root = self._session_root(session_id)
|
||
checkpoint = root / parent["checkpointPath"]
|
||
trial = self.storage.create_trial(
|
||
session_id,
|
||
parent["number"],
|
||
rung,
|
||
session["config"]["rungs"][rung],
|
||
parent["rewardConfig"],
|
||
parent["proposalId"],
|
||
f"trial-{parent['number']:03d}-rung-{rung}",
|
||
)
|
||
self._execute_trial(session, trial, checkpoint)
|
||
self._pause_after_step(session_id)
|
||
|
||
best = self._best_highest_rung(session_id)
|
||
if best is None:
|
||
raise TuningError("没有通过安全门槛的 trial")
|
||
preset_name = f"{session['config']['runName']}-{session_id[:8]}"
|
||
self.storage.save_preset(preset_name, session_id, best["id"], best["rewardConfig"])
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="succeeded",
|
||
best_trial_id=best["id"],
|
||
current_trial_id=None,
|
||
message="调参完成",
|
||
)
|
||
self.storage.audit(
|
||
session_id, "session_completed", {"bestTrialId": best["id"], "preset": preset_name}
|
||
)
|
||
except Exception as error:
|
||
state = self.storage.get_session(session_id)["state"]
|
||
if cancel.is_set() or state == "cancelled":
|
||
self.storage.update_session(
|
||
session_id, state="cancelled", message="调参已取消", current_trial_id=None
|
||
)
|
||
else:
|
||
self.storage.update_session(
|
||
session_id, state="failed", message=str(error), current_trial_id=None
|
||
)
|
||
self.storage.audit(session_id, "session_failed", {"error": str(error)[:1000]})
|
||
finally:
|
||
with self.lock:
|
||
self.workers.pop(session_id, None)
|
||
self.processes.pop(session_id, None)
|
||
|
||
def _wait_for_dispatch(self, session_id: str, cancel: threading.Event) -> None:
|
||
"""Wait at a scheduler boundary without consuming a one-shot token."""
|
||
with self.condition:
|
||
while not cancel.is_set():
|
||
session = self.storage.get_session(session_id)
|
||
control = self.storage.get_control(session_id)
|
||
if session["state"] == "paused":
|
||
self.condition.wait(timeout=1.0)
|
||
continue
|
||
if control["runPolicy"] == "continuous" or control["dispatchTokens"] > 0:
|
||
return
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="paused",
|
||
message="单步 Trial 已完成；等待下一个调度令牌",
|
||
)
|
||
self.storage.audit(session_id, "step_gate_waiting", {})
|
||
self.condition.wait(timeout=1.0)
|
||
raise TuningError("session 已取消")
|
||
|
||
def _claim_dispatch(self, session_id: str, cancel: threading.Event) -> None:
|
||
while not cancel.is_set():
|
||
self._wait_for_dispatch(session_id, cancel)
|
||
if self.storage.use_dispatch_token(session_id):
|
||
return
|
||
raise TuningError("session 已取消")
|
||
|
||
def _pause_after_step(self, session_id: str) -> None:
|
||
control = self.storage.get_control(session_id)
|
||
if control["runPolicy"] == "step" and control["dispatchTokens"] == 0:
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="paused",
|
||
message="单步 Trial 已完成；后续调度已暂停",
|
||
)
|
||
self.storage.audit(session_id, "step_trial_completed", {})
|
||
|
||
def approve(self, session_id: str, proposal_id: str, payload: Any) -> dict:
|
||
proposal = self.storage.get_proposal(proposal_id)
|
||
if proposal["sessionId"] != session_id:
|
||
raise TuningError("proposal 不属于该 session")
|
||
patch = proposal["patch"]
|
||
feedback = None
|
||
base = self.storage.get_trial(proposal["baseTrialId"])
|
||
constraints = self.storage.get_control(session_id)["constraints"]
|
||
if isinstance(payload, dict):
|
||
feedback = payload.get("feedback")
|
||
if "patch" in payload:
|
||
patch = payload["patch"]
|
||
if feedback is not None and (not isinstance(feedback, str) or len(feedback) > 2000):
|
||
raise TuningError("feedback 无效")
|
||
patch = validate_proposal(
|
||
patch,
|
||
base["rewardConfig"],
|
||
constraints,
|
||
self.storage.get_session(session_id)["config"]["taskId"],
|
||
)
|
||
if not self.storage.decide_proposal(proposal_id, "approved", feedback, patch):
|
||
raise TuningError("proposal 已处理")
|
||
self.storage.audit(
|
||
session_id,
|
||
"proposal_approved",
|
||
{"proposalId": proposal_id, "modified": patch != proposal["patch"]},
|
||
)
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def reject(self, session_id: str, proposal_id: str, payload: Any) -> dict:
|
||
feedback = payload.get("feedback", "") if isinstance(payload, dict) else ""
|
||
if not isinstance(feedback, str) or len(feedback) > 2000:
|
||
raise TuningError("feedback 无效")
|
||
proposal = self.storage.get_proposal(proposal_id)
|
||
if proposal["sessionId"] != session_id:
|
||
raise TuningError("proposal 不属于该 session")
|
||
if not self.storage.decide_proposal(proposal_id, "rejected", feedback):
|
||
raise TuningError("proposal 已处理")
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def set_mode(self, session_id: str, payload: Any) -> dict:
|
||
if not isinstance(payload, dict) or payload.get("mode") not in ("automatic", "approval"):
|
||
raise TuningError("mode 必须是 automatic 或 approval")
|
||
mode = payload["mode"]
|
||
with self.condition:
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in RUNNING_STATES | {"awaiting_approval", "paused"}:
|
||
raise TuningError("当前状态不能切换运行模式")
|
||
previous = session["mode"]
|
||
if previous == mode:
|
||
return self.detail(session_id)
|
||
self.storage.update_session(session_id, mode=mode)
|
||
approved_ids = []
|
||
if mode == "automatic":
|
||
constraints = self.storage.get_control(session_id)["constraints"]
|
||
for proposal in self.storage.list_proposals(session_id):
|
||
if proposal["state"] != "pending":
|
||
continue
|
||
base = self.storage.get_trial(proposal["baseTrialId"])
|
||
try:
|
||
validate_proposal(
|
||
proposal["patch"],
|
||
base["rewardConfig"],
|
||
constraints,
|
||
session["config"]["taskId"],
|
||
)
|
||
except Exception as error:
|
||
self.storage.decide_proposal(
|
||
proposal["id"], "rejected", f"参数护栏已变化：{error}"
|
||
)
|
||
continue
|
||
if self.storage.decide_proposal(
|
||
proposal["id"], "approved", "运行时切换为全自动模式"
|
||
):
|
||
approved_ids.append(proposal["id"])
|
||
if session["state"] == "awaiting_approval":
|
||
self.storage.update_session(
|
||
session_id, state="running", message="已切换为全自动模式，继续调参"
|
||
)
|
||
self.storage.audit(
|
||
session_id,
|
||
"session_mode_changed",
|
||
{"from": previous, "to": mode, "autoApprovedProposalIds": approved_ids},
|
||
)
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def set_constraints(self, session_id: str, payload: Any) -> dict:
|
||
if not isinstance(payload, dict) or set(payload) != {"revision", "constraints"}:
|
||
raise TuningError("参数护栏请求必须包含 revision 与 constraints")
|
||
revision = payload["revision"]
|
||
if isinstance(revision, bool) or not isinstance(revision, int) or revision < 0:
|
||
raise TuningError("constraints revision 必须是非负整数")
|
||
session = self.storage.get_session(session_id)
|
||
constraints = validate_constraints(payload["constraints"], session["config"]["taskId"])
|
||
if session["state"] not in ACTIVE_SESSION_STATES:
|
||
raise TuningError("终态 session 不能修改参数护栏")
|
||
control = self.storage.get_control(session_id)
|
||
base = None
|
||
if control["activeBaseTrialId"]:
|
||
base = self.storage.get_trial(control["activeBaseTrialId"])
|
||
elif session["currentTrialId"]:
|
||
base = self.storage.get_trial(session["currentTrialId"])
|
||
else:
|
||
base = self._best_highest_rung(session_id)
|
||
if base is not None:
|
||
validate_configuration_constraints(
|
||
base["rewardConfig"], constraints, session["config"]["taskId"]
|
||
)
|
||
try:
|
||
updated = self.storage.replace_constraints(session_id, revision, constraints)
|
||
except StorageConflict as error:
|
||
raise ResourceBusyError(str(error)) from error
|
||
|
||
rejected = []
|
||
for proposal in self.storage.list_proposals(session_id):
|
||
if proposal["state"] != "pending":
|
||
continue
|
||
proposal_base = self.storage.get_trial(proposal["baseTrialId"])
|
||
try:
|
||
validate_proposal(
|
||
proposal["patch"],
|
||
proposal_base["rewardConfig"],
|
||
constraints,
|
||
session["config"]["taskId"],
|
||
)
|
||
except Exception as error:
|
||
message = f"参数护栏 revision {updated['constraintsRevision']}：{error}"
|
||
if self.storage.decide_proposal(proposal["id"], "rejected", message):
|
||
rejected.append(proposal["id"])
|
||
self.storage.audit(
|
||
session_id,
|
||
"constraints_updated",
|
||
{
|
||
"revision": updated["constraintsRevision"],
|
||
"paths": sorted(constraints),
|
||
"rejectedProposalIds": rejected,
|
||
},
|
||
)
|
||
if rejected:
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def step(self, session_id: str, payload: Any) -> dict:
|
||
if not isinstance(payload, dict) or payload.get("count", 1) != 1:
|
||
raise TuningError("单步调度一次只能发放 1 个 Trial 令牌")
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in {"paused", "awaiting_approval"}:
|
||
raise TuningError("请先暂停或等待 Proposal 审批，再执行单步 Trial")
|
||
control = self.storage.get_control(session_id)
|
||
if control["dispatchTokens"] > 0:
|
||
raise TuningError("已有未消费的单步 Trial 令牌")
|
||
try:
|
||
self.storage.grant_dispatch_token(session_id)
|
||
except StorageConflict as error:
|
||
raise ResourceBusyError(str(error)) from error
|
||
if session["state"] == "paused":
|
||
has_pending = any(
|
||
proposal["state"] == "pending"
|
||
for proposal in self.storage.list_proposals(session_id)
|
||
)
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="awaiting_approval" if has_pending else "running",
|
||
message="单步令牌已就绪；请审批 Proposal"
|
||
if has_pending
|
||
else "已授权执行一个 Trial",
|
||
)
|
||
self.storage.audit(session_id, "step_token_granted", {"count": 1})
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def rollback(self, session_id: str, payload: Any) -> dict:
|
||
if not isinstance(payload, dict):
|
||
raise TuningError("rollback 请求体必须是对象")
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in {"paused", "awaiting_approval"}:
|
||
raise TuningError("回滚只能在安全暂停或等待审批时执行")
|
||
target_id = payload.get("trialId")
|
||
if payload.get("target") == "best":
|
||
target_id = session["bestTrialId"] or (self._best_highest_rung(session_id) or {}).get(
|
||
"id"
|
||
)
|
||
if not isinstance(target_id, str):
|
||
raise TuningError("rollback 必须指定 trialId 或 target=best")
|
||
target = self.storage.get_trial(target_id)
|
||
if target["sessionId"] != session_id:
|
||
raise TuningError("rollback Trial 不属于该 session")
|
||
if target["state"] != "completed" or not target["eligible"]:
|
||
raise TuningError("只能回滚到已完成且通过安全门槛的 Trial")
|
||
checkpoint = payload.get("checkpoint", False)
|
||
if not isinstance(checkpoint, bool):
|
||
raise TuningError("checkpoint 必须是布尔值")
|
||
if checkpoint:
|
||
path_value = target.get("checkpointPath")
|
||
if not path_value:
|
||
raise TuningError("目标 Trial 没有可用 checkpoint")
|
||
root = self._session_root(session_id)
|
||
path = (root / path_value).resolve()
|
||
if not path.is_relative_to(root) or not path.is_file():
|
||
raise TuningError("目标 checkpoint 不存在或路径非法")
|
||
|
||
superseded = []
|
||
for proposal in self.storage.list_proposals(session_id):
|
||
if proposal["state"] == "pending" and self.storage.decide_proposal(
|
||
proposal["id"], "rejected", f"由回滚到 Trial {target_id[:8]} 取代"
|
||
):
|
||
superseded.append(proposal["id"])
|
||
self.storage.set_active_base(session_id, target_id)
|
||
self.storage.reset_step_gate(session_id)
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="paused",
|
||
message=f"已回滚到 Trial {target['number']} / Rung {target['rung']}；等待单步或继续",
|
||
)
|
||
self.storage.audit(
|
||
session_id,
|
||
"rollback_selected",
|
||
{
|
||
"trialId": target_id,
|
||
"checkpoint": checkpoint,
|
||
"supersededProposalIds": superseded,
|
||
},
|
||
)
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def pause(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in RUNNING_STATES | {"awaiting_approval"}:
|
||
raise TuningError("当前状态不能暂停")
|
||
# 以持久化调度门记录暂停意图，而不是只依赖 session.state。当前 Trial
|
||
# 的训练/评估会继续完成，下一次 _claim_dispatch 必须等待显式继续或单步。
|
||
self.storage.reset_step_gate(session_id)
|
||
self.storage.update_session(
|
||
session_id, state="paused", message="已暂停后续调度；在途 Trial（如有）将完整结束"
|
||
)
|
||
return self.detail(session_id)
|
||
|
||
def resume(self, session_id: str) -> dict:
|
||
# Serialize state transition and worker registration against duplicate requests.
|
||
with self.condition:
|
||
return self._resume_locked(session_id)
|
||
|
||
def _resume_locked(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
if session["config"].get("pretrained") is not None:
|
||
if self.sources is None:
|
||
raise TuningError("基础策略快照服务未配置，无法恢复")
|
||
self.sources.verify(session["config"]["pretrained"])
|
||
if session["state"] == "paused":
|
||
self.storage.set_run_policy(session_id, "continuous")
|
||
has_pending = any(
|
||
proposal["state"] == "pending"
|
||
for proposal in self.storage.list_proposals(session_id)
|
||
)
|
||
self.storage.update_session(
|
||
session_id,
|
||
state="awaiting_approval" if has_pending else "running",
|
||
message="请审批待处理 Proposal" if has_pending else "连续调参已恢复",
|
||
)
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
elif session["state"] == "interrupted":
|
||
self.storage.set_run_policy(session_id, "continuous")
|
||
self.storage.update_session(session_id, state="queued", message="从最近完整结果恢复")
|
||
self._start_worker(session_id, resume=True)
|
||
else:
|
||
raise TuningError("当前状态不能恢复")
|
||
return self.detail(session_id)
|
||
|
||
def cancel(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in ACTIVE_SESSION_STATES:
|
||
return self.detail(session_id)
|
||
self.storage.update_session(session_id, state="cancelled", message="正在取消")
|
||
event = self.cancel_events.get(session_id)
|
||
if event:
|
||
event.set()
|
||
process = self.processes.get(session_id)
|
||
if process:
|
||
terminate_process(process)
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
return self.detail(session_id)
|
||
|
||
def metrics(
|
||
self,
|
||
session_id: str,
|
||
trial_id: str,
|
||
tags: list[str] | None,
|
||
max_points: int,
|
||
after_step: int | None = None,
|
||
) -> dict:
|
||
trial = self.storage.get_trial(trial_id)
|
||
if trial["sessionId"] != session_id:
|
||
raise TuningError("trial 不属于该 session")
|
||
series = self.storage.metrics(trial_id, tags, max_points, after_step)
|
||
next_step = max(
|
||
(point["step"] for item in series for point in item["points"]),
|
||
default=after_step,
|
||
)
|
||
return {"trialId": trial_id, "series": series, "nextStep": next_step}
|
||
|
||
def best_artifact(self, session_id: str) -> Path:
|
||
session = self.storage.get_session(session_id)
|
||
if not session["bestTrialId"]:
|
||
raise TuningError("尚无最佳策略")
|
||
trial = self.storage.get_trial(session["bestTrialId"])
|
||
if not trial["policyPath"]:
|
||
raise TuningError("最佳策略文件不存在")
|
||
root = self._session_root(session_id)
|
||
path = (root / trial["policyPath"]).resolve()
|
||
if not path.is_relative_to(root) or not path.is_file():
|
||
raise TuningError("最佳策略文件不存在")
|
||
return path
|
||
|
||
def preset_config(self, preset_id: str, task_id: str = FLAT_TASK) -> dict:
|
||
preset = self.storage.get_preset(preset_id)
|
||
if preset["taskId"] != task_id:
|
||
raise RewardConfigError("奖励 preset 来源任务与训练任务不匹配")
|
||
return validate_configuration(preset["rewardConfig"], task_id)
|
||
|
||
def test_agent(self) -> dict:
|
||
try:
|
||
return self.advisor.test_connection()
|
||
except Exception as error:
|
||
raise TuningError(f"Agent 连接测试失败：{error}") from error
|
||
|
||
def shutdown(self) -> None:
|
||
for session_id in list(self.workers):
|
||
with suppress(KeyError, TuningError):
|
||
self.cancel(session_id)
|
||
for worker in list(self.workers.values()):
|
||
worker.join(timeout=7)
|