| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for MoRIIOConnectorScheduler.{map,unmap}_request_id. |
| 4 | |
| 5 | vLLM's input_processor mutates a request_id after the connector has stored it |
| 6 | (e.g. ``cmpl-abc-0`` -> ``cmpl-abc-0-956053a4``), so an exact-match unmap can |
| 7 | miss. These tests cover the stable-``transfer_id`` fallback that handles that |
| 8 | case (and the clean exact-match and miss paths). |
| 9 | |
| 10 | The map/unmap methods touch only the two id<->id dicts, so we bind them to a |
| 11 | lightweight stand-in rather than constructing a full scheduler. |
| 12 | """ |
| 13 | |
| 14 | from types import SimpleNamespace |
| 15 | |
| 16 | from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( |
| 17 | MoRIIOConnectorScheduler, |
| 18 | ) |
| 19 | |
| 20 | _map = MoRIIOConnectorScheduler.map_request_id |
| 21 | _unmap = MoRIIOConnectorScheduler.unmap_request_id |
| 22 | |
| 23 | |
| 24 | def _sched(): |
| 25 | return SimpleNamespace( |
| 26 | transfer_id_to_request_id={}, |
| 27 | request_id_to_transfer_id={}, |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | def test_exact_match_unmap_clears_both_tables(): |
| 32 | s = _sched() |
| 33 | _map(s, "rid-1", "tid-1") |
| 34 | _unmap(s, "rid-1") |
| 35 | assert s.request_id_to_transfer_id == {} |
| 36 | assert s.transfer_id_to_request_id == {} |
| 37 | |
| 38 | |
| 39 | def test_unmap_via_transfer_id_after_request_id_mutation(): |
| 40 | s = _sched() |
| 41 | # Connector stores the ORIGINAL request_id at map time. |
| 42 | _map(s, "cmpl-abc-0", "tid-42") |
| 43 | # input_processor later mutates the rid; exact match would miss, so the |
| 44 | # stable transfer_id is used to find and clear the original entry. |
| 45 | _unmap(s, "cmpl-abc-0-956053a4", transfer_id="tid-42") |
| 46 | assert s.request_id_to_transfer_id == {} |
| 47 | assert s.transfer_id_to_request_id == {} |
| 48 | |
| 49 | |
| 50 | def test_unmap_miss_is_noop_and_does_not_raise(): |
| 51 | s = _sched() |
| 52 | _map(s, "rid-1", "tid-1") |
| 53 | # Unknown rid and unknown/None transfer_id: must not raise, must not touch |
| 54 | # the existing (unrelated) mapping. |
| 55 | _unmap(s, "rid-unknown", transfer_id=None) |
| 56 | _unmap(s, "rid-unknown", transfer_id="tid-unknown") |
| 57 | assert s.request_id_to_transfer_id == {"rid-1": "tid-1"} |
| 58 | assert s.transfer_id_to_request_id == {"tid-1": "rid-1"} |
| 59 | |
| 60 | |
| 61 | def test_exact_match_preferred_over_transfer_id_fallback(): |
| 62 | s = _sched() |
| 63 | _map(s, "rid-1", "tid-1") |
| 64 | # Exact rid is present, so it is removed directly even if a transfer_id is |
| 65 | # also supplied; no stale entries left behind. |
| 66 | _unmap(s, "rid-1", transfer_id="tid-1") |
| 67 | assert s.request_id_to_transfer_id == {} |
| 68 | assert s.transfer_id_to_request_id == {} |