From 0f8a6e4dafa5ac5bca4ccfdc405e5c2a807edcbe Mon Sep 17 00:00:00 2001 From: RYY177 <3346599702@qq.com> Date: Sat, 1 Aug 2026 06:33:05 +0000 Subject: [PATCH] feat(trainer): add standalone Trainer CLI --- loopai/skills/Trainer/README.md | 32 +++ loopai/skills/Trainer/cli.py | 449 ++++++++++++++++++++++++++++++++ setup.py | 1 + tests/test_trainer_cli.py | 349 +++++++++++++++++++++++++ 4 files changed, 831 insertions(+) create mode 100644 loopai/skills/Trainer/cli.py create mode 100644 tests/test_trainer_cli.py diff --git a/loopai/skills/Trainer/README.md b/loopai/skills/Trainer/README.md index 5dfb603..8b7d1d3 100644 --- a/loopai/skills/Trainer/README.md +++ b/loopai/skills/Trainer/README.md @@ -7,6 +7,38 @@ Trainer Skill 是 Dataflow-LoopAI 中负责模型训练的技能实现，支持 Trainer 不依赖 MCP。对调用方仍保持同步返回语义，但训练、进度持久化和结果收尾由独立本地 Worker 持有；Codex/API 会话提前结束不会中断 Worker。两条路径都先生成完整 YAML，用户确认后才启动训练。 +## 命令行入口 + +安装项目后可以通过 `loopai-trainer` 直接调用同一套 Trainer Skill。CLI +不会创建另一套训练实现，也不会改变 API、前端或独立 Worker 的行为。 + +```bash +pip install -e . + +# 1. 生成最终 YAML；从 JSON 结果的 data 中读取 config_yaml、 +# config_path、config_sha256 和 trainer_version_id。 +loopai-trainer prepare \ + --config ./starter.yaml \ + --task-id my-task + +# 2. 用户确认 YAML 后，使用 prepare 返回的同一个 version ID 启动训练。 +loopai-trainer run-prepared \ + --config ./starter.yaml \ + --task-id my-task \ + --version-id \ + --prepared-config \ + --sha256 + +# 查看状态、事件和分析结果。 +loopai-trainer status --task-id my-task --version-id +loopai-trainer events --task-id my-task --version-id +loopai-trainer analyze --task-id my-task +``` + +如果任务配置保存在数据库中，可同时传入 `--db-path`。`prepare` 与 +`run-prepared` 必须使用相同的任务配置来源；CLI 不提供绕过 YAML 确认的直接 +`run` 子命令。 + ## Verl GRPO 最小配置 ```python diff --git a/loopai/skills/Trainer/cli.py b/loopai/skills/Trainer/cli.py new file mode 100644 index 0000000..e57013c --- /dev/null +++ b/loopai/skills/Trainer/cli.py @@ -0,0 +1,449 @@ +"""Command-line entry point for the LoopAI Trainer skill.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any, Sequence + + +_FAILURE_STATUSES = {"cancelled", "error", "failed", "failure"} +_SAFE_RESULT_KEYS = { + "approval_required", + "code", + "config", + "config_path", + "config_sha256", + "config_yaml", + "data", + "error", + "errors", + "message", + "ok", + "status", + "success", + "task_id", + "trainer_output_dir", + "trainer_version_id", + "version_id", + "warnings", +} +_SECRET_KEYS = { + "api_key", + "apikey", + "authorization", + "credential", + "credentials", + "password", + "secret", + "token", +} +_SECRET_SUFFIXES = ( + "_api_key", + "_credential", + "_credentials", + "_password", + "_secret", + "_token", +) + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, set): + return sorted(value, key=str) + return str(value) + + +def _is_secret_key(key: Any) -> bool: + normalized = str(key).strip().lower().replace("-", "_") + return normalized in _SECRET_KEYS or normalized.endswith(_SECRET_SUFFIXES) + + +def _redact(value: Any) -> Any: + if isinstance(value, dict): + return { + str(key): "***" if _is_secret_key(key) else _redact(item) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [_redact(item) for item in value] + return value + + +def _print_json(payload: Any) -> None: + print( + json.dumps( + _redact(payload), + ensure_ascii=False, + sort_keys=True, + default=_json_default, + ) + ) + + +def _extract_trainer_result(state: Any) -> dict[str, Any]: + """Extract the public Trainer response without dumping the complete State.""" + if not isinstance(state, dict): + return { + "ok": True, + "status": "completed", + "message": "Trainer command completed.", + "data": None, + "error": None, + } + + trainer = state.get("trainer") + if isinstance(trainer, dict) and isinstance(trainer.get("trainer_result"), dict): + return trainer["trainer_result"] + + if isinstance(state.get("trainer_result"), dict): + return state["trainer_result"] + + safe_result = { + key: value + for key, value in state.items() + if key in _SAFE_RESULT_KEYS + } + if safe_result: + return safe_result + return { + "ok": True, + "status": "completed", + "message": "Trainer command completed.", + "data": None, + "error": None, + } + + +def _success_payload(message: str, data: Any) -> dict[str, Any]: + return { + "ok": True, + "status": "completed", + "message": message, + "data": data, + "error": None, + } + + +def _error_payload(exc: BaseException, message: str) -> dict[str, Any]: + return { + "ok": False, + "status": "failed", + "message": message, + "data": None, + "error": { + "type": type(exc).__name__, + "code": "UNEXPECTED_ERROR", + "detail": str(exc), + "recoverable": not isinstance(exc, KeyboardInterrupt), + }, + } + + +def _is_failure(result: Any) -> bool: + if not isinstance(result, dict): + return False + if result.get("ok") is False or result.get("success") is False: + return True + if str(result.get("status") or "").strip().lower() in _FAILURE_STATUSES: + return True + data = result.get("data") + return isinstance(data, dict) and _is_failure(data) + + +def _add_runtime_arguments( + parser: argparse.ArgumentParser, + *, + require_version: bool = False, +) -> None: + parser.add_argument( + "--config", + dest="config_path", + default=None, + help="LoopAI starter/state YAML used to initialize Trainer.", + ) + parser.add_argument( + "--thread-id", + default=None, + help="Optional LoopAI task/thread context ID; takes precedence over --task-id.", + ) + parser.add_argument( + "--task-id", + default=os.getenv("TASK_ID") or None, + help="Task ID override (defaults to TASK_ID).", + ) + parser.add_argument( + "--db-path", + default=os.getenv("DB_PATH") or None, + help="Runtime database path override (defaults to DB_PATH).", + ) + parser.add_argument( + "--output-dir", + default=os.getenv("OUTPUT_DIR") or None, + help="Output root override (defaults to OUTPUT_DIR or ./outputs).", + ) + parser.add_argument( + "--version-id", + default=os.getenv("VERSION_ID") or None, + required=require_version and not bool(os.getenv("VERSION_ID")), + help=( + "Trainer version ID. For run-prepared, pass trainer_version_id " + "returned by prepare (defaults to VERSION_ID)." + ), + ) + parser.add_argument( + "--trainer-output-dir", + default=os.getenv("TRAINER_OUTPUT_DIR") or None, + help="Explicit Trainer version directory override.", + ) + + +def _runtime_kwargs(args: argparse.Namespace) -> dict[str, Any]: + return { + key: value + for key, value in { + "task_id": getattr(args, "task_id", None), + "db_path": getattr(args, "db_path", None), + "output_dir": getattr(args, "output_dir", None), + "version_id": getattr(args, "version_id", None), + "trainer_output_dir": getattr(args, "trainer_output_dir", None), + }.items() + if value is not None + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="loopai-trainer", + description=( + "Prepare, approve, run, and inspect LoopAI Trainer jobs. " + "Training is deliberately split into prepare and run-prepared steps." + ), + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + prepare_parser = subparsers.add_parser( + "prepare", + help="Validate inputs and generate the exact training YAML for approval.", + ) + _add_runtime_arguments(prepare_parser) + + inspect_parser = subparsers.add_parser( + "inspect", + help="Print a prepared training YAML and its SHA256 approval digest.", + ) + inspect_parser.add_argument( + "--prepared-config", + required=True, + help="Prepared training YAML path.", + ) + + run_parser = subparsers.add_parser( + "run-prepared", + help="Run the exact prepared YAML after verifying its SHA256 digest.", + ) + run_parser.add_argument( + "--prepared-config", + required=True, + help="Prepared training YAML path.", + ) + run_parser.add_argument( + "--sha256", + required=True, + help="Expected SHA256 returned by prepare or inspect.", + ) + _add_runtime_arguments(run_parser, require_version=True) + + events_parser = subparsers.add_parser( + "events", + help="Read persisted Trainer events for a task.", + ) + events_parser.add_argument("--task-id", required=True) + events_parser.add_argument("--output-dir", default="./outputs") + events_parser.add_argument("--version-id", default=None) + events_parser.add_argument("--trainer-output-dir", default=None) + + status_parser = subparsers.add_parser( + "status", + help="Read persisted run state and the latest local training metric.", + ) + status_parser.add_argument("--task-id", default=None) + status_parser.add_argument("--version-id", default=None) + status_parser.add_argument("--output-dir", default="./outputs") + status_parser.add_argument("--trainer-output-dir", default=None) + + analyze_parser = subparsers.add_parser( + "analyze", + help="Analyze Trainer artifacts and select the best checkpoint.", + ) + analyze_parser.add_argument("--task-id", default=None) + analyze_parser.add_argument("--output-dir", default="./outputs") + analyze_parser.add_argument("--trainer-task-id", default=None) + analyze_parser.add_argument("--training-output-dir", default=None) + + guide_parser = subparsers.add_parser( + "guide", + help="Print Trainer configuration prefill guidance.", + ) + guide_parser.add_argument( + "--task-type", + default="sft", + help="Training task type, for example sft or grpo.", + ) + + return parser + + +def _dispatch(args: argparse.Namespace) -> tuple[dict[str, Any], int]: + # Lazy import keeps `loopai-trainer --help` free of training-side effects. + from loopai.skills import Trainer + + if args.command == "prepare": + state = Trainer.prepare( + thread_id=args.thread_id, + config_path=args.config_path, + **_runtime_kwargs(args), + ) + result = _extract_trainer_result(state) + return result, 1 if _is_failure(result) else 0 + + if args.command == "inspect": + data = Trainer.inspect_prepared_config(args.prepared_config) + return _success_payload("Prepared Trainer config inspected.", data), 0 + + if args.command == "run-prepared": + state = Trainer.run_prepared( + prepared_config_path=args.prepared_config, + expected_config_sha256=args.sha256, + thread_id=args.thread_id, + config_path=args.config_path, + **_runtime_kwargs(args), + ) + result = _extract_trainer_result(state) + return result, 1 if _is_failure(result) else 0 + + if args.command == "events": + data = Trainer.load_events( + task_id=args.task_id, + output_dir=args.output_dir, + version_id=args.version_id, + trainer_output_dir=args.trainer_output_dir, + ) + return _success_payload("Trainer events loaded.", data), 0 + + if args.command == "status": + data = _load_status( + task_id=args.task_id, + version_id=args.version_id, + output_dir=args.output_dir, + trainer_output_dir=args.trainer_output_dir, + ) + return _success_payload("Trainer status loaded.", data), 0 + + if args.command == "analyze": + result = Trainer.analyze_results( + task_id=args.task_id, + output_dir=args.output_dir, + trainer_task_id=args.trainer_task_id, + training_output_dir=args.training_output_dir, + ) + return result, 1 if _is_failure(result) else 0 + + if args.command == "guide": + data = Trainer.prefill_guide(task_type=args.task_type) + return _success_payload("Trainer prefill guide generated.", data), 0 + + raise ValueError(f"Unsupported Trainer command: {args.command}") + + +def _load_status( + *, + task_id: str | None, + version_id: str | None, + output_dir: str, + trainer_output_dir: str | None, +) -> dict[str, Any]: + if trainer_output_dir: + run_dir = Path(trainer_output_dir).expanduser().resolve() + else: + if not task_id or not version_id: + raise ValueError( + "status requires --trainer-output-dir or both --task-id and --version-id" + ) + run_dir = ( + Path(output_dir).expanduser() + / str(task_id) + / "trainer" + / str(version_id) + ).resolve() + + from loopai.skills.Trainer.results import load_live_training_metrics + from loopai.skills.Trainer.utils.persistent_worker import read_run_state, worker_paths + + paths = worker_paths(run_dir) + run_state = read_run_state(paths["state"]) + if not run_dir.is_dir() and not run_state: + raise FileNotFoundError(f"Trainer run directory does not exist: {run_dir}") + + metrics_task_info = None + latest_metric = None + try: + live_metrics = load_live_training_metrics(run_dir) + except FileNotFoundError: + live_metrics = {} + if isinstance(live_metrics, dict): + task_info = live_metrics.get("task_info") + if isinstance(task_info, dict): + metrics_task_info = task_info + records = live_metrics.get("metrics") + if isinstance(records, list): + metric_records = [ + item + for item in records + if isinstance(item, dict) and "log_line" not in item + ] + if metric_records: + latest_metric = metric_records[-1] + + return { + "run_dir": str(run_dir), + "run_state_path": str(paths["state"]), + "run_state": run_state, + "metrics_task_info": metrics_task_info, + "latest_metric": latest_metric, + } + + +def run(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(list(argv) if argv is not None else None) + try: + result, exit_code = _dispatch(args) + _print_json(result) + return exit_code + except KeyboardInterrupt as exc: + _print_json( + _error_payload( + exc, + ( + "Trainer CLI interrupted. A detached worker may still be running; " + "inspect persisted events or run_state.json before resubmitting." + ), + ) + ) + return 130 + except Exception as exc: + _print_json(_error_payload(exc, "Trainer CLI command failed.")) + return 1 + + +def main() -> None: + raise SystemExit(run(sys.argv[1:])) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 50fa07b..8952566 100644 --- a/setup.py +++ b/setup.py @@ -67,6 +67,7 @@ "loopai-obtainercli=loopai.skills.ObtainerCLI.cli:main", "loopai-judger=loopai.skills.Judger.cli:main", "loopai-analyzer=loopai.skills.Analyzer.cli:main", + "loopai-trainer=loopai.skills.Trainer.cli:main", ], }, python_requires=">=3.12", diff --git a/tests/test_trainer_cli.py b/tests/test_trainer_cli.py new file mode 100644 index 0000000..e6caddf --- /dev/null +++ b/tests/test_trainer_cli.py @@ -0,0 +1,349 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +from loopai.skills import Trainer +from loopai.skills.Trainer import cli + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _json_output(capsys) -> dict: + return json.loads(capsys.readouterr().out) + + +def test_setup_registers_trainer_console_script() -> None: + setup_text = (REPO_ROOT / "setup.py").read_text(encoding="utf-8") + + assert "loopai-trainer=loopai.skills.Trainer.cli:main" in setup_text + + +def test_prepare_forwards_runtime_args_and_only_prints_public_result( + tmp_path, + monkeypatch, + capsys, +) -> None: + captured = {} + public_result = { + "ok": True, + "status": "completed", + "message": "ready", + "data": { + "config_yaml": "model_name_or_path: /models/base\n", + "config_sha256": "abc123", + "trainer_version_id": "version-1", + }, + "error": None, + } + + def fake_prepare(**kwargs): + captured.update(kwargs) + return { + "system": {"starter_api_key": "must-not-leak"}, + "trainer": {"trainer_result": public_result}, + } + + monkeypatch.setattr(Trainer, "prepare", fake_prepare) + monkeypatch.delenv("TASK_ID", raising=False) + monkeypatch.delenv("DB_PATH", raising=False) + monkeypatch.delenv("OUTPUT_DIR", raising=False) + monkeypatch.delenv("VERSION_ID", raising=False) + monkeypatch.delenv("TRAINER_OUTPUT_DIR", raising=False) + + exit_code = cli.run( + [ + "prepare", + "--config", + str(tmp_path / "starter.yaml"), + "--thread-id", + "thread-1", + "--task-id", + "task-1", + "--db-path", + str(tmp_path / "db.sqlite3"), + "--output-dir", + str(tmp_path / "outputs"), + "--version-id", + "version-1", + "--trainer-output-dir", + str(tmp_path / "outputs" / "task-1" / "trainer" / "version-1"), + ] + ) + + payload = _json_output(capsys) + assert exit_code == 0 + assert payload == public_result + assert "must-not-leak" not in json.dumps(payload) + assert captured == { + "thread_id": "thread-1", + "config_path": str(tmp_path / "starter.yaml"), + "task_id": "task-1", + "db_path": str(tmp_path / "db.sqlite3"), + "output_dir": str(tmp_path / "outputs"), + "version_id": "version-1", + "trainer_output_dir": str( + tmp_path / "outputs" / "task-1" / "trainer" / "version-1" + ), + } + + +def test_run_prepared_requires_version_id(monkeypatch) -> None: + monkeypatch.delenv("VERSION_ID", raising=False) + + with pytest.raises(SystemExit) as exc_info: + cli.run( + [ + "run-prepared", + "--prepared-config", + "/tmp/training.yaml", + "--sha256", + "abc123", + ] + ) + + assert exc_info.value.code == 2 + + +def test_run_prepared_forwards_approval_and_returns_failure( + monkeypatch, + capsys, +) -> None: + captured = {} + failure = { + "ok": False, + "status": "failed", + "message": "training failed", + "data": None, + "error": {"detail": "synthetic failure"}, + } + + def fake_run_prepared(**kwargs): + captured.update(kwargs) + return { + "system": {"api_key": "must-not-leak"}, + "trainer": {"trainer_result": failure}, + } + + monkeypatch.setattr(Trainer, "run_prepared", fake_run_prepared) + monkeypatch.delenv("VERSION_ID", raising=False) + + exit_code = cli.run( + [ + "run-prepared", + "--prepared-config", + "/tmp/training.yaml", + "--sha256", + "abc123", + "--config", + "/tmp/starter.yaml", + "--task-id", + "task-1", + "--version-id", + "version-1", + ] + ) + + payload = _json_output(capsys) + assert exit_code == 1 + assert payload == failure + assert captured == { + "prepared_config_path": "/tmp/training.yaml", + "expected_config_sha256": "abc123", + "thread_id": None, + "config_path": "/tmp/starter.yaml", + "task_id": "task-1", + "version_id": "version-1", + } + + +def test_inspect_emits_structured_json(tmp_path, monkeypatch, capsys) -> None: + prepared_config = tmp_path / "training.yaml" + inspected = { + "config_path": str(prepared_config), + "config_yaml": "model_name_or_path: /models/base\n", + "config": {"model_name_or_path": "/models/base"}, + "config_sha256": "abc123", + } + monkeypatch.setattr( + Trainer, + "inspect_prepared_config", + lambda path: inspected if path == str(prepared_config) else None, + ) + + exit_code = cli.run( + ["inspect", "--prepared-config", str(prepared_config)] + ) + + payload = _json_output(capsys) + assert exit_code == 0 + assert payload["ok"] is True + assert payload["data"] == inspected + + +def test_events_forwards_filters(monkeypatch, capsys) -> None: + captured = {} + + def fake_load_events(**kwargs): + captured.update(kwargs) + return [{"status": "running", "version_id": "version-1"}] + + monkeypatch.setattr(Trainer, "load_events", fake_load_events) + + exit_code = cli.run( + [ + "events", + "--task-id", + "task-1", + "--output-dir", + "/tmp/outputs", + "--version-id", + "version-1", + "--trainer-output-dir", + "/tmp/run", + ] + ) + + payload = _json_output(capsys) + assert exit_code == 0 + assert payload["data"][0]["status"] == "running" + assert captured == { + "task_id": "task-1", + "output_dir": "/tmp/outputs", + "version_id": "version-1", + "trainer_output_dir": "/tmp/run", + } + + +def test_status_reads_persisted_state_and_latest_metric(tmp_path, capsys) -> None: + run_dir = tmp_path / "outputs" / "task-1" / "trainer" / "version-1" + metrics_dir = run_dir / "metrics" + metrics_dir.mkdir(parents=True) + (run_dir / "run_state.json").write_text( + json.dumps( + { + "status": "running", + "task_id": "task-1", + "version_id": "version-1", + "current_step": 2, + } + ), + encoding="utf-8", + ) + (metrics_dir / "metrics.json").write_text( + json.dumps( + { + "task_info": {"framework": "llamafactory"}, + "metrics": [ + {"step": 1, "loss": 2.0}, + {"step": 2, "loss": 1.0}, + ], + } + ), + encoding="utf-8", + ) + + exit_code = cli.run( + [ + "status", + "--task-id", + "task-1", + "--version-id", + "version-1", + "--output-dir", + str(tmp_path / "outputs"), + ] + ) + + payload = _json_output(capsys) + assert exit_code == 0 + assert payload["data"]["run_state"]["status"] == "running" + assert payload["data"]["latest_metric"] == {"step": 2, "loss": 1.0} + assert payload["data"]["metrics_task_info"]["framework"] == "llamafactory" + + +def test_status_requires_an_unambiguous_run(capsys) -> None: + exit_code = cli.run(["status"]) + + payload = _json_output(capsys) + assert exit_code == 1 + assert payload["error"]["type"] == "ValueError" + assert "--task-id and --version-id" in payload["error"]["detail"] + + +def test_analyze_preserves_structured_failure(monkeypatch, capsys) -> None: + failure = { + "ok": False, + "status": "failed", + "message": "not found", + "data": None, + "error": {"code": "NOT_FOUND"}, + } + monkeypatch.setattr(Trainer, "analyze_results", lambda **kwargs: failure) + + exit_code = cli.run( + [ + "analyze", + "--task-id", + "task-1", + "--training-output-dir", + "/tmp/training", + ] + ) + + assert exit_code == 1 + assert _json_output(capsys) == failure + + +def test_guide_redacts_secrets_without_redacting_tokenizer( + monkeypatch, + capsys, +) -> None: + monkeypatch.setattr( + Trainer, + "prefill_guide", + lambda **kwargs: { + "task_type": kwargs["task_type"], + "api_key": "secret-value", + "tokenizer": "qwen-tokenizer", + }, + ) + + exit_code = cli.run(["guide", "--task-type", "grpo"]) + + payload = _json_output(capsys) + assert exit_code == 0 + assert payload["data"]["api_key"] == "***" + assert payload["data"]["tokenizer"] == "qwen-tokenizer" + + +def test_unexpected_error_is_json_and_does_not_escape( + monkeypatch, + capsys, +) -> None: + def fail(**kwargs): + raise RuntimeError("synthetic CLI failure") + + monkeypatch.setattr(Trainer, "prepare", fail) + + exit_code = cli.run(["prepare"]) + + payload = _json_output(capsys) + assert exit_code == 1 + assert payload["ok"] is False + assert payload["error"]["type"] == "RuntimeError" + assert payload["error"]["detail"] == "synthetic CLI failure" + + +def test_main_exits_with_run_status(monkeypatch) -> None: + monkeypatch.setattr(cli, "run", lambda argv: 7) + monkeypatch.setattr(sys, "argv", ["loopai-trainer", "guide"]) + + with pytest.raises(SystemExit) as exc_info: + cli.main() + + assert exc_info.value.code == 7