716 lines
30 KiB
Python
716 lines
30 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 .advisor import AdvisorUnavailable, DeepSeekAdvisor
|
||
from .process import GpuLease, ResourceBusyError, terminate_process
|
||
from .schema import (
|
||
BASE_REWARD_CONFIGURATION,
|
||
merge_proposal,
|
||
validate_proposal,
|
||
)
|
||
from .scoring import DEFAULT_OBJECTIVE_WEIGHTS, score_evaluation, validate_objective_weights
|
||
from .storage import 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,
|
||
):
|
||
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] = {}
|
||
|
||
def capability(self) -> dict[str, Any]:
|
||
capability = self.advisor.capability()
|
||
capability.update({"ready": (self.trainer_root / "scripts" / "evaluate.py").is_file()})
|
||
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")
|
||
if payload.get("taskId", "Unitree-Go2-Flat") != "Unitree-Go2-Flat":
|
||
raise TuningError("第一版只支持 Unitree-Go2-Flat")
|
||
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, 4, 20)
|
||
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": "Unitree-Go2-Flat",
|
||
"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),
|
||
}
|
||
objective = validate_objective_weights(
|
||
payload.get("objectiveWeights", DEFAULT_OBJECTIVE_WEIGHTS)
|
||
)
|
||
fallback = payload.get("fallbackEnabled", False)
|
||
if not isinstance(fallback, bool):
|
||
raise TuningError("fallbackEnabled 必须是布尔值")
|
||
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)
|
||
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),
|
||
]
|
||
if resume_checkpoint is not None:
|
||
command.extend(("--resume-checkpoint", str(resume_checkpoint)))
|
||
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")
|
||
self._wait_if_paused(session_id, self.cancel_events[session_id])
|
||
|
||
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=(",", ":")),
|
||
]
|
||
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 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"],
|
||
)
|
||
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 _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": "服务端将验证固定 schema；最多四项修改",
|
||
"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) -> dict:
|
||
names = ("track_linear_velocity", "action_rate_l2", "body_orientation_l2", "foot_slip")
|
||
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)
|
||
|
||
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),
|
||
"rationale": f"Agent 不可用，显式 fallback：{error}",
|
||
"expectedImpact": {},
|
||
"confidence": 0.2,
|
||
"promptHash": None,
|
||
"usage": {},
|
||
"model": "optuna-fallback",
|
||
}
|
||
source = "fallback"
|
||
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 已取消")
|
||
|
||
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:
|
||
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],
|
||
deepcopy(BASE_REWARD_CONFIGURATION),
|
||
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"] or 0.0 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_if_paused(session_id, cancel)
|
||
base = self._best(session_id, rung=0) or completed_rung0[0]
|
||
proposal = self._request_proposal(
|
||
session, base["rewardConfig"], base["id"], next_number
|
||
)
|
||
if session["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._wait_if_paused(session_id, cancel)
|
||
reward_config = merge_proposal(base["rewardConfig"], proposal["patch"])
|
||
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 result["eligible"] and (result["score"] or -999) > 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):
|
||
self._wait_if_paused(session_id, cancel)
|
||
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 t: t["score"] or -999, 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 已取消")
|
||
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)
|
||
|
||
best = (
|
||
self._best(session_id, rung=2)
|
||
or self._best(session_id, rung=1)
|
||
or self._best(session_id, rung=0)
|
||
)
|
||
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_if_paused(self, session_id: str, cancel: threading.Event) -> None:
|
||
with self.condition:
|
||
while self.storage.get_session(session_id)["state"] == "paused" and not cancel.is_set():
|
||
self.condition.wait(timeout=1.0)
|
||
|
||
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
|
||
if isinstance(payload, dict):
|
||
feedback = payload.get("feedback")
|
||
if "patch" in payload:
|
||
base = self.storage.get_trial(proposal["baseTrialId"])
|
||
patch = validate_proposal(payload["patch"], base["rewardConfig"])
|
||
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 pause(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] not in RUNNING_STATES | {"awaiting_approval"}:
|
||
raise TuningError("当前状态不能暂停")
|
||
self.storage.update_session(
|
||
session_id, state="paused", message="已暂停后续调度；当前子进程将完成"
|
||
)
|
||
return self.detail(session_id)
|
||
|
||
def resume(self, session_id: str) -> dict:
|
||
session = self.storage.get_session(session_id)
|
||
if session["state"] == "paused":
|
||
self.storage.update_session(session_id, state="running", message="继续调参")
|
||
with self.condition:
|
||
self.condition.notify_all()
|
||
elif session["state"] == "interrupted":
|
||
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:
|
||
self.storage.get_session(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
|
||
) -> dict:
|
||
trial = self.storage.get_trial(trial_id)
|
||
if trial["sessionId"] != session_id:
|
||
raise TuningError("trial 不属于该 session")
|
||
return {"trialId": trial_id, "series": self.storage.metrics(trial_id, tags, max_points)}
|
||
|
||
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) -> dict:
|
||
return self.storage.get_preset(preset_id)["rewardConfig"]
|
||
|
||
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)
|