Hold the background warmup ladder while a request is in flight

mtplx/server/openai.py

@@ -20462,8 +20462,18 @@ def record_tokens(new_tokens: list[int]) -> None:
             else (getattr(out, "finish_reason", None) or "stop")
         )
         _record_request_metrics(state, dict(envelope))
-        state.last_request_at = time.time()
-        state.requests_completed += 1
+        if not bool((request_observability or {}).get("warmup")):
+            # Warming generations are not user requests: the same filter the
+            # dashboard gauge applies. Stamping this clock from a warm rung
+            # made the ladder defer against its own output (rung 512 finishes,
+            # rung 2560 then waits out a full idle grace) and made /health
+            # report seconds_since_last_request as if a user had just been
+            # served on a daemon nobody has touched. The counter moves with
+            # the clock so the two /health fields cannot disagree, and so the
+            # max-idle watchdog does not read warming as user traffic and
+            # hold the fans at performance.
+            state.last_request_at = time.time()
+            state.requests_completed += 1
         _dashboard_record_completion(state, envelope=envelope, stats=stats)
         last = {
             "text": out.text,
@@ -20691,10 +20701,42 @@ def _idle_grace_s(cls) -> float:
                 return cls.IDLE_GRACE_S
         return cls.IDLE_GRACE_S
 
+    def _foreground_busy(self) -> bool:
+        """True while a real request is generating or queued for the model.
+
+        Only sound at step admission. The warming generation itself calls
+        ``begin_foreground`` for the whole rung, so ``has_foreground()``
+        reads zero here but non-zero once the rung is running: a caller
+        that re-checked from inside a rung would see warmup's own work as
+        live traffic and defer the ladder forever. That is why the sibling
+        ``_ForegroundYield`` reads the scheduler queues instead — it runs
+        during the rung, this runs before it.
+        """
+        state = self.state
+        try:
+            if state.has_foreground():
+                return True
+        except BaseException:
+            pass
+        try:
+            return bool(_foreground_model_work_pending(state))
+        except BaseException:
+            return False
+
     def _foreground_quiet_for_s(self) -> float:
+        # last_request_at is stamped when a request COMPLETES, so a request
+        # that has arrived and is still generating leaves it untouched --
+        # 0.0 on a daemon that has not finished one yet, which read as
+        # "infinitely quiet" and admitted a warm rung against live traffic.
+        # That is the post-restart window an operator actually types into
+        # (a UI that restarts the engine on a config change is typed into
+        # immediately afterwards), so the very first request of a serve was
+        # the one most likely to be warmed over. In-flight and queued work
+        # is _foreground_busy's answer, and it holds the ladder whatever the
+        # grace is; this measures the completion stamp alone.
         last = float(getattr(self.state, "last_request_at", 0.0) or 0.0)
         if last <= 0.0:
-            # No request has ever landed (fresh boot): warm immediately.
+            # Nothing has ever completed: no traffic to stay clear of.
             return float("inf")
         return max(0.0, time.time() - last)
 
@@ -20730,10 +20772,15 @@ def _run_step_inner(self, index: int) -> None:
             self._finish()
             return
         grace = self._idle_grace_s()
-        quiet = self._foreground_quiet_for_s()
-        if quiet < grace:
-            # Recently-served traffic: hold the plan without burning GPU or
-            # the resubmit budget, and re-check when the grace can be met.
+        busy = self._foreground_busy()
+        quiet = 0.0 if busy else self._foreground_quiet_for_s()
+        if busy or quiet < grace:
+            # Live or recently-served traffic: hold the plan without burning
+            # GPU or the resubmit budget, and re-check when the grace can be
+            # met. Busy is its own branch, not a zero folded into the grace
+            # comparison: MTPLX_WARMUP_IDLE_GRACE_S=0 is how an operator says
+            # "do not wait between turns", and it must not also hand a warm
+            # rung a request that is still generating.
             step = self.steps[index]
             if step.get("state") in ("pending", "yielded", "waiting_idle"):
                 step["state"] = "waiting_idle"

tests/test_background_warmup.py

@@ -71,12 +71,17 @@ def make_state(scheduler: FakeScheduler | None = None, **args_overrides):
     )
     for key, value in args_overrides.items():
         setattr(args, key, value)
-    return SimpleNamespace(
+    state = SimpleNamespace(
         args=args,
         model_scheduler=scheduler or FakeScheduler(),
         runtime=SimpleNamespace(tokenizer=FakeTokenizer()),
         context_window=262144,
+        foreground_active=0,
     )
+    # Mirrors ServerState.has_foreground: non-zero while a request is being
+    # served, independent of the completion stamp in last_request_at.
+    state.has_foreground = lambda: state.foreground_active > 0
+    return state
 
 
 def test_background_warmup_enabled_env(monkeypatch):
@@ -115,6 +120,37 @@ def test_foreground_yield_shim_reads_scheduler_queues():
     assert server._ForegroundYield(state).is_set() is False
 
 
+def _deferral_probe(monkeypatch, state, scheduler, grace: str = "90"):
+    """Run one warmup plan with timers faked; report whether any rung
+    actually generated and what the first step's published state is."""
+    monkeypatch.setenv("MTPLX_WARMUP_LADDER", "16")
+    monkeypatch.setenv("MTPLX_WARMUP_IDLE_GRACE_S", grace)
+    monkeypatch.setattr(server, "_prewarm_gqa_packed_pipelines", lambda: True)
+    generations: list[int] = []
+    monkeypatch.setattr(
+        server,
+        "_run_generation",
+        lambda _state, prompt_ids, **kwargs: generations.append(len(prompt_ids))
+        or {"tok_s": 1.0},
+    )
+    timers: list[tuple[float, object, tuple]] = []
+
+    class FakeTimer:
+        def __init__(self, interval, fn, args=()):
+            timers.append((interval, fn, tuple(args)))
+            self.daemon = False
+
+        def start(self):
+            pass
+
+    monkeypatch.setattr(server.threading, "Timer", FakeTimer)
+    status_host: dict = {}
+    warming = server._BackgroundWarmup(state, status_host, [1, 2, 3])
+    warming.submit(0)
+    scheduler.drain()
+    return generations, status_host, timers
+
+
 def test_background_warmup_runs_all_steps_and_publishes_done(monkeypatch):
     scheduler = FakeScheduler()
     state = make_state(scheduler)
@@ -153,31 +189,7 @@ def test_background_warmup_defers_while_foreground_recent(monkeypatch):
     scheduler = FakeScheduler()
     state = make_state(scheduler)
     state.last_request_at = _time.time()  # a response just finished
-    monkeypatch.setenv("MTPLX_WARMUP_LADDER", "16")
-    monkeypatch.setenv("MTPLX_WARMUP_IDLE_GRACE_S", "90")
-    monkeypatch.setattr(server, "_prewarm_gqa_packed_pipelines", lambda: True)
-    generations: list[int] = []
-    monkeypatch.setattr(
-        server,
-        "_run_generation",
-        lambda _state, prompt_ids, **kwargs: generations.append(len(prompt_ids))
-        or {"tok_s": 1.0},
-    )
-    timers: list[tuple[float, object, tuple]] = []
-
-    class FakeTimer:
-        def __init__(self, interval, fn, args=()):
-            timers.append((interval, fn, tuple(args)))
-            self.daemon = False
-
-        def start(self):
-            pass
-
-    monkeypatch.setattr(server.threading, "Timer", FakeTimer)
-    status_host: dict = {}
-    warming = server._BackgroundWarmup(state, status_host, [1, 2, 3])
-    warming.submit(0)
-    scheduler.drain()
+    generations, status_host, timers = _deferral_probe(monkeypatch, state, scheduler)
     # No model work ran; the plan is waiting for idle, budget untouched.
     assert generations == []
     assert status_host["background"]["steps"][0]["state"] == "waiting_idle"
@@ -356,3 +368,78 @@ def test_dashboard_record_completion_skips_warmup_rows():
         stats={},
     )
     assert "lifetime" in calls and "rolling" in calls
+
+
+def test_background_warmup_defers_while_a_request_is_in_flight(monkeypatch):
+    """A request that has ARRIVED but not finished must hold the ladder.
+
+    last_request_at is stamped at completion, so a daemon that has not
+    completed a request yet still reads 0.0 while one is generating. That
+    read as "infinitely quiet" and let a warm rung run against live
+    traffic -- and because a UI that restarts the engine on a config change
+    is typed into immediately afterwards, the very first request of a serve
+    was the one most likely to be warmed over.
+    """
+    scheduler = FakeScheduler()
+    state = make_state(scheduler)
+    state.last_request_at = 0.0  # nothing has COMPLETED yet
+    state.foreground_active = 1  # ...but a real request is generating now
+    generations, status_host, timers = _deferral_probe(monkeypatch, state, scheduler)
+    assert generations == [], "warming generated while a request was in flight"
+    assert status_host["background"]["steps"][0]["state"] == "waiting_idle"
+    assert status_host["background"]["resubmits"] == 0
+    assert len(timers) == 1
+
+
+def test_background_warmup_holds_a_live_request_even_at_zero_grace(monkeypatch):
+    """MTPLX_WARMUP_IDLE_GRACE_S=0 drops the between-turns wait, not the
+    live-traffic guard.
+
+    Expressing "busy" as zero quiet made the hold collapse at grace 0
+    (``0.0 < 0.0`` is False), so the one knob an operator reaches for when
+    warm rungs are not running also handed them a request that was still
+    generating -- the exact starvation the guard exists to prevent.
+    """
+    scheduler = FakeScheduler()
+    state = make_state(scheduler)
+    state.last_request_at = 0.0
+    state.foreground_active = 1  # a real request is generating right now
+    generations, status_host, timers = _deferral_probe(
+        monkeypatch, state, scheduler, grace="0"
+    )
+    assert generations == [], "warming generated against a live request at grace 0"
+    assert status_host["background"]["steps"][0]["state"] == "waiting_idle"
+    assert status_host["background"]["resubmits"] == 0
+    assert len(timers) == 1
+    wait_s, fn, args = timers[0]
+    assert wait_s >= 1.0  # _defer_step floors the re-check, never a hot loop
+    # The request finishes and nothing else is queued: zero grace means the
+    # plan runs on the very next admission, with no wait to serve out.
+    state.foreground_active = 0
+    fn(*args)
+    scheduler.drain()
+    assert generations == [16]
+    assert status_host["background"]["state"] == "done"
+
+
+def test_background_warmup_defers_while_foreground_is_queued(monkeypatch):
+    """Foreground work queued on the scheduler counts as busy at admission
+    time, not only once it starts executing."""
+    scheduler = FakeScheduler()
+    scheduler.foreground_busy = True
+    state = make_state(scheduler)
+    state.last_request_at = 0.0
+    generations, status_host, _ = _deferral_probe(monkeypatch, state, scheduler)
+    assert generations == [], "warming generated while foreground work was queued"
+    assert status_host["background"]["steps"][0]["state"] == "waiting_idle"
+
+
+def test_background_warmup_still_warms_a_genuinely_idle_fresh_daemon(monkeypatch):
+    """The guard must not cost the case it exists for: nothing running and
+    nothing ever completed still warms immediately."""
+    scheduler = FakeScheduler()
+    state = make_state(scheduler)
+    state.last_request_at = 0.0
+    generations, status_host, _ = _deferral_probe(monkeypatch, state, scheduler)
+    assert generations == [16]
+    assert status_host["background"]["state"] == "done"