diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 7e4027425..16156b20d 100644 --- a/mtplx/server/openai.py +++ b/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" diff --git a/tests/test_background_warmup.py b/tests/test_background_warmup.py index 20ace8293..45d375510 100644 --- a/tests/test_background_warmup.py +++ b/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"