228 lines
7.6 KiB
Python
228 lines
7.6 KiB
Python
|
"""Batch manifest creation and provenance capture."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib
|
||
|
|
import importlib.metadata
|
||
|
|
import platform
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Mapping
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .hashing import file_sha256, stable_hash
|
||
|
|
from .io import atomic_write_json, utc_now
|
||
|
|
from .plan import validate_trial_plan
|
||
|
|
from .schema import SCHEMA_VERSION, STORAGE_FORMATS, ContractError, require_schema
|
||
|
|
|
||
|
|
|
||
|
|
_SOURCE_FILES = (
|
||
|
|
Path("pyproject.toml"),
|
||
|
|
Path("uv.lock"),
|
||
|
|
Path("code/config/master_7dof.urdf"),
|
||
|
|
Path("code/config/real_slave_7dof.urdf"),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _git_provenance(cwd: Path) -> dict[str, Any]:
|
||
|
|
try:
|
||
|
|
commit = subprocess.run(
|
||
|
|
["git", "rev-parse", "HEAD"],
|
||
|
|
cwd=cwd,
|
||
|
|
check=True,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=5,
|
||
|
|
).stdout.strip()
|
||
|
|
status = subprocess.run(
|
||
|
|
["git", "status", "--porcelain"],
|
||
|
|
cwd=cwd,
|
||
|
|
check=True,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=5,
|
||
|
|
).stdout
|
||
|
|
return {"available": True, "commit": commit, "dirty": bool(status.strip())}
|
||
|
|
except (OSError, subprocess.SubprocessError):
|
||
|
|
return {"available": False, "commit": None, "dirty": None}
|
||
|
|
|
||
|
|
|
||
|
|
def _source_file_hashes(source_root: Path) -> dict[str, dict[str, str]]:
|
||
|
|
"""Hash reproducibility-critical files without requiring every file."""
|
||
|
|
hashes: dict[str, dict[str, str]] = {}
|
||
|
|
for relative_path in _SOURCE_FILES:
|
||
|
|
path = source_root / relative_path
|
||
|
|
if path.is_file():
|
||
|
|
hashes[relative_path.as_posix()] = {"sha256": file_sha256(path)}
|
||
|
|
return hashes
|
||
|
|
|
||
|
|
|
||
|
|
def _optional_module_version(module_name: str) -> str | None:
|
||
|
|
"""Return an importable module's version, including distribution fallback."""
|
||
|
|
try:
|
||
|
|
module = importlib.import_module(module_name)
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
|
||
|
|
version = getattr(module, "__version__", None)
|
||
|
|
if version is not None:
|
||
|
|
return str(version)
|
||
|
|
|
||
|
|
try:
|
||
|
|
distributions = importlib.metadata.packages_distributions().get(
|
||
|
|
module_name, ()
|
||
|
|
)
|
||
|
|
for distribution in distributions:
|
||
|
|
try:
|
||
|
|
return importlib.metadata.version(distribution)
|
||
|
|
except importlib.metadata.PackageNotFoundError:
|
||
|
|
continue
|
||
|
|
except (ImportError, OSError):
|
||
|
|
pass
|
||
|
|
return "unknown"
|
||
|
|
|
||
|
|
|
||
|
|
def _runtime_provenance() -> dict[str, str]:
|
||
|
|
runtime = {
|
||
|
|
"python": sys.version.split()[0],
|
||
|
|
"numpy": np.__version__,
|
||
|
|
"platform": platform.platform(),
|
||
|
|
"processor": platform.processor(),
|
||
|
|
}
|
||
|
|
for module_name in ("scipy", "pinocchio"):
|
||
|
|
version = _optional_module_version(module_name)
|
||
|
|
if version is not None:
|
||
|
|
runtime[module_name] = version
|
||
|
|
return runtime
|
||
|
|
|
||
|
|
|
||
|
|
def _require_locked_git(
|
||
|
|
plan: Mapping[str, Any],
|
||
|
|
provenance: Mapping[str, Any],
|
||
|
|
source_root: Path,
|
||
|
|
) -> None:
|
||
|
|
"""Reject evidence-locked runs that cannot identify immutable source."""
|
||
|
|
if plan.get("split") != "locked":
|
||
|
|
return
|
||
|
|
if not provenance.get("available"):
|
||
|
|
raise ContractError(
|
||
|
|
"locked split requires available Git metadata; "
|
||
|
|
f"source_root={source_root} is not a usable Git worktree"
|
||
|
|
)
|
||
|
|
if provenance.get("dirty"):
|
||
|
|
raise ContractError(
|
||
|
|
"locked split requires a clean Git worktree; "
|
||
|
|
"commit or stash all source changes before running"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _require_locked_resume_match(
|
||
|
|
stored_manifest: Mapping[str, Any],
|
||
|
|
current_provenance: Mapping[str, Any],
|
||
|
|
source_root: Path,
|
||
|
|
) -> None:
|
||
|
|
"""Require a resumed locked batch to use its original immutable source."""
|
||
|
|
stored_source = stored_manifest.get("source")
|
||
|
|
if not isinstance(stored_source, Mapping):
|
||
|
|
raise ContractError("locked batch manifest is missing source provenance")
|
||
|
|
|
||
|
|
stored_commit = stored_source.get("commit")
|
||
|
|
current_commit = current_provenance.get("commit")
|
||
|
|
if not stored_commit or current_commit != stored_commit:
|
||
|
|
raise ContractError(
|
||
|
|
"locked batch Git commit mismatch; "
|
||
|
|
f"stored={stored_commit!r}, current={current_commit!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
stored_files = stored_source.get("files")
|
||
|
|
if not isinstance(stored_files, Mapping):
|
||
|
|
raise ContractError(
|
||
|
|
"locked batch manifest is missing reproducibility file hashes"
|
||
|
|
)
|
||
|
|
current_files = _source_file_hashes(source_root)
|
||
|
|
if dict(stored_files) != current_files:
|
||
|
|
changed_paths = sorted(
|
||
|
|
relative_path
|
||
|
|
for relative_path in set(stored_files) | set(current_files)
|
||
|
|
if stored_files.get(relative_path) != current_files.get(relative_path)
|
||
|
|
)
|
||
|
|
raise ContractError(
|
||
|
|
"locked batch source file hash mismatch for: "
|
||
|
|
+ ", ".join(changed_paths)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def build_batch_manifest(
|
||
|
|
plan: Mapping[str, Any],
|
||
|
|
*,
|
||
|
|
source_root: Path | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
validate_trial_plan(plan)
|
||
|
|
source_root = Path.cwd() if source_root is None else Path(source_root)
|
||
|
|
source = _git_provenance(source_root)
|
||
|
|
_require_locked_git(plan, source, source_root)
|
||
|
|
source["files"] = _source_file_hashes(source_root)
|
||
|
|
batch_id = f"batch-{str(plan['plan_hash'])[:16]}"
|
||
|
|
manifest = {
|
||
|
|
"kind": "batch_manifest",
|
||
|
|
"schema_version": SCHEMA_VERSION,
|
||
|
|
"batch_id": batch_id,
|
||
|
|
"study_id": plan["study_id"],
|
||
|
|
"split": plan["split"],
|
||
|
|
"plan_hash": plan["plan_hash"],
|
||
|
|
"created_utc": utc_now(),
|
||
|
|
"expected_pair_count": plan["pair_count"],
|
||
|
|
"expected_trial_count": plan["trial_count"],
|
||
|
|
"storage_formats": dict(STORAGE_FORMATS),
|
||
|
|
"runtime": _runtime_provenance(),
|
||
|
|
"source": source,
|
||
|
|
}
|
||
|
|
hash_basis = dict(manifest)
|
||
|
|
hash_basis.pop("created_utc")
|
||
|
|
manifest["manifest_hash"] = stable_hash(hash_basis, prefix="batch-manifest")
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_batch(
|
||
|
|
batch_dir: Path,
|
||
|
|
plan: Mapping[str, Any],
|
||
|
|
*,
|
||
|
|
source_root: Path | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Create or verify the immutable plan and batch manifest."""
|
||
|
|
validate_trial_plan(plan)
|
||
|
|
batch_dir = Path(batch_dir)
|
||
|
|
source_root = Path.cwd() if source_root is None else Path(source_root)
|
||
|
|
manifest_path = batch_dir / "batch_manifest.json"
|
||
|
|
plan_path = batch_dir / "plan.json"
|
||
|
|
if manifest_path.exists() or plan_path.exists():
|
||
|
|
if not manifest_path.exists() or not plan_path.exists():
|
||
|
|
raise ContractError("batch has only one of plan.json/batch_manifest.json")
|
||
|
|
current_source = _git_provenance(source_root)
|
||
|
|
_require_locked_git(plan, current_source, source_root)
|
||
|
|
from .plan import load_document
|
||
|
|
|
||
|
|
stored_plan = load_document(plan_path)
|
||
|
|
stored_manifest = load_document(manifest_path)
|
||
|
|
validate_trial_plan(stored_plan)
|
||
|
|
require_schema(stored_manifest, "batch_manifest")
|
||
|
|
if plan.get("split") == "locked":
|
||
|
|
_require_locked_resume_match(
|
||
|
|
stored_manifest,
|
||
|
|
current_source,
|
||
|
|
source_root,
|
||
|
|
)
|
||
|
|
if stored_plan["plan_hash"] != plan["plan_hash"]:
|
||
|
|
raise ContractError("existing batch was created from a different plan")
|
||
|
|
if stored_manifest.get("plan_hash") != plan["plan_hash"]:
|
||
|
|
raise ContractError("batch manifest plan_hash mismatch")
|
||
|
|
return stored_manifest
|
||
|
|
|
||
|
|
batch_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
manifest = build_batch_manifest(plan, source_root=source_root)
|
||
|
|
atomic_write_json(plan_path, plan)
|
||
|
|
atomic_write_json(manifest_path, manifest)
|
||
|
|
return manifest
|