add lightweight model training skill

2026-05-28 20:06:21 +08:00
parent 98b3e586de
commit 4665b9a632
3 changed files with 353 additions and 0 deletions
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import PurePosixPath
+from typing import Any
+
+
+DEFAULT_DATA_REPO = "git@git.n.xiaomi.com:ai-service/ai-planning.git"
+DEFAULT_BRANCH = "autoresearch-v1"
+DEFAULT_RECIPE = "zk_sft"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Render a lightweight SFT submission plan.")
+    parser.add_argument("--input", "-i", help="JSON input file. Reads stdin when omitted.")
+    args = parser.parse_args()
+    text = open(args.input, encoding="utf-8").read() if args.input else sys.stdin.read()
+    payload = json.loads(text or "{}")
+    result = render_plan(payload)
+    print(json.dumps(result, ensure_ascii=False, indent=2))
+    return 0 if result["ok"] else 1
+
+
+def render_plan(payload: dict[str, Any]) -> dict[str, Any]:
+    missing: list[str] = []
+    jupyter_url = clean(payload.get("jupyter_url"))
+    data_commit = clean(payload.get("data_commit") or payload.get("commit"))
+    data_path = clean(payload.get("data_path"))
+    if not jupyter_url:
+        missing.append("jupyter_url")
+    if not data_commit and not data_path:
+        missing.append("data_commit or data_path")
+
+    owner = clean(payload.get("owner")) or "wuyang6"
+    run_name = clean(payload.get("run_name")) or default_run_name(data_commit or data_path or "manual")
+    workspace = clean(payload.get("workspace")) or str(
+        PurePosixPath("/mnt/wangsenhao/autoresearch-zk-users") / owner / run_name
+    )
+    data_repo = clean(payload.get("data_repo")) or DEFAULT_DATA_REPO
+    data_branch = clean(payload.get("data_branch")) or DEFAULT_BRANCH
+    recipe = clean(payload.get("recipe")) or DEFAULT_RECIPE
+
+    commands = {
+        "prepare_workspace": [
+            f"mkdir -p {workspace}/{{scripts,assets,results,output}}",
+            f"cd {workspace}",
+        ],
+        "sync_data": [
+            f"[ -d ai-planning/.git ] || git clone -b {data_branch} {data_repo} ai-planning",
+            f"git -C ai-planning fetch origin {data_branch}",
+            f"git -C ai-planning checkout {data_branch}",
+            f"git -C ai-planning reset --hard {data_commit}" if data_commit else f"# use existing data_path: {data_path}",
+        ],
+        "submit_sft": [
+            "source ~/.cloudml-cli/.profile",
+            f"export AUTORESEARCH_CHAT_ROOT={workspace}",
+            'export AUTORESEARCH_ROOT="$AUTORESEARCH_CHAT_ROOT"',
+            f"cd {workspace}",
+            'eval "$(./scripts/resolve_run_ids.sh)"',
+            './scripts/submit_sft.sh "$SFT_RUNDIC"',
+        ],
+        "check_status": [
+            "cml custom_train describe <JOB_ID>",
+            f"test -f {workspace}/sft_output/_SUCCESS",
+        ],
+    }
+    return {
+        "ok": not missing,
+        "missing": missing,
+        "plan": {
+            "jupyter_url": jupyter_url,
+            "recipe": recipe,
+            "workspace": workspace,
+            "data_repo": data_repo,
+            "data_branch": data_branch,
+            "data_commit": data_commit,
+            "data_path": data_path,
+            "success_marker": f"{workspace}/sft_output/_SUCCESS",
+        },
+        "commands": commands,
+        "required_outputs": [
+            "CloudML JobID",
+            "CloudML task URL",
+            "SFT_RUNDIC",
+            "EVAL_RUNDIC",
+            "workspace",
+            "success_marker",
+        ],
+    }
+
+
+def clean(value: Any) -> str:
+    return str(value or "").strip()
+
+
+def default_run_name(seed: str) -> str:
+    slug = re.sub(r"[^a-zA-Z0-9_-]+", "_", seed).strip("_")[:32] or "manual"
+    return f"manual_sft_{slug}"
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())