File size: 7,119 Bytes
ed3aeeb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | #!/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
|