src/mcss/vision_probe/v2_runner.py
"""Stage-separated V2 discovery fitting, visible prediction, and sealed evaluation."""
from __future__ import annotations
import argparse
import hashlib
import json
import time
from datetime import UTC, datetime
from pathlib import Path
from types import SimpleNamespace
import numpy as np
import torch
from mcss.vision_probe.adaptation import FastState, Readout, masks, materialize, proposals
from mcss.vision_probe.experiment import _cached_extract, _module_hash
from mcss.vision_probe.features import FrozenDinoExtractor
from mcss.vision_probe.opportunity_selectors import (
enrich_candidates,
select_candidates,
visible_features,
)
from mcss.vision_probe.v2_access import PredictionAccessGuard
from mcss.vision_probe.v2_cycle import source_cycle
from mcss.vision_probe.v2_selectors import TRAJECTORIES, fit_v2, select_v2
ROOT = Path(__file__).resolve().parents[3]
V1_REPORT = ROOT / "docs/experiments/EXP-2D-20260926-opportunity-selector-v1"
V1_WORK = ROOT / "outputs/EXP-2D-20260926-opportunity-selector-v1"
DATA = ROOT / "data/vision_2d_probe/davis2017_trainval_480p/DAVIS"
WEIGHTS = ROOT / "data/vision_2d_probe/weights/dinov2_vits14_pretrain.pth"
METHODS = (
"OFF",
"discovery_fixed16",
"v1_rgb_selector",
"v1_visible_selector",
"GateOnly",
"CycleGate",
)
def sha(path):
return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def read(path):
return json.loads(Path(path).read_text())
def write(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value, sort_keys=True, indent=2, allow_nan=False) + "\n")
def read_rows(path):
return [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
def write_rows(path, rows):
Path(path).write_text(
"".join(json.dumps(r, sort_keys=True, allow_nan=False) + "\n" for r in rows)
)
def _context(paths):
paths = read(paths) if not isinstance(paths, dict) else paths
report, work = Path(paths["report"]), Path(paths["work"])
report.mkdir(parents=True, exist_ok=True)
work.mkdir(parents=True, exist_ok=True)
config = read(report / "config.json")
if config.get("eta", 5.0) != 5.0 or config.get("max_increment_norm", 0.15) != 0.15:
raise ValueError("V1 update formula and fixed hyperparameters must remain unchanged")
torch.set_num_threads(config.get("torch_num_threads", 4))
torch.manual_seed(config.get("seed", 20260926))
return report, work, config
def _v1_artifact():
return read(V1_REPORT / "selectors.json")["by_resolution"]["448"]
def _readout():
provenance = read(V1_REPORT / "readout_448_provenance.json")
if sha(V1_WORK / "readout_448.pt") != provenance["sha256"]:
raise ValueError("Frozen V1 readout hash mismatch")
return Readout(**torch.load(V1_WORK / "readout_448.pt", map_location="cpu", weights_only=True))
def _candidate_states(target, readout):
z = readout.project(target.features)
result = []
start = time.perf_counter()
for trajectory in TRAJECTORIES:
state = FastState.zero(64)
for step, action in enumerate(trajectory):
support, _ = masks(*z.shape[:2], step)
state = proposals(z, target.rgb, readout, state, support, 5.0, 0.15)[action]
adapted = readout.restore(target.features, z, materialize(z, state))
result.append((state, adapted))
return (
z,
result,
{
"proposal_calls": 32,
"increment_calculations": 64,
"candidate_materializations": 16,
"candidate_seconds": time.perf_counter() - start,
},
)
def _discovery_manifest(raw):
sequences = []
for sequence in sorted({r["sequence"] for r in raw}):
records = [r for r in raw if r["sequence"] == sequence]
first = records[0]
def frame(name, digest, sequence=sequence):
return {
"frame_id": name,
"image_path": str(DATA / "JPEGImages/480p" / sequence / name),
"mask_path": str(DATA / "Annotations/480p" / sequence / (Path(name).stem + ".png")),
"image_sha256": digest,
}
source = frame(first["source_frame"], first["source_image_sha256"])
targets = []
for slot in (1, 2):
item = next(r for r in records if r["target_slot"] == slot)
targets.append(
{**frame(item["target_frame"], item["target_image_sha256"]), "target_slot": slot}
)
sequences.append({"sequence": sequence, "source": source, "targets": targets})
return {"sequences": sequences}
def _cached_only(frame, cache_index):
digest = frame["image_sha256"]
candidates = cache_index.get(digest, [])
if len(candidates) != 1:
raise ValueError(f"Expected unique frozen V1 cached feature: {digest}")
tensor_path = candidates[0]
payload = torch.load(tensor_path, map_location="cpu", weights_only=True)
if tuple(payload["features"].shape) != (32, 32, 384):
raise ValueError("Frozen cache is not real 448 features")
return SimpleNamespace(**payload)
def discovery(paths):
report, work, config = _context(paths)
out = report / "discovery_rows.jsonl"
if out.exists():
raise FileExistsError("V2 discovery rows already exist")
raw = [
r
for r in read_rows(V1_REPORT / "discovery_raw.jsonl")
if r["resolution"] == 448 and r["split"] == "discovery"
]
artifact = _v1_artifact()
base = artifact["feature_names"]
manifest = _discovery_manifest(raw)
splits = read(V1_REPORT / "split_manifest.json")
if {s["sequence"] for s in manifest["sequences"]} != set(splits["discovery"]):
raise ValueError("Discovery identity mismatch")
index = {}
cache_hashes = {}
weight_hash, extractor_hash = sha(WEIGHTS), _module_hash()
for metadata in (V1_WORK / "cache/448").glob("*.json"):
info = read(metadata)
if info["weight_sha256"] != weight_hash or info["model_source_sha256"] != extractor_hash:
continue
if info["size"] == 448:
tensor = metadata.with_suffix(".pt")
index.setdefault(info["input_sha256"], []).append(tensor)
readout = _readout()
guard = PredictionAccessGuard(
manifest,
mode="discovery",
allowed_ids=splits["discovery"],
protected_ids=splits["reserve"] + splits["official_val_sealed"] + splits["validation"],
)
rows, costs = [], []
try:
with guard:
for seq in manifest["sequences"]:
source = _cached_only(seq["source"], index)
for target_frame in seq["targets"]:
target = _cached_only(target_frame, index)
for frame in (seq["source"], target_frame):
path = index[frame["image_sha256"]][0]
cache_hashes[str(path)] = sha(path)
mask = guard.read_source_mask(
seq["sequence"], seq["source"]["frame_id"], target_frame["frame_id"]
)
_, candidates, cost = _candidate_states(target, readout)
cycles = [
source_cycle(
source.features,
a,
mask,
seq["source"]["frame_id"],
target_frame["frame_id"],
)
for _, a in candidates
]
pair = [
r
for r in raw
if r["sequence"] == seq["sequence"]
and r["target_slot"] == target_frame["target_slot"]
]
lookup = {tuple(r["trajectory"]): r for r in pair}
if set(lookup) != set(TRAJECTORIES) or len(pair) != 16:
raise ValueError("RAW_INCOMPLETE discovery candidates")
for i, trajectory in enumerate(TRAJECTORIES):
old = lookup[trajectory]
if set(old["visible"]) != set(base):
raise ValueError("V1 visible schema differs from artifact")
rows.append(
{
"split": "discovery",
"sequence": seq["sequence"],
"target_slot": target_frame["target_slot"],
"trajectory": list(trajectory),
"J": old["J"],
"visible": {
**old["visible"],
"f_cycle": cycles[0]["L_cycle"] - cycles[i]["L_cycle"],
},
"L_cycle": cycles[i]["L_cycle"],
"cycle": cycles[i],
}
)
costs.append(
{
"sequence": seq["sequence"],
"target_slot": target_frame["target_slot"],
**cost,
"cycle_seconds": sum(c["cost"]["total_seconds"] for c in cycles),
"correspondence_seconds": sum(
c["cost"]["correspondence_seconds"] for c in cycles
),
"correspondence_calls": 32,
"cycle_calls": 16,
"backbone_calls": 0,
}
)
print(f"discovery {seq['sequence']}/{target_frame['target_slot']}", flush=True)
write_rows(out, rows)
write(
report / "discovery_feature_provenance.json",
{
"v1_raw_sha256": sha(V1_REPORT / "discovery_raw.jsonl"),
"v1_readout_sha256": sha(V1_WORK / "readout_448.pt"),
"cache_hashes": cache_hashes,
"base_feature_names": base,
"new_feature": "f_cycle",
"no_target_gt_reads": not guard.target_log,
},
)
write(report / "discovery_cost.json", costs)
finally:
write(
report / "discovery_access.json",
{
"source_reads": guard.source_log,
"target_reads": guard.target_log,
"denied": guard.denied_log,
},
)
def fit(paths):
report, work, config = _context(paths)
if (report / "selectors.json").exists():
raise FileExistsError("V2 selectors already fitted")
rows = read_rows(report / "discovery_rows.jsonl")
result = fit_v2(rows, _v1_artifact()["feature_names"], config)
write(report / "selectors.json", result["selectors"])
write(report / "nested_cv_results.json", result["nested_cv_results"])
# V1 discovery-fixed16 is also discovery-only and kept unchanged.
write(
report / "selector_lock.json",
{
"selectors_sha256": sha(report / "selectors.json"),
"config_sha256": sha(report / "config.json"),
"discovery_sha256": sha(report / "discovery_rows.jsonl"),
"v1_selectors_sha256": sha(V1_REPORT / "selectors.json"),
"readout_sha256": sha(V1_WORK / "readout_448.pt"),
"weights_sha256": sha(WEIGHTS),
"fit_split": "discovery",
"primary_method": "CycleGate",
"discovery_fixed16": _v1_artifact()["discovery_best_fixed"],
},
)
save_discovery_decisions(report)
def save_discovery_decisions(report):
"""Replay final frozen policies on discovery features, explicitly not OOF."""
report = Path(report)
rows = read_rows(report / "discovery_rows.jsonl")
selectors = read(report / "selectors.json")
v1 = _v1_artifact()
result = []
for sequence, slot in sorted({(r["sequence"], r["target_slot"]) for r in rows}):
group = [r for r in rows if r["sequence"] == sequence and r["target_slot"] == slot]
lookup = {tuple(r["trajectory"]): r for r in group}
group = [lookup[t] for t in TRAJECTORIES]
features = [r["visible"] for r in group]
paths = [list(t) for t in TRAJECTORIES]
base = [{n: f[n] for n in v1["feature_names"]} for f in features]
old = select_candidates(base, paths, v1)
choices = {
"OFF": 0,
"discovery_fixed16": paths.index(v1["discovery_best_fixed"]),
"v1_rgb_selector": old["rgb_selector"],
"v1_visible_selector": old["visible_selector"],
**{n: select_v2(features, paths, a) for n, a in selectors.items()},
}
result.append(
{
"sequence": sequence,
"target_slot": slot,
"choices": choices,
"trajectories": paths,
"scope": "discovery in-sample final policies, NOT outer OOF",
"selected_J": {n: group[i]["J"] for n, i in choices.items()},
}
)
write_rows(report / "discovery_decisions.jsonl", result)
def _verify_selectors(report):
lock = read(report / "selector_lock.json")
checks = {
"selectors_sha256": sha(report / "selectors.json"),
"config_sha256": sha(report / "config.json"),
"discovery_sha256": sha(report / "discovery_rows.jsonl"),
"v1_selectors_sha256": sha(V1_REPORT / "selectors.json"),
"readout_sha256": sha(V1_WORK / "readout_448.pt"),
"weights_sha256": sha(WEIGHTS),
}
if any(lock.get(k) != v for k, v in checks.items()):
raise ValueError("Selector lock mismatch")
return lock
def _source_hashes():
names = (
"adaptation",
"experiment",
"features",
"geometry",
"opportunity_metrics",
"opportunity_selectors",
"v2_selectors",
"v2_cycle",
"v2_access",
"v2_runner",
"v2_dataset",
)
files = [ROOT / f"src/mcss/vision_probe/{name}.py" for name in names]
files.extend((ROOT / "third_party/dinov2").rglob("*.py"))
files.append(ROOT / "scripts/run_vision_2d_v2.py")
return {str(p): sha(p) for p in sorted(files) if p.exists()}
def _guard(manifest):
splits = read(V1_REPORT / "split_manifest.json")
old = read_rows(V1_REPORT / "raw_results.jsonl")
hashes = {r[k] for r in old for k in ("source_image_sha256", "target_image_sha256")}
ids = [s for key in ("fit", "discovery", "validation", "reserve") for s in splits[key]]
return PredictionAccessGuard(
manifest,
historical_ids=ids,
historical_hashes=hashes,
protected_ids=splits["reserve"] + splits["official_val_sealed"],
protected_paths=[DATA],
)
def _prediction_checks(report, manifest_path, document):
lock = read(report / "prediction_lock.json")
checks = {
"selector_lock_sha256": sha(report / "selector_lock.json"),
"config_sha256": sha(report / "config.json"),
"split_sha256": sha(manifest_path),
"source_hashes": _source_hashes(),
}
if any(lock.get(k) != v for k, v in checks.items()):
raise ValueError("Frozen prediction/evaluation code or artifact changed")
manifest = read(manifest_path)
expected = {
(r["sequence"], t["target_slot"]) for r in manifest["sequences"] for t in r["targets"]
}
actual = {(r["sequence"], r["target_slot"]) for r in document["predictions"]}
if (
expected != actual
or len(actual) != len(document["predictions"])
or not document.get("complete")
):
raise ValueError("Incomplete prediction pair coverage")
if document.get("prediction_lock_sha256") != sha(report / "prediction_lock.json"):
raise ValueError("Prediction manifest lock binding mismatch")
for row in document["predictions"]:
if set(row["choices"]) != set(METHODS) or any(
type(i) is not int or not 0 <= i < 16 for i in row["choices"].values()
):
raise ValueError("Missing or invalid deployment decisions")
if row.get("trajectories") != [list(t) for t in TRAJECTORIES]:
raise ValueError("Missing candidate trajectories")
if row["selector_lock_sha256"] != lock["selector_lock_sha256"]:
raise ValueError("Per-pair selector binding mismatch")
path = Path(row["artifact_path"])
if sha(path) != row["artifact_sha256"]:
raise ValueError("Prediction tensor hash mismatch")
payload = torch.load(path, map_location="cpu", weights_only=True)
if payload["adapted"].shape != (16, 32, 32, 384) or payload["source_features"].shape != (
32,
32,
384,
):
raise ValueError("Prediction tensor candidate schema mismatch")
if payload["choices"] != row["choices"] or payload["trajectories"] != row["trajectories"]:
raise ValueError("Tensor/manifest decision mismatch")
_verify_selectors(report)
return lock
def freeze_implementation(paths):
report, _, _ = _context(paths)
_verify_selectors(report)
value = {
"source_hashes": _source_hashes(),
"selector_lock_sha256": sha(report / "selector_lock.json"),
"config_sha256": sha(report / "config.json"),
}
path = report / "implementation_lock.json"
if path.exists() and read(path) != value:
raise ValueError("Implementation snapshot differs; preserve and audit before replacing")
write(path, value)
return value
class _CountedExtractor:
def __init__(self, extractor):
self.extractor = extractor
self.calls = 0
def extract(self, path):
self.calls += 1
return self.extractor.extract(path)
def predict(paths, manifest_path):
report, work, config = _context(paths)
if (report / "prediction_lock.json").exists():
raise FileExistsError(
"Formal prediction already started; audit failures rather than overwrite"
)
selector_lock = _verify_selectors(report)
freeze_implementation(paths)
manifest = read(manifest_path)
guard = _guard(manifest)
selectors = read(report / "selectors.json")
v1 = _v1_artifact()
readout = _readout()
lock = {
"selector_lock_sha256": sha(report / "selector_lock.json"),
"config_sha256": sha(report / "config.json"),
"split_sha256": sha(manifest_path),
"source_hashes": _source_hashes(),
"expected_pairs": 2 * len(manifest["sequences"]),
"target_gt_reads_before_prediction": 0,
}
write(report / "prediction_lock.json", lock)
device = config.get("device", "cuda")
extractor = _CountedExtractor(
FrozenDinoExtractor(ROOT / "third_party/dinov2", WEIGHTS, device, 448)
)
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
records, costs = [], []
try:
with guard:
for seq in manifest["sequences"]:
started = time.perf_counter()
source_count_before = extractor.calls
source = _cached_extract(
extractor,
Path(seq["source"]["image_path"]),
work / "cache448",
sha(WEIGHTS),
_module_hash(),
448,
)
source_seconds = time.perf_counter() - started
source_calls = extractor.calls - source_count_before
for frame in seq["targets"]:
started = time.perf_counter()
target_count_before = extractor.calls
target = _cached_extract(
extractor,
Path(frame["image_path"]),
work / "cache448",
sha(WEIGHTS),
_module_hash(),
448,
)
backbone_seconds = time.perf_counter() - started
target_calls = extractor.calls - target_count_before
mask = guard.read_source_mask(
seq["sequence"], seq["source"]["frame_id"], frame["frame_id"]
)
z, candidates, cost = _candidate_states(target, readout)
start = time.perf_counter()
base = enrich_candidates(
[
visible_features(
source.features, target.features, a, target.rgb, readout, s, z
)
for s, a in candidates
],
[a for _, a in candidates],
)
visible_seconds = time.perf_counter() - start
cycles = [
source_cycle(
source.features, a, mask, seq["source"]["frame_id"], frame["frame_id"]
)
for _, a in candidates
]
features = [
{**b, "f_cycle": cycles[0]["L_cycle"] - cycles[i]["L_cycle"]}
for i, b in enumerate(base)
]
start = time.perf_counter()
paths_list = [list(t) for t in TRAJECTORIES]
v1choices = select_candidates(base, paths_list, v1)
choices = {
"OFF": 0,
"discovery_fixed16": paths_list.index(selector_lock["discovery_fixed16"]),
"v1_rgb_selector": v1choices["rgb_selector"],
"v1_visible_selector": v1choices["visible_selector"],
**{
name: select_v2(features, paths_list, artifact)
for name, artifact in selectors.items()
},
}
selector_seconds = time.perf_counter() - start
tensor_path = (
work / "predictions" / f"{seq['sequence']}_{frame['target_slot']}.pt"
).resolve()
tensor_path.parent.mkdir(exist_ok=True)
torch.save(
{
"source_features": source.features,
"source_mask": torch.from_numpy(mask.astype(np.int64)),
"adapted": torch.stack([a for _, a in candidates]),
"choices": choices,
"trajectories": paths_list,
},
tensor_path,
)
records.append(
{
"sequence": seq["sequence"],
"target_slot": frame["target_slot"],
"source_id": seq["source"]["frame_id"],
"target_id": frame["frame_id"],
"artifact_path": str(tensor_path),
"artifact_sha256": sha(tensor_path),
"choices": choices,
"trajectories": paths_list,
"visible": features,
"L_cycle": [c["L_cycle"] for c in cycles],
"selector_lock_sha256": lock["selector_lock_sha256"],
}
)
costs.append(
{
"sequence": seq["sequence"],
"target_slot": frame["target_slot"],
**cost,
"source_extraction_seconds_shared": source_seconds
if frame["target_slot"] == 1
else 0.0,
"target_extraction_seconds": backbone_seconds,
"backbone_source_calls_shared": source_calls
if frame["target_slot"] == 1
else 0,
"backbone_target_calls": target_calls,
"visible_seconds": visible_seconds,
"visible_readout_calls": 16,
"visible_matching_matrices": 16,
"cycle_calls": 16,
"cycle_correspondence_calls": 32,
"cycle_seconds": sum(c["cost"]["total_seconds"] for c in cycles),
"cycle_correspondence_seconds": sum(
c["cost"]["correspondence_seconds"] for c in cycles
),
"selector_seconds": selector_seconds,
"selector_calls": 4,
"cuda_peak_allocated_bytes": torch.cuda.max_memory_allocated()
if torch.cuda.is_available()
else 0,
}
)
print(f"predict {seq['sequence']}/{frame['target_slot']}", flush=True)
document = {
"complete": True,
"predictions": records,
"prediction_lock_sha256": sha(report / "prediction_lock.json"),
}
write(report / "predictions_manifest.json", document)
_prediction_checks(report, manifest_path, document)
receipt = guard.lock_predictions(
report / "predictions_manifest.json", sha(report / "predictions_manifest.json")
)
write(report / "predictions_lock_receipt.json", receipt)
write(
report / "cost_analysis.json",
{
"pairs": costs,
"scope": (
"All candidates trialed even when OFF selected; CPU adaptation/matching; "
"GPU backbone; no matched-budget efficiency claim"
),
},
)
finally:
write(
report / "prediction_access.json",
{
"source_reads": guard.source_log,
"target_reads": guard.target_log,
"denied": guard.denied_log,
},
)
def evaluate(paths, manifest_path):
report, work, config = _context(paths)
if (report / "raw_results.jsonl").exists() or (report / "evaluation_started.json").exists():
raise FileExistsError(
"Evaluation already started; explicit failure/exposure audit required"
)
document = read(report / "predictions_manifest.json")
_prediction_checks(report, manifest_path, document)
guard = _guard(read(manifest_path))
receipt = read(report / "predictions_lock_receipt.json")
if receipt["manifest_sha256"] != sha(report / "predictions_manifest.json"):
raise ValueError("Prediction manifest no longer matches sealed receipt")
guard.lock_predictions(report / "predictions_manifest.json", receipt["manifest_sha256"])
# Target retrieval is imported only after every prediction and lock has passed validation.
from mcss.vision_probe.opportunity_metrics import evaluate_pair
from mcss.vision_probe.v2_dataset import fetch_target_masks
write(
report / "evaluation_started.json",
{
"utc": datetime.now(UTC).isoformat(),
"predictions_manifest_sha256": receipt["manifest_sha256"],
"warning": "Any failure preserves exposure history; no unaudited restart",
},
)
rows = []
try:
fetch_receipt = fetch_target_masks(
read(manifest_path),
report / "predictions_manifest.json",
receipt["manifest_sha256"],
guard=guard,
)
write(report / "target_fetch_receipt.json", fetch_receipt)
with guard:
for pair in document["predictions"]:
payload = torch.load(pair["artifact_path"], map_location="cpu", weights_only=True)
mask = guard.read_target_mask(
pair["sequence"], pair["source_id"], pair["target_id"]
)
for i, trajectory in enumerate(TRAJECTORIES):
metric = evaluate_pair(
payload["source_features"],
payload["adapted"][i],
payload["source_mask"].numpy(),
mask,
)
rows.append(
{
"resolution": 448,
"split": "independent",
"sequence": pair["sequence"],
"target_slot": pair["target_slot"],
"source_frame": pair["source_id"],
"target_frame": pair["target_id"],
"trajectory": list(trajectory),
"visible": pair["visible"][i],
"L_cycle": pair["L_cycle"][i],
"selected_methods": [
name for name, index in pair["choices"].items() if index == i
],
**metric,
}
)
print(f"evaluate {pair['sequence']}/{pair['target_slot']}", flush=True)
write_rows(report / "raw_results.jsonl", rows)
finally:
write_rows(work / "evaluation_partial_rows.jsonl", rows)
write(
report / "evaluation_access.json",
{
"source_reads": guard.source_log,
"target_reads": guard.target_log,
"denied": guard.denied_log,
},
)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("stage", choices=("discovery", "fit", "predict", "evaluate"))
parser.add_argument("--paths", type=Path, default=Path("/tmp/mcss_v2_paths.json"))
parser.add_argument("--manifest", type=Path)
args = parser.parse_args()
if args.stage in ("predict", "evaluate"):
if args.manifest is None:
parser.error("--manifest required for confirmation stages")
globals()[args.stage](args.paths, args.manifest)
else:
globals()[args.stage](args.paths)
if __name__ == "__main__":
main()