scene-actor-extraction/scripts/optimizer/optimize.py

#!/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 anneal_sec:1:30 extinction_sec:1:15 \
      --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"))

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

_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")
    args = p.parse_args()

    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']}")

    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)
        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"ann={cfg['anneal_sec']:.0f} ext={cfg['extinction_sec']:.1f}  → "
                  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()