diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index 62aae5f15a4a..b997dc39a88f 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -246,13 +246,15 @@ class _BenchmarkGroupResult: @dataclass -class _BenchmarkStageExchange: - """One pending warm-up stage exchange (``_BenchmarkSynchronizer.stage_report``).""" +class _BenchmarkStageRound: + """Group view of one warm-up stage round (``_BenchmarkSynchronizer.stage_round``). - batch: int | None - deadline: float - reports: dict[int, bool] - identities: dict[int, bytes] + ``all_done`` once every rank has closed its build for the rung; ``ok`` is + then the group verdict (every rank's local outcome ok), else None. + """ + + all_done: bool + ok: bool | None @dataclass(frozen=True) @@ -559,9 +561,10 @@ class _BenchmarkSynchronizer: excluded from the measured iteration wall time. Small post-GO delivery and model-runner launch skew can remain because this operates at scheduler level. - The KV warm-up stage exchange (``stage_report`` / ``stage_poll``) is the - one non-blocking phase: ranks report and poll between idle steps, because - a peer may still be running the collective forward passes of its build. + The KV warm-up stage round (``stage_round``) is a blocking phase like the + others, but every rank enters it at the same point of every idle step, + before its forward pass, so no rank ever blocks here while a peer is + still running the collective forward passes of its build. """ MAX_SYNC_TIMEOUT_SECONDS = 10 @@ -573,6 +576,8 @@ class _BenchmarkSynchronizer: # between ranks. The phase precedes every measurement, so a long wait # costs startup time only. CAPACITY_TIMEOUT_SECONDS = 300 + # Stage status a rank reports in a warm-up round (``stage_round``). + STAGE_PHASES = ("none", "building", "done") def __init__( self, @@ -612,7 +617,6 @@ def __init__( self._socket.setsockopt(zmq.LINGER, 0) self._cleanup_complete = False self._flush_on_close = False - self._stage: _BenchmarkStageExchange | None = None @property def run_id(self) -> str | None: @@ -981,123 +985,115 @@ def synchronize_cleanup(self) -> None: self._recv_follower(deadline, benchmark_id, "cleanup_complete") self._cleanup_complete = True - def stage_report( - self, batch: int | None, ok: bool, *, timeout: float | None = None - ) -> None: - """Publish this rank's outcome for the warm-up stage of rung ``batch`` - without blocking; ``stage_poll`` then drives the exchange. - - A follower sends one ``stage_status``; rank 0 records its own. Both - keep polling from the scheduler's idle steps, so no rank ever blocks - in ``schedule()`` while a peer may still be running the collective - forward passes of its own build. ``timeout`` (default - ``timeout_seconds``) bounds the wait from this report; the caller - passes a longer one while peers may legitimately still be building. - - Isolation from the point protocol: the exchange runs only in the - DECODE_SWEEP window between the previous point's result commit (or - the grid commit) and the next point's boundary/READY, once this - rank's build for the rung has ended. Every rank derives the rung - from the negotiated plan, so every rank enters the exchange for the - same ``batch``, and no other protocol message is in flight in that - window: the only traffic is stage_status/stage_decision plus the + def stage_round( + self, seq: int, batch: int | None, phase: str, ok: bool | None + ) -> _BenchmarkStageRound: + """Exchange this rank's warm-up stage status with the group; blocking. + + Every attention-DP rank calls this once per idle DECODE_SWEEP step, at + the same point of the step and before its forward pass, so no rank + ever waits here while a peer is still inside a collective. ``seq`` is + this rank's count of rounds, ``phase`` is ``"none"`` (no stage + active), ``"building"`` or ``"done"``, and ``ok`` is the local outcome + once done. Rank 0 collects every follower's status and every rank + receives the same group view in the same step: ``all_done`` once every + rank is done, with the verdict ``ok`` = every local outcome ok. A + follower whose ``seq`` or ``batch`` disagrees with rank 0 has fallen + out of lockstep; that is a protocol error, never a wait. Past the + protocol timeout it raises TimeoutError; a protocol violation or a + peer abort raises RuntimeError, like the other phases. + + Isolation from the point protocol: rounds run only in the DECODE_SWEEP + idle steps between a point's result commit (or the grid commit) and + the next point's injection, so no other protocol message is in flight; + the only traffic is stage_round/stage_round_result plus the abort/error notices every phase honours. """ - if self._stage is not None: - raise RuntimeError("attention-DP warm-up stage exchange already pending") - wait = self.timeout_seconds if timeout is None else timeout - self._stage = _BenchmarkStageExchange( - batch=batch, - deadline=time.monotonic() + wait, - reports={self.dp_rank: ok}, - identities={}, - ) - if self.dp_rank != 0: - self._socket.send_json( - { - "type": "stage_status", - "benchmark_id": 0, - "dp_rank": self.dp_rank, - "batch": batch, - "ok": ok, - } + if phase not in self.STAGE_PHASES: + raise ValueError(f"invalid attention-DP warm-up stage phase: {phase!r}") + if (phase == "done") != isinstance(ok, bool): + raise ValueError( + f"attention-DP warm-up stage status {phase!r} with outcome {ok!r}" ) - - def stage_poll(self) -> bool | None: - """Advance the pending stage exchange without blocking. - - Returns the group verdict (every rank reported ok) once it is known, - else None so the caller yields the step and polls again. Past the - report deadline it raises TimeoutError; a protocol violation or a - peer abort raises RuntimeError, like the blocking phases. - """ - stage = self._stage - if stage is None: - raise RuntimeError("attention-DP warm-up stage poll without a report") + status = { + "type": "stage_round", + "benchmark_id": 0, + "dp_rank": self.dp_rank, + "seq": seq, + "batch": batch, + "phase": phase, + "ok": ok, + } if self.dp_rank == 0: - return self._coordinate_stage(stage) - return self._follow_stage(stage) + return self._coordinate_stage_round(status) + deadline = time.monotonic() + self.timeout_seconds + self._socket.send_json(status) + reply = self._recv_follower(deadline, 0, "stage_round_result") + all_done = reply.get("all_done") + verdict = reply.get("ok") + if ( + reply.get("seq") != seq + or not isinstance(all_done, bool) + or (all_done and not isinstance(verdict, bool)) + or (not all_done and verdict is not None) + ): + raise RuntimeError( + f"invalid attention-DP warm-up stage round result: {reply}" + ) + return _BenchmarkStageRound(all_done, verdict) - def _coordinate_stage(self, stage: _BenchmarkStageExchange) -> bool | None: + def _coordinate_stage_round(self, local: dict) -> _BenchmarkStageRound: + deadline = time.monotonic() + self.timeout_seconds + statuses = {self.dp_rank: local} + identities: dict[int, bytes] = {} try: - while len(stage.identities) < self.dp_size - 1: - if not self._socket.poll(0, zmq.POLLIN): - if time.monotonic() < stage.deadline: - return None - raise TimeoutError( - "timed out waiting for attention-DP warm-up stage reports " - f"for batch={stage.batch}; " - f"reported_ranks={sorted(stage.reports)}" - ) - identity, message = self._read_router(0) + while len(identities) < self.dp_size - 1: + identity, message = self._recv_router(deadline, 0) rank = message.get("dp_rank") + phase = message.get("phase") ok = message.get("ok") if ( - message.get("type") != "stage_status" + message.get("type") != "stage_round" or not isinstance(rank, int) or not 1 <= rank < self.dp_size - or rank in stage.identities + or rank in identities or identity != str(rank).encode() - or message.get("batch") != stage.batch - or not isinstance(ok, bool) + or phase not in self.STAGE_PHASES + or (phase == "done" and not isinstance(ok, bool)) + or (phase != "done" and ok is not None) + ): + raise RuntimeError( + f"invalid attention-DP warm-up stage status: {message}" + ) + if ( + message.get("seq") != local["seq"] + or message.get("batch") != local["batch"] ): raise RuntimeError( - f"invalid attention-DP warm-up stage report: {message}" + "attention-DP warm-up ranks are not in lockstep: rank 0 is at " + f"round {local['seq']} batch={local['batch']}, rank {rank} at " + f"round {message.get('seq')} batch={message.get('batch')}" ) - stage.identities[rank] = identity - stage.reports[rank] = ok - decision = all(stage.reports.values()) + identities[rank] = identity + statuses[rank] = message + all_done = all(status["phase"] == "done" for status in statuses.values()) + verdict = ( + all(status["ok"] for status in statuses.values()) if all_done else None + ) self._send_to_all( - stage.identities, + identities, { - "type": "stage_decision", + "type": "stage_round_result", "benchmark_id": 0, - "batch": stage.batch, - "ok": decision, + "seq": local["seq"], + "all_done": all_done, + "ok": verdict, }, ) except Exception as error: - self._stage = None self._notify_error(self._all_follower_identities(), str(error)) raise - self._stage = None - return decision - - def _follow_stage(self, stage: _BenchmarkStageExchange) -> bool | None: - if not self._socket.poll(0, zmq.POLLIN): - if time.monotonic() < stage.deadline: - return None - self._stage = None - raise TimeoutError( - "timed out waiting for attention-DP warm-up stage decision " - f"for batch={stage.batch}" - ) - self._stage = None - reply = self._read_follower(0, "stage_decision") - ok = reply.get("ok") - if reply.get("batch") != stage.batch or not isinstance(ok, bool): - raise RuntimeError(f"invalid attention-DP warm-up stage decision: {reply}") - return ok + return _BenchmarkStageRound(all_done, verdict) @staticmethod def _deadline_elapsed(deadline: float | None) -> bool: @@ -4666,9 +4662,13 @@ def _bench_step_prefill(self) -> SchedulerOutput | None: _KVWARM_DOWNLOAD_TIMEOUT_S = 60 _kvwarm_stage_t0: float | None _kvwarm_stage_batch: int | None - # Local outcome ``(batch, ok, detail)`` of the active stage while the - # attention-DP group verdict is pending (``_kvwarm_stage_await``). - _kvwarm_stage_reported: tuple[int | None, bool, dict] | None = None + # Local outcome ``(batch, ok, detail)`` of the active stage once this + # rank's build has closed, held until the group's round says every rank + # is done (``_kvwarm_stage_round``). + _kvwarm_stage_local: tuple[int | None, bool, dict] | None = None + # Rounds this rank has run; carried in every round so a rank that fell + # out of lockstep is detected instead of waited for. + _kvwarm_round_seq: int = 0 # Real-KV prefill seeding state: per-batch-size seed chains, the parked # point with its per-request KV and new-token lengths, which shot # ("warm" | "measure") comes next, and whether this point staged. @@ -5107,7 +5107,7 @@ def _kvwarm_prepare(self, mode: str) -> None: # The pool figure is the group's negotiated one (the smallest rank's; # local before negotiation), like the depth cap: the plan decides # which rung every rank builds and which points it warms, and the - # stage exchange (``_kvwarm_stage_outcome``) relies on every rank + # stage round (``_kvwarm_stage_round``) relies on every rank # agreeing on both. shadow_tail_blocks = self._kvwarm_shadow_tail_blocks(repeats) for batch, depth in list(plan.items()): @@ -5131,7 +5131,8 @@ def _kvwarm_prepare(self, mode: str) -> None: self._kvwarm_borrowed_ids: set = set() self._kvwarm_stage_batch = None self._kvwarm_building = False - self._kvwarm_stage_reported = None + self._kvwarm_stage_local = None + self._kvwarm_round_seq = 0 self._kvwarm_seq = 0 logger.info( "KVWARM: prepared %d stage plans over %d decode points", @@ -5177,9 +5178,12 @@ def _kvwarm_step_busy(self) -> bool: """DECODE_SWEEP phase: chain-fleet build/park/turnover. True = hand this step back to the real scheduler. - Under attention-DP a finished build first waits for the group's - verdict on the rung (``_kvwarm_stage_await``); those steps are idle - too, so the collective forward keeps running on every rank. + Under attention-DP every idle step runs one stage round + (``_kvwarm_stage_round``): each rank reports whether it is still + building, and the group settles the rung in the same step on every + rank, so the rung's first point is injected on the same step + everywhere and no rank ever blocks in the point's READY barrier while + a peer still runs a collective. Chains shed while their last step is still in flight leave their blocks behind the deferred-free fence; every shed branch then yields @@ -5198,9 +5202,21 @@ def _kvwarm_step_busy(self) -> bool: return False if self._bench_active_req_ids or self._bench_current_point is not None: return False - if self._kvwarm_stage_reported is not None: - # The rung's verdict is with the group: idle until it arrives. - return self._kvwarm_stage_await() + if self._kvwarm_building: + if self._bench_soft_timeout_elapsed() or getattr( + self, "_bench_stop_requested", False + ): + # Soft timeout mid-build: never finish this fleet. Close the + # stage as failed so the group leaves the rung the same way. + closed = self._kvwarm_stage_outcome(False, {"soft_timeout": True}) + else: + closed = self._kvwarm_monitor_build() + if self._bench_synchronizer is None: + # dp=1: the build's steps, and the step that closes it, stay + # idle steps for the real scheduler, as before. + return closed + if self._bench_synchronizer is not None and self._kvwarm_stage_round(): + return True grid = self._bench_grid nxt = grid[0] if grid and grid[0].point_type == "decode" else None if nxt is None: @@ -5211,11 +5227,6 @@ def _kvwarm_step_busy(self) -> bool: ): # Soft timeout: never build another fleet. Release the chains and # let the decode step reach the coordinated timeout boundary. - if self._kvwarm_building and self._bench_synchronizer is not None: - # Peers may already be waiting for this rank's stage report; - # abandon the build through the exchange so every rank leaves - # the rung the same way. - return self._kvwarm_stage_outcome(False, {"soft_timeout": True}) if self._kvwarm_chain_ids: self._kvwarm_shed_chains() return self._bench_frees_pending() @@ -5225,8 +5236,6 @@ def _kvwarm_step_busy(self) -> bool: if self._kvwarm_chain_ids: self._kvwarm_shed_chains() return self._bench_frees_pending() - if self._kvwarm_building: - return self._kvwarm_monitor_build() if self._kvwarm_stage_batch != nxt.batch_size: self._kvwarm_shed_chains() if self._bench_frees_pending(): @@ -5333,43 +5342,69 @@ def _kvwarm_stage_outcome(self, ok: bool, detail: dict) -> bool: A failed build releases its chains at once (a partial or unusable fleet only pins KV). Without a synchronizer the outcome is final and - settles here; under attention-DP it is reported to the group and the - stage waits in ``_kvwarm_stage_reported`` for the verdict, which - ``_kvwarm_stage_await`` applies. True either way: the step is idle. + settles here; under attention-DP it waits in ``_kvwarm_stage_local`` + for the round in which every rank reports done + (``_kvwarm_stage_round``). True either way: the step is idle. """ batch = self._kvwarm_stage_batch self._kvwarm_building = False if not ok: self._kvwarm_shed_chains() - synchronizer = self._bench_synchronizer - if synchronizer is None: + if self._bench_synchronizer is None: self._kvwarm_stage_settle(batch, ok, detail) return True - synchronizer.stage_report( - batch, ok, timeout=self._kvwarm_stage_sync_timeout(synchronizer) - ) - self._kvwarm_stage_reported = (batch, ok, detail) + self._kvwarm_stage_local = (batch, ok, detail) return True - def _kvwarm_stage_await(self) -> bool: - """Poll the group verdict for the reported stage; True (idle) while it - is pending. A group fallback zeroes the rung's plan on every rank, so - a rank whose own build succeeded sheds its chains too and the rung's - points take fake injection everywhere.""" + def _kvwarm_stage_round(self) -> bool: + """One stage round per idle step: exchange this rank's stage status + with the group and apply the group's view. + + True when the stage machinery owns this step (a build chunk for the + real scheduler, or an idle step while peers finish or the rung is + settled); False when nothing is pending and the shared-state decisions + of ``_kvwarm_step_busy`` may run. Every rank runs the round at the + same point of the same step and blocks in it, so the group view + arrives in the same step everywhere and the rung's first point is + injected on the same step by every rank; no rank waits here while a + peer is inside a forward pass. Attention-DP only: without a + synchronizer the local outcome settled in ``_kvwarm_stage_outcome``.""" + local = self._kvwarm_stage_local synchronizer = self._bench_synchronizer - reported = self._kvwarm_stage_reported - if synchronizer is None or reported is None: - # Nothing is awaiting a verdict (dp=1 settles locally). - return False - decision = synchronizer.stage_poll() - if decision is None: - return True - batch, _, detail = reported - self._kvwarm_stage_reported = None - if not decision: + if synchronizer is None: + raise RuntimeError("attention-DP warm-up round without a synchronizer") + if local is not None: + batch, ok, _ = local + phase = "done" + else: + batch = self._kvwarm_stage_batch + ok = None + phase = "building" if self._kvwarm_building else "none" + self._kvwarm_round_seq += 1 + view = synchronizer.stage_round(self._kvwarm_round_seq, batch, phase, ok) + if not view.all_done: + return phase != "none" + if local is None: + raise RuntimeError( + "attention-DP warm-up round settled a stage this rank never closed" + ) + verdict = view.ok + if not isinstance(verdict, bool): + raise RuntimeError( + "attention-DP warm-up round settled a stage without a verdict" + ) + batch, _, detail = local + self._kvwarm_stage_local = None + logger.info( + "KVWARM: stage batch=%s settled by the group (ok=%s) at round %d", + batch, + verdict, + self._kvwarm_round_seq, + ) + if not verdict: detail = {**detail, "group_fallback": True} self._kvwarm_shed_chains() - self._kvwarm_stage_settle(batch, decision, detail) + self._kvwarm_stage_settle(batch, verdict, detail) return True def _kvwarm_stage_settle(self, batch: int | None, ok: bool, detail: dict) -> None: @@ -5395,17 +5430,6 @@ def _kvwarm_stage_settle(self, batch: int | None, ok: bool, detail: dict) -> Non detail, ) - def _kvwarm_stage_sync_timeout(self, synchronizer: _BenchmarkSynchronizer) -> float: - """Wait budget for the stage exchange, counted from this rank's - report. Peers may still be building: a build runs to completion or, - at the latest, to the soft deadline, where the soft-timeout branch of - ``_kvwarm_step_busy`` abandons it through this same exchange. The - budget therefore reaches the soft deadline plus the protocol timeout - the blocking phases allow.""" - deadline = getattr(self, "_bench_deadline_monotonic", None) - remaining = 0.0 if deadline is None else max(0.0, deadline - time.monotonic()) - return remaining + synchronizer.timeout_seconds - def _kvwarm_stage_shadow_shortfall(self, batch: int | None) -> int: """Free blocks the pool lacks for the shadows of the most demanding point this stage will serve, with its chains parked (0 when every diff --git a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py index fee0c79edf4f..1f134ab85319 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -953,96 +953,198 @@ def _stage_pair(timeout=1): return rank0, rank1 -def _stage_verdict(synchronizer, budget=2.0): - """Drive ``stage_poll`` the way the scheduler does: one non-blocking poll - per idle step until the verdict arrives.""" - end = time.monotonic() + budget - while time.monotonic() < end: - verdict = synchronizer.stage_poll() - if verdict is not None: - return verdict - time.sleep(0.005) - raise AssertionError("no stage verdict within budget") - - -def test_benchmark_synchronizer_stage_exchange_agrees_when_every_rank_is_ok(): - rank0, rank1 = _stage_pair() - follower_result = {} +def _round(synchronizer, seq, batch, phase, ok): + """Run one blocking round on a thread so both ranks can be driven from the test.""" + result = {} - def follow(): - rank1.stage_report(8, True) - follower_result["verdict"] = _stage_verdict(rank1) + def run(): + try: + result["view"] = synchronizer.stage_round(seq, batch, phase, ok) + except Exception as error: # noqa: BLE001 + result["error"] = error - follower = threading.Thread(target=follow) - follower.start() + thread = threading.Thread(target=run) + thread.start() + return thread, result + + +def _join(thread, result, timeout=2.0): + thread.join(timeout=timeout) + assert not thread.is_alive(), "stage round did not complete" + if "error" in result: + raise result["error"] + return result["view"] + + +def test_benchmark_synchronizer_stage_round_reports_building_until_every_rank_is_done(): + rank0, rank1 = _stage_pair() + view_cls = instrumented_scheduler_module._BenchmarkStageRound try: - rank0.stage_report(8, True) - # Non-blocking: rank 0 keeps polling between idle steps. - assert _stage_verdict(rank0) is True - follower.join(timeout=2) - assert not follower.is_alive() - assert follower_result["verdict"] is True - # The exchange is closed on both sides once the verdict is out. - for synchronizer in (rank0, rank1): - with pytest.raises(RuntimeError, match="without a report"): - synchronizer.stage_poll() + t0, r0 = _round(rank0, 1, 8, "done", True) + t1, r1 = _round(rank1, 1, 8, "building", None) + assert _join(t0, r0) == view_cls(all_done=False, ok=None) + assert _join(t1, r1) == view_cls(all_done=False, ok=None) + t0, r0 = _round(rank0, 2, 8, "done", True) + t1, r1 = _round(rank1, 2, 8, "done", True) + assert _join(t0, r0) == view_cls(all_done=True, ok=True) + assert _join(t1, r1) == view_cls(all_done=True, ok=True) finally: rank1.close() rank0.close() @pytest.mark.parametrize("failing_rank", [0, 1]) -def test_benchmark_synchronizer_stage_exchange_fails_the_group_with_one_rank( - failing_rank, -): +def test_benchmark_synchronizer_stage_round_fails_the_group_with_one_rank(failing_rank): rank0, rank1 = _stage_pair() - follower_result = {} + try: + t0, r0 = _round(rank0, 1, 16, "done", failing_rank != 0) + t1, r1 = _round(rank1, 1, 16, "done", failing_rank != 1) + expected = instrumented_scheduler_module._BenchmarkStageRound( + all_done=True, ok=False + ) + assert _join(t0, r0) == expected + assert _join(t1, r1) == expected + finally: + rank1.close() + rank0.close() - def follow(): - rank1.stage_report(16, failing_rank != 1) - follower_result["verdict"] = _stage_verdict(rank1) - follower = threading.Thread(target=follow) - follower.start() +def test_benchmark_synchronizer_stage_round_blocks_the_follower_until_rank0_answers(): + """The follower's round returns only when the group view has arrived: a + late broadcast delays the follower's step instead of letting it idle + into the next collective while rank 0 has already moved on.""" + rank0, rank1 = _stage_pair() try: - rank0.stage_report(16, failing_rank != 0) - assert _stage_verdict(rank0) is False - follower.join(timeout=2) - assert not follower.is_alive() - assert follower_result["verdict"] is False + t1, r1 = _round(rank1, 1, 8, "done", True) + time.sleep(0.2) + assert t1.is_alive() # rank 0 has not run its round yet: still blocked + t0, r0 = _round(rank0, 1, 8, "done", True) + view = _join(t1, r1) + assert view.all_done is True and view.ok is True + assert _join(t0, r0) == view finally: rank1.close() rank0.close() -def test_benchmark_synchronizer_stage_exchange_times_out_without_follower_report(): +def test_benchmark_synchronizer_stage_round_times_out_without_a_peer(): rank0, rank1 = _stage_pair(timeout=0.05) try: - rank0.stage_report(8, True, timeout=0.05) - with pytest.raises(TimeoutError, match="stage reports.*batch=8"): - _stage_verdict(rank0) - # A second report is possible again (the failed exchange is closed), - # and the late follower learns about the failure instead of waiting. - rank1.stage_report(8, True) - with pytest.raises(RuntimeError, match="synchronization failed"): - _stage_verdict(rank1) - with pytest.raises(RuntimeError, match="already pending"): - rank0.stage_report(8, True) - rank0.stage_report(8, True) + # The follower alone: rank 0 never runs its round. + with pytest.raises(TimeoutError): + rank1.stage_round(1, 8, "done", True) + finally: + rank1.close() + rank0.close() + rank0, rank1 = _stage_pair(timeout=0.05) + try: + # Rank 0 alone: the follower never reports. + with pytest.raises(TimeoutError): + rank0.stage_round(1, 8, "done", True) + finally: + rank1.close() + rank0.close() + + +@pytest.mark.parametrize("field, value", [("seq", 2), ("batch", 4)]) +def test_benchmark_synchronizer_stage_round_rejects_a_rank_out_of_lockstep( + field, value +): + """A follower at another round or another rung is not waited for: the + ranks have lost lockstep and the group aborts instead of hanging.""" + rank0, rank1 = _stage_pair() + try: + follower = {"seq": 1, "batch": 8} + follower[field] = value + t1, r1 = _round(rank1, follower["seq"], follower["batch"], "done", True) + with pytest.raises(RuntimeError, match="not in lockstep"): + rank0.stage_round(1, 8, "done", True) + with pytest.raises(RuntimeError): + _join(t1, r1) + finally: + rank1.close() + rank0.close() + + +@pytest.mark.parametrize( + "bad", + [ + {"all_done": True, "ok": None}, + {"all_done": False, "ok": True}, + {"all_done": "yes", "ok": None}, + {"seq": 2, "all_done": True, "ok": True}, + ], +) +def test_benchmark_synchronizer_stage_round_follower_rejects_an_invalid_result(bad): + """A malformed ``stage_round_result`` from rank 0 is a protocol error for + the follower, never a hang: drive the follower directly so a raw reply + can be substituted for rank 0's real one.""" + rank0, rank1 = _stage_pair() + try: + t1, r1 = _round(rank1, 1, 8, "done", True) + ident, _ = rank0._read_router(0) + rank0._send_to_all( + {1: ident}, + {"type": "stage_round_result", "benchmark_id": 0, "seq": 1, **bad}, + ) + with pytest.raises( + RuntimeError, match="invalid attention-DP warm-up stage round result" + ): + _join(t1, r1) + finally: + rank1.close() + rank0.close() + + +@pytest.mark.parametrize( + "bad", + [ + {"phase": "done", "ok": None}, + {"phase": "building", "ok": True}, + {"dp_rank": 5}, + ], +) +def test_benchmark_synchronizer_stage_round_rank0_rejects_a_malformed_status(bad): + """A malformed raw ``stage_round`` status is rejected by rank 0 like any + other protocol violation, and the follower is notified instead of being + left to time out.""" + rank0, rank1 = _stage_pair() + try: + rank1._socket.send_json( + { + "type": "stage_round", + "benchmark_id": 0, + "dp_rank": 1, + "seq": 1, + "batch": 8, + "phase": "done", + "ok": True, + **bad, + } + ) + with pytest.raises( + RuntimeError, match="invalid attention-DP warm-up stage status" + ): + rank0.stage_round(1, 8, "done", True) + assert rank1._socket.poll(500) + reply = rank1._socket.recv_json() + assert reply["type"] == "error" finally: rank1.close() rank0.close() -def test_benchmark_synchronizer_stage_exchange_rejects_a_rung_mismatch(): +def test_benchmark_synchronizer_stage_round_rejects_an_invalid_phase(): rank0, rank1 = _stage_pair() try: - rank1.stage_report(4, True) - rank0.stage_report(8, True) - with pytest.raises(RuntimeError, match="invalid attention-DP warm-up stage"): - _stage_verdict(rank0) - with pytest.raises(RuntimeError, match="synchronization failed"): - _stage_verdict(rank1) + with pytest.raises( + ValueError, match="invalid attention-DP warm-up stage phase" + ): + rank0.stage_round(1, 8, "parked", None) + with pytest.raises(ValueError, match="with outcome"): + rank0.stage_round(1, 8, "done", None) + with pytest.raises(ValueError, match="with outcome"): + rank0.stage_round(1, 8, "building", True) finally: rank1.close() rank0.close() @@ -5493,10 +5595,10 @@ def test_kvwarm_chain_parks_only_after_in_flight_tokens_drain(): # negotiated envelope) but a chain can vanish, or the pool can fall short of # the shadows' tails, on one rank only. Zeroing the plan there alone would # send that rank to fake injection while its peers inject real KV: READY -# summaries differ and the sweep aborts. The outcome is therefore reported -# through the synchronizer's non-blocking stage exchange and applied only -# once the group verdict is in; meanwhile every step is handed to the real -# scheduler (the DP forward is collective, so no rank may block). +# summaries differ and the sweep aborts. The outcome is therefore settled +# through a round every rank runs at the same point of every idle +# decode-sweep step, before its forward: ranks block only together, and the +# group settles the rung in the same step everywhere. def _kvwarm_group_stage_stub(*, chains=("chain-a", "chain-b"), synchronizer=None): @@ -5504,7 +5606,8 @@ def _kvwarm_group_stage_stub(*, chains=("chain-a", "chain-b"), synchronizer=None stub = _kvwarm_busy_stub(point, chains=chains) stub._kvwarm_plan = {4: 128} # covers the point: max(63) + 1 + 3 <= 128 stub._bench_synchronizer = MagicMock() if synchronizer is None else synchronizer - stub._bench_synchronizer.timeout_seconds = 10.0 + if synchronizer is None: + stub._bench_synchronizer.timeout_seconds = 10.0 stub._kvwarm_building = True stub._kvwarm_stage_batch = 4 stub._kvwarm_stage_t0 = time.monotonic() @@ -5522,70 +5625,97 @@ def _kvwarm_group_stage_stub(*, chains=("chain-a", "chain-b"), synchronizer=None return stub, meta -def test_kvwarm_stage_verdict_is_taken_from_the_group(): +def _round_views(*views): + """side_effect for a mocked ``stage_round``: one group view per step.""" + return [instrumented_scheduler_module._BenchmarkStageRound(*view) for view in views] + + +def test_kvwarm_stage_is_settled_by_the_group_round_in_the_same_step(): stub, meta = _kvwarm_group_stage_stub() sync = stub._bench_synchronizer - sync.stage_poll.side_effect = [None, None, True] + # Step 1: this rank's build completes while a peer is still building. + # Step 2: every rank is done and the verdict is ok. + sync.stage_round.side_effect = _round_views((False, None), (True, True)) - # Build complete: the local outcome is reported, not applied. assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - sync.stage_report.assert_called_once() - (batch, ok), kwargs = sync.stage_report.call_args - # No soft deadline armed in the stub: the budget is the protocol timeout. - assert (batch, ok) == (4, True) and kwargs["timeout"] == 10.0 assert stub._kvwarm_building is False - assert stub._kvwarm_stage_reported[:2] == (4, True) + assert stub._kvwarm_stage_local[:2] == (4, True) assert meta["stages"] == [] - # Pending verdict: every step is an idle step for the real scheduler. assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - assert meta["stages"] == [] - # Verdict in: the stage is ready and the point flow resumes. - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - assert stub._kvwarm_stage_reported is None + assert stub._kvwarm_stage_local is None assert [entry["batch"] for entry in meta["stages"]] == [4] assert meta["stages"][0]["depth"] == 8 and "failed" not in meta["stages"][0] + assert sync.stage_round.call_args_list == [ + call(1, 4, "done", True), + call(2, 4, "done", True), + ] + # Step 3: nothing pending; the point flow resumes from shared state, and + # the round still runs (phase "none") so every rank stays in step. + sync.stage_round.side_effect = _round_views((False, None)) assert InstrumentedScheduler._kvwarm_step_busy(stub) is False + assert sync.stage_round.call_args_list[-1] == call(3, 4, "none", None) assert stub._kvwarm_plan[4] == 128 stub._kvwarm_shed_chains.assert_not_called() - assert sync.stage_poll.call_count == 3 + + +def test_kvwarm_round_keeps_a_building_rank_in_step(): + """A rank still building reports "building" and hands the step to the + real scheduler; the round runs on every step regardless.""" + stub, meta = _kvwarm_group_stage_stub() + stub.requests["chain-a"].num_computed_tokens = 2 # still building + sync = stub._bench_synchronizer + sync.stage_round.side_effect = _round_views((False, None), (False, None)) + + assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + assert stub._kvwarm_building is True and stub._kvwarm_stage_local is None + assert sync.stage_round.call_args_list == [ + call(1, 4, "building", None), + call(2, 4, "building", None), + ] + assert meta["stages"] == [] def test_kvwarm_group_fallback_zeroes_the_plan_and_sheds_a_healthy_fleet(): stub, meta = _kvwarm_group_stage_stub() sync = stub._bench_synchronizer - sync.stage_poll.side_effect = [None, False] + sync.stage_round.side_effect = _round_views((False, None), (True, False)) - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True # reported ok - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True # pending - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True # verdict: no + assert ( + InstrumentedScheduler._kvwarm_step_busy(stub) is True + ) # 1: done ok, peers pending + assert InstrumentedScheduler._kvwarm_step_busy(stub) is True # 2: the group says no assert stub._kvwarm_plan[4] == 0 stub._kvwarm_shed_chains.assert_called_once_with() (entry,) = meta["stages"] assert entry["batch"] == 4 and entry["failed"] is True assert entry["group_fallback"] is True # The rung's points now take fake injection: the shed blocks drain behind - # the fence first, then the fake path proceeds. - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + # the fence first, then the fake path proceeds; each step still rounds. + sync.stage_round.side_effect = _round_views((False, None), (False, None)) + assert InstrumentedScheduler._kvwarm_step_busy(stub) is True # 3: fence stub.deferred_frees.clear() - assert InstrumentedScheduler._kvwarm_step_busy(stub) is False + assert InstrumentedScheduler._kvwarm_step_busy(stub) is False # 4: fake path + assert sync.stage_round.call_args_list[-2:] == [ + call(3, None, "none", None), + call(4, None, "none", None), + ] stub._kvwarm_start_stage.assert_not_called() -def test_kvwarm_local_stage_failure_is_reported_before_it_is_applied(): +def test_kvwarm_local_stage_failure_is_reported_done_and_not_ok(): stub, meta = _kvwarm_group_stage_stub() del stub.requests["chain-a"] # vanished during the build sync = stub._bench_synchronizer - sync.stage_poll.side_effect = [None, False] + sync.stage_round.side_effect = _round_views((False, None), (True, False)) assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - # Survivors are released at once, but the plan waits for the group. + # Survivors are released at once, but the plan waits for the group; the + # status still names the rung the fleet was built for. stub._kvwarm_shed_chains.assert_called_once_with() - (batch, ok), _ = sync.stage_report.call_args - assert (batch, ok) == (4, False) - assert stub._kvwarm_plan[4] == 128 - assert meta["stages"] == [] - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + assert stub._kvwarm_stage_local[:2] == (4, False) + assert sync.stage_round.call_args_list == [call(1, 4, "done", False)] + assert stub._kvwarm_plan[4] == 128 and meta["stages"] == [] assert InstrumentedScheduler._kvwarm_step_busy(stub) is True assert stub._kvwarm_plan[4] == 0 assert meta["stages"] == [ @@ -5596,37 +5726,33 @@ def test_kvwarm_local_stage_failure_is_reported_before_it_is_applied(): def test_kvwarm_stage_pool_shortfall_fails_the_rung_for_the_group(): """The per-point pool check of injection would skip a point on one rank alone; under attention-DP the rung's worst case is checked once the - fleet is parked and goes into the shared verdict.""" + fleet is parked and goes into the group verdict through the round.""" stub, meta = _kvwarm_group_stage_stub() stub._kvwarm_stage_shadow_shortfall = lambda batch: 3 sync = stub._bench_synchronizer - sync.stage_poll.side_effect = [False] + sync.stage_round.side_effect = _round_views((True, False)) assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - (batch, ok), _ = sync.stage_report.call_args - assert (batch, ok) == (4, False) - stub._kvwarm_shed_chains.assert_called_once_with() - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + assert sync.stage_round.call_args_list == [call(1, 4, "done", False)] + # Once for the local failure, once more when the group falls back. + assert stub._kvwarm_shed_chains.call_count == 2 assert meta["stages"] == [ {"batch": 4, "failed": True, "pool_shortfall": 3, "group_fallback": True} ] -def test_kvwarm_soft_timeout_mid_build_abandons_the_stage_through_the_group(): - """A rank that reaches the soft timeout while its fleet is still - building must not walk off to the boundary handshake while a peer is - waiting for its stage report.""" +def test_kvwarm_soft_timeout_mid_build_abandons_the_stage_through_the_round(): + """A rank that reaches the soft timeout while its fleet is still building + closes the stage as failed and keeps rounding, so peers are never left + waiting for a report that will not come.""" stub, meta = _kvwarm_group_stage_stub() stub.requests["chain-a"].num_computed_tokens = 2 # still building stub._bench_deadline_monotonic = 0.0 # soft timeout elapsed sync = stub._bench_synchronizer - sync.stage_poll.side_effect = [False] + sync.stage_round.side_effect = _round_views((True, False)) assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - (batch, ok), _ = sync.stage_report.call_args - assert (batch, ok) == (4, False) - stub._kvwarm_shed_chains.assert_called_once_with() - assert InstrumentedScheduler._kvwarm_step_busy(stub) is True + assert sync.stage_round.call_args_list == [call(1, 4, "done", False)] assert meta["stages"] == [ {"batch": 4, "failed": True, "soft_timeout": True, "group_fallback": True} ] @@ -5639,12 +5765,60 @@ def test_kvwarm_stage_settles_locally_without_a_synchronizer(): stub._kvwarm_stage_shadow_shortfall = MagicMock() assert InstrumentedScheduler._kvwarm_step_busy(stub) is True - assert stub._kvwarm_stage_reported is None + assert stub._kvwarm_stage_local is None assert [entry["batch"] for entry in meta["stages"]] == [4] stub._kvwarm_stage_shadow_shortfall.assert_not_called() assert InstrumentedScheduler._kvwarm_step_busy(stub) is False +def test_kvwarm_round_delays_the_follower_step_until_the_group_view_arrives(): + """The delayed-decision regression: rank 0's group view is broadcast + later than any follower poll would have been. With a round the + follower's step cannot complete until the view arrives, so both ranks + settle the rung in the same step and resume the point flow on the same + step; neither can idle into a collective the other has already left.""" + rank0, rank1 = _stage_pair(timeout=5) + send_to_all = rank0._send_to_all + + def delayed_send(identities, message): + time.sleep(0.3) + send_to_all(identities, message) + + rank0._send_to_all = delayed_send + stubs = [_kvwarm_group_stage_stub(synchronizer=sync)[0] for sync in (rank0, rank1)] + sequences = {0: [], 1: []} + finished = {0: [], 1: []} + + def drive(rank, stub): + for _step in range(3): + sequences[rank].append(InstrumentedScheduler._kvwarm_step_busy(stub)) + finished[rank].append(time.monotonic()) + + threads = [ + threading.Thread(target=drive, args=(rank, stub)) + for rank, stub in enumerate(stubs) + ] + started = time.monotonic() + try: + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + assert not thread.is_alive() + # Step 1: both builds complete and the round settles the rung at once; + # steps 2 and 3: nothing pending, the point flow resumes. Identical. + assert sequences[0] == sequences[1] == [True, False, False] + for stub in stubs: + assert stub._kvwarm_stage_local is None and stub._kvwarm_plan[4] == 128 + # The follower's first step ended only after the delayed broadcast, + # within the same instant as rank 0's. + assert finished[1][0] - started >= 0.25 + assert abs(finished[1][0] - finished[0][0]) < 0.2 + finally: + rank1.close() + rank0.close() + + def test_kvwarm_stage_shadow_shortfall_takes_the_rung_worst_case(monkeypatch): monkeypatch.setenv("DYN_BENCH_GIANT_KV_REPEATS", "3") stub = InstrumentedScheduler.__new__(InstrumentedScheduler) @@ -5668,18 +5842,6 @@ def shortfall(injected, headroom): assert seen == [([63] * 4, 3), ([46] * 4, 3)] -def test_kvwarm_stage_sync_timeout_reaches_the_soft_deadline(): - stub = InstrumentedScheduler.__new__(InstrumentedScheduler) - synchronizer = SimpleNamespace(timeout_seconds=10.0) - stub._bench_deadline_monotonic = time.monotonic() + 100.0 - budget = InstrumentedScheduler._kvwarm_stage_sync_timeout(stub, synchronizer) - assert 109.0 < budget <= 110.0 - stub._bench_deadline_monotonic = 0.0 # already elapsed - assert InstrumentedScheduler._kvwarm_stage_sync_timeout(stub, synchronizer) == 10.0 - stub._bench_deadline_monotonic = None - assert InstrumentedScheduler._kvwarm_stage_sync_timeout(stub, synchronizer) == 10.0 - - # --------------------------------------------------------------------------- # Benchmark request retirement goes through the scheduler's abort path # ---------------------------------------------------------------------------