wuyang6
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
#!/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())
|