#!/usr/bin/env python3 """ optimize.py — Differential Evolution over pipeline thresholds, scored against X-Ray. Replaces the coarse grid sweep with scipy's differential_evolution over the continuous knob space. Each candidate config is a full-9-film replay (real KPN nodes) scored against Amazon X-Ray presence, micro-averaged. The gallery is loaded once per process (binding caches by path), so an evaluation is just N cheap replays. Objective: **maximize micro-F1** (DE minimizes, so we return -F1). NOTE: X-Ray recall is a face-vs-cast-in-scene ceiling (see [[xray-validation-results]]), so unconstrained F1 tends to push prob_threshold DOWN to recover unreachable recall — trading real precision for it. We therefore log precision/recall at every evaluation and print them at the optimum so the trade-off is visible and you can pick another operating point from the trajectory (--trajectory). Usage: python scripts/optimizer/optimize.py --manifest films.json \ --gallery gallery_arcface_w600k_r50.json \ --params prob_threshold:0.5:0.999 ownership_logodds:0.5:4 track_alpha:0:1 \ --popsize 20 --maxiter 25 --trajectory traj.json """ from __future__ import annotations import argparse import json import sys import time from pathlib import Path import numpy as np from scipy.optimize import differential_evolution REPO = Path(__file__).resolve().parent.parent.parent sys.path.insert(0, str(REPO / "scripts" / "optimizer")) sys.path.insert(0, str(REPO / "scripts" / "validation")) sys.path.insert(0, str(REPO / "scripts")) import json as _json import os import subprocess import tempfile import threading from concurrent.futures import ThreadPoolExecutor # Concurrent per-eval replays. Each replay is an isolated subprocess, so parallelism # is deadlock-safe; with 9 films/eval, 8 workers replays nearly all at once. Tune via # REPLAY_WORKERS (8 is the measured sweet spot on this 24GB GPU). REPLAY_WORKERS = int(os.environ.get("REPLAY_WORKERS", "8")) # DE-level parallelism: how many population candidates get evaluated concurrently # (each spawning its own REPLAY_WORKERS film subprocesses). Total concurrent GPU # replay processes ≈ DE_WORKERS × min(REPLAY_WORKERS, n_films). Threads, not # multiprocessing — each objective() call just waits on subprocess.run, so threads # share the GIL fine and avoid pickling the objective/gallery-key cache. DE_WORKERS = int(os.environ.get("DE_WORKERS", "1")) from second_score import score_seconds # noqa: E402 uniform per-second TPI/FPI scoring from sample_eval import load_gallery_keys # noqa: E402 from replay import dump_embedder_stamp # noqa: E402 from sae_stamp import EmbedderMismatch, verify_gallery_stamp # noqa: E402 _GAL_KEYS: dict = {} # gallery path → key set (fair-recall FN mask), loaded once _REPLAY_TIMEOUT = 45 # seconds per film; a wedged replay is killed, not left to hang REPLAY_CLI = str(Path(__file__).resolve().parent / "replay.py") def _gallery_keys(path): if path not in _GAL_KEYS: _GAL_KEYS[path] = load_gallery_keys(path) return _GAL_KEYS[path] def _replay_subprocess(dump, gallery, cfg, build_dir): """Run one replay in a SUBPROCESS with a timeout, returning its presence JSON. In-process replay intermittently DEADLOCKS at network teardown — a KPN worker stuck mid-rocBLAS GEMM inside the ROCm driver makes ~PyNode's jthread.join() hang forever (root-caused via gdb, 2026-07-15). Isolating each replay means a wedged GPU thread only kills that subprocess; the sweep continues. Returns None on timeout/failure (the caller drops that film from the average).""" with tempfile.NamedTemporaryFile("r", suffix=".json", delete=False) as tf: out = tf.name argv = [sys.executable, REPLAY_CLI, "--dump", dump, "--gallery", gallery, "--out", out, "--build-dir", build_dir] for k, v in cfg.items(): if isinstance(v, bool): # store_true flags: pass the flag, not a value if v: argv.append(f"--{k.replace('_', '-')}") else: argv += [f"--{k.replace('_', '-')}", str(v)] try: subprocess.run(argv, timeout=_REPLAY_TIMEOUT, capture_output=True, check=True) return _json.loads(Path(out).read_text()) except (subprocess.TimeoutExpired, subprocess.CalledProcessError, FileNotFoundError, ValueError) as e: print(f"[opt] replay failed for {Path(dump).name}: {type(e).__name__}", file=sys.stderr) return None finally: try: Path(out).unlink() except OSError: pass def evaluate(cfg, films, build_dir, step=None): """Objective = MACRO-mean over films of each film's duration-weighted per-scene F1. Each film's replay runs in a subprocess (timeout-guarded) to survive the intermittent ROCm teardown deadlock. If ANY film's replay times out, this evaluation is scored f1=0.0 (see below) rather than averaging over the survivors — a partial-coverage eval must never look better than a complete one, or DE will converge onto configs that make the hardest film time out. (An earlier version averaged over survivors, which silently rewarded truncation; the rep4 `mbf_full_noexp` winner was one such corrupted eval.) UNIFORM PER-SECOND scoring (second_score.py): every second of the film is sampled; GT(t) = the cast of the X-Ray scene containing t, Pred(t) = actors whose presence window covers t. Counts instances — TPI / FPI / FN — with FPI weighted 10× when the named actor isn't in the film's cast at all (a real misID vs a timing slip). FN counts only gallery-known actors (fair recall). Reports agreement_rate = mean per-second Jaccard (the "% of on-screen actors we agree with X-Ray about, over time"). Objective = macro-mean across films of the per-second weighted F1. expand_gallery: controlled by env SAE_EXPAND (default on). Set SAE_EXPAND=0 to run the no-expansion arm — the overnight matrix tests both to quantify what expansion buys. The 9 films' replays run CONCURRENTLY (REPLAY_WORKERS) — each is an isolated subprocess, so parallelism is safe (a wedged one only kills itself).""" if os.environ.get("SAE_EXPAND", "1") == "1": cfg = {**cfg, "expand_gallery": True} def _one(film): pj = _replay_subprocess(film["dump"], film.get("gallery"), cfg, build_dir) if pj is None: return None return score_seconds(pj, film["xray"], gallery_keys=_gallery_keys(film.get("gallery"))) with ThreadPoolExecutor(max_workers=REPLAY_WORKERS) as ex: per_film = [m for m in ex.map(_one, films) if m is not None] n = len(per_film) n_expected = len(films) # Incomplete coverage (a replay timed out) is scored as a failure, not # averaged over survivors: dropping the hardest film would otherwise inflate # the score and let DE reward exactly the configs that cause timeouts. We # still record the real survivor counts so a truncated eval is diagnosable # in the trajectory (f1=0.0, films_scored < films_expected). if n < n_expected: agg = {"precision": 0.0, "recall": 0.0, "f1": 0.0, "agreement": 0.0, "TPI": sum(m["TPI"] for m in per_film), "FPI": sum(m["FPI"] for m in per_film), "FPI_misid": sum(m["FPI_misid"] for m in per_film), "FN": sum(m["FN"] for m in per_film)} agg["films_scored"] = n agg["films_expected"] = n_expected return agg return {"precision": sum(m["precision"] for m in per_film) / n, "recall": sum(m["recall"] for m in per_film) / n, "f1": sum(m["f1"] for m in per_film) / n, "agreement": sum(m["agreement_rate"] for m in per_film) / n, "TPI": sum(m["TPI"] for m in per_film), "FPI": sum(m["FPI"] for m in per_film), "FPI_misid": sum(m["FPI_misid"] for m in per_film), "FN": sum(m["FN"] for m in per_film), "films_scored": n, "films_expected": n_expected} def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--manifest", required=True) p.add_argument("--gallery", help="default gallery if not per-film") p.add_argument("--params", nargs="+", required=True, help="knob:lo:hi (e.g. prob_threshold:0.5:0.999). Int knobs kept float, rounded in cfg.") p.add_argument("--build-dir", default=str(REPO / "build")) p.add_argument("--step", type=float, default=5.0) p.add_argument("--popsize", type=int, default=20) p.add_argument("--maxiter", type=int, default=25) p.add_argument("--seed", type=int, default=0) p.add_argument("--trajectory", help="write every evaluation here (JSON lines)") p.add_argument("--out", help="write best config + metrics") # TRACES: GR-004 | SR-001 p.add_argument("--require-gallery-stamp", action="store_true", help="unprovable gallery/dump model binding is a hard error, " "not a warning (also via SAE_REQUIRE_GALLERY_STAMP=1)") args = p.parse_args() if args.require_gallery_stamp: # Set the env var rather than threading a flag through cfg: replays run as # subprocesses and inherit it, so strictness cannot be lost in the handoff. os.environ["SAE_REQUIRE_GALLERY_STAMP"] = "1" films = json.loads(Path(args.manifest).read_text()) for f in films: f.setdefault("gallery", args.gallery) if not Path(f["dump"]).exists(): sys.exit(f"[opt] missing dump for {f['name']}: {f['dump']}") # TRACES: GR-004 | SR-001 # every (dump, gallery) pair is checked ONCE here, # before the first evaluation. A DE sweep is thousands of replays; discovering # a cross-model pair at the end (or never) means every number it produced was # noise. Each replay subprocess re-checks its own pair anyway. for f in films: try: verify_gallery_stamp(f["gallery"], stamp=dump_embedder_stamp(f["dump"]), embedder_desc=f"embedding dump {Path(f['dump']).name}", require_stamp=args.require_gallery_stamp) except EmbedderMismatch as e: sys.exit(f"[opt] {f['name']}: {e}") names, bounds = [], [] int_knobs = {"track_max_frames_missing", "cut_inactive_max_frames"} for spec in args.params: k, lo, hi = spec.split(":") names.append(k); bounds.append((float(lo), float(hi))) print(f"[opt] DE over {names} bounds={bounds}", file=sys.stderr) print(f"[opt] {len(films)} films, popsize={args.popsize}, maxiter={args.maxiter}", file=sys.stderr) traj = [] evals = [0] t0 = time.time() traj_lock = threading.Lock() def vec_to_cfg(x): cfg = {} for k, v in zip(names, x): cfg[k] = int(round(v)) if k in int_knobs else float(v) # The expansion band is [lo, hi]; independent DE bounds can invert it, # and an inverted band admits nothing (track_gallery.hpp). Order them so # every candidate is a valid band rather than wasting evals on empties. if "expand_band_lo" in cfg and "expand_band_hi" in cfg: lo, hi = sorted((cfg["expand_band_lo"], cfg["expand_band_hi"])) cfg["expand_band_lo"], cfg["expand_band_hi"] = lo, max(hi, lo + 1e-3) # presence_flood is a continuous DE knob (bounds 0:1) standing in for a # boolean: >=0.5 selects flood-fill presence. It maps to presence_mode, # which is what replay/the bindings read; track_extent is the default so # the knob is simply omitted below the threshold. if "presence_flood" in cfg: flood = cfg.pop("presence_flood") >= 0.5 if flood: cfg["presence_mode"] = "flood" return cfg def objective(x): cfg = vec_to_cfg(x) m = evaluate(cfg, films, args.build_dir, args.step) with traj_lock: evals[0] += 1 rec = {"eval": evals[0], "config": cfg, **m, "t": round(time.time() - t0, 1)} traj.append(rec) print(f"[opt] eval {evals[0]:3d} thr={cfg['prob_threshold']:.2f} " f"own={cfg.get('ownership_logodds', float('nan')):.2f} → " f"F1={m['f1']*100:.1f}% P={m['precision']*100:.1f}% R={m['recall']*100:.1f}% " f"agree={m.get('agreement', 0)*100:.1f}% misID={m.get('FPI_misid', 0)}", file=sys.stderr) if args.trajectory: with open(args.trajectory, "a") as tf: tf.write(json.dumps(rec) + "\n") return -m["f1"] de_kwargs = dict( popsize=args.popsize, maxiter=args.maxiter, seed=args.seed, polish=False, tol=1e-4, mutation=(0.5, 1.0), recombination=0.7, init="sobol") if DE_WORKERS > 1: pool = ThreadPoolExecutor(max_workers=DE_WORKERS) de_kwargs["workers"] = pool.map result = differential_evolution(objective, bounds, **de_kwargs) best_cfg = vec_to_cfg(result.x) best = evaluate(best_cfg, films, args.build_dir, args.step) print("\n══ DE optimum (by F1) ═══════════════════════════") print(f" config : {best_cfg}") print(f" F1 : {best['f1']*100:.2f}%") print(f" precision: {best['precision']*100:.2f}% recall: {best['recall']*100:.2f}%") print(f" TP/FP/FN: {best['TP']}/{best['FP']}/{best['FN']}") print(f" evaluations: {evals[0]} time: {time.time()-t0:.0f}s") # Also surface the highest-precision config seen (the ship-safe operating point). if traj: hp = max(traj, key=lambda r: (r["precision"], r["recall"])) print("\n── highest-precision config seen (ship-safe) ──") print(f" config : {hp['config']}") print(f" P={hp['precision']*100:.2f}% R={hp['recall']*100:.2f}% F1={hp['f1']*100:.2f}%") if args.out: Path(args.out).write_text(json.dumps( {"best_by_f1": {"config": best_cfg, **best}, "n_evals": evals[0]}, indent=2)) if __name__ == "__main__": main()