From f0e876444fdd9510576dde2bc38d4e8e9ae3adb3 Mon Sep 17 00:00:00 2001 From: kliuae Date: Wed, 28 Oct 2026 17:20:44 +0000 Subject: [PATCH] wait in short chunks for mori write Signed-off-by: kliuae Co-authored-by: Claude Opus --- .../kv_connector/unit/test_moriio_tp_ack.py | 15 +++-- .../unit/test_moriio_transfer_completion.py | 43 +++++++++++++- .../kv_connector/v1/moriio/moriio_engine.py | 56 ++++++++++++------- 3 files changed, 88 insertions(+), 26 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index fe6b116476e8..64447424a7b5 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -655,10 +655,14 @@ def test_step_barrier_timeout_fails_closed_while_reads_in_flight(has_mamba, fake moriio_engine.StatusCode is None, reason="mori batched wait unavailable" ) @pytest.mark.parametrize("has_mamba", [False, True]) -def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamba): - engine = SimpleNamespace( - wait_all=Mock(return_value=moriio_engine.StatusCode.IN_PROGRESS) - ) +def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight( + has_mamba, fake_clock +): + def wait_all(statuses, timeout_ms): + fake_clock.sleep(timeout_ms / 1000) + return moriio_engine.StatusCode.IN_PROGRESS + + engine = SimpleNamespace(wait_all=Mock(side_effect=wait_all)) worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) worker._has_mamba = has_mamba worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) @@ -667,7 +671,8 @@ def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamb with pytest.raises(TransferError, match="timed out"): worker._await_reads_issued_this_step() - engine.wait_all.assert_called_once() + assert fake_clock.now >= 0.05 + assert all(c.kwargs["timeout_ms"] <= 10 for c in engine.wait_all.call_args_list) def test_step_barrier_noop_when_no_reads_posted(): diff --git a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py index 8fa59e5fb5cc..f2d93fafe73a 100644 --- a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py +++ b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py @@ -131,7 +131,7 @@ def test_wait_blocks_until_terminal(): timer.cancel() -@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 50), (0.0001, 1)]) +@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 10), (0.004, 4), (0.0001, 1)]) def test_native_wait_receives_all_statuses_and_a_positive_deadline( timeout, expected_ms ): @@ -186,5 +186,44 @@ def test_native_wait_drains_only_the_unscoped_status_snapshot(): wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) wrapper.transfer_status.extend(statuses) wrapper.waiting_for_transfer_complete() - engine.wait_all.assert_called_once_with(statuses, timeout_ms=50) + engine.wait_all.assert_called_once_with(statuses, timeout_ms=10) assert wrapper.transfer_status == [] + + +def _wait_all_that_completes(status, on_call, result="IN_PROGRESS"): + def wait_all(statuses, timeout_ms): + if wait_all_mock.call_count >= on_call: + status.state = "success" + return getattr(moriio_engine.StatusCode, result) + + wait_all_mock = Mock(side_effect=wait_all) + return wait_all_mock + + +@pytest.mark.parametrize("on_call", [1, 3]) +def test_native_wait_returns_once_the_statuses_succeed(on_call): + """MoRI's xgmi wait_all can return IN_PROGRESS at its deadline although + every status has succeeded.""" + late = _pending() + wait_all = _wait_all_that_completes(late, on_call) + wrapper = MoRIIOWrapper( + moriio_engine=SimpleNamespace(wait_all=wait_all), transfer_timeout=30.0 + ) + wrapper.waiting_for_transfer_complete([_ok(), late]) + assert wait_all.call_count == on_call + assert all(c.kwargs["timeout_ms"] == 10 for c in wait_all.call_args_list) + + +@pytest.mark.parametrize("fail_fast", [True, False]) +def test_native_wait_after_a_failure_waits_for_in_flight_statuses_unless_fail_fast( + fail_fast, +): + late = _pending() + wait_all = _wait_all_that_completes(late, on_call=3, result="ERR_RDMA_OP") + wrapper = MoRIIOWrapper( + moriio_engine=SimpleNamespace(wait_all=wait_all), transfer_timeout=30.0 + ) + with pytest.raises(TransferError, match="SQ full"): + wrapper.waiting_for_transfer_complete([_bad(), late], fail_fast=fail_fast) + assert wait_all.call_count == (1 if fail_fast else 3) + assert late.Succeeded() is not fail_fast diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 61e2c8d38b37..d43d6b2c5199 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -74,6 +74,7 @@ _MAX_TERMINAL_TRANSFER_IDS = 4096 +_NATIVE_WAIT_SLICE_S = 0.01 WriteGeometryKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype] @@ -911,25 +912,42 @@ def waiting_for_transfer_complete( self._poll_transfers_until_done(transfers_to_wait, fail_fast) return - # wait_all waits on every status until the deadline; a failure does not - # end it early, so fail_fast only applies to the polling fallback. - # Zero means a nonblocking progress call in Mori, not an immediate - # deadline. Round up so even sub-millisecond waits remain blocking. - timeout_ms = ceil(self._transfer_timeout * 1000) - result = wait_all(transfers_to_wait, timeout_ms=timeout_ms) - if result != StatusCode.SUCCESS: - details = [ - f"{status.Message()} (code={status.Code()})" - for status in transfers_to_wait - if status.Failed() - ] - if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS): - details.append( - f"RDMA transfer timed out after {self._transfer_timeout:g}s" - ) - raise TransferError( - f"MoRIIO batch wait failed ({result}): " + "; ".join(details) - ) + # On MoRI's xgmi backend a bounded wait_all polls each status once and + # then sleeps until its deadline. + remaining = self._transfer_timeout + deadline = time.monotonic() + remaining + while True: + # Zero means a nonblocking progress call in Mori, not an immediate + # deadline. Round up so even sub-millisecond waits remain blocking. + slice_ms = ceil(min(remaining, _NATIVE_WAIT_SLICE_S) * 1000) + result = wait_all(transfers_to_wait, timeout_ms=slice_ms) + if result == StatusCode.SUCCESS: + return + state = self.poll_transfer_batch(transfers_to_wait) + if state is TransferBatchState.DONE: + return + if state is TransferBatchState.FAILED: + if fail_fast or all( + status.Succeeded() or status.Failed() + for status in transfers_to_wait + ): + break + elif result not in (StatusCode.INIT, StatusCode.IN_PROGRESS): + break + remaining = deadline - time.monotonic() + if remaining <= 0: + break + + details = [ + f"{status.Message()} (code={status.Code()})" + for status in transfers_to_wait + if status.Failed() + ] + if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS): + details.append(f"RDMA transfer timed out after {self._transfer_timeout:g}s") + raise TransferError( + f"MoRIIO batch wait failed ({result}): " + "; ".join(details) + ) def _poll_transfers_until_done( self, transfers_to_wait: list[Any], fail_fast: bool = True