zk-data-agent/skills/model-training-lite/scripts/render_sft_submission_plan.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path
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]:
    defaults = load_local_defaults()
    missing: list[str] = []
    jupyter_url = clean(payload.get("jupyter_url")) or clean(defaults.get("jupyter_url"))
    jupyter_auth_type = clean(payload.get("jupyter_auth_type") or payload.get("auth_type")) or clean(defaults.get("auth_type"))
    jupyter_password = clean(payload.get("jupyter_password") or payload.get("password")) or clean(defaults.get("password"))
    jupyter_token = clean(payload.get("jupyter_token") or payload.get("token")) or clean(defaults.get("token"))
    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 jupyter_url and not (jupyter_password or jupyter_token or clean(payload.get("jupyter_cookie")) or clean(defaults.get("cookie"))):
        missing.append("jupyter auth")

    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
    run_eval = bool(payload.get("run_eval") or payload.get("eval_after_train"))
    model_old = clean(payload.get("model_old")) or "/mnt/wangsenhao/verl_zk/qwen4b_cispo_wokl_add_bvt_2/global_step_5/actor/huggingface"
    wf_root = clean(payload.get("workflow_root")) or "/mnt/xiaoai-zk-model-train-tj5/workflow5"
    warnings: list[str] = []
    if not data_commit and not data_path:
        warnings.append("data_commit 未提供：计划将使用 data_branch 当前 HEAD，复现性弱于固定 commit。")

    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}" if data_path else
                f"git -C ai-planning reset --hard origin/{data_branch}"
            ),
            "git -C ai-planning rev-parse HEAD",
        ],
        "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",
        ],
        "submit_eval": [
            "source ~/.cloudml-cli/.profile",
            f"export AUTORESEARCH_CHAT_ROOT={workspace}",
            'export AUTORESEARCH_ROOT="$AUTORESEARCH_CHAT_ROOT"',
            f"export WF_ROOT={wf_root}",
            f"cd {workspace}",
            'eval "$(./scripts/resolve_run_ids.sh)"',
            f'./scripts/submit_cml_eval.sh "$EVAL_RUNDIC" "{workspace}/sft_output" "{model_old}"',
        ],
        "analyze_eval": [
            f"python {workspace}/scripts/summarize_eval_result.py --workflow-root {wf_root} --run-dic <EVAL_RUNDIC> --output {workspace}/results/eval_summary_<EVAL_RUNDIC>.md",
        ],
    }
    return {
        "ok": not missing,
        "missing": missing,
        "plan": {
            "jupyter_url": jupyter_url,
            "jupyter_auth": {
                "type": jupyter_auth_type or ("token" if jupyter_token else "password" if jupyter_password else ""),
                "configured": bool(jupyter_password or jupyter_token or clean(payload.get("jupyter_cookie")) or clean(defaults.get("cookie"))),
                "source": auth_source(payload, defaults),
            },
            "recipe": recipe,
            "workspace": workspace,
            "data_repo": data_repo,
            "data_branch": data_branch,
            "data_commit": data_commit,
            "data_path": data_path,
            "run_eval": run_eval,
            "model_old": model_old,
            "workflow_root": wf_root,
            "success_marker": f"{workspace}/sft_output/_SUCCESS",
        },
        "commands": commands,
        "warnings": warnings,
        "required_outputs": [
            "CloudML JobID",
            "CloudML task URL",
            "SFT_RUNDIC",
            "EVAL_RUNDIC",
            "workspace",
            "success_marker",
            "eval workflow id when run_eval=true",
            "eval summary path when analysis is requested",
        ],
    }


def clean(value: Any) -> str:
    return str(value or "").strip()


def load_local_defaults() -> dict[str, Any]:
    defaults: dict[str, Any] = {}
    skill_dir = Path(__file__).resolve().parents[1]
    shared_config_path = skill_dir / "jupyter_defaults.json"
    if shared_config_path.exists():
        try:
            shared_defaults = json.loads(shared_config_path.read_text(encoding="utf-8"))
            if isinstance(shared_defaults, dict):
                defaults.update({key: value for key, value in shared_defaults.items() if value})
                if any(clean(shared_defaults.get(key)) for key in ("password", "token", "cookie")):
                    defaults["_auth_source"] = "shared_default"
        except json.JSONDecodeError as exc:
            raise SystemExit(f"invalid shared defaults JSON: {shared_config_path}: {exc}") from exc
    env_defaults = {
        "jupyter_url": os.environ.get("MODEL_TRAINING_LITE_JUPYTER_URL"),
        "password": os.environ.get("MODEL_TRAINING_LITE_JUPYTER_PASSWORD"),
        "token": os.environ.get("MODEL_TRAINING_LITE_JUPYTER_TOKEN"),
        "cookie": os.environ.get("MODEL_TRAINING_LITE_JUPYTER_COOKIE"),
        "auth_type": os.environ.get("MODEL_TRAINING_LITE_JUPYTER_AUTH_TYPE"),
    }
    defaults.update({key: value for key, value in env_defaults.items() if value})
    if any(clean(env_defaults.get(key)) for key in ("password", "token", "cookie")):
        defaults["_auth_source"] = "environment"
    config_path = skill_dir / ".local" / "jupyter_defaults.json"
    if config_path.exists():
        try:
            file_defaults = json.loads(config_path.read_text(encoding="utf-8"))
            if isinstance(file_defaults, dict):
                defaults.update({key: value for key, value in file_defaults.items() if value})
                if any(clean(file_defaults.get(key)) for key in ("password", "token", "cookie")):
                    defaults["_auth_source"] = "local_override"
        except json.JSONDecodeError as exc:
            raise SystemExit(f"invalid local defaults JSON: {config_path}: {exc}") from exc
    return defaults


def auth_source(payload: dict[str, Any], defaults: dict[str, Any]) -> str:
    if any(clean(payload.get(key)) for key in ("jupyter_password", "password", "jupyter_token", "token", "jupyter_cookie")):
        return "input"
    if any(clean(defaults.get(key)) for key in ("password", "token", "cookie")):
        return clean(defaults.get("_auth_source")) or "default"
    return ""


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())