| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the scheduler-driven heartbeat / lease-renewal system.""" |
| 4 | |
| 5 | import time |
| 6 | from unittest.mock import MagicMock |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.v1.outputs import KVConnectorOutput |
| 11 | |
| 12 | from .utils import create_request, make_nixl_scheduler |
| 13 | |
| 14 | _ENGINE_A = "my-engine-id" |
| 15 | |
| 16 | |
| 17 | def _sched(kv_lease_duration: int = 30): |
| 18 | return make_nixl_scheduler(heartbeat=True, kv_lease_duration=kv_lease_duration) |
| 19 | |
| 20 | |
| 21 | def _req(request_id: int = 1): |
| 22 | return create_request(request_id=request_id, do_remote_prefill=True) |
| 23 | |
| 24 | |
| 25 | def _worker_stub(): |
| 26 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( |
| 27 | NixlConnectorWorker, |
| 28 | ) |
| 29 | |
| 30 | w = object.__new__(NixlConnectorWorker) |
| 31 | w._reqs_to_send = {} |
| 32 | w._lease_extension = 20 |
| 33 | w._reqs_to_process = set() |
| 34 | w.consumer_notification_counts_by_req = {} |
| 35 | w.expected_consumer_notifications_by_req = {} |
| 36 | w.xfer_stats = MagicMock() |
| 37 | return w |
| 38 | |
| 39 | |
| 40 | # =================================================================== |
| 41 | # Scheduler: on_new_request |
| 42 | # =================================================================== |
| 43 | |
| 44 | |
| 45 | def test_on_new_request_tracks_and_groups(): |
| 46 | """Add two reqs to same engine, one to another; verify grouping.""" |
| 47 | s = _sched() |
| 48 | s.on_new_request(_req(1)) |
| 49 | s.on_new_request(_req(2)) |
| 50 | |
| 51 | assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-1", "prefill-2"} |
| 52 | info = s._heartbeat_by_engine[_ENGINE_A] |
| 53 | assert (info.host, info.port, info.tp_size) == ("my-host", 1234, 1) |
| 54 | assert s._heartbeat_req_engine["id-1"] == (_ENGINE_A, "prefill-1") |
| 55 | |
| 56 | # Different engine. |
| 57 | r3 = _req(3) |
| 58 | r3.kv_transfer_params["remote_engine_id"] = "engine-b" |
| 59 | s.on_new_request(r3) |
| 60 | assert len(s._heartbeat_by_engine) == 2 |
| 61 | |
| 62 | |
| 63 | @pytest.mark.parametrize( |
| 64 | "make_req", |
| 65 | [ |
| 66 | lambda: create_request(request_id=2, do_remote_decode=True), |
| 67 | lambda: create_request(request_id=3), # no kv_transfer_params |
| 68 | ], |
| 69 | ids=["decode", "plain"], |
| 70 | ) |
| 71 | def test_on_new_request_ignores_non_prefill(make_req): |
| 72 | s = _sched() |
| 73 | s.on_new_request(make_req()) |
| 74 | assert len(s._heartbeat_by_engine) == 0 |
| 75 | |
| 76 | |
| 77 | # =================================================================== |
| 78 | # Scheduler: _stop_heartbeat |
| 79 | # =================================================================== |
| 80 | |
| 81 | |
| 82 | def test_stop_heartbeat_partial_and_full(): |
| 83 | """Stop one of two reqs on same engine, then stop the other.""" |
| 84 | s = _sched() |
| 85 | s.on_new_request(_req(1)) |
| 86 | s.on_new_request(_req(2)) |
| 87 | |
| 88 | s._stop_heartbeat("id-1") |
| 89 | assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-2"} |
| 90 | assert "id-1" not in s._heartbeat_req_engine |
| 91 | |
| 92 | s._stop_heartbeat("id-2") |
| 93 | assert len(s._heartbeat_by_engine) == 0 |
| 94 | assert len(s._heartbeat_req_engine) == 0 |
| 95 | |
| 96 | |
| 97 | # =================================================================== |
| 98 | # Scheduler: build_connector_meta throttling |
| 99 | # =================================================================== |
| 100 | |
| 101 | |
| 102 | def test_build_connector_meta_heartbeat_throttling(): |
| 103 | # kv_lease_duration=30 => _heartbeat_interval = 30 // 6 = 5 |
| 104 | s = _sched(kv_lease_duration=30) |
| 105 | s.on_new_request(_req(1)) |
| 106 | |
| 107 | # Ensure the first call triggers by placing last_heartbeat far in the past. |
| 108 | s._last_heartbeat_time = time.perf_counter() - 10 |
| 109 | meta1 = s.build_connector_meta(MagicMock()) |
| 110 | assert _ENGINE_A in meta1.heartbeat_by_engine |
| 111 | |
| 112 | # Immediate second call is throttled (< 5s since last). |
| 113 | meta2 = s.build_connector_meta(MagicMock()) |
| 114 | assert len(meta2.heartbeat_by_engine) == 0 |
| 115 | |
| 116 | |
| 117 | # =================================================================== |
| 118 | # Scheduler: cleanup paths (update_connector_output / request_finished) |
| 119 | # =================================================================== |
| 120 | |
| 121 | |
| 122 | def test_update_connector_output_stops_heartbeat(): |
| 123 | s = _sched() |
| 124 | s.on_new_request(_req(1)) |
| 125 | |
| 126 | s.update_connector_output( |
| 127 | KVConnectorOutput( |
| 128 | finished_sending=None, |
| 129 | finished_recving={"id-1"}, |
| 130 | invalid_block_ids=set(), |
| 131 | ) |
| 132 | ) |
| 133 | |
| 134 | assert len(s._heartbeat_by_engine) == 0 |
| 135 | assert len(s._heartbeat_req_engine) == 0 |
| 136 | |
| 137 | |
| 138 | def test_request_finished_stops_heartbeat(): |
| 139 | s = _sched() |
| 140 | r = _req(1) |
| 141 | s.on_new_request(r) |
| 142 | |
| 143 | # Simulate update_state_after_alloc having consumed do_remote_prefill. |
| 144 | r.kv_transfer_params["do_remote_prefill"] = False |
| 145 | s.request_finished(r, block_ids=()) |
| 146 | |
| 147 | assert len(s._heartbeat_by_engine) == 0 |
| 148 | assert len(s._heartbeat_req_engine) == 0 |
| 149 | |
| 150 | |
| 151 | # =================================================================== |
| 152 | # Worker: _handle_heartbeat |
| 153 | # =================================================================== |
| 154 | |
| 155 | |
| 156 | def test_handle_heartbeat(): |
| 157 | w = _worker_stub() |
| 158 | far_future = time.perf_counter() + 99999 |
| 159 | w._reqs_to_send = {"req-a": 100.0, "req-b": far_future} |
| 160 | |
| 161 | before = time.perf_counter() |
| 162 | w._handle_heartbeat("req-a,req-b,req-unknown") |
| 163 | |
| 164 | # req-a: pushed forward to ~now+20. |
| 165 | assert w._reqs_to_send["req-a"] >= before + 20 |
| 166 | # req-b: already far out, max() keeps it. |
| 167 | assert w._reqs_to_send["req-b"] >= far_future |
| 168 | # req-unknown: not added. |
| 169 | assert "req-unknown" not in w._reqs_to_send |
| 170 | |
| 171 | |
| 172 | def test_reaper_reclaims_expired_leases_behind_heartbeated_head(monkeypatch): |
| 173 | """A heartbeated head entry must not strand expired leases behind it. |
| 174 | |
| 175 | Mirrors the stdlib repro in https://github.com/vllm-project/vllm/issues/58222: |
| 176 | A is inserted first and kept alive via heartbeats; B and C expire and |
| 177 | must still be reaped. The old early-exit scan returned nothing here. |
| 178 | """ |
| 179 | w = _worker_stub() |
| 180 | w._reqs_to_process.update(("A", "B", "C")) |
| 181 | # Insertion order is A, B, C. Heartbeats will keep A at the head. |
| 182 | w._reqs_to_send = {"A": 30.0, "B": 32.0, "C": 34.0} |
| 183 | |
| 184 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker |
| 185 | |
| 186 | now = {"t": 0.0} |
| 187 | monkeypatch.setattr(base_worker.time, "perf_counter", lambda: now["t"]) |
| 188 | for t in (20.0, 40.0, 60.0, 80.0): |
| 189 | now["t"] = t |
| 190 | w._handle_heartbeat("A") |
| 191 | |
| 192 | # A was renewed in place: still the dict head, expiry 100. |
| 193 | assert next(iter(w._reqs_to_send)) == "A" |
| 194 | assert w._reqs_to_send["A"] == pytest.approx(100.0) |
| 195 | |
| 196 | done: set[str] = set() |
| 197 | now["t"] = 50.0 |
| 198 | w._reap_expired_send_leases(done) |
| 199 | |
| 200 | assert done == {"B", "C"} |
| 201 | assert list(w._reqs_to_send) == ["A"] |
| 202 | assert w._reqs_to_process == {"A"} |
| 203 | assert w.xfer_stats.record_kv_expired_req.call_count == 2 |
| 204 | |
| 205 | # now == expires is expired (now >= expires): A is due at 100. |
| 206 | done_eq: set[str] = set() |
| 207 | now["t"] = 100.0 |
| 208 | w._reap_expired_send_leases(done_eq) |
| 209 | assert done_eq == {"A"} |
| 210 | assert not w._reqs_to_send |
| 211 | assert not w._reqs_to_process |
| 212 | assert w.xfer_stats.record_kv_expired_req.call_count == 3 |
| 213 | |
| 214 | |
| 215 | def test_reaper_reclaims_when_now_equals_expiry(monkeypatch): |
| 216 | """The reaper must reclaim on the now == expires boundary.""" |
| 217 | w = _worker_stub() |
| 218 | w._reqs_to_process.add("eq") |
| 219 | w._reqs_to_send = {"eq": 50.0} |
| 220 | |
| 221 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker |
| 222 | |
| 223 | monkeypatch.setattr(base_worker.time, "perf_counter", lambda: 50.0) |
| 224 | |
| 225 | done: set[str] = set() |
| 226 | w._reap_expired_send_leases(done) |
| 227 | |
| 228 | assert done == {"eq"} |
| 229 | assert not w._reqs_to_send |
| 230 | assert not w._reqs_to_process |
| 231 | w.xfer_stats.record_kv_expired_req.assert_called_once() |
| 232 | |
| 233 | |
| 234 | def test_reaper_reclaims_shorter_lease_behind_later_deadline(monkeypatch): |
| 235 | """A later-inserted longer TTL must not block an earlier-expiring entry. |
| 236 | |
| 237 | Delete+reinsert on heartbeat cannot restore expiry order in this case |
| 238 | (issue #58222), so the reaper must scan the whole map. |
| 239 | """ |
| 240 | w = _worker_stub() |
| 241 | w._reqs_to_process.update(("X", "Y")) |
| 242 | w._reqs_to_send = {"X": 55.0, "Y": 46.0} |
| 243 | |
| 244 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker |
| 245 | |
| 246 | monkeypatch.setattr(base_worker.time, "perf_counter", lambda: 50.0) |
| 247 | |
| 248 | done: set[str] = set() |
| 249 | w._reap_expired_send_leases(done) |
| 250 | |
| 251 | assert done == {"Y"} |
| 252 | assert list(w._reqs_to_send) == ["X"] |
| 253 | assert w._reqs_to_process == {"X"} |