#!/usr/bin/env python3 """Run independent, resumable stages for one registered model pair.""" from __future__ import annotations import argparse import json import os import subprocess import sys import traceback from copy import deepcopy from pathlib import Path from typing import Any from pipeline_common import ( REPO_ROOT, atomic_write_json, basic_validate_config, canonical_json_sha256, expand, file_record, load_json, resolve_path, run_id, safe_slug, shell_join, tool_versions, utc_now, ) def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", required=True, type=Path) parser.add_argument("--stage", action="append", dest="stages", help="stage id to run; repeatable") parser.add_argument("--dry-run", action="store_true", help="record resolved commands without executing them") parser.add_argument("--force", action="store_true", help="rerun a matching successful stage") parser.add_argument( "--resume-failed", action="store_true", help="run failed/blocked/missing stages and reuse matching successful stages", ) parser.add_argument("--result-out", type=Path, help="override aggregate result path") return parser.parse_args(argv) def load_and_validate(config_path: Path) -> tuple[dict[str, Any], str]: config = load_json(config_path) errors = basic_validate_config(config) try: import jsonschema # type: ignore schema = load_json(REPO_ROOT / "schemas" / "model_config.schema.json") validator = jsonschema.Draft202012Validator(schema) errors.extend(error.message for error in validator.iter_errors(config)) except ImportError: pass if errors: raise ValueError("invalid model configuration:\n- " + "\n- ".join(sorted(set(errors)))) return config, canonical_json_sha256(config) def read_previous(stage_root: Path) -> dict[str, Any] | None: latest = stage_root / "latest_result.json" if not latest.is_file(): return None try: value = load_json(latest) return value if isinstance(value, dict) else None except (OSError, json.JSONDecodeError): return None def same_inputs(previous: dict[str, Any], current: list[dict[str, Any]]) -> bool: def signature(records: list[dict[str, Any]]) -> list[tuple[str, bool, str | None, int | None]]: return [(r["path"], r["exists"], r["sha256"], r["bytes"]) for r in records] return signature(previous.get("inputs", [])) == signature(current) def stage_config_sha256(config: dict[str, Any], stage: dict[str, Any]) -> str: variant = stage["variant"] artifact = config["artifacts"].get(variant) if variant in {"fp32", "public_quantized"} else None return canonical_json_sha256({ "schema_version": config["schema_version"], "model_id": config["model"]["model_id"], "random_seed": config["model"].get("random_seed", 0), "artifact": artifact, "stage": stage, }) def previous_stage_matches( previous: dict[str, Any] | None, stage_sha: str, argv: list[str], stage: dict[str, Any] ) -> bool: if previous is None: return False if previous.get("stage_config_sha256") == stage_sha: return True # Migration path for T00 results written before stage-level hashes existed. return ( "stage_config_sha256" not in previous and previous.get("stage_id") == stage["id"] and previous.get("stage") == stage["stage"] and previous.get("variant") == stage["variant"] and previous.get("command_argv") == argv and previous.get("options", {}) == stage.get("options", {}) and previous.get("patch") == stage.get("patch") ) def summarize_error(stderr_path: Path, exception_text: str = "") -> str: text = exception_text if stderr_path.is_file(): try: lines = stderr_path.read_text(encoding="utf-8", errors="replace").splitlines() if lines: text = "\n".join(lines[-20:]) except OSError: pass return text[-4000:] def load_validation_report(path: Path | None) -> dict[str, Any]: if path is None: return {} if not path.is_file(): return {"validation_report_exists": False, "validation_report_path": str(path)} try: value = load_json(path) except (OSError, json.JSONDecodeError) as error: return { "validation_report_exists": True, "validation_report_path": str(path), "validation_report_parse_error": str(error), } return { "validation_report_exists": True, "validation_report_path": str(path), "stage_report": value, } def make_result_base( config: dict[str, Any], stage: dict[str, Any], stage_id: str, config_sha: str, stage_sha: str, argv: list[str], cwd: Path, inputs: list[dict[str, Any]], stdout_path: Path, stderr_path: Path, ) -> dict[str, Any]: variant = stage["variant"] artifact = config["artifacts"].get(variant, {}) if variant in {"fp32", "public_quantized"} else {} return { "model_id": config["model"]["model_id"], "artifact_id": stage.get("artifact_id", artifact.get("artifact_id", f"{config['model']['model_id']}-{variant}")), "variant": variant, "stage": stage["stage"], "stage_id": stage_id, "status": "RUNNING", "failure_code": None, "source_artifact": artifact.get("local_path"), "source_checksum": artifact.get("sha256"), "artifact": None, "command": shell_join(argv), "command_argv": argv, "working_directory": str(cwd), "inputs": inputs, "outputs": [], "tool_versions": tool_versions(), "options": stage.get("options", {}), "random_seed": config["model"].get("random_seed", 0), "started_at": utc_now(), "ended_at": utc_now(), "duration_sec": 0.0, "exit_code": None, "error_summary": "", "stdout_log": str(stdout_path), "stderr_log": str(stderr_path), "patch": stage.get("patch"), "validation": {}, "config_sha256": config_sha, "stage_config_sha256": stage_sha, "reused_from_previous_run": False, } def run_stage( config: dict[str, Any], stage: dict[str, Any], variables: dict[str, str], config_sha: str, current_run_id: str, dry_run: bool, force: bool, selected_ids: set[str] | None, statuses: dict[str, str], ) -> dict[str, Any]: stage_id = stage["id"] stage_root = Path(variables["model_dir"]) / "logs" / stage_id invocation_dir = stage_root / current_run_id invocation_dir.mkdir(parents=True, exist_ok=True) stdout_path = invocation_dir / "stdout.log" stderr_path = invocation_dir / "stderr.log" cwd = resolve_path(stage.get("working_directory", "{repo_root}"), variables) argv = [expand(str(part), variables) for part in stage["command"]] inputs = [file_record(resolve_path(value, variables, cwd)) for value in stage.get("inputs", [])] outputs_paths = [resolve_path(value, variables, cwd) for value in stage.get("outputs", [])] previous = read_previous(stage_root) stage_sha = stage_config_sha256(config, stage) result = make_result_base(config, stage, stage_id, config_sha, stage_sha, argv, cwd, inputs, stdout_path, stderr_path) if selected_ids is not None and stage_id not in selected_ids: if previous is not None: reused = deepcopy(previous) reused["reused_result_config_sha256"] = previous.get("config_sha256") reused["config_sha256"] = config_sha reused["stage_config_sha256"] = stage_sha reused["reused_from_previous_run"] = True atomic_write_json(invocation_dir / "run_result.json", reused) atomic_write_json(stage_root / "latest_result.json", reused) return reused result.update(status="SKIPPED", ended_at=utc_now(), error_summary="stage not selected") return result dependencies = stage.get("requires", []) dependency_success = {"PASS", "PASS_WITH_PATCH", "QUEUED"} if dry_run else {"PASS", "PASS_WITH_PATCH"} unavailable = [dependency for dependency in dependencies if statuses.get(dependency) not in dependency_success] if unavailable: result.update( status="BLOCKED", failure_code=stage["failure_code_on_error"], ended_at=utc_now(), error_summary=f"required stages not successful: {', '.join(unavailable)}", validation={"dependencies_satisfied": False, "unavailable_dependencies": unavailable}, ) atomic_write_json(invocation_dir / "run_result.json", result) atomic_write_json(stage_root / "latest_result.json", result) return result if dry_run: result.update( status="QUEUED", ended_at=utc_now(), validation={"dry_run": True, "resolved_outputs": [str(path) for path in outputs_paths]}, ) atomic_write_json(invocation_dir / "run_result.json", result) return result if ( not force and previous is not None and previous.get("status") in {"PASS", "PASS_WITH_PATCH"} and previous_stage_matches(previous, stage_sha, argv, stage) and same_inputs(previous, inputs) and all(record["exists"] for record in previous.get("outputs", [])) ): reused = deepcopy(previous) reused["reused_result_config_sha256"] = previous.get("config_sha256") reused["config_sha256"] = config_sha reused["stage_config_sha256"] = stage_sha reused["reused_from_previous_run"] = True atomic_write_json(invocation_dir / "run_result.json", reused) atomic_write_json(stage_root / "latest_result.json", reused) return reused conflicting_outputs = [path for path in outputs_paths if path.exists()] previous_matches = previous_stage_matches(previous, stage_sha, argv, stage) if conflicting_outputs and not previous_matches: result.update( status="FAIL", failure_code="FAIL_ENVIRONMENT", ended_at=utc_now(), error_summary="refusing to overwrite output from a different or unknown configuration", outputs=[file_record(path) for path in outputs_paths], validation={"overwrite_guard": False, "conflicting_outputs": [str(path) for path in conflicting_outputs]}, ) atomic_write_json(invocation_dir / "run_result.json", result) atomic_write_json(stage_root / "latest_result.json", result) return result import time started_monotonic = time.monotonic() exit_code: int | None = None exception_text = "" cwd.mkdir(parents=True, exist_ok=True) for path in outputs_paths: path.parent.mkdir(parents=True, exist_ok=True) try: with stdout_path.open("wb") as stdout_handle, stderr_path.open("wb") as stderr_handle: completed = subprocess.run( argv, cwd=cwd, stdout=stdout_handle, stderr=stderr_handle, timeout=stage["timeout_sec"], check=False, env={**os.environ, "PYTHONHASHSEED": str(config["model"].get("random_seed", 0))}, ) exit_code = completed.returncode except subprocess.TimeoutExpired as error: exit_code = 124 exception_text = f"timeout after {stage['timeout_sec']} seconds: {error}" except (OSError, ValueError) as error: exit_code = 127 exception_text = f"execution failed: {error}" duration = time.monotonic() - started_monotonic output_records = [file_record(path) for path in outputs_paths] outputs_exist = all(record["exists"] for record in output_records) validation_path = resolve_path(stage["validation_report"], variables, cwd) if stage.get("validation_report") else None validation = { "outputs_exist": outputs_exist, "expected_output_count": len(output_records), "quantization_preserved": None, "runtime_output_match": None, **load_validation_report(validation_path), } passed = exit_code == 0 and outputs_exist result.update( status="PASS_WITH_PATCH" if passed and stage.get("patch") else ("PASS" if passed else "FAIL"), failure_code=None if passed else stage["failure_code_on_error"], artifact=str(outputs_paths[0]) if passed and outputs_paths else None, outputs=output_records, ended_at=utc_now(), duration_sec=round(duration, 6), exit_code=exit_code, error_summary="" if passed else summarize_error(stderr_path, exception_text or "stage failed or output missing"), validation=validation, ) atomic_write_json(invocation_dir / "run_result.json", result) atomic_write_json(stage_root / "latest_result.json", result) return result def overall_status(results: list[dict[str, Any]]) -> str: relevant = [result["status"] for result in results if result["status"] != "SKIPPED"] if relevant and all(status in {"PASS", "PASS_WITH_PATCH"} for status in relevant): return "PASS" if relevant and all(status == "QUEUED" for status in relevant): return "QUEUED" if any(status in {"PASS", "PASS_WITH_PATCH"} for status in relevant): return "PARTIAL" if any(status == "BLOCKED" for status in relevant) and not any(status == "FAIL" for status in relevant): return "BLOCKED" return "FAIL" def execute(args: argparse.Namespace) -> tuple[dict[str, Any], int]: config_path = args.config.resolve() config, config_sha = load_and_validate(config_path) model = config["model"] model_dir = REPO_ROOT / "models" / safe_slug(model["task"]) / model["model_id"] for relative in ( "config", "source/fp32", "source/quantized", "baseline", "onnx/fp32", "onnx/quantized", "tflite/fp32", "tflite/quantized", "mlir/fp32", "mlir/quantized", "graphs", "analysis", "logs", ): (model_dir / relative).mkdir(parents=True, exist_ok=True) variables = { "repo_root": str(REPO_ROOT), "model_dir": str(model_dir), "python": sys.executable, "config_dir": str(config_path.parent), } current_run_id = run_id() started_at = utc_now() selected = set(args.stages) if args.stages else None known_ids = {stage["id"] for stage in config["stages"]} if selected and (unknown := selected - known_ids): raise ValueError(f"unknown stage ids: {', '.join(sorted(unknown))}") statuses: dict[str, str] = {} results: list[dict[str, Any]] = [] for stage in config["stages"]: result = run_stage( config, stage, variables, config_sha, current_run_id, args.dry_run, args.force, selected, statuses, ) statuses[stage["id"]] = result["status"] results.append(result) aggregate = { "schema_version": "1.0", "model_id": model["model_id"], "config_path": str(config_path), "config_sha256": config_sha, "run_id": current_run_id, "started_at": started_at, "ended_at": utc_now(), "overall_status": overall_status(results), "stages": results, } output_path = args.result_out.resolve() if args.result_out else model_dir / ("dry_run_result.json" if args.dry_run else "run_result.json") atomic_write_json(output_path, aggregate) return aggregate, 0 if aggregate["overall_status"] in {"PASS", "QUEUED"} else 1 def main(argv: list[str] | None = None) -> int: args = parse_args(argv) try: aggregate, exit_code = execute(args) print(json.dumps({ "model_id": aggregate["model_id"], "run_id": aggregate["run_id"], "overall_status": aggregate["overall_status"], }, sort_keys=True)) return exit_code except Exception as error: # preserve a useful CLI failure while avoiding false PASS traceback.print_exc(file=sys.stderr) print(f"run_model failed: {error}", file=sys.stderr) return 2 if __name__ == "__main__": raise SystemExit(main())