v0.1.8: two-Spark pipeline + kept-warm serving

- Secondary-Spark models now work in air-gapped mode: the LiteLLM proxy is
  dual-homed onto the default bridge (docker network connect) to reach the
  secondary's published vLLM port; grader containers stay on the --internal
  network with zero egress. The head-only enforcement is replaced by a
  secondary-configured check.
- Extraction runs in parallel with grading when the extractor's model and
  every grader model in the wave sit on different Sparks (separate GPUs).
- Keep-warm: single-wave jobs no longer tear the wave down between decks
  (was a ~6-min 31B reload per deck); a kept-warm wave that fails preflight
  is restarted once. Adjudicator reuses the live wave when its model is
  already serving instead of cycling the shared proxy.
- clear_resident_containers (preJobStopContainers) now stops names on every
  configured Spark; health() reports containers on both Sparks.
- Verified with a mocked dry-run of the full job loop (3 decks: one
  bring-up, zero mid-job teardowns, parallel overlap, stale-wave restart).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Jonathan Kirkwood
2026-07-31 15:01:10 -05:00
co-authored by Claude Fable 5
parent 7ec5222834
commit 506e6c79bd
7 changed files with 158 additions and 46 deletions
+87 -18
View File
@@ -94,6 +94,15 @@ def _token(s: str) -> str:
return t or "deck"
def _sparks_disjoint(cfg: dict, extractor_alias: str, wpanel: list[dict]) -> bool:
"""True when the extractor's model and every grader model in the wave sit on
different Sparks (separate GPUs) — extraction can then overlap grading."""
catalog = {m["alias"]: m for m in (cfg.get("models") or [])}
ex_spark = (catalog.get(extractor_alias) or {}).get("spark") or "primary"
return all(((catalog.get(r["model"]) or {}).get("spark") or "primary") != ex_spark
for r in wpanel)
def _composite(record) -> float | None:
"""Best-effort composite lookup on the scoring record (shape owned by scoring.py)."""
if not isinstance(record, dict):
@@ -125,6 +134,7 @@ class JobRunner:
self.period = None
self.decks: list[dict] = []
self.last_report_path = None
self._live_aliases: set[str] | None = None # kept-warm wave's model aliases
self._thread = None
self._last_sig = None
self._last_done_sig = None
@@ -250,6 +260,7 @@ class JobRunner:
self.period = None
self.decks = []
self.panel = []
self._live_aliases = None
remote_root = f"{cfg['remoteWorkDir'].rstrip('/')}/jobs/{job_id}"
self.log(f"=== Grading job {job_id} begins ===")
@@ -294,15 +305,16 @@ class JobRunner:
adjudicate, adj_model = False, ""
needed_all = needed | ({adj_model} if adjudicate and adj_model else set())
# Air-gapped mode can't route to second-Spark models (internal net).
if cfg.get("networkMode") == "airgapped":
cat = {m["alias"]: m for m in models}
offenders = [a for a in needed_all if cat.get(a, {}).get("spark") == "secondary"]
if offenders:
raise RuntimeError(
"air-gapped mode requires all models on the head Spark, but these are "
f"on the secondary: {', '.join(sorted(offenders))}. Move them to the "
"primary Spark or switch to local-services mode.")
# Second-Spark models need the secondary configured. (In air-gapped
# mode the proxy is dual-homed onto the bridge to reach them — the
# grader containers themselves stay on the zero-egress internal net.)
cat = {m["alias"]: m for m in models}
on_secondary = [a for a in needed_all if cat.get(a, {}).get("spark") == "secondary"]
if on_secondary and not (cfg.get("useBothSparks") and cfg.get("secondarySparkHost")):
raise RuntimeError(
f"model(s) {', '.join(sorted(on_secondary))} are assigned to the "
"secondary Spark, but no secondary Spark is configured — "
"run Configure Sparks (enable both Sparks) or move them to primary.")
# 3. Infra once per job.
serving.clear_resident_containers(cfg, self.log)
@@ -350,6 +362,7 @@ class JobRunner:
sc.run(sc.head(cfg), f"rm -rf {remote_root}", timeout=120)
self.log("[runner] wiped deck text from the Spark")
serving.tear_down_all(cfg, self.log)
self._live_aliases = None
# 7. Move the graded originals aside (the company folders stay).
self._drain_inbox(job_id, units)
@@ -376,6 +389,7 @@ class JobRunner:
serving.tear_down_all(cfg, self.log)
except Exception:
pass
self._live_aliases = None
self._persist()
# ------------------------------------------------------------- one deck
@@ -416,19 +430,63 @@ class JobRunner:
self.log(f"[runner] wave {i}/{len(waves)}: models={sorted(aliases)} "
f"graders={[r['name'] for r in wpanel]}"
f"{' +extractor' if extractor_model in aliases else ''}")
serving.bring_up_wave(cfg, wave, hf, self.log)
try:
# Single-wave jobs keep the wave serving across decks — reloading a
# 31B from disk between decks costs minutes for the same model set.
reusable = len(waves) == 1
reused = reusable and self._live_aliases == aliases
if reused:
self.log("[runner] wave already serving from the previous deck — reusing")
else:
serving.bring_up_wave(cfg, wave, hf, self.log)
self._await_serving(cfg, wave)
preflight.check_wave(cfg, wave, self.log)
if extractor_model in aliases:
er = gr_mod.run_extractor(cfg, remote_deck, extractor_model, self.log)
if not er.get("report"):
raise RuntimeError("extractor produced no extraction.json")
if wpanel:
try:
try:
preflight.check_wave(cfg, wave, self.log)
except Exception as e:
if not reused:
raise
# The kept-warm wave went stale (e.g. a vLLM crashed between
# decks) — restart it once and re-check.
self.log(f"[runner] kept-warm wave failed preflight ({e}); restarting it")
self._live_aliases = None
serving.tear_down_wave(cfg, wave, self.log)
serving.bring_up_wave(cfg, wave, hf, self.log)
self._await_serving(cfg, wave)
preflight.check_wave(cfg, wave, self.log)
extract_here = extractor_model in aliases
if (extract_here and wpanel
and _sparks_disjoint(cfg, extractor_model, wpanel)):
# Extractor and graders sit on different Sparks (separate
# GPUs) — run them concurrently.
self.log("[runner] extractor and graders are on different Sparks — "
"running them in parallel")
holder: dict = {}
t = threading.Thread(
target=lambda: holder.update(
gr_mod.run_extractor(cfg, remote_deck, extractor_model,
self.log)),
daemon=True)
t.start()
res = gr_mod.run_wave_graders(cfg, remote_deck, wpanel, self.log)
self._mark_panel(res)
t.join()
if not holder.get("report"):
raise RuntimeError("extractor produced no extraction.json")
else:
if extract_here:
er = gr_mod.run_extractor(cfg, remote_deck, extractor_model,
self.log)
if not er.get("report"):
raise RuntimeError("extractor produced no extraction.json")
if wpanel:
res = gr_mod.run_wave_graders(cfg, remote_deck, wpanel, self.log)
self._mark_panel(res)
finally:
serving.tear_down_wave(cfg, wave, self.log)
if reusable:
self._live_aliases = aliases # leave it serving for the next deck
else:
serving.tear_down_wave(cfg, wave, self.log)
self._live_aliases = None
# --- pull the panel outputs + validate --------------------------------
local_out = os.path.join(local_deck, "out")
@@ -466,13 +524,24 @@ class JobRunner:
try:
adj_model = adj_mod.pick_model(cfg)
for wave in serving.plan_waves(cfg, {adj_model}):
if self._live_aliases and \
{m["alias"] for m in wave} <= self._live_aliases:
# The adjudicator's model is already serving on the
# kept-warm wave — use it; cycling the containers here
# would tear down the proxy the next deck reuses.
self.log("[runner] adjudicator model already serving — reusing wave")
adj_mod.run_adjudication(cfg, remote_deck, self.log)
continue
serving.bring_up_wave(cfg, wave, hf, self.log)
try:
self._await_serving(cfg, wave)
preflight.check_wave(cfg, wave, self.log)
adj_mod.run_adjudication(cfg, remote_deck, self.log)
finally:
# Cycling a different model set invalidates the kept-warm
# wave (this replaces/removes the shared proxy).
serving.tear_down_wave(cfg, wave, self.log)
self._live_aliases = None
local_adj = os.path.join(local_deck, "adjudicator-out")
sc.pull_dir(head, f"{remote_deck}/adjudicator-out", local_adj)
apath = os.path.join(local_adj, "ADJUDICATION.md")