scripts/pipeline_common.py

#!/usr/bin/env python3
"""Shared, dependency-free helpers for the reproducible model pipeline."""

from __future__ import annotations

import hashlib
import json
import os
import platform
import re
import shlex
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[1]
STATUS_VALUES = {
    "TODO", "QUEUED", "RUNNING", "PASS", "PASS_WITH_PATCH", "PARTIAL",
    "BLOCKED", "FAIL", "SKIPPED", "DISCOVERY_ONLY",
}
FAILURE_CODES = {
    "FAIL_SOURCE", "FAIL_ENVIRONMENT", "FAIL_BASELINE",
    "FAIL_PUBLIC_QUANTIZED_ARTIFACT", "FAIL_PUBLIC_QUANTIZED_LOAD",
    "FAIL_BASELINE_PAIR_MISMATCH", "FAIL_EXPORT", "FAIL_UNSUPPORTED_OP",
    "FAIL_DYNAMIC_SHAPE", "FAIL_CONTROL_FLOW", "FAIL_QUANTIZATION_PRESERVATION",
    "FAIL_NUMERICAL_MISMATCH", "FAIL_RUNTIME", "FAIL_MLIR_IMPORT",
    "FAIL_MLIR_QUANT_LEGALIZATION", "FAIL_MLIR_LOWERING", "FAIL_CODEGEN",
    "FAIL_ANALYSIS", "OVER_8MIB",
}


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")


def run_id() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")


def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(chunk_size):
            digest.update(chunk)
    return digest.hexdigest()


def canonical_json_sha256(value: Any) -> str:
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def atomic_write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        mode="w", encoding="utf-8", dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False
    ) as handle:
        json.dump(value, handle, indent=2, sort_keys=True, ensure_ascii=False)
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def load_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def safe_slug(value: str) -> str:
    slug = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip()).strip("_.-")
    if not slug:
        raise ValueError(f"cannot derive a safe path component from {value!r}")
    return slug.lower()


def expand(value: str, variables: dict[str, str]) -> str:
    try:
        return value.format_map(variables)
    except KeyError as error:
        raise ValueError(f"unknown placeholder {error.args[0]!r} in {value!r}") from error


def resolve_path(value: str, variables: dict[str, str], base: Path | None = None) -> Path:
    expanded = Path(expand(value, variables))
    if expanded.is_absolute():
        return expanded.resolve()
    return ((base or REPO_ROOT) / expanded).resolve()


def file_record(path: Path) -> dict[str, Any]:
    exists = path.is_file()
    return {
        "path": str(path),
        "exists": exists,
        "sha256": sha256_file(path) if exists else None,
        "bytes": path.stat().st_size if exists else None,
    }


def shell_join(argv: list[str]) -> str:
    return shlex.join(argv)


def basic_validate_config(config: Any) -> list[str]:
    """Validate safety-critical structure without third-party jsonschema."""
    errors: list[str] = []
    if not isinstance(config, dict):
        return ["configuration root must be an object"]
    for key in ("schema_version", "model", "artifacts", "stages"):
        if key not in config:
            errors.append(f"missing required key: {key}")
    if config.get("schema_version") != "1.0":
        errors.append("schema_version must be '1.0'")
    model = config.get("model", {})
    for key in (
        "model_id", "model_name", "task", "architecture_family", "source_framework",
        "source_repository", "license", "public_quantized_available",
        "paired_fp32_available", "pair_compatibility", "dataset", "input_shape",
        "eligibility", "priority",
    ):
        if key not in model:
            errors.append(f"model missing required key: {key}")
    if model.get("eligibility") not in {
        "ELIGIBLE", "DISCOVERY_ONLY", "BLOCKED_SOURCE", "BLOCKED_BASELINE_PAIR"
    }:
        errors.append("model.eligibility is invalid")
    artifacts = config.get("artifacts", {})
    for variant in ("fp32", "public_quantized"):
        artifact = artifacts.get(variant)
        if not isinstance(artifact, dict):
            errors.append(f"artifacts.{variant} must be an object")
            continue
        for key in ("artifact_id", "source_url", "local_path", "format", "sha256", "license"):
            if key not in artifact:
                errors.append(f"artifacts.{variant} missing required key: {key}")
    stages = config.get("stages")
    if not isinstance(stages, list) or not stages:
        errors.append("stages must be a non-empty array")
        return errors
    seen: set[str] = set()
    for index, stage in enumerate(stages):
        prefix = f"stages[{index}]"
        if not isinstance(stage, dict):
            errors.append(f"{prefix} must be an object")
            continue
        for key in ("id", "stage", "variant", "command", "inputs", "outputs", "timeout_sec", "failure_code_on_error"):
            if key not in stage:
                errors.append(f"{prefix} missing required key: {key}")
        stage_id = stage.get("id")
        if stage_id in seen:
            errors.append(f"duplicate stage id: {stage_id}")
        if isinstance(stage_id, str):
            seen.add(stage_id)
        if stage.get("variant") not in {"fp32", "public_quantized", "pair", "common"}:
            errors.append(f"{prefix}.variant is invalid")
        if not isinstance(stage.get("command"), list) or not stage.get("command"):
            errors.append(f"{prefix}.command must be a non-empty array")
        if stage.get("failure_code_on_error") not in FAILURE_CODES:
            errors.append(f"{prefix}.failure_code_on_error is invalid")
        if not isinstance(stage.get("timeout_sec"), int) or stage.get("timeout_sec", 0) < 1:
            errors.append(f"{prefix}.timeout_sec must be a positive integer")
    for index, stage in enumerate(stages):
        for dependency in stage.get("requires", []):
            if dependency not in seen:
                errors.append(f"stages[{index}] references unknown dependency {dependency!r}")
    return errors


def tool_versions() -> dict[str, Any]:
    result: dict[str, Any] = {
        "python": platform.python_version(),
        "python_implementation": platform.python_implementation(),
        "platform": platform.platform(),
        "executable": sys.executable,
    }
    version_file = REPO_ROOT / "environment" / "tool_versions.json"
    if version_file.is_file():
        try:
            result["captured_environment"] = load_json(version_file)
        except (OSError, json.JSONDecodeError) as error:
            result["captured_environment_error"] = str(error)
    return result