| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for MoRIIO router-authoritative DP-rank routing. |
| 4 | |
| 5 | These lock in the redesign that addresses the PR #45043 review feedback |
| 6 | (njhill): the connector must NOT self-derive a DP rank, and the prefill leg |
| 7 | must *return* the rank it actually ran on via ``request_finished`` so the |
| 8 | decode->prefill notify is routed by pure propagation (READ / serial WRITE), |
| 9 | not by an independent hash on each side. |
| 10 | |
| 11 | Following test_moriio_unmap.py, we bind the unbound method to a lightweight |
| 12 | stand-in rather than constructing a full scheduler. |
| 13 | """ |
| 14 | |
| 15 | from types import SimpleNamespace |
| 16 | |
| 17 | from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( |
| 18 | MoRIIOConnectorScheduler, |
| 19 | ) |
| 20 | from vllm.v1.request import RequestStatus |
| 21 | |
| 22 | _request_finished = MoRIIOConnectorScheduler.request_finished |
| 23 | _clip_blocks = MoRIIOConnectorScheduler.get_exchange_clipped_blocks |
| 24 | |
| 25 | |
| 26 | def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): |
| 27 | """Minimal prefill-side (producer) scheduler stand-in for request_finished.""" |
| 28 | s = SimpleNamespace( |
| 29 | is_producer=True, |
| 30 | _is_hma_required=False, |
| 31 | _has_mamba=False, |
| 32 | blocks_per_sw=[0], |
| 33 | _global_dp_rank=global_dp_rank, |
| 34 | engine_id="engine-abc", |
| 35 | host_ip="10.0.0.1", |
| 36 | handshake_port=6300, |
| 37 | side_notify_port=7300, |
| 38 | vllm_config=SimpleNamespace( |
| 39 | parallel_config=SimpleNamespace( |
| 40 | data_parallel_size=dp_size, |
| 41 | data_parallel_size_local=dp_size_local, |
| 42 | tensor_parallel_size=1, |
| 43 | ) |
| 44 | ), |
| 45 | ) |
| 46 | s.get_exchange_clipped_blocks = _clip_blocks.__get__(s) |
| 47 | return s |
| 48 | |
| 49 | |
| 50 | def _req( |
| 51 | request_id: str, params: dict | None, status=RequestStatus.FINISHED_LENGTH_CAPPED |
| 52 | ): |
| 53 | return SimpleNamespace( |
| 54 | request_id=request_id, |
| 55 | kv_transfer_params=params, |
| 56 | status=status, |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | def test_request_finished_returns_prefill_rank_for_propagation(): |
| 61 | # 2P2D: global rank 11 lives on pod 1 (ranks 8..15), local size 8. |
| 62 | s = _producer(global_dp_rank=11, dp_size=16, dp_size_local=8) |
| 63 | req = _req("r1", {"do_remote_decode": True, "transfer_id": "t1"}) |
| 64 | |
| 65 | delay_free, params = _request_finished(s, req, block_ids=[]) |
| 66 | |
| 67 | # The prefill leg returns the rank IT actually ran on, marked authoritative, |
| 68 | # so the router forwards it to decode and the notify propagates (no hashing). |
| 69 | assert params is not None |
| 70 | assert params["remote_dp_rank"] == 11 |
| 71 | assert params["remote_dp_rank_override"] is True |
| 72 | assert params["remote_dp_size"] == 16 |
| 73 | assert params["remote_dp_size_local"] == 8 |
| 74 | assert params["transfer_id"] == "t1" |
| 75 | # Empty block_ids => nothing deferred, blocks freed immediately. |
| 76 | assert delay_free is False |
| 77 | |
| 78 | |
| 79 | def test_request_finished_rank_tracks_the_actual_rank(): |
| 80 | # A different rank must be echoed verbatim -- the value is the producer's |
| 81 | # own rank, never a recomputation of the request id. |
| 82 | s = _producer(global_dp_rank=4, dp_size=8, dp_size_local=8) |
| 83 | req = _req("r2", {"do_remote_decode": True, "transfer_id": "t2"}) |
| 84 | |
| 85 | _, params = _request_finished(s, req, block_ids=[]) |
| 86 | |
| 87 | assert params["remote_dp_rank"] == 4 |
| 88 | assert params["remote_dp_size_local"] == 8 |
| 89 | |
| 90 | |
| 91 | def test_request_finished_no_kv_params_is_noop(): |
| 92 | s = _producer(global_dp_rank=0, dp_size=16, dp_size_local=8) |
| 93 | req = _req("r3", None) |
| 94 | |
| 95 | delay_free, params = _request_finished(s, req, block_ids=[]) |
| 96 | |
| 97 | assert delay_free is False |
| 98 | assert params is None |