dtourolle
add7e22053
Two failures surfaced scaling the DE sweep up. The replay sink prints a per-second progress line with an explicit flush; under subprocess.run(capture_output=True) those thousands of writes fill a fixed OS pipe buffer that nothing drains until exit, so the long films blocked on write to stderr and looked like hangs. Discard the child's stdout/stderr (DEVNULL) — it was captured and thrown away anyway; the long films then finish in seconds. Separately, the per-film timeout is only a backstop for the rare, intermittent ROCm GEMM wedge (a wedged replay hangs forever and must be killed so the sweep continues), not a performance bound. It had been set huge, which let a single flake stall the whole sweep; set it to a sane 180s (overridable via REPLAY_TIMEOUT) — well above a healthy replay, short enough to reap a wedge quickly.
315 lines
15 KiB
Python
315 lines
15 KiB
Python
#!/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
|
||
# Seconds per film before a replay is killed. Its ONLY job is to escape the rare,
|
||
# intermittent ROCm GEMM wedge (github ROCT-Thunk #56): a wedged replay hangs
|
||
# forever and would otherwise stall the whole sweep, so it must be killed and that
|
||
# film dropped (the eval is then scored as incomplete → F1=0, and DE moves on). It
|
||
# is NOT a performance bound. A healthy replay finishes in ~15-30s even for the
|
||
# long films with stderr discarded, so 180s is comfortably above any real run yet
|
||
# short enough that a wedge is reaped quickly rather than after half an hour.
|
||
# Raise via REPLAY_TIMEOUT if a legitimately slow config is being killed.
|
||
_REPLAY_TIMEOUT = int(os.environ.get("REPLAY_TIMEOUT", "180"))
|
||
|
||
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:
|
||
# Discard the child's stdout/stderr rather than capture it. replay's sink
|
||
# prints a per-second "[result_sink] t=Ns" progress line with an explicit
|
||
# flush; on a long film that is thousands of writes, and under
|
||
# subprocess.run(capture_output=True) they accumulate in a fixed OS pipe
|
||
# buffer that nothing drains until the process exits. On the long films
|
||
# (Valerian, Sound of Metal) under DE concurrency the buffer fills and the
|
||
# C++ process BLOCKS on write to stderr — indistinguishable from a hang, so
|
||
# it hit the timeout and scored F1=0. DEVNULL never fills, so the process
|
||
# runs to completion. (Any real error is still surfaced by check=True.)
|
||
subprocess.run(argv, timeout=_REPLAY_TIMEOUT, check=True,
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||
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()
|