| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for P2PSecondaryTierManager. |
| 4 | |
| 5 | Tests the manager's job routing, session lifecycle, and result collection |
| 6 | using fake transport and session objects. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | import time |
| 12 | import uuid |
| 13 | from types import SimpleNamespace |
| 14 | |
| 15 | import numpy as np |
| 16 | import pytest |
| 17 | |
| 18 | from vllm.utils.hashing import sha256 |
| 19 | from vllm.v1.core.kv_cache_utils import DEFAULT_NONE_HASH_SEED, init_none_hash |
| 20 | from vllm.v1.kv_offload.base import ( |
| 21 | LookupResult, |
| 22 | OffloadPolicy, |
| 23 | ReqContext, |
| 24 | ScheduleEndContext, |
| 25 | ) |
| 26 | from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob |
| 27 | from vllm.v1.kv_offload.tiering.p2p import manager as manager_module |
| 28 | from vllm.v1.kv_offload.tiering.p2p.manager import ( |
| 29 | _UNBOUND_STORE_TIMEOUT_S, |
| 30 | P2PSecondaryTierManager, |
| 31 | _annotate_req_context, |
| 32 | ) |
| 33 | from vllm.v1.kv_offload.tiering.p2p.session import ( |
| 34 | LoadResult, |
| 35 | SessionCloseResult, |
| 36 | SessionPollResult, |
| 37 | StoreResult, |
| 38 | ) |
| 39 | |
| 40 | # --------------------------------------------------------------------------- |
| 41 | # Helpers |
| 42 | # --------------------------------------------------------------------------- |
| 43 | |
| 44 | |
| 45 | def _remote_prefiller_kv_params( |
| 46 | remote_host: str = "10.0.0.1", |
| 47 | remote_port: int = 8000, |
| 48 | kv_request_id: str = "req-1", |
| 49 | ) -> dict: |
| 50 | """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries |
| 51 | kv_request_id + remote_host + remote_port.""" |
| 52 | return { |
| 53 | "remote_prefiller": { |
| 54 | "kv_request_id": kv_request_id, |
| 55 | "remote_host": remote_host, |
| 56 | "remote_port": remote_port, |
| 57 | }, |
| 58 | } |
| 59 | |
| 60 | |
| 61 | def _remote_kv_source_kv_params( |
| 62 | remote_host: str = "10.0.0.1", |
| 63 | remote_port: int = 8000, |
| 64 | kv_request_id: str = "req-1", |
| 65 | ) -> dict: |
| 66 | """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has |
| 67 | the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" |
| 68 | return { |
| 69 | "remote_kv_source": { |
| 70 | "kv_request_id": kv_request_id, |
| 71 | "remote_host": remote_host, |
| 72 | "remote_port": remote_port, |
| 73 | }, |
| 74 | } |
| 75 | |
| 76 | |
| 77 | def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: |
| 78 | """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries |
| 79 | kv_request_id only.""" |
| 80 | return {"remote_decoder": {"kv_request_id": kv_request_id}} |
| 81 | |
| 82 | |
| 83 | def _req_context(kv_params: dict | None = None) -> ReqContext: |
| 84 | ctx = ReqContext(req_id="test", kv_transfer_params=kv_params) |
| 85 | # Mirror on_new_request: parse the P2P routing state once and cache it, |
| 86 | # so lookup/submit_*/on_request_finished can read it back via get_state. |
| 87 | _annotate_req_context(ctx) |
| 88 | return ctx |
| 89 | |
| 90 | |
| 91 | def _job_metadata( |
| 92 | job_id: int, |
| 93 | keys: list[bytes] | None = None, |
| 94 | chunk_ids: list[int] | None = None, |
| 95 | kv_params: dict | None = None, |
| 96 | ) -> TransferJob: |
| 97 | if keys is None: |
| 98 | keys = [b"key1"] |
| 99 | if chunk_ids is None: |
| 100 | chunk_ids = list(range(len(keys))) |
| 101 | return TransferJob( |
| 102 | job_id=job_id, |
| 103 | keys=keys, |
| 104 | chunk_ids=np.array(chunk_ids), |
| 105 | is_promotion=False, |
| 106 | req_context=_req_context(kv_params), |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | def _make_manager() -> P2PSecondaryTierManager: |
| 111 | """Create a manager with stubbed __init__.""" |
| 112 | mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager) |
| 113 | mgr._local_id = "127.0.0.1:7777" |
| 114 | mgr._hash_seed = "0" |
| 115 | mgr._finished_jobs = [] |
| 116 | mgr._failed_req_ids = set() |
| 117 | mgr._sessions = {} |
| 118 | mgr._kv_to_session = {} |
| 119 | mgr._unbound_stores = {} |
| 120 | mgr._reaped_stores = {} |
| 121 | mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S |
| 122 | mgr._failed_serve_ctxs = [] |
| 123 | return mgr |
| 124 | |
| 125 | |
| 126 | def _init_offloading_spec() -> SimpleNamespace: |
| 127 | """Minimal offloading_spec for driving the real __init__.""" |
| 128 | return SimpleNamespace( |
| 129 | config=SimpleNamespace(parallel=SimpleNamespace(data_parallel_index=0)), |
| 130 | blocks_per_chunk=1, |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | # --------------------------------------------------------------------------- |
| 135 | # Tests for __init__ hash seed resolution |
| 136 | # --------------------------------------------------------------------------- |
| 137 | |
| 138 | |
| 139 | class TestInitHashSeed: |
| 140 | def _build(self, monkeypatch) -> P2PSecondaryTierManager: |
| 141 | monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) |
| 142 | monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) |
| 143 | monkeypatch.setattr( |
| 144 | manager_module.FileMapper, |
| 145 | "from_offloading_spec", |
| 146 | lambda **k: SimpleNamespace(get_run_config=lambda: {}), |
| 147 | ) |
| 148 | return P2PSecondaryTierManager( |
| 149 | offloading_spec=_init_offloading_spec(), |
| 150 | primary_kv_view=memoryview(bytearray(16)), |
| 151 | ) |
| 152 | |
| 153 | def test_missing_pythonhashseed_uses_default(self, monkeypatch): |
| 154 | """P2P falls back to the deterministic default seed when unset.""" |
| 155 | monkeypatch.delenv("PYTHONHASHSEED", raising=False) |
| 156 | mgr = self._build(monkeypatch) |
| 157 | init_none_hash(sha256) |
| 158 | assert mgr._get_hash_seed() == DEFAULT_NONE_HASH_SEED |
| 159 | |
| 160 | def test_pythonhashseed_set_succeeds(self, monkeypatch): |
| 161 | """With PYTHONHASHSEED set, the handshake advertises it.""" |
| 162 | monkeypatch.setenv("PYTHONHASHSEED", "12345") |
| 163 | mgr = self._build(monkeypatch) |
| 164 | init_none_hash(sha256) |
| 165 | assert mgr._get_hash_seed() == "12345" |
| 166 | |
| 167 | def test_seed_resolved_after_init_none_hash(self, monkeypatch): |
| 168 | """The seed is read lazily, not at construction time. |
| 169 | |
| 170 | This tier is built before init_none_hash runs, and a non-cryptographic |
| 171 | hash algorithm seeds NONE_HASH randomly, so resolving in __init__ would |
| 172 | advertise a value that does not match the NONE_HASH actually in use. |
| 173 | """ |
| 174 | monkeypatch.delenv("PYTHONHASHSEED", raising=False) |
| 175 | mgr = self._build(monkeypatch) |
| 176 | assert mgr._hash_seed is None |
| 177 | monkeypatch.setattr( |
| 178 | manager_module, "get_none_hash_seed", lambda: "random-seed-abc" |
| 179 | ) |
| 180 | assert mgr._get_hash_seed() == "random-seed-abc" |
| 181 | |
| 182 | |
| 183 | class TestBackpressureRejected: |
| 184 | def _patch_transports(self, monkeypatch) -> None: |
| 185 | monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) |
| 186 | monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) |
| 187 | monkeypatch.setattr( |
| 188 | manager_module.FileMapper, |
| 189 | "from_offloading_spec", |
| 190 | lambda **k: SimpleNamespace(get_run_config=lambda: {}), |
| 191 | ) |
| 192 | |
| 193 | def test_explicit_backpressure_detector_rejected(self, monkeypatch): |
| 194 | """An explicitly configured detector must fail fast for P2P. |
| 195 | |
| 196 | The generic store-latency detector cannot fail fast in PD mode and |
| 197 | rendezvous time makes store latency an unreliable pressure signal, |
| 198 | so P2P rejects it rather than degrading silently. |
| 199 | """ |
| 200 | self._patch_transports(monkeypatch) |
| 201 | with pytest.raises(ValueError, match="not supported for the P2P"): |
| 202 | P2PSecondaryTierManager( |
| 203 | offloading_spec=_init_offloading_spec(), |
| 204 | primary_kv_view=memoryview(bytearray(16)), |
| 205 | backpressure_detector=object(), |
| 206 | ) |
| 207 | |
| 208 | def test_no_detector_constructs_normally(self, monkeypatch): |
| 209 | self._patch_transports(monkeypatch) |
| 210 | mgr = P2PSecondaryTierManager( |
| 211 | offloading_spec=_init_offloading_spec(), |
| 212 | primary_kv_view=memoryview(bytearray(16)), |
| 213 | ) |
| 214 | assert mgr.bp_detector is None |
| 215 | |
| 216 | |
| 217 | # --------------------------------------------------------------------------- |
| 218 | # Tests for _peer_id_from_params |
| 219 | # --------------------------------------------------------------------------- |
| 220 | |
| 221 | |
| 222 | class TestPeerIdFromParams: |
| 223 | def test_valid_params(self): |
| 224 | result = manager_module._peer_id_from_params( |
| 225 | {"remote_host": "10.0.0.1", "remote_port": 8000} |
| 226 | ) |
| 227 | assert result == "10.0.0.1:8000" |
| 228 | |
| 229 | def test_missing_host(self): |
| 230 | result = manager_module._peer_id_from_params({"remote_port": 8000}) |
| 231 | assert result is None |
| 232 | |
| 233 | def test_missing_port(self): |
| 234 | result = manager_module._peer_id_from_params({"remote_host": "10.0.0.1"}) |
| 235 | assert result is None |
| 236 | |
| 237 | def test_empty_dict(self): |
| 238 | result = manager_module._peer_id_from_params({}) |
| 239 | assert result is None |
| 240 | |
| 241 | |
| 242 | # --------------------------------------------------------------------------- |
| 243 | # Tests for lookup |
| 244 | # --------------------------------------------------------------------------- |
| 245 | |
| 246 | |
| 247 | class TestLookup: |
| 248 | def test_lookup_returns_miss_without_kv_params(self): |
| 249 | mgr = _make_manager() |
| 250 | ctx = _req_context(kv_params=None) |
| 251 | assert mgr.lookup(b"key", ctx) is LookupResult.MISS |
| 252 | |
| 253 | def test_lookup_returns_miss_without_required_fields(self): |
| 254 | mgr = _make_manager() |
| 255 | ctx = _req_context(kv_params={"remote_prefiller": {"remote_host": "x"}}) |
| 256 | assert mgr.lookup(b"key", ctx) is LookupResult.MISS |
| 257 | |
| 258 | def test_lookup_returns_hit_for_valid_request(self): |
| 259 | mgr = _make_manager() |
| 260 | ctx = _req_context(kv_params=_remote_prefiller_kv_params()) |
| 261 | assert mgr.lookup(b"key", ctx) is LookupResult.HIT |
| 262 | |
| 263 | def test_lookup_returns_miss_for_failed_request(self): |
| 264 | mgr = _make_manager() |
| 265 | mgr._failed_req_ids.add("req-1") |
| 266 | ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) |
| 267 | assert mgr.lookup(b"key", ctx) is LookupResult.MISS |
| 268 | |
| 269 | def test_lookup_returns_hit_for_different_request_id(self): |
| 270 | mgr = _make_manager() |
| 271 | mgr._failed_req_ids.add("req-1") |
| 272 | ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-2")) |
| 273 | assert mgr.lookup(b"key", ctx) is LookupResult.HIT |
| 274 | |
| 275 | def test_lookup_returns_miss_without_prefill_key(self): |
| 276 | """No ``remote_prefiller`` sub-dict means the request was not routed for |
| 277 | remote prefill — local prefill should run instead, so lookup() |
| 278 | returns MISS even when a stale ``remote_decoder`` block is present.""" |
| 279 | mgr = _make_manager() |
| 280 | ctx = _req_context(kv_params=_remote_decoder_kv_params()) |
| 281 | assert mgr.lookup(b"key", ctx) is LookupResult.MISS |
| 282 | |
| 283 | |
| 284 | # --------------------------------------------------------------------------- |
| 285 | # Tests for on_new_request offload policy |
| 286 | # --------------------------------------------------------------------------- |
| 287 | |
| 288 | |
| 289 | @pytest.mark.parametrize( |
| 290 | "kv_params,expected", |
| 291 | [ |
| 292 | (_remote_decoder_kv_params(), OffloadPolicy.REQUEST_LEVEL), |
| 293 | (None, OffloadPolicy.CHUNK_LEVEL), |
| 294 | (_remote_prefiller_kv_params(), OffloadPolicy.CHUNK_LEVEL), |
| 295 | ({"remote_decoder": {}}, OffloadPolicy.CHUNK_LEVEL), |
| 296 | ], |
| 297 | ids=["producer", "plain", "consumer", "producer_no_id"], |
| 298 | ) |
| 299 | def test_on_new_request_policy(monkeypatch, kv_params, expected): |
| 300 | """Only a producer leg carrying a kv_request_id widens to REQUEST_LEVEL.""" |
| 301 | mgr = _make_manager() |
| 302 | monkeypatch.setattr(mgr, "_get_or_create_session", lambda peer_id: None) |
| 303 | assert mgr.on_new_request(_req_context(kv_params=kv_params)).policy is expected |
| 304 | |
| 305 | |
| 306 | # --------------------------------------------------------------------------- |
| 307 | # Tests for serve_external_requests |
| 308 | # --------------------------------------------------------------------------- |
| 309 | |
| 310 | |
| 311 | class _RecordingParent: |
| 312 | """Minimal ParentManager stub recording on_request_finished calls.""" |
| 313 | |
| 314 | def __init__(self) -> None: |
| 315 | self.finished: list[str] = [] |
| 316 | |
| 317 | def on_new_request(self, ctx): |
| 318 | from vllm.v1.kv_offload.base import RequestOffloadingContext |
| 319 | |
| 320 | return RequestOffloadingContext() |
| 321 | |
| 322 | def lookup(self, key, ctx): |
| 323 | return LookupResult.MISS |
| 324 | |
| 325 | def create_store_job(self, keys, ctx): |
| 326 | raise AssertionError("unreachable") |
| 327 | |
| 328 | def on_request_finished(self, ctx) -> None: |
| 329 | self.finished.append(ctx.req_id) |
| 330 | |
| 331 | |
| 332 | class _RecordingSession: |
| 333 | """Fake P2PSession that records the parent it was served with.""" |
| 334 | |
| 335 | def __init__(self) -> None: |
| 336 | self.served_with: list[object] = [] |
| 337 | |
| 338 | def serve_external_requests(self, parent) -> None: |
| 339 | self.served_with.append(parent) |
| 340 | |
| 341 | |
| 342 | class TestServeExternalRequests: |
| 343 | def test_flushes_failed_serve_ctxs_then_serves_each_session(self): |
| 344 | """serve_external_requests releases the failed serves left by |
| 345 | reaped sessions via parent.on_request_finished (clearing the |
| 346 | queue), then delegates to every live session with the same parent.""" |
| 347 | mgr = _make_manager() |
| 348 | ctx = ReqContext(req_id="p2p:peer:req-1:lu1") |
| 349 | mgr._failed_serve_ctxs = [ctx] |
| 350 | sess_a = _RecordingSession() |
| 351 | sess_b = _RecordingSession() |
| 352 | mgr._sessions = {"a": sess_a, "b": sess_b} |
| 353 | |
| 354 | parent = _RecordingParent() |
| 355 | mgr.serve_external_requests(parent) |
| 356 | |
| 357 | # Failed serve released and queue cleared. |
| 358 | assert parent.finished == ["p2p:peer:req-1:lu1"] |
| 359 | assert mgr._failed_serve_ctxs == [] |
| 360 | # Every live session served with the same parent handle. |
| 361 | assert sess_a.served_with == [parent] |
| 362 | assert sess_b.served_with == [parent] |
| 363 | |
| 364 | def test_no_failed_serves_still_serves_sessions(self): |
| 365 | mgr = _make_manager() |
| 366 | sess = _RecordingSession() |
| 367 | mgr._sessions = {"a": sess} |
| 368 | |
| 369 | parent = _RecordingParent() |
| 370 | mgr.serve_external_requests(parent) |
| 371 | |
| 372 | assert parent.finished == [] |
| 373 | assert sess.served_with == [parent] |
| 374 | |
| 375 | |
| 376 | # --------------------------------------------------------------------------- |
| 377 | # Tests for submit_store |
| 378 | # --------------------------------------------------------------------------- |
| 379 | |
| 380 | |
| 381 | class TestSubmitStore: |
| 382 | def test_no_decode_succeeds_immediately(self): |
| 383 | """Without a ``remote_decoder`` block, job succeeds immediately.""" |
| 384 | mgr = _make_manager() |
| 385 | job = _job_metadata(job_id=1, kv_params={}) |
| 386 | mgr.submit_store(job) |
| 387 | assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] |
| 388 | |
| 389 | def test_missing_kv_request_id_fails(self): |
| 390 | """Missing kv_request_id inside ``remote_decoder`` fails the job.""" |
| 391 | mgr = _make_manager() |
| 392 | params: dict = {"remote_decoder": {}} |
| 393 | job = _job_metadata(job_id=1, kv_params=params) |
| 394 | mgr.submit_store(job) |
| 395 | assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] |
| 396 | |
| 397 | def test_no_binding_yet_parks_in_unbound_stores(self): |
| 398 | """submit_store without a bound session buffers the batch keyed |
| 399 | by kv_request_id; no session is pre-created (peer_id is unknown |
| 400 | to the producer at store time).""" |
| 401 | mgr = _make_manager() |
| 402 | job = _job_metadata( |
| 403 | job_id=1, |
| 404 | keys=[b"k1", b"k2"], |
| 405 | chunk_ids=[3, 4], |
| 406 | kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), |
| 407 | ) |
| 408 | mgr.submit_store(job) |
| 409 | |
| 410 | assert mgr._sessions == {} |
| 411 | assert mgr._finished_jobs == [] |
| 412 | batches = mgr._unbound_stores["req-1"] |
| 413 | assert len(batches) == 1 |
| 414 | assert ( |
| 415 | batches[0].job_id, |
| 416 | list(batches[0].keys), |
| 417 | list(batches[0].block_ids), |
| 418 | ) == ( |
| 419 | 1, |
| 420 | [b"k1", b"k2"], |
| 421 | [3, 4], |
| 422 | ) |
| 423 | |
| 424 | def test_routes_to_bound_session(self): |
| 425 | """If a session has already received FetchMsg for this |
| 426 | kv_request_id (so _kv_to_session is populated), submit_store |
| 427 | forwards directly to that session rather than re-buffering.""" |
| 428 | mgr = _make_manager() |
| 429 | bound = _FakeSession(peer_id="10.0.0.1:8000", connected=True) |
| 430 | mgr._kv_to_session["req-1"] = bound |
| 431 | # Note: _sessions is intentionally untouched — this test isolates |
| 432 | # the kv_request_id → session fast path. |
| 433 | job = _job_metadata( |
| 434 | job_id=7, |
| 435 | keys=[b"k1", b"k2"], |
| 436 | chunk_ids=[3, 4], |
| 437 | kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), |
| 438 | ) |
| 439 | mgr.submit_store(job) |
| 440 | |
| 441 | assert len(bound.stores_added) == 1 |
| 442 | kv_req_id, keys, _, job_id = bound.stores_added[0] |
| 443 | assert (kv_req_id, keys, job_id) == ("req-1", [b"k1", b"k2"], 7) |
| 444 | assert mgr._unbound_stores == {} |
| 445 | assert mgr._finished_jobs == [] |
| 446 | |
| 447 | def test_extra_top_level_keys_are_ignored(self): |
| 448 | """Producer-side kv_transfer_params should not pre-create a |
| 449 | session even when a stale caller still passes a top-level |
| 450 | ``remote_host``/``remote_port`` next to ``remote_decoder``.""" |
| 451 | mgr = _make_manager() |
| 452 | params = _remote_decoder_kv_params() |
| 453 | params["remote_host"] = "stale" |
| 454 | params["remote_port"] = 12345 |
| 455 | job = _job_metadata(job_id=1, kv_params=params) |
| 456 | mgr.submit_store(job) |
| 457 | # No session pre-created, no peer-keyed state. |
| 458 | assert mgr._sessions == {} |
| 459 | assert "req-1" in mgr._unbound_stores |
| 460 | |
| 461 | |
| 462 | # --------------------------------------------------------------------------- |
| 463 | # Tests for submit_load |
| 464 | # --------------------------------------------------------------------------- |
| 465 | |
| 466 | |
| 467 | class TestSubmitLoad: |
| 468 | def test_missing_params_fails(self): |
| 469 | """Missing required kv_params fields fails the job.""" |
| 470 | mgr = _make_manager() |
| 471 | job = _job_metadata(job_id=1, kv_params={}) |
| 472 | mgr.submit_load(job) |
| 473 | assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] |
| 474 | |
| 475 | def test_empty_keys_succeeds_immediately(self): |
| 476 | """Empty key list succeeds immediately.""" |
| 477 | mgr = _make_manager() |
| 478 | job = _job_metadata( |
| 479 | job_id=1, keys=[], chunk_ids=[], kv_params=_remote_prefiller_kv_params() |
| 480 | ) |
| 481 | mgr.submit_load(job) |
| 482 | assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] |
| 483 | |
| 484 | def test_no_session_fails(self): |
| 485 | """No session for peer fails and marks request failed.""" |
| 486 | mgr = _make_manager() |
| 487 | job = _job_metadata(job_id=1, kv_params=_remote_prefiller_kv_params()) |
| 488 | mgr.submit_load(job) |
| 489 | assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] |
| 490 | assert "req-1" in mgr._failed_req_ids |
| 491 | |
| 492 | def test_happy_path_with_active_session(self): |
| 493 | """When the peer's session exists, submit_load forwards to |
| 494 | session.request_blocks and does NOT add a finished result yet.""" |
| 495 | mgr = _make_manager() |
| 496 | peer_id = "10.0.0.1:8000" |
| 497 | existing = _FakeSession(peer_id=peer_id, connected=True) |
| 498 | mgr._sessions[peer_id] = existing |
| 499 | job = _job_metadata( |
| 500 | job_id=42, |
| 501 | keys=[b"k1", b"k2"], |
| 502 | chunk_ids=[5, 6], |
| 503 | kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"), |
| 504 | ) |
| 505 | mgr.submit_load(job) |
| 506 | |
| 507 | assert existing.requests == [(42, "req-42")] |
| 508 | assert mgr._finished_jobs == [] |
| 509 | assert "req-42" not in mgr._failed_req_ids |
| 510 | |
| 511 | def test_missing_consumer_flag_fails(self): |
| 512 | """Peer fields present but neither do_remote_prefill nor |
| 513 | do_p2p_fetch is set — submit_load fails the job rather than |
| 514 | emit a stray FetchMsg.""" |
| 515 | mgr = _make_manager() |
| 516 | params = { |
| 517 | "remote_host": "10.0.0.1", |
| 518 | "remote_port": 8000, |
| 519 | "kv_request_id": "req-1", |
| 520 | } |
| 521 | job = _job_metadata(job_id=1, kv_params=params) |
| 522 | mgr.submit_load(job) |
| 523 | assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] |
| 524 | |
| 525 | |
| 526 | # --------------------------------------------------------------------------- |
| 527 | # Tests for on_request_finished |
| 528 | # --------------------------------------------------------------------------- |
| 529 | |
| 530 | |
| 531 | class TestOnRequestFinished: |
| 532 | def _make_with_failed(self) -> P2PSecondaryTierManager: |
| 533 | mgr = _make_manager() |
| 534 | mgr._failed_req_ids = {"req-1"} |
| 535 | return mgr |
| 536 | |
| 537 | def test_prunes_failed_req_ids(self): |
| 538 | mgr = self._make_with_failed() |
| 539 | ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) |
| 540 | mgr.on_request_finished(ctx) |
| 541 | assert "req-1" not in mgr._failed_req_ids |
| 542 | |
| 543 | def test_no_kv_params_does_nothing(self): |
| 544 | mgr = self._make_with_failed() |
| 545 | ctx = _req_context(kv_params=None) |
| 546 | mgr.on_request_finished(ctx) |
| 547 | assert "req-1" in mgr._failed_req_ids |
| 548 | |
| 549 | def test_no_kv_request_id_does_nothing(self): |
| 550 | mgr = self._make_with_failed() |
| 551 | ctx = _req_context(kv_params={"remote_host": "x", "remote_port": 1}) |
| 552 | mgr.on_request_finished(ctx) |
| 553 | assert "req-1" in mgr._failed_req_ids |
| 554 | |
| 555 | def test_decoder_side_calls_session_finish_request(self): |
| 556 | """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id |
| 557 | because the consumer addresses the producer it loaded from. The |
| 558 | session's finish_request cancels the client-role load.""" |
| 559 | mgr = _make_manager() |
| 560 | peer_id = "10.0.0.1:8000" |
| 561 | session = _FakeSession(peer_id=peer_id) |
| 562 | mgr._sessions[peer_id] = session |
| 563 | ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) |
| 564 | mgr.on_request_finished(ctx) |
| 565 | assert session.finishes == ["req-1"] |
| 566 | |
| 567 | def test_p2p_consumer_side_calls_session_finish_request(self): |
| 568 | """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id |
| 569 | so the session drops any pending lookups (cancel_lookups) and |
| 570 | cancels any inbound load.""" |
| 571 | mgr = _make_manager() |
| 572 | peer_id = "10.0.0.1:8000" |
| 573 | session = _FakeSession(peer_id=peer_id) |
| 574 | mgr._sessions[peer_id] = session |
| 575 | ctx = _req_context(kv_params=_remote_kv_source_kv_params(kv_request_id="req-1")) |
| 576 | mgr.on_request_finished(ctx) |
| 577 | assert session.finishes == ["req-1"] |
| 578 | |
| 579 | def test_prefiller_bound_id_routes_via_kv_to_session(self): |
| 580 | """Prefiller-side finish for an id whose session is already bound |
| 581 | (FetchMsg received) routes via _kv_to_session and pops the entry.""" |
| 582 | mgr = _make_manager() |
| 583 | bound = _FakeSession(peer_id="some-peer:1", connected=True) |
| 584 | mgr._kv_to_session["req-1"] = bound |
| 585 | ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) |
| 586 | mgr.on_request_finished(ctx) |
| 587 | assert bound.finishes == ["req-1"] |
| 588 | assert "req-1" not in mgr._kv_to_session |
| 589 | |
| 590 | def test_prefiller_unbound_id_leaves_batches_parked(self): |
| 591 | """Prefiller-side finish for an id with parked unbound batches |
| 592 | and no session binding is a no-op on `_unbound_stores`. The |
| 593 | parked batches survive until a peer fetches them or the |
| 594 | `_reap_unbound_stores` timeout fires — `on_request_finished` |
| 595 | must not evict them.""" |
| 596 | from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch |
| 597 | |
| 598 | mgr = _make_manager() |
| 599 | mgr._unbound_stores["req-1"] = [ |
| 600 | _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]), |
| 601 | _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]), |
| 602 | ] |
| 603 | ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) |
| 604 | mgr.on_request_finished(ctx) |
| 605 | assert "req-1" in mgr._unbound_stores |
| 606 | assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11] |
| 607 | outcomes = {(r.job_id, r.success) for r in mgr._finished_jobs} |
| 608 | assert (10, False) not in outcomes |
| 609 | assert (11, False) not in outcomes |
| 610 | |
| 611 | |
| 612 | # --------------------------------------------------------------------------- |
| 613 | # Tests for get_finished_jobs |
| 614 | # --------------------------------------------------------------------------- |
| 615 | |
| 616 | |
| 617 | class _FakeServerHalf: |
| 618 | def __init__(self) -> None: |
| 619 | self._inflight: dict[int, object] = {} |
| 620 | |
| 621 | @property |
| 622 | def has_inflight_transfers(self) -> bool: |
| 623 | return bool(self._inflight) |
| 624 | |
| 625 | |
| 626 | class _FakeClientHalf: |
| 627 | def __init__(self) -> None: |
| 628 | self._inbound: dict[int, object] = {} |
| 629 | |
| 630 | @property |
| 631 | def has_active_loads(self) -> bool: |
| 632 | return bool(self._inbound) |
| 633 | |
| 634 | |
| 635 | class _FakeSession: |
| 636 | """Fake bidirectional session that returns canned poll() results.""" |
| 637 | |
| 638 | def __init__( |
| 639 | self, |
| 640 | peer_id: str = "fake:1", |
| 641 | alive: bool = True, |
| 642 | connected: bool = True, |
| 643 | loads: list[LoadResult] | None = None, |
| 644 | stores: list[StoreResult] | None = None, |
| 645 | new_fetch_ids: list[str] | None = None, |
| 646 | close_jobs: list[int] | None = None, |
| 647 | close_req_ids: list[str] | None = None, |
| 648 | close_stores: list[int] | None = None, |
| 649 | close_failed_serves: list[ReqContext] | None = None, |
| 650 | ) -> None: |
| 651 | self.peer_id = peer_id |
| 652 | self.alive = alive |
| 653 | self.connected = connected |
| 654 | self.ready = True |
| 655 | self._loads = loads or [] |
| 656 | self._stores = stores or [] |
| 657 | self._new_fetch_ids = new_fetch_ids or [] |
| 658 | self._close_jobs = close_jobs or [] |
| 659 | self._close_req_ids = close_req_ids or [] |
| 660 | self._close_stores = close_stores or [] |
| 661 | self._close_failed_serves = close_failed_serves or [] |
| 662 | self.requests: list[tuple[int, str]] = [] |
| 663 | self.stores_added: list[tuple[str, list, object, int]] = [] |
| 664 | self.attached: list[object] = [] |
| 665 | self.finishes: list[str] = [] |
| 666 | # Mirror P2PSession._server._inflight (transfer_id → handle) and |
| 667 | # P2PSession._client.has_active_loads for the shutdown-drain and |
| 668 | # drain_jobs paths. Tests populate _server._inflight when needed. |
| 669 | self._server = _FakeServerHalf() |
| 670 | self._client = _FakeClientHalf() |
| 671 | |
| 672 | @property |
| 673 | def has_pending_work(self) -> bool: |
| 674 | return self._client.has_active_loads or self._server.has_inflight_transfers |
| 675 | |
| 676 | def poll(self): |
| 677 | result = SessionPollResult( |
| 678 | loads=self._loads, |
| 679 | stores=self._stores, |
| 680 | new_fetch_ids=self._new_fetch_ids, |
| 681 | ) |
| 682 | self._loads = [] |
| 683 | self._stores = [] |
| 684 | self._new_fetch_ids = [] |
| 685 | return result |
| 686 | |
| 687 | def request_blocks(self, job_id, kv_request_id, keys, block_ids): |
| 688 | self.requests.append((job_id, kv_request_id)) |
| 689 | |
| 690 | def add_stored_blocks(self, kv_request_id, keys, block_ids, job_id): |
| 691 | self.stores_added.append((kv_request_id, list(keys), block_ids, job_id)) |
| 692 | |
| 693 | def attach_connection(self, conn): |
| 694 | self.attached.append(conn) |
| 695 | self.connected = True |
| 696 | |
| 697 | def finish_request(self, kv_request_id): |
| 698 | self.finishes.append(kv_request_id) |
| 699 | |
| 700 | def close(self): |
| 701 | return SessionCloseResult( |
| 702 | failed_jobs=self._close_jobs, |
| 703 | failed_req_ids=self._close_req_ids, |
| 704 | failed_stores=self._close_stores, |
| 705 | failed_serves=self._close_failed_serves, |
| 706 | ) |
| 707 | |
| 708 | |
| 709 | class TestGetFinished: |
| 710 | def _make(self) -> P2PSecondaryTierManager: |
| 711 | mgr = _make_manager() |
| 712 | mgr._finished_jobs = [ |
| 713 | JobResult(job_id=1, success=True), |
| 714 | JobResult(job_id=2, success=False), |
| 715 | ] |
| 716 | |
| 717 | class FakeControl: |
| 718 | def poll(self): |
| 719 | return [] |
| 720 | |
| 721 | mgr._control = FakeControl() |
| 722 | mgr._data = None |
| 723 | return mgr |
| 724 | |
| 725 | def test_drains_finished_jobs(self): |
| 726 | """get_finished_jobs returns and clears accumulated results.""" |
| 727 | mgr = self._make() |
| 728 | results = list(mgr.get_finished_jobs()) |
| 729 | assert len(results) == 2 |
| 730 | assert JobResult(job_id=1, success=True) in results |
| 731 | assert JobResult(job_id=2, success=False) in results |
| 732 | # Second call returns empty. |
| 733 | assert list(mgr.get_finished_jobs()) == [] |
| 734 | |
| 735 | def test_reaps_dead_sessions(self): |
| 736 | """Dead connected sessions are removed and their pending jobs failed.""" |
| 737 | |
| 738 | class FakeData: |
| 739 | def remove_remote_peer(self, pid): |
| 740 | pass |
| 741 | |
| 742 | mgr = self._make() |
| 743 | mgr._data = FakeData() |
| 744 | dead = _FakeSession( |
| 745 | peer_id="dead:1234", |
| 746 | alive=False, |
| 747 | connected=True, |
| 748 | close_jobs=[20], |
| 749 | close_req_ids=["req-load"], |
| 750 | close_stores=[10, 11], |
| 751 | ) |
| 752 | mgr._sessions["dead:1234"] = dead |
| 753 | |
| 754 | results = list(mgr.get_finished_jobs()) |
| 755 | # 2 baseline + 1 failed load + 2 failed stores |
| 756 | assert len(results) == 5 |
| 757 | assert JobResult(job_id=10, success=False) in results |
| 758 | assert JobResult(job_id=11, success=False) in results |
| 759 | assert JobResult(job_id=20, success=False) in results |
| 760 | assert "dead:1234" not in mgr._sessions |
| 761 | assert "req-load" in mgr._failed_req_ids |
| 762 | |
| 763 | def test_reap_fails_probes(self): |
| 764 | """A reaped session's in-flight lookups land in _failed_req_ids so |
| 765 | the consumer's lookup() returns MISS instead of RETRY forever.""" |
| 766 | |
| 767 | class FakeData: |
| 768 | def remove_remote_peer(self, pid): |
| 769 | pass |
| 770 | |
| 771 | mgr = self._make() |
| 772 | mgr._data = FakeData() |
| 773 | dead = _FakeSession( |
| 774 | peer_id="dead:1234", |
| 775 | alive=False, |
| 776 | connected=True, |
| 777 | close_req_ids=["req-probe-1", "req-probe-2"], |
| 778 | ) |
| 779 | mgr._sessions["dead:1234"] = dead |
| 780 | |
| 781 | list(mgr.get_finished_jobs()) |
| 782 | assert "dead:1234" not in mgr._sessions |
| 783 | assert "req-probe-1" in mgr._failed_req_ids |
| 784 | assert "req-probe-2" in mgr._failed_req_ids |
| 785 | |
| 786 | def test_unbound_store_kept_within_timeout(self): |
| 787 | """Recently-parked unbound stores stay across a poll.""" |
| 788 | mgr = self._make() |
| 789 | from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch |
| 790 | |
| 791 | mgr._unbound_stores["req-fresh"] = [ |
| 792 | _UnboundStoreBatch(job_id=99, keys=[b"k"], block_ids=[0]) |
| 793 | ] |
| 794 | list(mgr.get_finished_jobs()) |
| 795 | assert "req-fresh" in mgr._unbound_stores |
| 796 | |
| 797 | def test_unbound_store_reaped_after_timeout(self): |
| 798 | """Unbound stores past the reap deadline surface as failed and their |
| 799 | kv_request_id lands in _reaped_stores so a late FetchMsg is rejected |
| 800 | instead of parking demand nothing can satisfy.""" |
| 801 | from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch |
| 802 | |
| 803 | mgr = self._make() |
| 804 | stale = _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]) |
| 805 | # Backdate the submission so the head batch is past the deadline. |
| 806 | stale.submitted_at = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S - 1.0 |
| 807 | mgr._unbound_stores["req-stale"] = [ |
| 808 | stale, |
| 809 | _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]), |
| 810 | ] |
| 811 | |
| 812 | results = list(mgr.get_finished_jobs()) |
| 813 | |
| 814 | assert "req-stale" not in mgr._unbound_stores |
| 815 | # 2 baseline + 2 buffered stores |
| 816 | assert JobResult(job_id=10, success=False) in results |
| 817 | assert JobResult(job_id=11, success=False) in results |
| 818 | assert "req-stale" in mgr._reaped_stores |
| 819 | |
| 820 | def test_submit_store_parks_unbound_batch(self): |
| 821 | """submit_store on an unbound id appends a batch with a fresh |
| 822 | submitted_at stamp so the unbound-store sweep can age it out.""" |
| 823 | mgr = _make_manager() |
| 824 | job = _job_metadata( |
| 825 | job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") |
| 826 | ) |
| 827 | before = time.monotonic() |
| 828 | mgr.submit_store(job) |
| 829 | after = time.monotonic() |
| 830 | batches = mgr._unbound_stores["req-1"] |
| 831 | assert len(batches) == 1 |
| 832 | assert before <= batches[0].submitted_at <= after |
| 833 | |
| 834 | |
| 835 | # --------------------------------------------------------------------------- |
| 836 | # has_pending_work |
| 837 | # --------------------------------------------------------------------------- |
| 838 | |
| 839 | |
| 840 | class TestHasPendingWork: |
| 841 | """has_pending_work() must always return True so the engine keeps |
| 842 | ticking the offload pipeline — that's the only thread driving |
| 843 | _control.poll() (incoming peer connects) and session.poll() |
| 844 | (incoming fetch messages on existing sessions).""" |
| 845 | |
| 846 | def test_returns_true_unconditionally_to_keep_engine_ticking(self): |
| 847 | """Even with no sessions and no jobs, has_pending_work() returns |
| 848 | True so the engine keeps calling get_finished_jobs(), which is |
| 849 | what drives _control.poll() for inbound peer connects.""" |
| 850 | mgr = _make_manager() |
| 851 | assert mgr.has_pending_work() is True |
| 852 | |
| 853 | def test_returns_true_even_when_sessions_present(self): |
| 854 | """The result is the same regardless of session state — there is |
| 855 | no 'idle' branch.""" |
| 856 | mgr = _make_manager() |
| 857 | mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") |
| 858 | assert mgr.has_pending_work() is True |
| 859 | |
| 860 | |
| 861 | # --------------------------------------------------------------------------- |
| 862 | # Shutdown drain |
| 863 | # --------------------------------------------------------------------------- |
| 864 | |
| 865 | |
| 866 | class _ShutdownFakeData: |
| 867 | """Fake DataTransport that records cancel/poll/close calls and |
| 868 | drives the wait-cancel loop with a scriptable `still` queue.""" |
| 869 | |
| 870 | def __init__(self, still_queue: list[list[int]] | None = None) -> None: |
| 871 | # Each list in still_queue is the set of ids the next |
| 872 | # cancel(mode="wait") should report as still inflight. The last |
| 873 | # entry repeats once exhausted. |
| 874 | self._still_queue = list(still_queue) if still_queue else [[]] |
| 875 | self.cancel_calls: list[tuple[list[int], str]] = [] |
| 876 | self.poll_calls: int = 0 |
| 877 | self.close_calls: int = 0 |
| 878 | |
| 879 | def cancel(self, transfer_ids, mode: str = "immediate") -> list[int]: |
| 880 | ids = list(transfer_ids) |
| 881 | self.cancel_calls.append((ids, mode)) |
| 882 | if mode == "wait": |
| 883 | if len(self._still_queue) > 1: |
| 884 | return list(self._still_queue.pop(0)) |
| 885 | return list(self._still_queue[0]) |
| 886 | return [] |
| 887 | |
| 888 | def poll(self, peer_id=None): |
| 889 | self.poll_calls += 1 |
| 890 | |
| 891 | class _Empty: |
| 892 | done: list[int] = [] |
| 893 | failed: list[int] = [] |
| 894 | |
| 895 | return _Empty() |
| 896 | |
| 897 | def close(self) -> None: |
| 898 | self.close_calls += 1 |
| 899 | |
| 900 | |
| 901 | class _ShutdownFakeControl: |
| 902 | def __init__(self) -> None: |
| 903 | self.close_calls: int = 0 |
| 904 | |
| 905 | def close(self) -> None: |
| 906 | self.close_calls += 1 |
| 907 | |
| 908 | |
| 909 | class TestShutdownDrain: |
| 910 | """shutdown() drains inflight transfers via cancel(mode='wait') |
| 911 | before calling _data.close(), with a 3s deadline fallback to |
| 912 | cancel(mode='immediate').""" |
| 913 | |
| 914 | def _prep( |
| 915 | self, |
| 916 | still_queue: list[list[int]] | None = None, |
| 917 | inflight_ids: list[int] | None = None, |
| 918 | ) -> tuple[P2PSecondaryTierManager, _ShutdownFakeData, _ShutdownFakeControl]: |
| 919 | mgr = _make_manager() |
| 920 | data = _ShutdownFakeData(still_queue=still_queue) |
| 921 | control = _ShutdownFakeControl() |
| 922 | mgr._data = data |
| 923 | mgr._control = control |
| 924 | if inflight_ids: |
| 925 | session = _FakeSession(peer_id="peer:1", connected=True) |
| 926 | session._server._inflight = {tid: object() for tid in inflight_ids} |
| 927 | mgr._sessions["peer:1"] = session |
| 928 | return mgr, data, control |
| 929 | |
| 930 | def test_shutdown_drains_inflight_via_wait_cancel(self): |
| 931 | # First cancel(wait) returns the input still inflight; second returns []. |
| 932 | mgr, data, control = self._prep( |
| 933 | still_queue=[[42, 43], []], |
| 934 | inflight_ids=[42, 43], |
| 935 | ) |
| 936 | mgr.shutdown() |
| 937 | |
| 938 | wait_calls = [c for c in data.cancel_calls if c[1] == "wait"] |
| 939 | immediate_calls = [c for c in data.cancel_calls if c[1] == "immediate"] |
| 940 | assert len(wait_calls) >= 1 |
| 941 | assert wait_calls[0][0] == [42, 43] |
| 942 | assert immediate_calls == [] |
| 943 | # poll() was driven between cancel attempts. |
| 944 | assert data.poll_calls >= 1 |
| 945 | # _data and _control were closed exactly once each, after the drain. |
| 946 | assert data.close_calls == 1 |
| 947 | assert control.close_calls == 1 |
| 948 | |
| 949 | def test_shutdown_force_cancels_after_timeout(self, monkeypatch): |
| 950 | # Drain never completes — wait-cancel keeps returning the inflight set. |
| 951 | # Use a synthetic clock so the test does not depend on real wallclock |
| 952 | # being able to advance in <50ms on a loaded CI node: |
| 953 | # call 1 (deadline calc): 100.0 -> deadline = 100.05 |
| 954 | # call 2 (loop predicate): 100.0 -> enters loop, one wait-cancel |
| 955 | # call 3 (loop predicate): 100.06 -> exits, force-cancel runs |
| 956 | monkeypatch.setattr(manager_module, "_SHUTDOWN_DRAIN_TIMEOUT_S", 0.05) |
| 957 | monkeypatch.setattr(manager_module, "_DRAIN_SLEEP_S", 0.0) |
| 958 | times = iter([100.0, 100.0, 100.06]) |
| 959 | # Patch via a fake module on `manager_module.time` so we do not mutate |
| 960 | # the global `time` module — other code in the process (e.g. the |
| 961 | # buildkite test collector's pytest_runtest_logreport hook) calls |
| 962 | # time.monotonic() before monkeypatch teardown. |
| 963 | fake_time = SimpleNamespace(monotonic=lambda: next(times), sleep=time.sleep) |
| 964 | monkeypatch.setattr(manager_module, "time", fake_time) |
| 965 | |
| 966 | mgr, data, control = self._prep( |
| 967 | still_queue=[[42]], |
| 968 | inflight_ids=[42], |
| 969 | ) |
| 970 | mgr.shutdown() |
| 971 | |
| 972 | wait_calls = [c for c in data.cancel_calls if c[1] == "wait"] |
| 973 | immediate_calls = [c for c in data.cancel_calls if c[1] == "immediate"] |
| 974 | assert len(wait_calls) == 1 |
| 975 | assert wait_calls[0][0] == [42] |
| 976 | assert immediate_calls == [([42], "immediate")] |
| 977 | assert data.close_calls == 1 |
| 978 | assert control.close_calls == 1 |
| 979 | |
| 980 | def test_shutdown_no_inflight_skips_drain(self): |
| 981 | mgr, data, control = self._prep() |
| 982 | mgr.shutdown() |
| 983 | |
| 984 | assert data.cancel_calls == [] |
| 985 | assert data.poll_calls == 0 |
| 986 | assert data.close_calls == 1 |
| 987 | assert control.close_calls == 1 |
| 988 | |
| 989 | |
| 990 | # --------------------------------------------------------------------------- |
| 991 | # Bidirectional regression test — both managers act as client AND server |
| 992 | # toward each other on the same peer_id. This is the case the unification |
| 993 | # is meant to fix. |
| 994 | # --------------------------------------------------------------------------- |
| 995 | |
| 996 | |
| 997 | class _LoopbackControl: |
| 998 | """In-memory control transport that pairs two managers head-to-head. |
| 999 | |
| 1000 | Each manager hands its outbound message buffer to the other's inbound |
| 1001 | queue on poll(). connect() returns a connection whose send() writes |
| 1002 | into the peer's inbound side; recv() reads what the peer's connect-or- |
| 1003 | poll path delivered for us. |
| 1004 | """ |
| 1005 | |
| 1006 | def __init__(self, local_id: str) -> None: |
| 1007 | self._local_id = local_id |
| 1008 | self._peer: _LoopbackControl | None = None |
| 1009 | self._inbound_outbound: dict[str, _LoopbackConnection] = {} |
| 1010 | # Pending inbound for a peer that has not yet been registered. |
| 1011 | self._pending: list[tuple[str, dict]] = [] |
| 1012 | |
| 1013 | def pair(self, peer: _LoopbackControl) -> None: |
| 1014 | self._peer = peer |
| 1015 | peer._peer = self |
| 1016 | |
| 1017 | def connect(self, peer_id: str): |
| 1018 | if peer_id in self._inbound_outbound: |
| 1019 | raise AssertionError(f"already connected to {peer_id}") |
| 1020 | conn = _LoopbackConnection(self, peer_id) |
| 1021 | self._inbound_outbound[peer_id] = conn |
| 1022 | return conn |
| 1023 | |
| 1024 | def poll(self): |
| 1025 | # Drain whatever the peer has sent toward us. |
| 1026 | new = [] |
| 1027 | if self._peer is not None: |
| 1028 | for pid, msg in self._peer._drain_outbound_to(self._local_id): |
| 1029 | conn = self._inbound_outbound.get(pid) |
| 1030 | if conn is None: |
| 1031 | conn = _LoopbackConnection(self, pid) |
| 1032 | self._inbound_outbound[pid] = conn |
| 1033 | new.append(conn) |
| 1034 | conn._inbox.append(msg) |
| 1035 | return new |
| 1036 | |
| 1037 | def _drain_outbound_to(self, peer_local_id: str): |
| 1038 | # Peer is calling our poll → return all messages we've sent toward |
| 1039 | # peer_local_id (which is the peer's own local_id). |
| 1040 | out: list[tuple[str, dict]] = [] |
| 1041 | for pid, conn in list(self._inbound_outbound.items()): |
| 1042 | # Each outgoing message goes to peer_local_id and is tagged |
| 1043 | # with the sender's local_id (i.e., self._local_id). |
| 1044 | for msg in conn._outbox: |
| 1045 | out.append((self._local_id, msg)) |
| 1046 | conn._outbox.clear() |
| 1047 | return out |
| 1048 | |
| 1049 | def close(self): |
| 1050 | for conn in self._inbound_outbound.values(): |
| 1051 | conn.close() |
| 1052 | self._inbound_outbound.clear() |
| 1053 | |
| 1054 | |
| 1055 | class _LoopbackConnection: |
| 1056 | def __init__(self, transport: _LoopbackControl, peer_id: str) -> None: |
| 1057 | self._transport = transport |
| 1058 | self.peer_id = peer_id |
| 1059 | self._inbox: list[dict] = [] |
| 1060 | self._outbox: list[dict] = [] |
| 1061 | self._closed = False |
| 1062 | |
| 1063 | @property |
| 1064 | def alive(self) -> bool: |
| 1065 | return not self._closed |
| 1066 | |
| 1067 | def send(self, msg: dict) -> None: |
| 1068 | if self._closed: |
| 1069 | raise RuntimeError("send on closed conn") |
| 1070 | self._outbox.append(msg) |
| 1071 | |
| 1072 | def recv(self) -> list[dict]: |
| 1073 | msgs = self._inbox |
| 1074 | self._inbox = [] |
| 1075 | return msgs |
| 1076 | |
| 1077 | def mark_dead(self) -> None: |
| 1078 | self._closed = True |
| 1079 | |
| 1080 | def close(self) -> None: |
| 1081 | self._closed = True |
| 1082 | |
| 1083 | |
| 1084 | class _FakeData: |
| 1085 | """Minimal NIXL fake that lets matched transfers complete on the next poll.""" |
| 1086 | |
| 1087 | def __init__(self, local_id: str) -> None: |
| 1088 | self._local_id = local_id |
| 1089 | self.block_len = 4096 |
| 1090 | self.base_addr = 0x1000 |
| 1091 | self.num_blocks = 16 |
| 1092 | self.config_fingerprint = "" |
| 1093 | self._remote_peers: dict[str, dict] = {} |
| 1094 | self._inflight_done: list[int] = [] |
| 1095 | self._next_id = 0 |
| 1096 | |
| 1097 | def get_agent_metadata(self) -> bytes: |
| 1098 | return f"meta-{self._local_id}".encode() |
| 1099 | |
| 1100 | def add_remote_peer( |
| 1101 | self, peer_id, agent_metadata, base_addr, num_blocks, block_len |
| 1102 | ) -> None: |
| 1103 | self._remote_peers[peer_id] = { |
| 1104 | "agent_metadata": agent_metadata, |
| 1105 | "base_addr": base_addr, |
| 1106 | "num_blocks": num_blocks, |
| 1107 | "block_len": block_len, |
| 1108 | } |
| 1109 | |
| 1110 | def remove_remote_peer(self, peer_id: str) -> None: |
| 1111 | self._remote_peers.pop(peer_id, None) |
| 1112 | |
| 1113 | def write_blocks(self, peer_id, local_idxs, remote_idxs): |
| 1114 | if peer_id not in self._remote_peers: |
| 1115 | return None |
| 1116 | tid = self._next_id |
| 1117 | self._next_id += 1 |
| 1118 | self._inflight_done.append(tid) |
| 1119 | return tid |
| 1120 | |
| 1121 | def poll(self, peer_id=None): |
| 1122 | from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult |
| 1123 | |
| 1124 | done = self._inflight_done[:] |
| 1125 | self._inflight_done.clear() |
| 1126 | return PollResult(done=done, failed=[]) |
| 1127 | |
| 1128 | def cancel(self, transfer_ids) -> None: |
| 1129 | pass |
| 1130 | |
| 1131 | def close(self) -> None: |
| 1132 | pass |
| 1133 | |
| 1134 | |
| 1135 | def _build_paired_managers() -> tuple[P2PSecondaryTierManager, P2PSecondaryTierManager]: |
| 1136 | """Two managers each acting as both client and server toward the other. |
| 1137 | |
| 1138 | Wires _LoopbackControl pair + per-side _FakeData so transfers complete |
| 1139 | on the next poll. The test drives polling by calling get_finished_jobs(), |
| 1140 | which invokes _poll_once synchronously on the calling thread. |
| 1141 | """ |
| 1142 | mgr_a = _make_manager() |
| 1143 | mgr_b = _make_manager() |
| 1144 | mgr_a._local_id = "A:1" |
| 1145 | mgr_b._local_id = "B:2" |
| 1146 | |
| 1147 | ctrl_a = _LoopbackControl(mgr_a._local_id) |
| 1148 | ctrl_b = _LoopbackControl(mgr_b._local_id) |
| 1149 | ctrl_a.pair(ctrl_b) |
| 1150 | |
| 1151 | mgr_a._control = ctrl_a |
| 1152 | mgr_b._control = ctrl_b |
| 1153 | mgr_a._data = _FakeData(mgr_a._local_id) |
| 1154 | mgr_b._data = _FakeData(mgr_b._local_id) |
| 1155 | |
| 1156 | return mgr_a, mgr_b |
| 1157 | |
| 1158 | |
| 1159 | class TestBidirectionalManager: |
| 1160 | """Two managers each load FROM and serve TO the other over a single peer.""" |
| 1161 | |
| 1162 | def test_both_loads_succeed(self): |
| 1163 | mgr_a, mgr_b = _build_paired_managers() |
| 1164 | |
| 1165 | a_loads_kv = "req-AtoB-load" # A loads, B serves |
| 1166 | b_loads_kv = "req-BtoA-load" # B loads, A serves |
| 1167 | |
| 1168 | a_decoder_params = { |
| 1169 | "remote_prefiller": { |
| 1170 | "kv_request_id": a_loads_kv, |
| 1171 | "remote_host": "B", |
| 1172 | "remote_port": 2, |
| 1173 | }, |
| 1174 | } |
| 1175 | b_decoder_params = { |
| 1176 | "remote_prefiller": { |
| 1177 | "kv_request_id": b_loads_kv, |
| 1178 | "remote_host": "A", |
| 1179 | "remote_port": 1, |
| 1180 | }, |
| 1181 | } |
| 1182 | a_prefiller_params = {"remote_decoder": {"kv_request_id": b_loads_kv}} |
| 1183 | b_prefiller_params = {"remote_decoder": {"kv_request_id": a_loads_kv}} |
| 1184 | |
| 1185 | # 1. Both sides open client-role sessions toward the peer. |
| 1186 | mgr_a.on_new_request(_req_context(a_decoder_params)) |
| 1187 | mgr_b.on_new_request(_req_context(b_decoder_params)) |
| 1188 | |
| 1189 | # 2. Both sides store the blocks the peer will fetch. |
| 1190 | mgr_a.submit_store( |
| 1191 | _job_metadata( |
| 1192 | job_id=100, |
| 1193 | keys=[b"a-block"], |
| 1194 | chunk_ids=[0], |
| 1195 | kv_params=a_prefiller_params, |
| 1196 | ) |
| 1197 | ) |
| 1198 | mgr_b.submit_store( |
| 1199 | _job_metadata( |
| 1200 | job_id=200, |
| 1201 | keys=[b"b-block"], |
| 1202 | chunk_ids=[0], |
| 1203 | kv_params=b_prefiller_params, |
| 1204 | ) |
| 1205 | ) |
| 1206 | |
| 1207 | # 3. Both sides submit loads. |
| 1208 | mgr_a.submit_load( |
| 1209 | _job_metadata( |
| 1210 | job_id=101, |
| 1211 | keys=[b"b-block"], |
| 1212 | chunk_ids=[0], |
| 1213 | kv_params=a_decoder_params, |
| 1214 | ) |
| 1215 | ) |
| 1216 | mgr_b.submit_load( |
| 1217 | _job_metadata( |
| 1218 | job_id=201, |
| 1219 | keys=[b"a-block"], |
| 1220 | chunk_ids=[0], |
| 1221 | kv_params=b_decoder_params, |
| 1222 | ) |
| 1223 | ) |
| 1224 | |
| 1225 | # 4. Drive several poll iterations on each side. Each |
| 1226 | # get_finished_jobs() call invokes _poll_once synchronously. |
| 1227 | all_a: list[JobResult] = [] |
| 1228 | all_b: list[JobResult] = [] |
| 1229 | for _ in range(8): |
| 1230 | all_a.extend(list(mgr_a.get_finished_jobs())) |
| 1231 | all_b.extend(list(mgr_b.get_finished_jobs())) |
| 1232 | |
| 1233 | # Both load jobs and both store jobs must complete successfully. |
| 1234 | a_ok = {r.job_id for r in all_a if r.success} |
| 1235 | b_ok = {r.job_id for r in all_b if r.success} |
| 1236 | # A: load job 101 + store job 100 |
| 1237 | assert 101 in a_ok, f"A loads succeeded: {a_ok}" |
| 1238 | assert 100 in a_ok, f"A stores succeeded: {a_ok}" |
| 1239 | # B: load job 201 + store job 200 |
| 1240 | assert 201 in b_ok, f"B loads succeeded: {b_ok}" |
| 1241 | assert 200 in b_ok, f"B stores succeeded: {b_ok}" |
| 1242 | |
| 1243 | def test_late_fetch_after_reap_fails_immediately(self): |
| 1244 | """A fetch for a reaped kv_request_id fails in one round trip. |
| 1245 | |
| 1246 | The producer drops parked blocks once unbound_store_timeout_s |
| 1247 | expires. A consumer whose fetch arrives after that must learn on the |
| 1248 | next tick via TransferDoneMsg(success=False) — not stall for |
| 1249 | _LOAD_TIMEOUT_S and then take the abort path. |
| 1250 | """ |
| 1251 | from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S |
| 1252 | |
| 1253 | mgr_a, mgr_b = _build_paired_managers() |
| 1254 | kv_id = "req-late" |
| 1255 | |
| 1256 | # Record every message type leaving the consumer, so the assertion |
| 1257 | # below can prove no abort was ever needed. |
| 1258 | sent_from_a: list[str] = [] |
| 1259 | drain = mgr_a._control._drain_outbound_to |
| 1260 | |
| 1261 | def recording_drain(peer_local_id: str): |
| 1262 | out = drain(peer_local_id) |
| 1263 | sent_from_a.extend(msg.get("type") for _, msg in out) |
| 1264 | return out |
| 1265 | |
| 1266 | mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign] |
| 1267 | |
| 1268 | # Producer parks the blocks, then reaps them before any fetch lands. |
| 1269 | mgr_b.submit_store( |
| 1270 | _job_metadata( |
| 1271 | job_id=200, |
| 1272 | keys=[b"b-block"], |
| 1273 | chunk_ids=[0], |
| 1274 | kv_params={"remote_decoder": {"kv_request_id": kv_id}}, |
| 1275 | ) |
| 1276 | ) |
| 1277 | mgr_b._unbound_store_timeout_s = 0.0 |
| 1278 | reap_results = list(mgr_b.get_finished_jobs()) |
| 1279 | assert JobResult(job_id=200, success=False) in reap_results |
| 1280 | assert kv_id in mgr_b._reaped_stores |
| 1281 | |
| 1282 | # Consumer now asks for the blocks that no longer exist. |
| 1283 | consumer_params = { |
| 1284 | "remote_prefiller": { |
| 1285 | "kv_request_id": kv_id, |
| 1286 | "remote_host": "B", |
| 1287 | "remote_port": 2, |
| 1288 | }, |
| 1289 | } |
| 1290 | mgr_a.on_new_request(_req_context(consumer_params)) |
| 1291 | mgr_a.submit_load( |
| 1292 | _job_metadata( |
| 1293 | job_id=101, |
| 1294 | keys=[b"b-block"], |
| 1295 | chunk_ids=[0], |
| 1296 | kv_params=consumer_params, |
| 1297 | ) |
| 1298 | ) |
| 1299 | |
| 1300 | started = time.monotonic() |
| 1301 | all_a: list[JobResult] = [] |
| 1302 | for _ in range(8): |
| 1303 | all_a.extend(list(mgr_a.get_finished_jobs())) |
| 1304 | list(mgr_b.get_finished_jobs()) |
| 1305 | elapsed = time.monotonic() - started |
| 1306 | |
| 1307 | assert JobResult(job_id=101, success=False) in all_a, all_a |
| 1308 | assert kv_id in mgr_a._failed_req_ids |
| 1309 | assert "abort_fetch" not in sent_from_a, sent_from_a |
| 1310 | assert elapsed < _LOAD_TIMEOUT_S |
| 1311 | |
| 1312 | |
| 1313 | # --------------------------------------------------------------------------- |
| 1314 | # _accept_new_peers — duplicate connection rejection |
| 1315 | # --------------------------------------------------------------------------- |
| 1316 | |
| 1317 | |
| 1318 | class _RecordingConn: |
| 1319 | """Inbound connection stub that records close()/peer_id only.""" |
| 1320 | |
| 1321 | def __init__(self, peer_id: str) -> None: |
| 1322 | self.peer_id = peer_id |
| 1323 | self.close_calls: int = 0 |
| 1324 | |
| 1325 | def close(self) -> None: |
| 1326 | self.close_calls += 1 |
| 1327 | |
| 1328 | |
| 1329 | class TestAcceptNewPeers: |
| 1330 | """A second inbound from an already-connected peer is rejected and the |
| 1331 | new conn is closed; the existing session is left untouched.""" |
| 1332 | |
| 1333 | def test_duplicate_connection_is_closed_and_existing_session_untouched(self): |
| 1334 | mgr = _make_manager() |
| 1335 | peer_id = "10.0.0.1:8000" |
| 1336 | existing = _FakeSession(peer_id=peer_id, connected=True) |
| 1337 | mgr._sessions[peer_id] = existing |
| 1338 | |
| 1339 | new_conn = _RecordingConn(peer_id) |
| 1340 | mgr._accept_new_peers([new_conn]) |
| 1341 | |
| 1342 | # Manager swallowed the ValueError and closed the duplicate conn. |
| 1343 | assert new_conn.close_calls == 1 |
| 1344 | # Existing session was NOT re-attached. |
| 1345 | assert existing.attached == [] |
| 1346 | # Session map unchanged. |
| 1347 | assert mgr._sessions[peer_id] is existing |
| 1348 | |
| 1349 | def test_creates_session_for_new_peer(self): |
| 1350 | """An inbound conn from a peer with no existing session creates |
| 1351 | a fresh connected session and registers it under conn.peer_id. |
| 1352 | The prefiller has no pre-created pending session anymore — the |
| 1353 | first signal of a peer's existence is its inbound connection.""" |
| 1354 | |
| 1355 | class FakeData: |
| 1356 | block_len = 4096 |
| 1357 | base_addr = 0x1000 |
| 1358 | num_blocks = 16 |
| 1359 | config_fingerprint = "" |
| 1360 | |
| 1361 | def get_agent_metadata(self): |
| 1362 | return b"meta" |
| 1363 | |
| 1364 | def add_remote_peer(self, *args, **kwargs): |
| 1365 | pass |
| 1366 | |
| 1367 | mgr = _make_manager() |
| 1368 | mgr._data = FakeData() |
| 1369 | peer_id = "10.0.0.1:8000" |
| 1370 | |
| 1371 | # Real ControlConnection-shaped fake: send/close/peer_id only. |
| 1372 | sent: list[dict] = [] |
| 1373 | |
| 1374 | class _Conn: |
| 1375 | def __init__(self, pid: str) -> None: |
| 1376 | self.peer_id = pid |
| 1377 | self.alive = True |
| 1378 | |
| 1379 | def send(self, msg: dict) -> None: |
| 1380 | sent.append(msg) |
| 1381 | |
| 1382 | def close(self) -> None: |
| 1383 | self.alive = False |
| 1384 | |
| 1385 | mgr._accept_new_peers([_Conn(peer_id)]) |
| 1386 | |
| 1387 | assert peer_id in mgr._sessions |
| 1388 | assert mgr._sessions[peer_id].connected is True |
| 1389 | # Session sent its ConnectMsg on the new connection. |
| 1390 | assert any(m for m in sent) |
| 1391 | |
| 1392 | |
| 1393 | # --------------------------------------------------------------------------- |
| 1394 | # _poll_once orchestration |
| 1395 | # --------------------------------------------------------------------------- |
| 1396 | |
| 1397 | |
| 1398 | class TestPollOnce: |
| 1399 | """_poll_once must drain control, accept new peers, poll every session, |
| 1400 | surface results, and reap dead sessions — in that order.""" |
| 1401 | |
| 1402 | def test_orchestrates_accept_poll_and_reap(self): |
| 1403 | mgr = _make_manager() |
| 1404 | |
| 1405 | # Alive session whose poll() returns one load + one store. |
| 1406 | peer_alive = "10.0.0.2:9000" |
| 1407 | alive = _FakeSession( |
| 1408 | peer_id=peer_alive, |
| 1409 | alive=True, |
| 1410 | connected=True, |
| 1411 | loads=[LoadResult(job_id=11, kv_request_id="req-11", success=True)], |
| 1412 | stores=[StoreResult(job_id=22, success=True)], |
| 1413 | ) |
| 1414 | mgr._sessions[peer_alive] = alive |
| 1415 | |
| 1416 | # Dead session whose pending close() jobs surface as failures. |
| 1417 | peer_dead = "10.0.0.3:9999" |
| 1418 | dead = _FakeSession( |
| 1419 | peer_id=peer_dead, |
| 1420 | alive=False, |
| 1421 | connected=True, |
| 1422 | close_jobs=[33], |
| 1423 | close_req_ids=["req-33"], |
| 1424 | close_stores=[44], |
| 1425 | ) |
| 1426 | mgr._sessions[peer_dead] = dead |
| 1427 | |
| 1428 | class _Ctrl: |
| 1429 | def poll(self_inner): |
| 1430 | return [] |
| 1431 | |
| 1432 | class _Data: |
| 1433 | def remove_remote_peer(self_inner, pid): |
| 1434 | pass |
| 1435 | |
| 1436 | mgr._control = _Ctrl() |
| 1437 | mgr._data = _Data() |
| 1438 | |
| 1439 | mgr._poll_once() |
| 1440 | |
| 1441 | # Every session was polled — alive's results landed. |
| 1442 | # Dead session was reaped — its close() failures landed. |
| 1443 | finished = mgr._finished_jobs |
| 1444 | ok = {(r.job_id, r.success) for r in finished} |
| 1445 | assert (11, True) in ok # alive load result |
| 1446 | assert (22, True) in ok # alive store result |
| 1447 | assert (33, False) in ok # dead session's pending load |
| 1448 | assert (44, False) in ok # dead session's pending store |
| 1449 | assert "req-33" in mgr._failed_req_ids |
| 1450 | assert peer_dead not in mgr._sessions |
| 1451 | assert peer_alive in mgr._sessions |
| 1452 | |
| 1453 | def test_new_fetch_id_binds_and_replays_unbound_batches(self): |
| 1454 | """When session.poll() reports a kv_request_id whose FetchMsg |
| 1455 | arrived this tick, the manager binds it to that session and |
| 1456 | replays every parked submit_store batch via add_stored_blocks.""" |
| 1457 | from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch |
| 1458 | |
| 1459 | mgr = _make_manager() |
| 1460 | peer = "10.0.0.1:8000" |
| 1461 | sess = _FakeSession( |
| 1462 | peer_id=peer, |
| 1463 | alive=True, |
| 1464 | connected=True, |
| 1465 | new_fetch_ids=["req-1"], |
| 1466 | ) |
| 1467 | mgr._sessions[peer] = sess |
| 1468 | mgr._unbound_stores["req-1"] = [ |
| 1469 | _UnboundStoreBatch(job_id=5, keys=[b"k1"], block_ids=[0]), |
| 1470 | _UnboundStoreBatch(job_id=6, keys=[b"k2"], block_ids=[1]), |
| 1471 | ] |
| 1472 | |
| 1473 | class _Ctrl: |
| 1474 | def poll(self_inner): |
| 1475 | return [] |
| 1476 | |
| 1477 | mgr._control = _Ctrl() |
| 1478 | mgr._poll_once() |
| 1479 | |
| 1480 | assert mgr._kv_to_session["req-1"] is sess |
| 1481 | assert "req-1" not in mgr._unbound_stores |
| 1482 | replayed = [ |
| 1483 | (kv_req_id, list(keys), job_id) |
| 1484 | for kv_req_id, keys, _, job_id in sess.stores_added |
| 1485 | ] |
| 1486 | assert replayed == [("req-1", [b"k1"], 5), ("req-1", [b"k2"], 6)] |
| 1487 | |
| 1488 | def test_new_fetch_id_with_no_unbound_still_binds(self): |
| 1489 | """A FetchMsg for a kv_request_id with no parked batches still |
| 1490 | records the binding so subsequent submit_stores route fast.""" |
| 1491 | mgr = _make_manager() |
| 1492 | peer = "10.0.0.1:8000" |
| 1493 | sess = _FakeSession( |
| 1494 | peer_id=peer, |
| 1495 | alive=True, |
| 1496 | connected=True, |
| 1497 | new_fetch_ids=["req-fast"], |
| 1498 | ) |
| 1499 | mgr._sessions[peer] = sess |
| 1500 | |
| 1501 | class _Ctrl: |
| 1502 | def poll(self_inner): |
| 1503 | return [] |
| 1504 | |
| 1505 | mgr._control = _Ctrl() |
| 1506 | mgr._poll_once() |
| 1507 | |
| 1508 | assert mgr._kv_to_session["req-fast"] is sess |
| 1509 | assert sess.stores_added == [] |
| 1510 | |
| 1511 | def test_failed_load_records_kv_request_id(self): |
| 1512 | """A LoadResult(success=False) from session.poll() must add its |
| 1513 | kv_request_id to _failed_req_ids so future lookups return MISS.""" |
| 1514 | mgr = _make_manager() |
| 1515 | peer = "10.0.0.1:8000" |
| 1516 | sess = _FakeSession( |
| 1517 | peer_id=peer, |
| 1518 | alive=True, |
| 1519 | connected=True, |
| 1520 | loads=[LoadResult(job_id=5, kv_request_id="req-5", success=False)], |
| 1521 | ) |
| 1522 | mgr._sessions[peer] = sess |
| 1523 | |
| 1524 | class _Ctrl: |
| 1525 | def poll(self_inner): |
| 1526 | return [] |
| 1527 | |
| 1528 | mgr._control = _Ctrl() |
| 1529 | |
| 1530 | mgr._poll_once() |
| 1531 | |
| 1532 | assert mgr._finished_jobs == [JobResult(job_id=5, success=False)] |
| 1533 | assert "req-5" in mgr._failed_req_ids |
| 1534 | |
| 1535 | |
| 1536 | # --------------------------------------------------------------------------- |
| 1537 | # drain_jobs |
| 1538 | # --------------------------------------------------------------------------- |
| 1539 | |
| 1540 | |
| 1541 | class _DrainCtrl: |
| 1542 | """Trivial control fake whose poll() returns an empty list.""" |
| 1543 | |
| 1544 | def poll(self): |
| 1545 | return [] |
| 1546 | |
| 1547 | |
| 1548 | class TestDrainJobs: |
| 1549 | def test_returns_immediately_when_quiescent(self): |
| 1550 | """No sessions and no inflight: drain_jobs returns without sleeping.""" |
| 1551 | mgr = _make_manager() |
| 1552 | mgr._control = _DrainCtrl() |
| 1553 | |
| 1554 | sleeps: list[float] = [] |
| 1555 | # If drain_jobs sleeps when nothing is pending, that's a regression. |
| 1556 | import vllm.v1.kv_offload.tiering.p2p.manager as m |
| 1557 | |
| 1558 | original_sleep = m.time.sleep |
| 1559 | m.time.sleep = lambda s: sleeps.append(s) |
| 1560 | try: |
| 1561 | mgr.drain_jobs() |
| 1562 | finally: |
| 1563 | m.time.sleep = original_sleep |
| 1564 | |
| 1565 | assert sleeps == [] |
| 1566 | |
| 1567 | def test_returns_when_session_has_no_inflight_or_inbound(self): |
| 1568 | """A session with empty _inbound and _inflight does not block drain.""" |
| 1569 | mgr = _make_manager() |
| 1570 | mgr._control = _DrainCtrl() |
| 1571 | mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") |
| 1572 | # Should return on the first iteration. |
| 1573 | mgr.drain_jobs() |
| 1574 | |
| 1575 | def test_logs_warning_after_5s_then_completes(self, monkeypatch): |
| 1576 | """A session that stays inflight past 5s triggers the warning, and |
| 1577 | once it clears the loop returns.""" |
| 1578 | mgr = _make_manager() |
| 1579 | mgr._control = _DrainCtrl() |
| 1580 | sess = _FakeSession(peer_id="peer:1") |
| 1581 | sess._server._inflight = {1: object()} # non-empty |
| 1582 | mgr._sessions["peer:1"] = sess |
| 1583 | |
| 1584 | # Synthetic monotonic clock: 100.0 for the start stamp, then 106.0 |
| 1585 | # for the first elapsed-check (past the 5s warning threshold), then |
| 1586 | # steady at 106.0 for any later checks. |
| 1587 | clock = iter([100.0, 106.0]) |
| 1588 | |
| 1589 | def fake_monotonic() -> float: |
| 1590 | try: |
| 1591 | return next(clock) |
| 1592 | except StopIteration: |
| 1593 | return 106.0 |
| 1594 | |
| 1595 | monkeypatch.setattr(manager_module, "_DRAIN_SLEEP_S", 0.0) |
| 1596 | |
| 1597 | # Spy on the warning logger directly — vllm's logger does not |
| 1598 | # propagate to root, so caplog can't see it. |
| 1599 | warnings: list[str] = [] |
| 1600 | |
| 1601 | def record_warning(msg, *args, **_kwargs): |
| 1602 | warnings.append(msg % args if args else msg) |
| 1603 | |
| 1604 | monkeypatch.setattr(manager_module.logger, "warning", record_warning) |
| 1605 | |
| 1606 | # Clear inflight after the first sleep so drain can exit on the |
| 1607 | # next iteration's `pending` check. |
| 1608 | n_sleeps = 0 |
| 1609 | |
| 1610 | def clearing_sleep(_s): |
| 1611 | nonlocal n_sleeps |
| 1612 | n_sleeps += 1 |
| 1613 | sess._server._inflight = {} |
| 1614 | |
| 1615 | # Patch via a fake module on `manager_module.time` so we do not mutate |
| 1616 | # the global `time` module — other code in the process (e.g. the |
| 1617 | # buildkite test collector's pytest_runtest_logreport hook) calls |
| 1618 | # time.monotonic() before monkeypatch teardown. |
| 1619 | fake_time = SimpleNamespace(monotonic=fake_monotonic, sleep=clearing_sleep) |
| 1620 | monkeypatch.setattr(manager_module, "time", fake_time) |
| 1621 | |
| 1622 | mgr.drain_jobs() |
| 1623 | |
| 1624 | assert any("still draining after 5s" in w for w in warnings), warnings |
| 1625 | |
| 1626 | |
| 1627 | # --------------------------------------------------------------------------- |
| 1628 | # on_schedule_end |
| 1629 | # --------------------------------------------------------------------------- |
| 1630 | |
| 1631 | |
| 1632 | class TestOnScheduleEnd: |
| 1633 | def test_is_noop(self): |
| 1634 | """on_schedule_end is a documented no-op; just confirm it doesn't |
| 1635 | raise and doesn't mutate state.""" |
| 1636 | mgr = _make_manager() |
| 1637 | before_sessions = dict(mgr._sessions) |
| 1638 | before_jobs = list(mgr._finished_jobs) |
| 1639 | assert ( |
| 1640 | mgr.on_schedule_end( |
| 1641 | ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) |
| 1642 | ) |
| 1643 | is None |
| 1644 | ) |
| 1645 | assert mgr._sessions == before_sessions |
| 1646 | assert mgr._finished_jobs == before_jobs |
| 1647 | |
| 1648 | |
| 1649 | # --------------------------------------------------------------------------- |
| 1650 | # Connection death mid-transfer (real P2PSession via paired managers) |
| 1651 | # --------------------------------------------------------------------------- |
| 1652 | |
| 1653 | |
| 1654 | class TestConnectionDeathMidTransfer: |
| 1655 | """When a peer's control connection dies while a load is in flight, |
| 1656 | the load surfaces as failed and its kv_request_id lands in |
| 1657 | _failed_req_ids so future lookups route to local prefill. The |
| 1658 | prefiller-side store no longer travels through the session at store |
| 1659 | time (it's parked in _unbound_stores keyed by kv_request_id), so its |
| 1660 | cleanup on connection death is via on_request_finished or the |
| 1661 | unbound-store timeout — covered separately below.""" |
| 1662 | |
| 1663 | def test_dead_connection_with_pending_work_surfaces_failures(self): |
| 1664 | mgr_a, mgr_b = _build_paired_managers() |
| 1665 | |
| 1666 | a_decoder_params = { |
| 1667 | "remote_prefiller": { |
| 1668 | "kv_request_id": "req-load", |
| 1669 | "remote_host": "B", |
| 1670 | "remote_port": 2, |
| 1671 | }, |
| 1672 | } |
| 1673 | a_prefiller_params = {"remote_decoder": {"kv_request_id": "req-store"}} |
| 1674 | |
| 1675 | # Open the outbound session A->B and submit one load + one store. |
| 1676 | mgr_a.on_new_request(_req_context(a_decoder_params)) |
| 1677 | mgr_a.submit_store( |
| 1678 | _job_metadata( |
| 1679 | job_id=900, |
| 1680 | keys=[b"a-block"], |
| 1681 | chunk_ids=[0], |
| 1682 | kv_params=a_prefiller_params, |
| 1683 | ) |
| 1684 | ) |
| 1685 | mgr_a.submit_load( |
| 1686 | _job_metadata( |
| 1687 | job_id=901, |
| 1688 | keys=[b"b-block"], |
| 1689 | chunk_ids=[0], |
| 1690 | kv_params=a_decoder_params, |
| 1691 | ) |
| 1692 | ) |
| 1693 | |
| 1694 | # Drain anything the loopback can deliver synchronously, but stop |
| 1695 | # before the remote side has had time to complete the transfers. |
| 1696 | list(mgr_a.get_finished_jobs()) |
| 1697 | |
| 1698 | # Sanity: store 900 is parked in unbound_stores, not in any |
| 1699 | # session — the producer no longer learns the peer at store time. |
| 1700 | assert "req-store" in mgr_a._unbound_stores |
| 1701 | |
| 1702 | # Kill the control connection out from under the session. |
| 1703 | peer_id = "B:2" |
| 1704 | sess = mgr_a._sessions[peer_id] |
| 1705 | assert sess._conn is not None |
| 1706 | sess._conn.mark_dead() |
| 1707 | |
| 1708 | # Reap surfaces the load (which lived inside the session) as |
| 1709 | # failed. The store 900 is not session-scoped — it survives the |
| 1710 | # session reap and waits for the unbound-store timeout. |
| 1711 | results: list[JobResult] = [] |
| 1712 | for _ in range(3): |
| 1713 | results.extend(list(mgr_a.get_finished_jobs())) |
| 1714 | |
| 1715 | outcomes = {(r.job_id, r.success) for r in results} |
| 1716 | assert (901, False) in outcomes, f"load should fail: {outcomes}" |
| 1717 | assert (900, False) not in outcomes, f"store should still be parked: {outcomes}" |
| 1718 | assert "req-load" in mgr_a._failed_req_ids |
| 1719 | # Session removed. |
| 1720 | assert peer_id not in mgr_a._sessions |
| 1721 | # Store batch is still parked. |
| 1722 | assert "req-store" in mgr_a._unbound_stores |
| 1723 | |
| 1724 | # The engine signals the producer's request is done. That is a |
| 1725 | # no-op for the parked batch — `on_request_finished` does not |
| 1726 | # evict unbound stores; only `_reap_unbound_stores` does, after |
| 1727 | # the unbound-store timeout. Job 900 stays unfinished here. |
| 1728 | mgr_a.on_request_finished(_req_context(a_prefiller_params)) |
| 1729 | finishes = {(r.job_id, r.success) for r in mgr_a._finished_jobs} |
| 1730 | assert (900, False) not in finishes |
| 1731 | assert (900, True) not in finishes |
| 1732 | assert "req-store" in mgr_a._unbound_stores |
| 1733 | |
| 1734 | |
| 1735 | # --------------------------------------------------------------------------- |
| 1736 | # Tests for host/port resolution in __init__ (env-var defaults) |
| 1737 | # --------------------------------------------------------------------------- |
| 1738 | |
| 1739 | |
| 1740 | class TestBindHostPortDefaults: |
| 1741 | """host/port fall back to VLLM_P2P_SIDE_CHANNEL_* when not in config.""" |
| 1742 | |
| 1743 | @staticmethod |
| 1744 | def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager: |
| 1745 | """Build a manager with the transports/file-mapper stubbed out. |
| 1746 | |
| 1747 | The host is used verbatim (no resolution). The transport constructor |
| 1748 | args are recorded on ``mgr._test_calls`` so tests can assert the ZMQ |
| 1749 | identity (``host:port``) stays decoupled from the NIXL agent name |
| 1750 | (a uuid). |
| 1751 | """ |
| 1752 | monkeypatch.setenv("PYTHONHASHSEED", "0") |
| 1753 | monkeypatch.setattr( |
| 1754 | manager_module, |
| 1755 | "FileMapper", |
| 1756 | SimpleNamespace( |
| 1757 | from_offloading_spec=lambda **_: SimpleNamespace( |
| 1758 | get_run_config=lambda: {} |
| 1759 | ) |
| 1760 | ), |
| 1761 | ) |
| 1762 | calls: dict = {} |
| 1763 | monkeypatch.setattr( |
| 1764 | manager_module, |
| 1765 | "NixlTransport", |
| 1766 | lambda agent_name, *a, **k: ( |
| 1767 | calls.update(nixl_name=agent_name) or SimpleNamespace() |
| 1768 | ), |
| 1769 | ) |
| 1770 | monkeypatch.setattr( |
| 1771 | manager_module, |
| 1772 | "ZmqTransport", |
| 1773 | lambda local_id, host, port, *a, **k: ( |
| 1774 | calls.update(zmq_id=local_id, zmq_host=host, zmq_port=port) |
| 1775 | or SimpleNamespace() |
| 1776 | ), |
| 1777 | ) |
| 1778 | spec = SimpleNamespace( |
| 1779 | blocks_per_chunk=1, |
| 1780 | config=SimpleNamespace( |
| 1781 | parallel=SimpleNamespace(data_parallel_index=dp_index) |
| 1782 | ), |
| 1783 | ) |
| 1784 | mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs) |
| 1785 | mgr._test_calls = calls |
| 1786 | return mgr |
| 1787 | |
| 1788 | def test_defaults_from_env_unset(self, monkeypatch): |
| 1789 | monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) |
| 1790 | monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) |
| 1791 | mgr = self._construct(monkeypatch) |
| 1792 | # localhost default is used verbatim as the dial-back identity. |
| 1793 | assert mgr._local_id == "localhost:5710" |
| 1794 | |
| 1795 | def test_env_override(self, monkeypatch): |
| 1796 | monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3") |
| 1797 | monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799") |
| 1798 | mgr = self._construct(monkeypatch) |
| 1799 | assert mgr._local_id == "10.1.2.3:5799" |
| 1800 | |
| 1801 | def test_explicit_config_wins(self, monkeypatch): |
| 1802 | monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3") |
| 1803 | monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799") |
| 1804 | mgr = self._construct(monkeypatch, host="192.0.2.5", port=6001) |
| 1805 | assert mgr._local_id == "192.0.2.5:6001" |
| 1806 | |
| 1807 | def test_dp_index_offsets_default_port(self, monkeypatch): |
| 1808 | monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) |
| 1809 | monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) |
| 1810 | mgr = self._construct(monkeypatch, dp_index=2) |
| 1811 | assert mgr._local_id == "localhost:5712" |
| 1812 | |
| 1813 | def test_dp_index_offsets_explicit_port(self, monkeypatch): |
| 1814 | mgr = self._construct(monkeypatch, dp_index=1, host="192.0.2.5", port=6000) |
| 1815 | assert mgr._local_id == "192.0.2.5:6001" |
| 1816 | |
| 1817 | @pytest.mark.parametrize( |
| 1818 | "bind_host", ["localhost", "127.0.0.1", "::1", "0.0.0.0", "::", "192.0.2.5"] |
| 1819 | ) |
| 1820 | def test_host_used_verbatim(self, monkeypatch, bind_host): |
| 1821 | # The host is never rewritten — loopbacks and wildcards alike become |
| 1822 | # the dial-back identity verbatim (mirrors the NIXL connector). |
| 1823 | mgr = self._construct(monkeypatch, host=bind_host, port=5710) |
| 1824 | assert mgr._local_id == f"{bind_host}:5710" |
| 1825 | |
| 1826 | def test_nixl_name_decoupled_from_identity(self, monkeypatch): |
| 1827 | # The ZMQ identity is the verbatim host:port; the NIXL agent name is a |
| 1828 | # uuid, distinct from the identity and never used as an address. |
| 1829 | mgr = self._construct(monkeypatch, host="127.0.0.1", port=5710) |
| 1830 | assert mgr._test_calls["zmq_host"] == "127.0.0.1" |
| 1831 | assert mgr._test_calls["zmq_port"] == 5710 |
| 1832 | assert mgr._test_calls["zmq_id"] == "127.0.0.1:5710" |
| 1833 | assert mgr._local_id == "127.0.0.1:5710" |
| 1834 | # nixl_name is a valid uuid4 and not the host:port identity. |
| 1835 | nixl_name = mgr._test_calls["nixl_name"] |
| 1836 | assert nixl_name != mgr._local_id |
| 1837 | assert uuid.UUID(nixl_name).version == 4 |
| 1838 | |
| 1839 | def test_same_host_port_gets_distinct_nixl_names(self, monkeypatch): |
| 1840 | # The original collision: two peers sharing a host:port must still get |
| 1841 | # distinct NIXL agent names so add_remote_agent doesn't reject a remote |
| 1842 | # whose name equals the local. The per-process uuid guarantees this. |
| 1843 | mgr_a = self._construct(monkeypatch, host="localhost", port=5710) |
| 1844 | mgr_b = self._construct(monkeypatch, host="localhost", port=5710) |
| 1845 | assert mgr_a._local_id == mgr_b._local_id == "localhost:5710" |
| 1846 | assert mgr_a._nixl_agent_name != mgr_b._nixl_agent_name |