| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ProducerXfer, ProducerSession, ConsumerXfer, ConsumerSession.""" |
| 4 | |
| 5 | import time |
| 6 | from unittest.mock import MagicMock |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( |
| 11 | XferAck, |
| 12 | XferReq, |
| 13 | XferStatus, |
| 14 | ) |
| 15 | from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( |
| 16 | EmbeddingCache, |
| 17 | ) |
| 18 | from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( |
| 19 | _CONSUMER_QUARANTINE_TIMEOUT_S, |
| 20 | ConsumerSession, |
| 21 | ConsumerXfer, |
| 22 | ProducerSession, |
| 23 | ProducerXfer, |
| 24 | XferState, |
| 25 | ) |
| 26 | |
| 27 | # ── helpers ─────────────────────────────────────────────────────────────────── |
| 28 | |
| 29 | |
| 30 | def _make_data(xfer_state: str = "PROC") -> MagicMock: |
| 31 | data = MagicMock() |
| 32 | data.check_xfer_state.return_value = xfer_state |
| 33 | data.post_read.return_value = 99 |
| 34 | return data |
| 35 | |
| 36 | |
| 37 | def _ok_ack( |
| 38 | mm_hash: str = "h1", |
| 39 | session_id: str = "sess-1", |
| 40 | src_indices: list | None = None, |
| 41 | ) -> XferAck: |
| 42 | return XferAck( |
| 43 | mm_hash=mm_hash, |
| 44 | status=XferStatus.OK, |
| 45 | session_id=session_id, |
| 46 | src_block_indices=src_indices or [0, 1], |
| 47 | agent_metadata=b"meta", |
| 48 | mem_descriptor=b"desc", |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def _xfer( |
| 53 | data: MagicMock | None = None, |
| 54 | deadline: float | None = None, |
| 55 | consumer_session_id: str = "sess-1", |
| 56 | ) -> ConsumerXfer: |
| 57 | return ConsumerXfer( |
| 58 | mm_hash="h1", |
| 59 | block_indices=[10, 11], |
| 60 | addr=("host", 1234), |
| 61 | deadline=deadline if deadline is not None else time.monotonic() + 60, |
| 62 | data=data or _make_data(), |
| 63 | consumer_session_id=consumer_session_id, |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def _started( |
| 68 | data: MagicMock | None = None, deadline: float | None = None |
| 69 | ) -> ConsumerXfer: |
| 70 | """ConsumerXfer that has already received an OK XferAck.""" |
| 71 | x = _xfer(data, deadline) |
| 72 | x.handle_ack(_ok_ack(session_id="sess-1", src_indices=[5, 6]), agent_name="agent-1") |
| 73 | return x |
| 74 | |
| 75 | |
| 76 | # ── ProducerXfer ────────────────────────────────────────────────────────────── |
| 77 | |
| 78 | |
| 79 | def test_producer_xfer_not_expired_before_deadline(): |
| 80 | x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() + 60) |
| 81 | assert not x.is_expired() |
| 82 | |
| 83 | |
| 84 | def test_producer_xfer_expired_after_deadline(): |
| 85 | x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() - 1) |
| 86 | assert x.is_expired() |
| 87 | |
| 88 | |
| 89 | # ── ProducerSession ─────────────────────────────────────────────────────────── |
| 90 | |
| 91 | |
| 92 | def _xfer_req(mm_hash: str = "h1", session_id: str = "sess-1") -> XferReq: |
| 93 | return XferReq( |
| 94 | mm_hash=mm_hash, compatibility_hash="hash-abc", session_id=session_id |
| 95 | ) |
| 96 | |
| 97 | |
| 98 | def _make_producer_session(cache: EmbeddingCache | None = None) -> ProducerSession: |
| 99 | data = _make_data() |
| 100 | data.get_agent_metadata.return_value = b"meta" |
| 101 | data.get_mem_descriptor.return_value = b"desc" |
| 102 | return ProducerSession( |
| 103 | transport=MagicMock(), |
| 104 | data=data, |
| 105 | cache=cache or EmbeddingCache(num_blocks=8), |
| 106 | compat_hash="hash-abc", |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | def test_producer_session_grant_pins_ready_entry_and_returns_blocks(): |
| 111 | cache = EmbeddingCache(num_blocks=8) |
| 112 | cache.alloc("h1", 2) |
| 113 | cache.mark_ready("h1") |
| 114 | s = _make_producer_session(cache) |
| 115 | |
| 116 | ack = s._grant_or_nack(_xfer_req(mm_hash="h1")) |
| 117 | |
| 118 | assert ack.status == XferStatus.OK |
| 119 | entry = cache.get("h1") |
| 120 | assert not entry.evictable # pinned by the grant |
| 121 | assert list(ack.src_block_indices) == list(entry.block_ids) |
| 122 | |
| 123 | |
| 124 | def test_producer_session_grant_nacks_missing_entry(): |
| 125 | s = _make_producer_session() |
| 126 | ack = s._grant_or_nack(_xfer_req(mm_hash="nope")) |
| 127 | assert ack.status == XferStatus.NACK_MISSING |
| 128 | |
| 129 | |
| 130 | def test_producer_session_grant_nacks_not_ready_entry_as_retryable(): |
| 131 | """An announced entry whose save is still in flight is distinguished from a |
| 132 | miss, so the consumer retries instead of falling back to local compute.""" |
| 133 | cache = EmbeddingCache(num_blocks=8) |
| 134 | cache.alloc("h1", 2) # not ready |
| 135 | s = _make_producer_session(cache) |
| 136 | ack = s._grant_or_nack(_xfer_req(mm_hash="h1")) |
| 137 | assert ack.status == XferStatus.NACK_NOT_READY |
| 138 | |
| 139 | |
| 140 | def test_producer_session_grant_nacks_version_mismatch(): |
| 141 | cache = EmbeddingCache(num_blocks=8) |
| 142 | cache.alloc("h1", 2) |
| 143 | cache.mark_ready("h1") |
| 144 | s = _make_producer_session(cache) |
| 145 | req = XferReq( |
| 146 | mm_hash="h1", |
| 147 | compatibility_hash="hash-abc", |
| 148 | session_id="sess-1", |
| 149 | connector_version=-1, |
| 150 | ) |
| 151 | |
| 152 | ack = s._grant_or_nack(req) |
| 153 | |
| 154 | assert ack.status == XferStatus.NACK_VERSION |
| 155 | entry = cache.get("h1") |
| 156 | assert entry.evictable # nothing was pinned |
| 157 | |
| 158 | |
| 159 | def test_producer_session_grant_nacks_compat_hash_mismatch(): |
| 160 | cache = EmbeddingCache(num_blocks=8) |
| 161 | cache.alloc("h1", 2) |
| 162 | cache.mark_ready("h1") |
| 163 | s = _make_producer_session(cache) |
| 164 | req = XferReq(mm_hash="h1", compatibility_hash="wrong-hash", session_id="sess-1") |
| 165 | |
| 166 | ack = s._grant_or_nack(req) |
| 167 | |
| 168 | assert ack.status == XferStatus.NACK_INCOMPAT |
| 169 | entry = cache.get("h1") |
| 170 | assert entry.evictable # nothing was pinned |
| 171 | |
| 172 | |
| 173 | def test_producer_session_notif_unpins_via_cache(): |
| 174 | cache = EmbeddingCache(num_blocks=8) |
| 175 | cache.alloc("h1", 2) |
| 176 | cache.mark_ready("h1") |
| 177 | s = _make_producer_session(cache) |
| 178 | s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1")) |
| 179 | entry = cache.get("h1") |
| 180 | assert not entry.evictable # pinned |
| 181 | |
| 182 | s._data.get_new_notifs.return_value = {"agent": [b"sess-1:h1"]} |
| 183 | s._drain_notifs() |
| 184 | |
| 185 | assert entry.evictable # unpinned via cache |
| 186 | assert s._active_xfers == {} |
| 187 | |
| 188 | |
| 189 | def test_producer_session_sweep_timeouts_unpins_via_cache(): |
| 190 | cache = EmbeddingCache(num_blocks=8) |
| 191 | cache.alloc("h1", 2) |
| 192 | cache.mark_ready("h1") |
| 193 | s = _make_producer_session(cache) |
| 194 | s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1")) |
| 195 | key = next(iter(s._active_xfers)) |
| 196 | s._active_xfers[key].deadline = time.monotonic() - 1 |
| 197 | |
| 198 | s._sweep_timeouts() |
| 199 | |
| 200 | assert cache.get("h1").evictable |
| 201 | assert s._active_xfers == {} |
| 202 | |
| 203 | |
| 204 | # ── ConsumerXfer.handle_ack ─────────────────────────────────────────────────── |
| 205 | |
| 206 | |
| 207 | def test_consumer_xfer_handle_ack_ok_starts_read(): |
| 208 | data = _make_data() |
| 209 | x = _xfer(data) |
| 210 | assert x.handle_ack(_ok_ack(), agent_name="agent-1") is True |
| 211 | assert x.transfer_handle == 99 |
| 212 | |
| 213 | |
| 214 | def test_consumer_xfer_handle_ack_notif_msg_is_session_id_colon_mm_hash(): |
| 215 | data = _make_data() |
| 216 | x = _xfer(data, consumer_session_id="my-sess") |
| 217 | x.handle_ack(_ok_ack(mm_hash="h1"), agent_name="agent-1") |
| 218 | _, kwargs = data.post_read.call_args |
| 219 | assert kwargs["notif_msg"] == b"my-sess:h1" |
| 220 | |
| 221 | |
| 222 | def test_consumer_xfer_handle_ack_forwards_correct_args(): |
| 223 | data = _make_data() |
| 224 | x = _xfer(data) |
| 225 | x.handle_ack(_ok_ack(src_indices=[5, 6]), agent_name="agent-1") |
| 226 | args, _ = data.post_read.call_args |
| 227 | assert args[0] == [10, 11] # local block_indices |
| 228 | assert args[1] == "agent-1" # agent_name (not a raw handle) |
| 229 | assert args[2] == [5, 6] # remote src_block_indices |
| 230 | |
| 231 | |
| 232 | def test_consumer_xfer_handle_ack_resets_deadline(): |
| 233 | original = time.monotonic() + 2 |
| 234 | x = _xfer(deadline=original) |
| 235 | x.handle_ack(_ok_ack(), agent_name="agent-1") |
| 236 | assert x.deadline > original |
| 237 | |
| 238 | |
| 239 | def test_consumer_xfer_handle_ack_nack_returns_false(): |
| 240 | data = _make_data() |
| 241 | x = _xfer(data) |
| 242 | assert ( |
| 243 | x.handle_ack(XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING), "a") |
| 244 | is False |
| 245 | ) |
| 246 | assert x.transfer_handle is None |
| 247 | data.post_read.assert_not_called() |
| 248 | |
| 249 | |
| 250 | # ── ConsumerXfer.poll — WAITING_ACK ────────────────────────────────────────── |
| 251 | |
| 252 | |
| 253 | def test_consumer_xfer_poll_waiting_ack_before_deadline(): |
| 254 | assert _xfer().poll(time.monotonic()) == XferState.WAITING_ACK |
| 255 | |
| 256 | |
| 257 | def test_consumer_xfer_poll_waiting_ack_timeout(): |
| 258 | assert ( |
| 259 | _xfer(deadline=time.monotonic() - 1).poll(time.monotonic()) |
| 260 | == XferState.ACK_TIMEOUT |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | # ── ConsumerXfer.poll — READING ─────────────────────────────────────────────── |
| 265 | |
| 266 | |
| 267 | def test_consumer_xfer_poll_reading_proc(): |
| 268 | assert _started().poll(time.monotonic()) == XferState.READING |
| 269 | |
| 270 | |
| 271 | def test_consumer_xfer_poll_reading_done_releases_handle(): |
| 272 | data = _make_data(xfer_state="DONE") |
| 273 | x = _started(data) |
| 274 | assert x.poll(time.monotonic()) == XferState.DONE |
| 275 | data.release_xfer_handle.assert_called_once_with(99) |
| 276 | assert x.transfer_handle is None |
| 277 | |
| 278 | |
| 279 | def test_consumer_xfer_poll_reading_unexpected_state_fails(): |
| 280 | data = _make_data(xfer_state="ERR") |
| 281 | x = _started(data) |
| 282 | assert x.poll(time.monotonic()) == XferState.READ_FAILED |
| 283 | data.release_xfer_handle.assert_called_once_with(99) |
| 284 | |
| 285 | |
| 286 | def test_consumer_xfer_poll_reading_exception_fails_gracefully(): |
| 287 | data = _make_data() |
| 288 | data.check_xfer_state.side_effect = RuntimeError("nixl crash") |
| 289 | x = _started(data) |
| 290 | assert x.poll(time.monotonic()) == XferState.READ_FAILED |
| 291 | |
| 292 | |
| 293 | def test_consumer_xfer_poll_timeout_quarantines_without_releasing(): |
| 294 | data = _make_data(xfer_state="PROC") |
| 295 | x = _started(data) |
| 296 | assert x.poll(time.monotonic() + 9999) == XferState.QUARANTINED |
| 297 | data.release_xfer_handle.assert_not_called() |
| 298 | assert x.transfer_handle == 99 |
| 299 | |
| 300 | |
| 301 | # ── ConsumerXfer.poll — QUARANTINED ─────────────────────────────────────────── |
| 302 | |
| 303 | |
| 304 | def test_consumer_xfer_quarantined_proc_stays(): |
| 305 | data = _make_data(xfer_state="PROC") |
| 306 | x = _started(data) |
| 307 | x.poll(time.monotonic() + 9999) # → QUARANTINED |
| 308 | assert x.poll(time.monotonic()) == XferState.QUARANTINED |
| 309 | data.release_xfer_handle.assert_not_called() |
| 310 | |
| 311 | |
| 312 | def test_consumer_xfer_quarantined_terminal_settles(): |
| 313 | data = _make_data(xfer_state="PROC") |
| 314 | x = _started(data) |
| 315 | x.poll(time.monotonic() + 9999) # → QUARANTINED |
| 316 | data.check_xfer_state.return_value = "DONE" |
| 317 | assert x.poll(time.monotonic()) == XferState.SETTLED |
| 318 | data.release_xfer_handle.assert_called_once_with(99) |
| 319 | assert x.transfer_handle is None |
| 320 | |
| 321 | |
| 322 | def test_consumer_xfer_quarantined_exception_settles(): |
| 323 | data = _make_data(xfer_state="PROC") |
| 324 | x = _started(data) |
| 325 | x.poll(time.monotonic() + 9999) |
| 326 | data.check_xfer_state.side_effect = RuntimeError("gone") |
| 327 | assert x.poll(time.monotonic()) == XferState.SETTLED |
| 328 | |
| 329 | |
| 330 | def test_consumer_xfer_quarantined_timeout_evicts_blocks(): |
| 331 | """Quarantine timeout should release handle and return READ_FAILED.""" |
| 332 | data = _make_data(xfer_state="PROC") |
| 333 | x = _started(data) |
| 334 | # First poll transitions to QUARANTINED (deadline extended by 60s) |
| 335 | x.poll(time.monotonic() + 9999) |
| 336 | # Second poll with expired quarantine deadline should evict |
| 337 | expired_time = time.monotonic() + 9999 + _CONSUMER_QUARANTINE_TIMEOUT_S + 1 |
| 338 | assert x.poll(expired_time) == XferState.READ_FAILED |
| 339 | data.release_xfer_handle.assert_called_once_with(99) |
| 340 | assert x.transfer_handle is None |
| 341 | |
| 342 | |
| 343 | # ── ConsumerXfer.cancel / release ──────────────────────────────────────────── |
| 344 | |
| 345 | |
| 346 | def test_consumer_xfer_cancel_valid_in_waiting_ack(): |
| 347 | _xfer().cancel() # must not raise |
| 348 | |
| 349 | |
| 350 | def test_consumer_xfer_cancel_raises_when_read_started(): |
| 351 | with pytest.raises(AssertionError): |
| 352 | _started().cancel() |
| 353 | |
| 354 | |
| 355 | def test_consumer_xfer_release_releases_handle(): |
| 356 | data = _make_data() |
| 357 | x = _started(data) |
| 358 | x.release() |
| 359 | data.release_xfer_handle.assert_called_once_with(99) |
| 360 | |
| 361 | |
| 362 | def test_consumer_xfer_release_noop_when_no_handle(): |
| 363 | data = _make_data() |
| 364 | _xfer(data).release() |
| 365 | data.release_xfer_handle.assert_not_called() |
| 366 | |
| 367 | |
| 368 | # ── ConsumerSession ─────────────────────────────────────────────────────────── |
| 369 | |
| 370 | |
| 371 | def _make_consumer_session(data=None): |
| 372 | zmq_conn = MagicMock() |
| 373 | zmq_conn.recv.return_value = [] |
| 374 | return ConsumerSession( |
| 375 | addr=("host", 1234), |
| 376 | zmq_conn=zmq_conn, |
| 377 | transport=MagicMock(), |
| 378 | data=data or _make_data(), |
| 379 | compat_hash="hash-abc", |
| 380 | ) |
| 381 | |
| 382 | |
| 383 | def test_consumer_session_has_stable_session_id(): |
| 384 | s = _make_consumer_session() |
| 385 | assert isinstance(s._session_id, str) and len(s._session_id) > 0 |
| 386 | assert s._session_id == s._session_id # stable |
| 387 | |
| 388 | |
| 389 | def test_consumer_session_start_xfer_sends_req_with_session_id(): |
| 390 | import msgspec |
| 391 | |
| 392 | s = _make_consumer_session() |
| 393 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) |
| 394 | raw = s._zmq.send.call_args[0][0] |
| 395 | req = msgspec.msgpack.decode(raw, type=XferReq) |
| 396 | assert req.mm_hash == "h1" |
| 397 | assert req.session_id == s._session_id |
| 398 | |
| 399 | |
| 400 | def test_consumer_session_poll_done_goes_to_completed(): |
| 401 | import msgspec |
| 402 | |
| 403 | data = _make_data(xfer_state="PROC") |
| 404 | s = _make_consumer_session(data) |
| 405 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) |
| 406 | |
| 407 | ack = XferAck( |
| 408 | mm_hash="h1", |
| 409 | status=XferStatus.OK, |
| 410 | session_id=s._session_id, |
| 411 | src_block_indices=[5, 6], |
| 412 | agent_metadata=b"meta", |
| 413 | mem_descriptor=b"desc", |
| 414 | ) |
| 415 | data.add_remote_peer.return_value = "agent-1" |
| 416 | raw_ack = msgspec.msgpack.encode(ack) |
| 417 | |
| 418 | data.check_xfer_state.return_value = "DONE" |
| 419 | s.poll([raw_ack], time.monotonic()) |
| 420 | |
| 421 | results = s.take_results() |
| 422 | assert "h1" in results.completed |
| 423 | assert not results.tombstoned |
| 424 | |
| 425 | |
| 426 | def test_consumer_session_poll_nack_goes_to_tombstoned(): |
| 427 | import msgspec |
| 428 | |
| 429 | s = _make_consumer_session() |
| 430 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) |
| 431 | nack = XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING) |
| 432 | s.poll([msgspec.msgpack.encode(nack)], time.monotonic()) |
| 433 | results = s.take_results() |
| 434 | assert "h1" in results.tombstoned |
| 435 | assert not results.completed |
| 436 | |
| 437 | |
| 438 | def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): |
| 439 | """A not-ready NACK must not tombstone: tombstoning admits the request with |
| 440 | no embedding, which is fatal once the media has been rewritten away.""" |
| 441 | import msgspec |
| 442 | |
| 443 | s = _make_consumer_session() |
| 444 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) |
| 445 | nack = XferAck(mm_hash="h1", status=XferStatus.NACK_NOT_READY) |
| 446 | s.poll([msgspec.msgpack.encode(nack)], time.monotonic()) |
| 447 | results = s.take_results() |
| 448 | assert "h1" in results.retryable |
| 449 | assert not results.tombstoned |
| 450 | assert not results.completed |
| 451 | |
| 452 | |
| 453 | def test_consumer_session_ack_timeout_reports_retry_once(): |
| 454 | """A delayed producer ACK is retryable, and its result is drained once.""" |
| 455 | s = _make_consumer_session() |
| 456 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) |
| 457 | s.poll([], time.monotonic()) |
| 458 | r1 = s.take_results() |
| 459 | assert "h1" in r1.retryable |
| 460 | assert not r1.tombstoned |
| 461 | r2 = s.take_results() |
| 462 | assert not r2.retryable # cleared after first take |
| 463 | |
| 464 | |
| 465 | def test_consumer_session_on_peer_down_cancels_waiting_ack(): |
| 466 | s = _make_consumer_session() |
| 467 | s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 60) |
| 468 | s.on_peer_down() |
| 469 | results = s.take_results() |
| 470 | assert "h1" in results.retryable |
| 471 | assert not results.tombstoned |