scripts/netron_capture_common.py

#!/usr/bin/env python3
"""Shared provenance and file helpers for the T80 Netron capture batch."""

from __future__ import annotations

import csv
import hashlib
import json
import os
import tempfile
from collections.abc import Iterable
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[1]
VARIANTS = ("fp32", "public_quantized")
FORMATS = ("onnx",)
STAGE_COLUMNS = {
    ("onnx", "fp32"): "s2_fp32_onnx",
    ("onnx", "public_quantized"): "s3_public_quantized_onnx",
}
STAGE_PRIORITY = {
    ("onnx", "fp32"): ("validate_fp32_onnx",),
    ("onnx", "public_quantized"): ("validate_quantized_onnx",),
}
PRODUCTION_STAGE = {
    ("onnx", "fp32"): "convert_fp32_onnx",
    ("onnx", "public_quantized"): "convert_quantized_onnx",
}


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def relative(path: Path, root: Path) -> str:
    resolved = path.resolve()
    try:
        return str(resolved.relative_to(root.resolve()))
    except ValueError:
        return str(resolved)


def resolve(root: Path, value: str | Path) -> Path:
    path = Path(value)
    if not path.is_absolute():
        return root / path
    if path.exists():
        return path
    # Immutable run records contain the absolute repository root used at
    # execution time. Rebase only a known repository-owned suffix after the
    # workspace is moved; checksums still guard the selected artifact bytes.
    anchors = ("models", "configs", "environment", "reports", "results", "logs", "research")
    for anchor in anchors:
        if anchor in path.parts:
            index = path.parts.index(anchor)
            return root.joinpath(*path.parts[index:])
    return path


def load_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, 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 atomic_csv(path: Path, rows: Iterable[dict[str, Any]], fieldnames: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", newline="", encoding="utf-8", dir=path.parent, delete=False) as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)
        temporary = Path(handle.name)
    os.replace(temporary, path)


def task_group(task: str) -> str:
    normalized = task.strip().lower().replace(" ", "_")
    if normalized == "anomaly_detection":
        return "Anomaly Detection"
    if normalized == "object_detection":
        return "Detection"
    if normalized == "semantic_segmentation":
        return "Segmentation"
    if normalized == "keyword_spotting":
        return "Speech / KWS"
    if normalized == "vision_classification":
        return "Vision"
    if normalized.startswith("language_modeling"):
        return "Language Model"
    return task


def _model_directory(root: Path, model_id: str) -> Path:
    matches = sorted(path for path in (root / "models").glob(f"**/{model_id}") if path.is_dir())
    if len(matches) != 1:
        raise ValueError(f"expected one model directory for {model_id}, found {matches}")
    return matches[0]


def _run_results(model_dir: Path) -> list[tuple[Path, dict[str, Any]]]:
    results: list[tuple[Path, dict[str, Any]]] = []
    for path in sorted(model_dir.glob("*run_result.json")):
        if "dry_run" in path.name:
            continue
        value = json.loads(path.read_text(encoding="utf-8"))
        if isinstance(value.get("stages"), list):
            results.append((path, value))
    return results


def _stage_input(stage: dict[str, Any], fmt: str) -> dict[str, Any] | None:
    matches = []
    for value in stage.get("inputs", []):
        path = str(value.get("path", ""))
        if Path(path).suffix.lower() != f".{fmt}":
            continue
        if fmt == "onnx" and ".inferred.onnx" in path.lower():
            continue
        matches.append(value)
    if not matches:
        return None
    matches.sort(key=lambda value: (not bool(value.get("exists")), str(value.get("path", ""))))
    return matches[0]


def _validation_evidence(
    run_results: list[tuple[Path, dict[str, Any]]], fmt: str, variant: str
) -> tuple[Path, dict[str, Any], dict[str, Any]] | None:
    for stage_id in STAGE_PRIORITY[(fmt, variant)]:
        candidates: list[tuple[Path, dict[str, Any], dict[str, Any]]] = []
        for result_path, result in run_results:
            for stage in result["stages"]:
                if stage.get("stage_id") != stage_id:
                    continue
                artifact_input = _stage_input(stage, fmt)
                if artifact_input is not None:
                    candidates.append((result_path, stage, artifact_input))
        if candidates:
            candidates.sort(
                key=lambda value: (
                    not bool(value[2].get("exists")),
                    value[1].get("status") not in {"PASS", "PASS_WITH_PATCH"},
                    value[0].name != "run_result.json",
                    str(value[0]),
                )
            )
            return candidates[0]
    return None


def _production_evidence(
    run_results: list[tuple[Path, dict[str, Any]]], fmt: str, variant: str
) -> tuple[Path, dict[str, Any]] | None:
    stage_id = PRODUCTION_STAGE[(fmt, variant)]
    candidates: list[tuple[Path, dict[str, Any]]] = []
    for result_path, result in run_results:
        for stage in result["stages"]:
            if stage.get("stage_id") == stage_id:
                candidates.append((result_path, stage))
    if not candidates:
        return None
    candidates.sort(
        key=lambda value: (
            value[1].get("status") not in {"PASS", "PASS_WITH_PATCH"},
            value[0].name != "run_result.json",
            str(value[0]),
        )
    )
    return candidates[0]


def discover_slots(root: Path = REPO_ROOT) -> list[dict[str, Any]]:
    """Return the 21 x 2 ONNX artifact inventory used for Netron export."""
    root = root.resolve()
    registry = [
        row for row in load_csv(root / "model_registry.csv") if row.get("eligibility") == "ELIGIBLE"
    ]
    conversion = {
        row["model_id"]: row for row in load_csv(root / "reports/conversion/pipeline_status.csv")
    }
    slots: list[dict[str, Any]] = []
    for registry_row in sorted(registry, key=lambda row: row["model_id"]):
        model_id = registry_row["model_id"]
        model_dir = _model_directory(root, model_id)
        results = _run_results(model_dir)
        for variant in VARIANTS:
            for fmt in FORMATS:
                evidence = _validation_evidence(results, fmt, variant)
                production_evidence = _production_evidence(results, fmt, variant)
                stage_result_path: Path | None = None
                stage: dict[str, Any] = {}
                artifact_input: dict[str, Any] = {}
                if evidence is not None:
                    stage_result_path, stage, artifact_input = evidence
                production_result_path: Path | None = None
                production_stage: dict[str, Any] = {}
                if production_evidence is not None:
                    production_result_path, production_stage = production_evidence
                raw_path = str(artifact_input.get("path", ""))
                artifact_path = resolve(root, raw_path) if raw_path else None
                exists = bool(artifact_path and artifact_path.is_file())
                recorded_sha = str(artifact_input.get("sha256") or "")
                current_sha = sha256(artifact_path) if exists and artifact_path else ""
                checksum_match = bool(exists and recorded_sha and current_sha == recorded_sha)
                if exists and checksum_match:
                    artifact_status = "AVAILABLE"
                elif exists:
                    artifact_status = "BLOCKED_CHECKSUM_MISMATCH"
                else:
                    artifact_status = "NOT_AVAILABLE"
                # ONNX is the common format used by the ONNX-MLIR pipeline, so
                # the Netron pair view uses the two ONNX variants directly.
                canonical = bool(exists and fmt == "onnx")
                output_dir = model_dir / "graphs/netron" / variant
                output_png = output_dir / f"{fmt}_netron.png"
                metadata_json = output_dir / f"{fmt}_netron.metadata.json"
                matrix = conversion[model_id]
                slots.append(
                    {
                        "model_id": model_id,
                        "task": registry_row["task"],
                        "task_group": task_group(registry_row["task"]),
                        "architecture_family": registry_row["architecture_family"],
                        "variant": variant,
                        "format": fmt,
                        "pipeline_stage": STAGE_COLUMNS[(fmt, variant)],
                        "pipeline_stage_status": matrix[STAGE_COLUMNS[(fmt, variant)]],
                        "validation_stage_id": stage.get("stage_id", "NOT_FOUND"),
                        "validation_stage_status": stage.get("status", "NOT_FOUND"),
                        "validation_failure_code": stage.get("failure_code") or "",
                        "validation_exit_code": stage.get("exit_code", ""),
                        "validation_command": stage.get("command", ""),
                        "validation_stdout_log": stage.get("stdout_log", ""),
                        "validation_stderr_log": stage.get("stderr_log", ""),
                        "production_stage_id": production_stage.get("stage_id", "NOT_FOUND"),
                        "production_stage_status": production_stage.get("status", "NOT_FOUND"),
                        "production_failure_code": production_stage.get("failure_code") or "",
                        "production_exit_code": production_stage.get("exit_code", ""),
                        "production_command": production_stage.get("command", ""),
                        "production_stdout_log": production_stage.get("stdout_log", ""),
                        "production_stderr_log": production_stage.get("stderr_log", ""),
                        "production_run_result": relative(production_result_path, root) if production_result_path else "",
                        "production_run_result_sha256": sha256(production_result_path) if production_result_path else "",
                        "source_run_result": relative(stage_result_path, root) if stage_result_path else "",
                        "source_run_result_sha256": sha256(stage_result_path) if stage_result_path else "",
                        "source_artifact": relative(artifact_path, root) if artifact_path else raw_path,
                        "source_artifact_exists": exists,
                        "source_artifact_bytes": artifact_path.stat().st_size if exists and artifact_path else 0,
                        "recorded_source_sha256": recorded_sha,
                        "current_source_sha256": current_sha,
                        "source_checksum_match": checksum_match,
                        "artifact_status": artifact_status,
                        "canonical_s7_selected": canonical,
                        "canonical_s7_analysis_format": "onnx",
                        "canonical_s7_source_sha256": current_sha if fmt == "onnx" else "",
                        "output_png": relative(output_png, root),
                        "metadata_json": relative(metadata_json, root),
                    }
                )
    return slots


INVENTORY_FIELDS = [
    "model_id",
    "task",
    "task_group",
    "architecture_family",
    "variant",
    "format",
    "pipeline_stage",
    "pipeline_stage_status",
    "validation_stage_id",
    "validation_stage_status",
    "validation_failure_code",
    "validation_exit_code",
    "validation_command",
    "validation_stdout_log",
    "validation_stderr_log",
    "production_stage_id",
    "production_stage_status",
    "production_failure_code",
    "production_exit_code",
    "production_command",
    "production_stdout_log",
    "production_stderr_log",
    "production_run_result",
    "production_run_result_sha256",
    "source_run_result",
    "source_run_result_sha256",
    "source_artifact",
    "source_artifact_exists",
    "source_artifact_bytes",
    "recorded_source_sha256",
    "current_source_sha256",
    "source_checksum_match",
    "artifact_status",
    "canonical_s7_selected",
    "canonical_s7_analysis_format",
    "canonical_s7_source_sha256",
    "output_png",
    "metadata_json",
]