[EC Connector] P2P NIXL + CPU EC Connector (#47941)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
omerpaz95 committed
Sep 4, 2026 at 06:23 UTC
8a0a7ee40a5915dd7330c8f6e6c58ca565566471
28 files changed
+3521
-72
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+16
-6
@@ -297,16 +297,26 @@ async def fanout_encoder_primer(
297
detail=f"Encoder request failed: {detail}",
298
)
299
300
- # The encoder reports each item's cache key and grid here.
300
+ # The encoder reports each mm_hash's metadata (e.g. the grid) here,
301
+ # keyed by the same uuid this proxy assigned above.
302
try:
303
params = (await r.json()).get("ec_transfer_params") or {}
303
- reported = params.get("ec_items") or []
304
except Exception:
305
logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
306
- reported = []
307
- if reported and idx in item_uuids:
308
- # One item per encoder request, so the first entry is this item's.
309
- item_meta[idx] = {**reported[0], "mm_hash": item_uuids[idx]}
306
+ params = {}
307
+ if idx in item_uuids:
308
+ reported = params.get(item_uuids[idx])
309
+ if reported:
310
+ metadata = reported.get("metadata") or {}
311
+ if metadata:
312
+ item_meta[idx] = {**metadata, "mm_hash": item_uuids[idx]}
313
+ # Whatever the encoder reported alongside `metadata` is the
314
+ # connector's own handle on the published embedding (for NIXL,
315
+ # peer_host/peer_port/size_bytes). The decoder's connector
316
+ # looks it up by mm_hash on the request, so carry it through.
317
+ orig_request.setdefault("ec_transfer_params", {})[item_uuids[idx]] = (
318
+ reported
319
+ )
320
321
logger.info(
322
"[%s] All %d encoder requests completed successfully", req_id, len(mm_items)
tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py
+54
@@ -333,3 +333,57 @@ def test_concurrent_alloc_evict():
333
t.join()
334
335
assert not errors, errors
336
+
337
+
338
+# ── discard ──────────────────────────────────────────────────────────────────
339
+
340
+
341
+def test_discard_not_ready_returns_blocks():
342
+ cache = EmbeddingCache(num_blocks=4)
343
+ entry = cache.alloc("h", 3)
344
+ assert entry is not None
345
+ cache.discard("h")
346
+ assert cache.get("h") is None
347
+ # All blocks reclaimed: a fresh 4-block alloc now fits.
348
+ assert cache.alloc("h2", 4) is not None
349
+
350
+
351
+def test_discard_rejects_ready_entry():
352
+ cache = EmbeddingCache(num_blocks=4)
353
+ cache.alloc("h", 1)
354
+ cache.mark_ready("h")
355
+ with pytest.raises(AssertionError):
356
+ cache.discard("h")
357
+
358
+
359
+def test_discard_rejects_absent_key():
360
+ cache = EmbeddingCache(num_blocks=4)
361
+ with pytest.raises(KeyError):
362
+ cache.discard("nope")
363
+
364
+
365
+# ── pin_if_ready ─────────────────────────────────────────────────────────────
366
+
367
+
368
+def test_pin_if_ready_pins_and_returns_entry():
369
+ cache = EmbeddingCache(num_blocks=4)
370
+ cache.alloc("h", 2)
371
+ cache.mark_ready("h")
372
+ entry = cache.pin_if_ready("h")
373
+ assert entry is not None and entry.ready and len(entry.block_ids) == 2
374
+ assert cache.alloc("h2", 4) is None # pinned ⇒ not evictable
375
+
376
+
377
+def test_pin_if_ready_none_when_absent():
378
+ cache = EmbeddingCache(num_blocks=4)
379
+ assert cache.pin_if_ready("nope") is None
380
+
381
+
382
+def test_pin_if_ready_returns_unpinned_entry_when_not_ready():
383
+ """The producer needs "present but not ready" told apart from "absent":
384
+ the first is NACKed as retryable, the second as a miss."""
385
+ cache = EmbeddingCache(num_blocks=4)
386
+ cache.alloc("h", 1) # not-ready
387
+ entry = cache.pin_if_ready("h")
388
+ assert entry is not None and not entry.ready
389
+ assert cache.alloc("h2", 4) is None # not-ready ⇒ still not evictable
tests/v1/ec_connector/unit/cpu/test_connector.py
+27
-3
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from unittest.mock import MagicMock
4
5
+import pytest
6
+
7
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
8
from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector
9
from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
@@ -34,14 +36,36 @@ def test_worker_role_builds_only_worker(monkeypatch):
36
assert c.connector_scheduler is None
37
38
37
-def test_request_finished_inherited_noop(monkeypatch):
39
+@pytest.mark.parametrize("role", [ECConnectorRole.SCHEDULER, ECConnectorRole.WORKER])
40
+def test_v1_model_runner_is_rejected(monkeypatch, role):
41
+ """The V1 runner never reports save completions, so refuse to start."""
42
monkeypatch.setattr(
43
ECCPUConnector, "_make_scheduler", lambda self, cfg: MagicMock()
44
)
41
- c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
42
- assert c.request_finished(MagicMock()) == (False, None)
45
+ monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: MagicMock())
46
+ cfg = _cfg()
47
+ cfg.use_v2_model_runner = False
48
+
49
+ with pytest.raises(ValueError, match="V2 model runner"):
50
+ ECCPUConnector(cfg, role)
51
52
53
def test_factory_registered():
54
cls = ECConnectorFactory._registry["ECCPUConnector"]()
55
assert cls is ECCPUConnector
56
+
57
+
58
+def test_request_finished_forwards_to_scheduler(monkeypatch):
59
+ from unittest.mock import MagicMock
60
+
61
+ import vllm.distributed.ec_transfer.ec_connector.cpu.connector as conn_mod # noqa: F401
62
+ from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
63
+ from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector
64
+
65
+ fake_sched = MagicMock()
66
+ fake_sched.request_finished.return_value = (False, {"h1": {"peer_port": 1}})
67
+ monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched)
68
+ c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
69
+ req = MagicMock()
70
+ assert c.request_finished(req) == (False, {"h1": {"peer_port": 1}})
71
+ fake_sched.request_finished.assert_called_once_with(req)
tests/v1/ec_connector/unit/cpu/worker/test_worker.py
+19
@@ -64,6 +64,14 @@ _requires_cuda_alike = pytest.mark.skipif(
64
reason="stalling the compute stream requires a CUDA-like platform",
65
)
66
67
+_requires_swap_blocks_batch = pytest.mark.skipif(
68
+ not hasattr(torch.ops._C_cache_ops, "swap_blocks_batch"),
69
+ reason=(
70
+ "installed vllm C++ extension predates the swap_blocks_batch op "
71
+ "flush_saves/start_load_caches use; rebuild the extension to run this"
72
+ ),
73
+)
74
+
75
# Cycles to stall the compute stream for. Must comfortably outlast the host-side
76
# work the lifetime tests do while a copy is deliberately held back — ~0.5 s on
77
# a current datacenter GPU.
@@ -194,6 +202,7 @@ def _warm_up_stream_pools(worker: ECCPUWorker) -> None:
202
# ── save_caches ──────────────────────────────────────────────────────────────
203
204
205
+@_requires_swap_blocks_batch
206
@_requires_accelerator
207
@pytest.mark.parametrize(
208
"n_elements,block_ids",
@@ -326,6 +335,7 @@ def test_save_caches_raises_when_allocated_blocks_too_small(make_worker):
335
worker.save_caches({"h": src}, "h", _meta(saves={"h": [0, 1]}))
336
337
338
+@_requires_swap_blocks_batch
339
@_requires_accelerator
340
def test_save_caches_batches_multiple_hashes(make_worker):
341
"""Multiple save_caches calls are batched into a single flush."""
@@ -401,6 +411,7 @@ def test_save_descriptors_hold_high_bit_source_addresses(make_worker):
411
# ── start_load_caches ────────────────────────────────────────────────────────
412
413
414
+@_requires_swap_blocks_batch
415
@_requires_accelerator
416
def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker):
417
"""Single batched load across all hashes with correct byte→dtype→shape."""
@@ -465,6 +476,7 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker):
476
assert encoder_cache == {}
477
478
479
+@_requires_swap_blocks_batch
480
@_requires_accelerator
481
def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):
482
"""Every hash in ``meta.loads`` is copied from mmap, including one whose
@@ -497,6 +509,7 @@ def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):
509
assert torch.equal(new.cpu(), src_orig[1:])
510
511
512
+@_requires_swap_blocks_batch
513
@_requires_accelerator
514
@pytest.mark.parametrize(
515
"tp_rank,pcp_rank",
@@ -526,6 +539,7 @@ def test_start_load_caches_works_on_all_ranks(make_worker, tp_rank, pcp_rank):
539
# ── round-trip ───────────────────────────────────────────────────────────────
540
541
542
+@_requires_swap_blocks_batch
543
@_requires_accelerator
544
def test_save_then_load_round_trips_bytes(make_worker):
545
"""Full producer→mmap→consumer byte path in one shot."""
@@ -552,6 +566,7 @@ def test_save_then_load_round_trips_bytes(make_worker):
566
# ── memory lifetime across in-flight copies ─────────────────────────────────
567
568
569
+@_requires_swap_blocks_batch
570
@_requires_cuda_alike
571
def test_save_survives_encoder_cache_free_before_copy_runs(make_worker):
572
"""The bytes handed to ``save_caches`` must reach the mmap intact even when
@@ -612,6 +627,7 @@ def test_save_survives_encoder_cache_free_before_copy_runs(make_worker):
627
)
628
629
630
+@_requires_swap_blocks_batch
631
@_requires_cuda_alike
632
def test_load_buffer_survives_eviction_while_consumer_read_is_queued(make_worker):
633
"""A loaded encoder cache entry must keep its bytes for a consumer already
@@ -662,6 +678,7 @@ def test_load_buffer_survives_eviction_while_consumer_read_is_queued(make_worker
678
# ── buffer recycling ────────────────────────────────────────────────────────
679
680
681
+@_requires_swap_blocks_batch
682
@_requires_accelerator
683
def test_buffer_pool_is_reused_across_save_steps(make_worker):
684
"""Once a save copy completes its descriptor buffers return to the pool and
@@ -689,6 +706,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker):
706
assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id
707
708
709
+@_requires_swap_blocks_batch
710
@_requires_accelerator
711
def test_buffer_pool_is_reused_across_load_steps(make_worker):
712
"""Once a load copy completes its descriptor buffers return to the pool and
@@ -793,6 +811,7 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):
811
# ── e2e: scheduler + worker pipeline ────────────────────────────────────────
812
813
814
+@_requires_swap_blocks_batch
815
@_requires_accelerator
816
def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):
817
"""Full pipeline: scheduler allocates blocks, worker saves a GPU tensor to
tests/v1/ec_connector/unit/test_control.py
new
+285
@@ -0,0 +1,285 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for control-layer ABCs, ZmqClientConnection, and ECSession."""
4
+
5
+from unittest.mock import MagicMock, patch
6
+
7
+import pytest
8
+import zmq
9
+
10
+from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import (
11
+ ControlConnection,
12
+)
13
+
14
+
15
+class _ConcreteConn(ControlConnection):
16
+ """Minimal concrete subclass for ABC-instantiation tests."""
17
+
18
+ def __init__(self):
19
+ self._alive = True
20
+ self._sent = []
21
+ self._inbox = []
22
+
23
+ @property
24
+ def alive(self) -> bool:
25
+ return self._alive
26
+
27
+ def send(self, msg: bytes) -> None:
28
+ self._sent.append(msg)
29
+
30
+ def recv(self) -> list[bytes]:
31
+ msgs, self._inbox = self._inbox, []
32
+ return msgs
33
+
34
+ def mark_dead(self) -> None:
35
+ self._alive = False
36
+
37
+ def close(self) -> None:
38
+ self._alive = False
39
+
40
+
41
+def test_control_connection_abc_cannot_be_instantiated_directly():
42
+ with pytest.raises(TypeError):
43
+ ControlConnection() # type: ignore[abstract]
44
+
45
+
46
+def test_concrete_connection_send_recv():
47
+ conn = _ConcreteConn()
48
+ assert conn.alive
49
+ conn.send(b"hello")
50
+ conn._inbox.append(b"world")
51
+ assert conn.recv() == [b"world"]
52
+ assert conn._sent == [b"hello"]
53
+
54
+
55
+def test_mark_dead_sets_alive_false():
56
+ conn = _ConcreteConn()
57
+ conn.mark_dead()
58
+ assert not conn.alive
59
+
60
+
61
+def test_close_sets_alive_false():
62
+ conn = _ConcreteConn()
63
+ conn.close()
64
+ assert not conn.alive
65
+
66
+
67
+# ── ZmqClientConnection ───────────────────────────────────────────────────────
68
+
69
+
70
+from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( # noqa: E402
71
+ ZmqClientConnection,
72
+ ZmqClientTransport,
73
+)
74
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr # noqa: E402
75
+
76
+
77
+def _conn(messages: list[list[bytes]], monitor=None) -> ZmqClientConnection:
78
+ """Build a ZmqClientConnection whose dealer yields the given frame-lists."""
79
+ dealer = MagicMock()
80
+ dealer.recv_multipart.side_effect = [*messages, zmq.Again()]
81
+ return ZmqClientConnection(dealer=dealer, monitor=monitor)
82
+
83
+
84
+def test_zmq_client_connection_recv_extracts_payloads_from_envelopes():
85
+ conn = _conn([[b"", b"msg1"], [b"", b"msg2"]])
86
+ assert conn.recv() == [b"msg1", b"msg2"]
87
+
88
+
89
+def test_zmq_client_connection_recv_skips_frames_without_empty_delimiter():
90
+ # A frame that does not start with b"" is malformed and must be dropped.
91
+ conn = _conn([[b"", b"ok"], [b"no-delimiter", b"bad"]])
92
+ assert conn.recv() == [b"ok"]
93
+
94
+
95
+def test_zmq_client_connection_recv_breaks_on_non_again_exception():
96
+ # A socket error should stop the drain and return whatever was collected.
97
+ dealer = MagicMock()
98
+ dealer.recv_multipart.side_effect = [[b"", b"first"], OSError("broken")]
99
+ conn = ZmqClientConnection(dealer=dealer, monitor=None)
100
+ assert conn.recv() == [b"first"]
101
+
102
+
103
+def test_zmq_client_connection_send_prefixes_empty_delimiter():
104
+ dealer = MagicMock()
105
+ conn = ZmqClientConnection(dealer=dealer, monitor=None)
106
+ conn.send(b"payload")
107
+ dealer.send_multipart.assert_called_once_with([b"", b"payload"])
108
+
109
+
110
+def test_zmq_client_connection_mark_dead_does_not_close_socket():
111
+ # mark_dead signals liveness only; resource teardown is close()'s job.
112
+ dealer = MagicMock()
113
+ conn = ZmqClientConnection(dealer=dealer, monitor=None)
114
+ conn.mark_dead()
115
+ assert not conn.alive
116
+ dealer.close.assert_not_called()
117
+
118
+
119
+def test_zmq_client_connection_close_without_monitor():
120
+ dealer = MagicMock()
121
+ conn = ZmqClientConnection(dealer=dealer, monitor=None)
122
+ conn.close()
123
+ dealer.close.assert_called_once_with(linger=0)
124
+ assert not conn.alive
125
+
126
+
127
+def test_zmq_client_connection_close_with_monitor_disables_then_closes_both():
128
+ # The disable_monitor → monitor.close → dealer.close sequence must be
129
+ # respected: closing the dealer before disabling the monitor can cause
130
+ # a zmq assertion in some versions.
131
+ dealer = MagicMock()
132
+ monitor = MagicMock()
133
+ conn = ZmqClientConnection(dealer=dealer, monitor=monitor)
134
+ conn.close()
135
+ dealer.disable_monitor.assert_called_once()
136
+ monitor.close.assert_called_once_with(linger=0)
137
+ dealer.close.assert_called_once_with(linger=0)
138
+ assert not conn.alive
139
+
140
+
141
+def test_zmq_client_connection_close_is_idempotent():
142
+ dealer = MagicMock()
143
+ conn = ZmqClientConnection(dealer=dealer, monitor=None)
144
+ conn.close()
145
+ conn.close() # must not raise
146
+ assert not conn.alive
147
+
148
+
149
+# ── ZmqClientTransport ────────────────────────────────────────────────────────
150
+
151
+
152
+def _make_transport() -> tuple[ZmqClientTransport, MagicMock]:
153
+ """Return a ZmqClientTransport with a mocked ZMQ context."""
154
+ ctx = MagicMock()
155
+ ctx.socket.return_value = MagicMock()
156
+ with (
157
+ patch(
158
+ "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context",
159
+ return_value=ctx,
160
+ ),
161
+ patch(
162
+ "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket",
163
+ return_value=MagicMock(),
164
+ ),
165
+ patch(
166
+ "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path",
167
+ return_value="tcp://host:0",
168
+ ),
169
+ ):
170
+ transport = ZmqClientTransport()
171
+ transport._ctx = ctx
172
+ return transport, ctx
173
+
174
+
175
+def _inject_conn(transport: ZmqClientTransport, addr: PeerAddr) -> ZmqClientConnection:
176
+ """Inject a ZmqClientConnection directly into the transport pool."""
177
+ conn = ZmqClientConnection(dealer=MagicMock(), monitor=MagicMock())
178
+ transport._connections[addr] = conn
179
+ return conn
180
+
181
+
182
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
183
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
184
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
185
+def test_zmq_client_transport_connect_caches_connection(
186
+ mock_ctx_cls, mock_path, mock_sock
187
+):
188
+ mock_ctx_cls.return_value = MagicMock()
189
+ mock_ctx_cls.return_value.socket.return_value = MagicMock()
190
+ mock_sock.return_value = MagicMock()
191
+ transport = ZmqClientTransport()
192
+ addr: PeerAddr = ("host", 1234)
193
+ assert transport.connect(addr) is transport.connect(addr)
194
+
195
+
196
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
197
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
198
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
199
+def test_zmq_client_transport_connect_creates_independent_conns_per_addr(
200
+ mock_ctx_cls, mock_path, mock_sock
201
+):
202
+ mock_ctx_cls.return_value = MagicMock()
203
+ mock_ctx_cls.return_value.socket.return_value = MagicMock()
204
+ mock_sock.return_value = MagicMock()
205
+ transport = ZmqClientTransport()
206
+ conn_a = transport.connect(("host-a", 1111))
207
+ conn_b = transport.connect(("host-b", 2222))
208
+ assert conn_a is not conn_b
209
+
210
+
211
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
212
+def test_zmq_client_transport_poll_dead_detects_disconnected_peer(mock_recv_mon):
213
+ transport, _ = _make_transport()
214
+ addr: PeerAddr = ("host", 1234)
215
+ conn = _inject_conn(transport, addr)
216
+ mock_recv_mon.side_effect = [{"event": zmq.EVENT_DISCONNECTED}, zmq.Again()]
217
+
218
+ dead = transport.poll_dead()
219
+
220
+ assert addr in dead
221
+ assert not conn.alive
222
+
223
+
224
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
225
+def test_zmq_client_transport_poll_dead_skips_conn_without_monitor(mock_recv_mon):
226
+ transport, _ = _make_transport()
227
+ addr: PeerAddr = ("host", 9999)
228
+ conn = ZmqClientConnection(dealer=MagicMock(), monitor=None)
229
+ transport._connections[addr] = conn
230
+
231
+ dead = transport.poll_dead()
232
+
233
+ assert dead == []
234
+ mock_recv_mon.assert_not_called()
235
+
236
+
237
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
238
+def test_zmq_client_transport_poll_dead_ignores_non_disconnect_events(mock_recv_mon):
239
+ transport, _ = _make_transport()
240
+ addr: PeerAddr = ("host", 1234)
241
+ conn = _inject_conn(transport, addr)
242
+ # A CONNECTED event should not trigger a dead notification.
243
+ mock_recv_mon.side_effect = [{"event": zmq.EVENT_CONNECTED}, zmq.Again()]
244
+
245
+ dead = transport.poll_dead()
246
+
247
+ assert dead == []
248
+ assert conn.alive
249
+
250
+
251
+def test_zmq_client_transport_remove_closes_and_evicts_conn():
252
+ transport, _ = _make_transport()
253
+ addr: PeerAddr = ("host", 5678)
254
+ conn = _inject_conn(transport, addr)
255
+
256
+ removed = transport.remove(addr)
257
+
258
+ assert removed is conn
259
+ assert addr not in transport._connections
260
+ conn.dealer.close.assert_called()
261
+
262
+
263
+def test_zmq_client_transport_remove_returns_none_for_unknown_addr():
264
+ transport, _ = _make_transport()
265
+ assert transport.remove(("ghost", 0)) is None
266
+
267
+
268
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
269
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
270
+@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
271
+def test_zmq_client_transport_close_closes_all_connections(
272
+ mock_ctx_cls, mock_path, mock_sock
273
+):
274
+ mock_ctx = MagicMock()
275
+ mock_ctx_cls.return_value = mock_ctx
276
+ mock_ctx.socket.return_value = MagicMock()
277
+ dealer_a, dealer_b = MagicMock(), MagicMock()
278
+ mock_sock.side_effect = [dealer_a, dealer_b]
279
+ transport = ZmqClientTransport()
280
+ transport.connect(("a", 1))
281
+ transport.connect(("b", 2))
282
+ transport.close()
283
+ dealer_a.close.assert_called()
284
+ dealer_b.close.assert_called()
285
+ assert transport._connections == {}
tests/v1/ec_connector/unit/test_data.py
new
+184
@@ -0,0 +1,184 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for DataTransport ABC and NixlDataTransport."""
4
+
5
+from unittest.mock import MagicMock, patch
6
+
7
+import pytest
8
+
9
+from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
10
+ DataTransport,
11
+)
12
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
13
+ serialize_mem_descriptor,
14
+)
15
+
16
+_NIXL_PATH = "vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl"
17
+
18
+
19
+# ── helpers ───────────────────────────────────────────────────────────────────
20
+
21
+
22
+def _make_nixl_mock(prep_side_effect=None) -> MagicMock:
23
+ nixl = MagicMock()
24
+ nixl.get_agent_metadata.return_value = b"agent-meta"
25
+ nixl.get_reg_descs.return_value = MagicMock()
26
+ nixl.get_xfer_descs.return_value = MagicMock()
27
+ nixl.prep_xfer_dlist.side_effect = prep_side_effect or [77, 88]
28
+ nixl.add_remote_agent.return_value = "remote-agent-1"
29
+ nixl.get_new_notifs.return_value = {}
30
+ nixl.check_xfer_state.return_value = "PROC"
31
+ return nixl
32
+
33
+
34
+def _make_transport(nixl_mock=None):
35
+ from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
36
+ NixlDataTransport,
37
+ )
38
+
39
+ mock = nixl_mock or _make_nixl_mock()
40
+ with (
41
+ patch(f"{_NIXL_PATH}.NixlWrapper", return_value=mock),
42
+ patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
43
+ ):
44
+ t = NixlDataTransport("agent", 0, 4, 64, 256)
45
+ t._nixl = mock
46
+ return t
47
+
48
+
49
+# ── ABC ───────────────────────────────────────────────────────────────────────
50
+
51
+
52
+def test_data_transport_abc_cannot_be_instantiated():
53
+ with pytest.raises(TypeError):
54
+ DataTransport() # type: ignore[abstract]
55
+
56
+
57
+# ── NixlDataTransport construction ───────────────────────────────────────────
58
+
59
+
60
+def test_nixl_data_transport_registers_memory_on_init():
61
+ nixl = _make_nixl_mock()
62
+ _make_transport(nixl)
63
+ nixl.register_memory.assert_called_once()
64
+
65
+
66
+def test_nixl_data_transport_unwinds_registration_if_init_fails():
67
+ """A failure after register_memory must not leave the region pinned.
68
+
69
+ __init__ propagates, so the caller never gets a transport to deregister.
70
+ """
71
+ nixl = _make_nixl_mock(prep_side_effect=RuntimeError("no dlist"))
72
+ from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
73
+ NixlDataTransport,
74
+ )
75
+
76
+ with (
77
+ patch(f"{_NIXL_PATH}.NixlWrapper", return_value=nixl),
78
+ patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
79
+ pytest.raises(RuntimeError, match="no dlist"),
80
+ ):
81
+ NixlDataTransport("agent", 0, 4, 64, 256)
82
+
83
+ nixl.deregister_memory.assert_called_once_with(nixl.get_reg_descs.return_value)
84
+
85
+
86
+def test_nixl_data_transport_raises_if_nixl_unavailable():
87
+ from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
88
+ NixlDataTransport,
89
+ )
90
+
91
+ with (
92
+ patch(f"{_NIXL_PATH}.NixlWrapper", None),
93
+ patch(f"{_NIXL_PATH}.nixl_agent_config", None),
94
+ pytest.raises(RuntimeError, match="requires NIXL"),
95
+ ):
96
+ NixlDataTransport("a", 0, 1, 64, 64)
97
+
98
+
99
+# ── public accessors ──────────────────────────────────────────────────────────
100
+
101
+
102
+def test_get_agent_metadata_returns_nixl_blob():
103
+ assert _make_transport().get_agent_metadata() == b"agent-meta"
104
+
105
+
106
+def test_get_mem_descriptor_roundtrips_block_descs():
107
+ t = _make_transport()
108
+ raw = t.get_mem_descriptor()
109
+ assert isinstance(raw, bytes) and len(raw) > 0
110
+
111
+
112
+# ── add_remote_peer ───────────────────────────────────────────────────────────
113
+
114
+
115
+def test_add_remote_peer_returns_agent_name_and_stores_handle_internally():
116
+ nixl = _make_nixl_mock(prep_side_effect=[77, 88])
117
+ t = _make_transport(nixl)
118
+ mem_desc = serialize_mem_descriptor([(0, 64, 0)])
119
+ agent_name = t.add_remote_peer(b"meta", mem_desc)
120
+ assert agent_name == "remote-agent-1"
121
+ # Dlist handle stored internally; accessible via _peer_handles.
122
+ assert t._peer_handles["remote-agent-1"] == 88
123
+
124
+
125
+def test_add_remote_peer_registers_remote_agent():
126
+ nixl = _make_nixl_mock(prep_side_effect=[77, 88])
127
+ t = _make_transport(nixl)
128
+ mem_desc = serialize_mem_descriptor([(0, 64, 0)])
129
+ t.add_remote_peer(b"fresh-meta", mem_desc)
130
+ nixl.add_remote_agent.assert_called_once_with(b"fresh-meta")
131
+
132
+
133
+# ── remove_remote_peer ────────────────────────────────────────────────────────
134
+
135
+
136
+def test_remove_remote_peer_delegates_to_nixl():
137
+ t = _make_transport()
138
+ t.remove_remote_peer("agent-x")
139
+ t._nixl.remove_remote_agent.assert_called_once_with("agent-x")
140
+
141
+
142
+def test_remove_remote_peer_swallows_exceptions():
143
+ t = _make_transport()
144
+ t._nixl.remove_remote_agent.side_effect = RuntimeError("gone")
145
+ t.remove_remote_peer("agent-x") # must not raise
146
+
147
+
148
+# ── post_read ─────────────────────────────────────────────────────────────────
149
+
150
+
151
+def test_post_read_raises_on_index_count_mismatch():
152
+ nixl = _make_nixl_mock(prep_side_effect=[77, 88])
153
+ t = _make_transport(nixl)
154
+ mem_desc = serialize_mem_descriptor([(0, 64, 0)])
155
+ agent = t.add_remote_peer(b"meta", mem_desc)
156
+ with pytest.raises(ValueError, match="mismatch"):
157
+ t.post_read([0, 1], agent, [5], notif_msg=b"sid")
158
+
159
+
160
+def test_post_read_calls_make_prepped_xfer_and_transfer():
161
+ nixl = _make_nixl_mock(prep_side_effect=[77, 88])
162
+ t = _make_transport(nixl)
163
+ mem_desc = serialize_mem_descriptor([(0, 64, 0)])
164
+ agent = t.add_remote_peer(b"meta", mem_desc)
165
+ nixl.make_prepped_xfer.return_value = 999
166
+ handle = t.post_read([0, 1], agent, [5, 6], notif_msg=b"sid")
167
+ nixl.make_prepped_xfer.assert_called_once()
168
+ nixl.transfer.assert_called_once_with(999)
169
+ assert handle == 999
170
+
171
+
172
+# ── release_xfer_handle / deregister ─────────────────────────────────────────
173
+
174
+
175
+def test_release_xfer_handle_swallows_exceptions():
176
+ t = _make_transport()
177
+ t._nixl.release_xfer_handle.side_effect = RuntimeError("stale")
178
+ t.release_xfer_handle(42) # must not raise
179
+
180
+
181
+def test_deregister_swallows_exceptions():
182
+ t = _make_transport()
183
+ t._nixl.deregister_memory.side_effect = RuntimeError("already gone")
184
+ t.deregister() # must not raise
tests/v1/ec_connector/unit/test_ec_transfer_params.py
+12
-1
@@ -7,7 +7,8 @@ No running engine required.
7
8
from unittest.mock import MagicMock
9
10
-from tests.v1.core.utils import create_scheduler
10
+import pytest
11
+
12
from vllm.entrypoints.openai.chat_completion.protocol import (
13
ChatCompletionRequest,
14
)
@@ -15,6 +16,16 @@ from vllm.outputs import CompletionOutput, RequestOutput
16
from vllm.sampling_params import SamplingParams
17
from vllm.v1.request import Request, RequestStatus
18
19
+# tests.v1.core.utils transitively imports the example KV connector, which
20
+# requires the compiled CUDA flash-attention extensions; skip cleanly rather
21
+# than erroring at collection when those extensions aren't built.
22
+_core_utils = pytest.importorskip(
23
+ "tests.v1.core.utils",
24
+ reason="requires compiled CUDA flash-attention extensions",
25
+ exc_type=ImportError,
26
+)
27
+create_scheduler = _core_utils.create_scheduler
28
+
29
EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}}
30
31
tests/v1/ec_connector/unit/test_protocol.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import msgspec
4
+
5
+from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
6
+ EC_CONNECTOR_VERSION,
7
+ XferAck,
8
+ XferReq,
9
+ XferStatus,
10
+ compute_ec_compatibility_hash,
11
+)
12
+
13
+
14
+def test_xferreq_roundtrip():
15
+ req = XferReq(mm_hash="h1", compatibility_hash="c1", session_id="s1")
16
+ data = msgspec.msgpack.encode(req)
17
+ out = msgspec.msgpack.decode(data, type=XferReq)
18
+ assert out.mm_hash == "h1"
19
+ assert out.compatibility_hash == "c1"
20
+ assert out.connector_version == EC_CONNECTOR_VERSION
21
+
22
+
23
+def test_xferack_roundtrip_ok():
24
+ ack = XferAck(
25
+ mm_hash="h1",
26
+ status=XferStatus.OK,
27
+ src_block_indices=[1, 2],
28
+ agent_metadata=b"meta",
29
+ mem_descriptor=b"desc",
30
+ )
31
+ out = msgspec.msgpack.decode(msgspec.msgpack.encode(ack), type=XferAck)
32
+ assert out.status == XferStatus.OK
33
+ assert out.src_block_indices == [1, 2]
34
+
35
+
36
+def test_compat_hash_deterministic_and_sensitive():
37
+ a = compute_ec_compatibility_hash("0.1", "m", "float16", 64)
38
+ b = compute_ec_compatibility_hash("0.1", "m", "float16", 64)
39
+ c = compute_ec_compatibility_hash("0.1", "m", "float16", 128)
40
+ assert a == b
41
+ assert a != c
tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py
new
+350
@@ -0,0 +1,350 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import uuid
4
+
5
+import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
6
+from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
7
+from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
8
+ ECSharedRegion,
9
+)
10
+from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
11
+
12
+_N, _BS, _HID, _ES = 16, 64, 32, 2
13
+
14
+
15
+class _Pos:
16
+ def __init__(self, offset, length):
17
+ self.offset, self.length = offset, length
18
+
19
+
20
+class _Feature:
21
+ def __init__(self, mm_hash, length=1):
22
+ self.mm_hash = mm_hash
23
+ self.identifier = mm_hash
24
+ self.mm_position = _Pos(0, length)
25
+
26
+
27
+class _Request:
28
+ def __init__(self, features, params=None, req_id="r1"):
29
+ self.mm_features = features
30
+ self.request_id = req_id
31
+ self.ec_transfer_params = params
32
+
33
+
34
+class _FakeResults:
35
+ def __init__(self):
36
+ self.completed = set()
37
+ self.tombstoned = set()
38
+ self.quarantined = set()
39
+ self.retryable = set()
40
+ self.settled = []
41
+
42
+
43
+class _FakeSession:
44
+ """Records started transfers; lets the test flip them to completed."""
45
+
46
+ def __init__(self):
47
+ self.started = []
48
+ self._results = _FakeResults()
49
+
50
+ def start_xfer(self, mm_hash, block_indices, deadline):
51
+ self.started.append(mm_hash)
52
+
53
+ def poll(self, messages, now):
54
+ pass
55
+
56
+ def take_results(self):
57
+ r, self._results = self._results, _FakeResults()
58
+ return r
59
+
60
+ def close(self):
61
+ pass
62
+
63
+
64
+class _FakeTransport:
65
+ """Stand-in for ZmqClientTransport so _poll_step runs without ZMQ."""
66
+
67
+ def poll(self):
68
+ return {}
69
+
70
+ def poll_dead(self):
71
+ return []
72
+
73
+ def close(self):
74
+ pass
75
+
76
+
77
+def _consumer_sched(monkeypatch):
78
+ def _region(cfg):
79
+ return ECSharedRegion(
80
+ engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
81
+ )
82
+
83
+ monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region)
84
+
85
+ # ec_enable_nixl defaults to False, so this builds gate-off; the fields
86
+ # below flip on the NIXL consumer state directly.
87
+ s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_consumer"))
88
+ # Turn on NIXL consumer state without constructing real transports.
89
+ s._nixl_enabled = True
90
+ s._transport = _FakeTransport()
91
+ s._data = None
92
+ s._compat_hash = "c"
93
+ # _setup_nixl normally computes these from model_config; set them
94
+ # directly since this helper builds gate-off then flips fields on.
95
+ s._hidden_dim = _HID
96
+ s._element_size = _ES
97
+ return s
98
+
99
+
100
+def _params(mm_hash, length):
101
+ return {
102
+ mm_hash: {
103
+ "peer_host": "h",
104
+ "peer_port": 1,
105
+ "size_bytes": length * _HID * _ES,
106
+ }
107
+ }
108
+
109
+
110
+def test_new_remote_read_defers_then_completes(monkeypatch):
111
+ s = _consumer_sched(monkeypatch)
112
+ fake = _FakeSession()
113
+
114
+ # Route _start_xfer to our fake session instead of real ZMQ/NIXL, but
115
+ # still reserve a real not-ready cache entry so mark_ready works.
116
+ def _fake_start(mm_hash, info, size):
117
+ entry = s._cache.alloc(mm_hash, 1)
118
+ assert entry is not None
119
+ fake.started.append(mm_hash)
120
+ return True
121
+
122
+ monkeypatch.setattr(s, "_start_xfer", _fake_start)
123
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
124
+
125
+ # Step 1: unseen remote item -> read started, request deferred.
126
+ assert s.ensure_cache_available(req, 0) is False
127
+ assert "h1" in s._in_flight
128
+ assert "h1" in fake.started
129
+ s.build_connector_meta(scheduler_output=None) # re-arms _first_in_batch
130
+
131
+ # Session reports completion this step.
132
+ fake._results.completed.add("h1")
133
+ s._sessions[("h", 1)] = fake
134
+
135
+ # Step 2: poll drains completion -> entry marked ready, request admitted.
136
+ assert s.ensure_cache_available(req, 0) is True
137
+ assert "h1" not in s._in_flight
138
+ assert "h1" in s._step_completed
139
+ entry = s._cache.get("h1")
140
+ assert entry is not None and entry.ready
141
+ # The scheduler now allocates encoder input; the ready entry loads through
142
+ # the same local path as a natively cached encoding.
143
+ s.update_state_after_alloc(req, 0)
144
+ meta = s.build_connector_meta(scheduler_output=None)
145
+ assert "h1" in meta.loads
146
+ assert "h1" not in s._step_completed # cleared by promote
147
+
148
+ # Step 3: still cached -> admitted directly, no new transfer.
149
+ assert s.ensure_cache_available(req, 0) is True
150
+ s.shutdown()
151
+
152
+
153
+def test_no_params_is_ready(monkeypatch):
154
+ s = _consumer_sched(monkeypatch)
155
+ assert s.ensure_cache_available(_Request([_Feature("x")], params=None), 0) is True
156
+ s.shutdown()
157
+
158
+
159
+def test_tombstoned_read_discards_and_blocks_retry(monkeypatch):
160
+ s = _consumer_sched(monkeypatch)
161
+ fake = _FakeSession()
162
+
163
+ def _fake_start(mm_hash, info, size):
164
+ entry = s._cache.alloc(mm_hash, 1)
165
+ assert entry is not None
166
+ fake.started.append(mm_hash)
167
+ return True
168
+
169
+ monkeypatch.setattr(s, "_start_xfer", _fake_start)
170
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
171
+
172
+ # Step 1: start the read.
173
+ assert s.ensure_cache_available(req, 0) is False
174
+ s.build_connector_meta(scheduler_output=None)
175
+
176
+ # Producer rejected the read -> tombstoned.
177
+ fake._results.tombstoned.add("h1")
178
+ s._sessions[("h", 1)] = fake
179
+
180
+ # Step 2: poll discards the entry and records a tombstone; the tombstone
181
+ # is consumed by admit so the request proceeds to local compute.
182
+ assert s.ensure_cache_available(req, 0) is True
183
+ assert s._cache.get("h1") is None
184
+ assert "h1" not in s._in_flight
185
+ s.build_connector_meta(scheduler_output=None)
186
+
187
+ # Step 3: tombstone was consumed, so a fresh transfer starts again.
188
+ assert s.ensure_cache_available(req, 0) is False
189
+ assert "h1" in s._in_flight
190
+ s.shutdown()
191
+
192
+
193
+def test_retryable_read_re_requests_without_admitting(monkeypatch):
194
+ """A not-ready NACK defers the request and re-requests the read.
195
+
196
+ Unlike a tombstone it never admits the request, because an admitted request
197
+ whose media was rewritten away upstream has no embedding and no way to
198
+ recompute one.
199
+ """
200
+ s = _consumer_sched(monkeypatch)
201
+ fake = _FakeSession()
202
+
203
+ def _fake_start(mm_hash, info, size):
204
+ entry = s._cache.alloc(mm_hash, 1)
205
+ assert entry is not None
206
+ fake.started.append(mm_hash)
207
+ return True
208
+
209
+ monkeypatch.setattr(s, "_start_xfer", _fake_start)
210
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
211
+
212
+ assert s.ensure_cache_available(req, 0) is False
213
+ s.build_connector_meta(scheduler_output=None)
214
+
215
+ # Producer's save has not landed: retryable, not tombstoned.
216
+ fake._results.retryable.add("h1")
217
+ s._sessions[("h", 1)] = fake
218
+
219
+ # The entry is released and the read re-issued in the same step, so the
220
+ # request stays deferred rather than being admitted unfetched.
221
+ assert s.ensure_cache_available(req, 0) is False
222
+ assert "h1" in s._in_flight
223
+ assert "h1" not in s._tombstones
224
+ assert fake.started.count("h1") == 2
225
+ s.shutdown()
226
+
227
+
228
+def test_size_mismatch_skips_transfer(monkeypatch):
229
+ s = _consumer_sched(monkeypatch)
230
+ # Advertised size disagrees with pos.length * hidden_dim * element_size.
231
+ bad = {"h1": {"peer_host": "h", "peer_port": 1, "size_bytes": 999}}
232
+ req = _Request([_Feature("h1", 1)], params=bad)
233
+ assert s.ensure_cache_available(req, 0) is True
234
+ assert "h1" not in s._in_flight
235
+ assert s._cache.get("h1") is None
236
+ s.shutdown()
237
+
238
+
239
+def test_metadata_only_entry_admits_without_transfer(monkeypatch):
240
+ # A producer that reported placeholder metadata but had no cache entry
241
+ # to transfer omits peer_host/peer_port/size_bytes entirely; the consumer
242
+ # must fall back to local compute rather than treating that as a size
243
+ # mismatch.
244
+ s = _consumer_sched(monkeypatch)
245
+ metadata_only = {"h1": {"image_grid_thw": [1, 2, 3]}}
246
+ req = _Request([_Feature("h1", 1)], params=metadata_only)
247
+ assert s.ensure_cache_available(req, 0) is True
248
+ assert "h1" not in s._in_flight
249
+ assert s._cache.get("h1") is None
250
+ s.shutdown()
251
+
252
+
253
+def test_orphan_not_ready_entry_falls_back_no_realloc(monkeypatch):
254
+ # Post-quarantine, post-tombstone-consumed orphan: a not-ready cache
255
+ # entry exists for the mm_hash, but the hash is tracked in none of
256
+ # _in_flight / _step_completed / _tombstones. Its blocks are held by a
257
+ # quarantined/settling DMA and must not be re-allocated.
258
+ s = _consumer_sched(monkeypatch)
259
+ entry = s._cache.alloc("h1", 1)
260
+ assert entry is not None and not entry.ready
261
+ assert "h1" not in s._in_flight
262
+ assert "h1" not in s._step_completed
263
+ assert "h1" not in s._tombstones
264
+
265
+ # Spy on alloc: record calls, delegate to the real allocator (which
266
+ # asserts on a duplicate key). A re-alloc for "h1" is the bug.
267
+ real_alloc = s._cache.alloc
268
+ calls: list[str] = []
269
+
270
+ def _spy_alloc(key, n):
271
+ calls.append(key)
272
+ return real_alloc(key, n)
273
+
274
+ monkeypatch.setattr(s._cache, "alloc", _spy_alloc)
275
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
276
+
277
+ # Must not raise; request falls back to local compute (admitted True),
278
+ # the orphan entry is untouched, and alloc is never called again.
279
+ assert s.ensure_cache_available(req, 0) is True
280
+ assert calls == []
281
+ assert s._cache.get("h1") is entry
282
+ assert not entry.ready
283
+ s.shutdown()
284
+
285
+
286
+def test_alloc_failure_falls_back_to_local(monkeypatch):
287
+ s = _consumer_sched(monkeypatch)
288
+
289
+ # Force the cache to reject the allocation.
290
+ monkeypatch.setattr(s._cache, "alloc", lambda key, n: None)
291
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
292
+
293
+ # _start_xfer returns False -> request admitted for local recompute.
294
+ assert s.ensure_cache_available(req, 0) is True
295
+ assert "h1" not in s._in_flight
296
+ s.shutdown()
297
+
298
+
299
+def test_already_computed_feature_admits_without_transfer(monkeypatch):
300
+ # A feature entirely within num_computed_tokens is skipped by admit; no
301
+ # cache lookup/alloc or transfer should be triggered for it.
302
+ s = _consumer_sched(monkeypatch)
303
+ calls: list[str] = []
304
+ real_alloc = s._cache.alloc
305
+
306
+ def _spy_alloc(key, n):
307
+ calls.append(key)
308
+ return real_alloc(key, n)
309
+
310
+ monkeypatch.setattr(s._cache, "alloc", _spy_alloc)
311
+ # Feature spans [0, 4); with num_computed_tokens=4 it is fully computed.
312
+ req = _Request([_Feature("h1", 4)], params=_params("h1", 4))
313
+
314
+ assert s.ensure_cache_available(req, 4) is True
315
+ assert "h1" not in s._in_flight
316
+ assert calls == []
317
+ s.shutdown()
318
+
319
+
320
+def test_local_ready_entry_admits_without_transfer(monkeypatch):
321
+ # A READY entry already present in the cache is a local hit: admit
322
+ # returns True and no new transfer is started for that mm_hash.
323
+ s = _consumer_sched(monkeypatch)
324
+ entry = s._cache.alloc("h1", 1)
325
+ assert entry is not None
326
+ s._cache.mark_ready("h1")
327
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
328
+
329
+ assert s.ensure_cache_available(req, 0) is True
330
+ assert "h1" not in s._in_flight
331
+ s.shutdown()
332
+
333
+
334
+def test_in_flight_hash_defers_without_second_transfer(monkeypatch):
335
+ # A mm_hash already tracked in _in_flight must defer the request rather
336
+ # than starting a duplicate transfer.
337
+ s = _consumer_sched(monkeypatch)
338
+ started: list[str] = []
339
+
340
+ def _spy_start_xfer(mm_hash, info, size):
341
+ started.append(mm_hash)
342
+ return True
343
+
344
+ monkeypatch.setattr(s, "_start_xfer", _spy_start_xfer)
345
+ s._in_flight.add("h1")
346
+ req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
347
+
348
+ assert s.ensure_cache_available(req, 0) is False
349
+ assert started == []
350
+ s.shutdown()
tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py
new
+78
@@ -0,0 +1,78 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import uuid
4
+
5
+import pytest
6
+
7
+import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
8
+from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
9
+from vllm.config.ec_transfer import ECTransferConfig
10
+from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
11
+ ECSharedRegion,
12
+)
13
+from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
14
+from vllm.distributed.nixl_utils import NixlWrapper
15
+
16
+_N = 16
17
+_BS = 64
18
+
19
+
20
+def _region() -> ECSharedRegion:
21
+ return ECSharedRegion(
22
+ engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
23
+ )
24
+
25
+
26
+def test_ec_enable_nixl_absent_from_extra_config_by_default():
27
+ cfg = ECTransferConfig()
28
+ assert cfg.get_from_extra_config("ec_enable_nixl", False) is False
29
+
30
+
31
+def test_ec_enable_nixl_read_from_extra_config():
32
+ cfg = ECTransferConfig(
33
+ ec_connector="ECCPUConnector",
34
+ ec_role="ec_both",
35
+ ec_connector_extra_config={"ec_enable_nixl": True},
36
+ )
37
+ assert cfg.get_from_extra_config("ec_enable_nixl", False) is True
38
+
39
+
40
+def test_gate_off_builds_no_nixl(monkeypatch):
41
+ monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
42
+ # ec_enable_nixl defaults to False.
43
+ s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both"))
44
+ assert s._nixl_enabled is False
45
+ assert getattr(s, "_data", None) is None
46
+ assert getattr(s, "_producer_session", None) is None
47
+ s.shutdown()
48
+
49
+
50
+def test_string_false_in_extra_config_leaves_nixl_off(monkeypatch):
51
+ """Extra config is not type coerced, and bool("false") is True."""
52
+ monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
53
+ cfg = create_ec_vllm_config(ec_role="ec_both")
54
+ cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = "false"
55
+ s = ECCPUScheduler(cfg)
56
+ assert s._nixl_enabled is False
57
+ s.shutdown()
58
+
59
+
60
+@pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package")
61
+def test_gate_on_wires_data_transport_and_producer_session(monkeypatch):
62
+ # Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind
63
+ # in ProducerSession.start() cannot collide with another test/process.
64
+ monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1")
65
+ monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", "0")
66
+ monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
67
+ cfg = create_ec_vllm_config(ec_role="ec_both")
68
+ cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
69
+ s = ECCPUScheduler(cfg)
70
+ try:
71
+ assert s._nixl_enabled is True
72
+ assert s._data is not None
73
+ assert s._compat_hash is not None
74
+ assert s._producer_session is not None
75
+ assert s._peer_host == "127.0.0.1"
76
+ assert s._transport is not None # is_ec_consumer is also True
77
+ finally:
78
+ s.shutdown()
tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
new
+113
@@ -0,0 +1,113 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import uuid
4
+
5
+import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
6
+from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
7
+from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
8
+ ECSharedRegion,
9
+)
10
+from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
11
+
12
+_N, _BS = 16, 64
13
+
14
+
15
+class _Pos:
16
+ def __init__(self, offset, length):
17
+ self.offset, self.length = offset, length
18
+
19
+
20
+class _Feature:
21
+ def __init__(self, mm_hash, length=1):
22
+ self.mm_hash = mm_hash
23
+ self.identifier = mm_hash
24
+ self.mm_position = _Pos(0, length)
25
+ self.data = None
26
+ self.modality = "image"
27
+
28
+
29
+class _Request:
30
+ def __init__(self, features, req_id="r1"):
31
+ self.mm_features = features
32
+ self.request_id = req_id
33
+
34
+
35
+def _sched_gate_off(monkeypatch):
36
+ def _region(cfg):
37
+ return ECSharedRegion(
38
+ engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
39
+ )
40
+
41
+ monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region)
42
+ return ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both"))
43
+
44
+
45
+def test_request_finished_gate_off_returns_none(monkeypatch):
46
+ s = _sched_gate_off(monkeypatch)
47
+ assert s.request_finished(_Request([_Feature("h1")])) == (False, None)
48
+ s.shutdown()
49
+
50
+
51
+def test_request_finished_producer_emits_params(monkeypatch):
52
+ s = _sched_gate_off(monkeypatch)
53
+ # Simulate NIXL-enabled producer bookkeeping without building real NIXL.
54
+ s._nixl_enabled = True
55
+ s._peer_host, s._peer_port = "1.2.3.4", 5601
56
+ # _setup_nixl normally computes these from model_config; set them
57
+ # directly since this test builds gate-off then flips fields on.
58
+ s._hidden_dim, s._element_size = 32, 2
59
+ # feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks.
60
+ entry = s._cache.alloc("h1", 2)
61
+ assert entry is not None
62
+ s._cache.mark_ready("h1")
63
+
64
+ delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
65
+ assert delay is False
66
+ assert params == {
67
+ "h1": {
68
+ "metadata": {},
69
+ "peer_host": "1.2.3.4",
70
+ "peer_port": 5601,
71
+ "size_bytes": 2 * 32 * 2,
72
+ }
73
+ }
74
+ s.shutdown()
75
+
76
+
77
+def test_request_finished_announces_not_ready_entry(monkeypatch):
78
+ """A save whose GPU->mmap copy hasn't been confirmed complete yet is
79
+ still announced: the entry can't be evicted before it's ready, and a read
80
+ arriving too early is NACKed NACK_NOT_READY for the consumer to retry —
81
+ protecting it is not this method's job."""
82
+ s = _sched_gate_off(monkeypatch)
83
+ s._nixl_enabled = True
84
+ s._peer_host, s._peer_port = "1.2.3.4", 5601
85
+ s._hidden_dim, s._element_size = 32, 2
86
+ s._cache.alloc("h1", 2) # allocated but not marked ready
87
+
88
+ delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
89
+ assert delay is False
90
+ assert params == {
91
+ "h1": {
92
+ "metadata": {},
93
+ "peer_host": "1.2.3.4",
94
+ "peer_port": 5601,
95
+ "size_bytes": 2 * 32 * 2,
96
+ }
97
+ }
98
+ s.shutdown()
99
+
100
+
101
+def test_request_finished_skips_unallocated_entry(monkeypatch):
102
+ s = _sched_gate_off(monkeypatch)
103
+ s._nixl_enabled = True
104
+ s._peer_host, s._peer_port = "1.2.3.4", 5601
105
+ s._hidden_dim, s._element_size = 32, 2
106
+ # No alloc() for "h1" at all — e.g. the cache was full at save time.
107
+
108
+ delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
109
+ assert delay is False
110
+ # The item's placeholder metadata is still reported even though there's
111
+ # no cache entry to transfer (empty "metadata": no fields, no transfer).
112
+ assert params == {"h1": {"metadata": {}}}
113
+ s.shutdown()
tests/v1/ec_connector/unit/test_session.py
new
+470
@@ -0,0 +1,470 @@
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
+
333
+ data = _make_data(xfer_state="PROC")
334
+ x = _started(data)
335
+ # First poll transitions to QUARANTINED (deadline extended by 60s)
336
+ x.poll(time.monotonic() + 9999)
337
+ # Second poll with expired quarantine deadline should evict
338
+ expired_time = time.monotonic() + 9999 + _CONSUMER_QUARANTINE_TIMEOUT_S + 1
339
+ assert x.poll(expired_time) == XferState.READ_FAILED
340
+ data.release_xfer_handle.assert_called_once_with(99)
341
+ assert x.transfer_handle is None
342
+
343
+
344
+# ── ConsumerXfer.cancel / release ────────────────────────────────────────────
345
+
346
+
347
+def test_consumer_xfer_cancel_valid_in_waiting_ack():
348
+ _xfer().cancel() # must not raise
349
+
350
+
351
+def test_consumer_xfer_cancel_raises_when_read_started():
352
+ with pytest.raises(AssertionError):
353
+ _started().cancel()
354
+
355
+
356
+def test_consumer_xfer_release_releases_handle():
357
+ data = _make_data()
358
+ x = _started(data)
359
+ x.release()
360
+ data.release_xfer_handle.assert_called_once_with(99)
361
+
362
+
363
+def test_consumer_xfer_release_noop_when_no_handle():
364
+ data = _make_data()
365
+ _xfer(data).release()
366
+ data.release_xfer_handle.assert_not_called()
367
+
368
+
369
+# ── ConsumerSession ───────────────────────────────────────────────────────────
370
+
371
+
372
+def _make_consumer_session(data=None):
373
+ zmq_conn = MagicMock()
374
+ zmq_conn.recv.return_value = []
375
+ return ConsumerSession(
376
+ addr=("host", 1234),
377
+ zmq_conn=zmq_conn,
378
+ transport=MagicMock(),
379
+ data=data or _make_data(),
380
+ compat_hash="hash-abc",
381
+ )
382
+
383
+
384
+def test_consumer_session_has_stable_session_id():
385
+ s = _make_consumer_session()
386
+ assert isinstance(s._session_id, str) and len(s._session_id) > 0
387
+ assert s._session_id == s._session_id # stable
388
+
389
+
390
+def test_consumer_session_start_xfer_sends_req_with_session_id():
391
+ import msgspec
392
+
393
+ s = _make_consumer_session()
394
+ s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
395
+ raw = s._zmq.send.call_args[0][0]
396
+ req = msgspec.msgpack.decode(raw, type=XferReq)
397
+ assert req.mm_hash == "h1"
398
+ assert req.session_id == s._session_id
399
+
400
+
401
+def test_consumer_session_poll_done_goes_to_completed():
402
+ import msgspec
403
+
404
+ data = _make_data(xfer_state="PROC")
405
+ s = _make_consumer_session(data)
406
+ s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
407
+
408
+ ack = XferAck(
409
+ mm_hash="h1",
410
+ status=XferStatus.OK,
411
+ session_id=s._session_id,
412
+ src_block_indices=[5, 6],
413
+ agent_metadata=b"meta",
414
+ mem_descriptor=b"desc",
415
+ )
416
+ data.add_remote_peer.return_value = "agent-1"
417
+ raw_ack = msgspec.msgpack.encode(ack)
418
+
419
+ data.check_xfer_state.return_value = "DONE"
420
+ s.poll([raw_ack], time.monotonic())
421
+
422
+ results = s.take_results()
423
+ assert "h1" in results.completed
424
+ assert not results.tombstoned
425
+
426
+
427
+def test_consumer_session_poll_nack_goes_to_tombstoned():
428
+ import msgspec
429
+
430
+ s = _make_consumer_session()
431
+ s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
432
+ nack = XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING)
433
+ s.poll([msgspec.msgpack.encode(nack)], time.monotonic())
434
+ results = s.take_results()
435
+ assert "h1" in results.tombstoned
436
+ assert not results.completed
437
+
438
+
439
+def test_consumer_session_poll_not_ready_nack_goes_to_retryable():
440
+ """A not-ready NACK must not tombstone: tombstoning admits the request with
441
+ no embedding, which is fatal once the media has been rewritten away."""
442
+ import msgspec
443
+
444
+ s = _make_consumer_session()
445
+ s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
446
+ nack = XferAck(mm_hash="h1", status=XferStatus.NACK_NOT_READY)
447
+ s.poll([msgspec.msgpack.encode(nack)], time.monotonic())
448
+ results = s.take_results()
449
+ assert "h1" in results.retryable
450
+ assert not results.tombstoned
451
+ assert not results.completed
452
+
453
+
454
+def test_consumer_session_take_results_clears_state():
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.tombstoned
460
+ r2 = s.take_results()
461
+ assert not r2.tombstoned # cleared after first take
462
+
463
+
464
+def test_consumer_session_on_peer_down_cancels_waiting_ack():
465
+ s = _make_consumer_session()
466
+ s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 60)
467
+ s.on_peer_down()
468
+ results = s.take_results()
469
+ assert "h1" in results.retryable
470
+ assert not results.tombstoned
tests/v1/ec_connector/unit/test_utils.py
new
+54
@@ -0,0 +1,54 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for ECCPUConnector scheduler utilities."""
4
+
5
+import msgspec
6
+import pytest
7
+
8
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
9
+ build_block_descs,
10
+ deserialize_mem_descriptor,
11
+ serialize_mem_descriptor,
12
+)
13
+
14
+# ── build_block_descs ────────────────────────────────────────────────────────
15
+
16
+
17
+@pytest.mark.parametrize("device_id,expected_dev", [(7, 7), (None, 0)])
18
+def test_build_block_descs(device_id, expected_dev):
19
+ kwargs = dict(base_ptr=1000, num_blocks=4, block_size_bytes=256)
20
+ if device_id is not None:
21
+ kwargs["device_id"] = device_id
22
+ descs = build_block_descs(**kwargs)
23
+ assert len(descs) == 4
24
+ for i, (addr, size, dev) in enumerate(descs):
25
+ assert addr == 1000 + i * 256
26
+ assert size == 256
27
+ assert dev == expected_dev
28
+
29
+
30
+def test_build_block_descs_zero_blocks_returns_empty():
31
+ assert build_block_descs(base_ptr=100, num_blocks=0, block_size_bytes=64) == []
32
+
33
+
34
+# ── serialize / deserialize_mem_descriptor ───────────────────────────────────
35
+
36
+
37
+def test_mem_descriptor_roundtrip():
38
+ descs = [(100, 64, 0), (164, 64, 0), (228, 64, 1)]
39
+ assert deserialize_mem_descriptor(serialize_mem_descriptor(descs)) == descs
40
+
41
+
42
+@pytest.mark.parametrize(
43
+ "bad_value",
44
+ [
45
+ [(1, 2)], # 2-tuple instead of 3
46
+ [("a", "b", "c")], # strings, not ints
47
+ ],
48
+)
49
+def test_mem_descriptor_rejects_malformed_payload(bad_value):
50
+ """Malformed descriptor lists must fail to decode."""
51
+ encoder = msgspec.msgpack.Encoder()
52
+ bad_payload = encoder.encode(bad_value)
53
+ with pytest.raises((msgspec.DecodeError, msgspec.ValidationError)):
54
+ deserialize_mem_descriptor(bad_payload)
vllm/distributed/ec_transfer/ec_connector/cpu/connector.py
+13
-1
@@ -8,7 +8,7 @@ GPU copy plumbing. An ec_both instance reuses encoder outputs it has already
8
offloaded to CPU instead of recomputing them.
9
"""
10
11
-from typing import TYPE_CHECKING
11
+from typing import TYPE_CHECKING, Any
12
13
import torch
14
@@ -37,6 +37,12 @@ class ECCPUConnector(ECConnectorBase):
37
def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole) -> None:
38
super().__init__(vllm_config=vllm_config, role=role)
39
40
+ if not vllm_config.use_v2_model_runner:
41
+ raise ValueError(
42
+ "ECCPUConnector requires the V2 model runner. "
43
+ "Set VLLM_USE_V2_MODEL_RUNNER=1."
44
+ )
45
+
46
self.connector_worker = None
47
self.connector_scheduler = None
48
@@ -108,6 +114,12 @@ class ECCPUConnector(ECConnectorBase):
114
assert self.connector_scheduler is not None
115
return self.connector_scheduler.build_connector_meta(scheduler_output)
116
117
+ def request_finished(
118
+ self, request: "Request"
119
+ ) -> tuple[bool, "dict[str, Any] | None"]:
120
+ assert self.connector_scheduler is not None
121
+ return self.connector_scheduler.request_finished(request)
122
+
123
def get_finished(
124
self, finished_req_ids: set[str]
125
) -> tuple[set[str] | None, set[str] | None]:
vllm/distributed/ec_transfer/ec_connector/cpu/control/__init__.py
vllm/distributed/ec_transfer/ec_connector/cpu/control/base.py
new
+37
@@ -0,0 +1,37 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Abstract base classes for the EC control plane.
4
+
5
+ControlConnection wraps a single peer channel (one DEALER socket).
6
+"""
7
+
8
+from abc import ABC, abstractmethod
9
+
10
+
11
+class ControlConnection(ABC):
12
+ """Bidirectional message channel to a single remote peer.
13
+
14
+ All I/O is non-blocking: send enqueues without blocking; recv drains
15
+ all buffered inbound messages since the last call.
16
+ """
17
+
18
+ @property
19
+ @abstractmethod
20
+ def alive(self) -> bool:
21
+ """True if the connection is still usable."""
22
+
23
+ @abstractmethod
24
+ def send(self, msg: bytes) -> None:
25
+ """Enqueue msg for delivery. Must not block."""
26
+
27
+ @abstractmethod
28
+ def recv(self) -> list[bytes]:
29
+ """Drain and return all buffered inbound messages."""
30
+
31
+ @abstractmethod
32
+ def mark_dead(self) -> None:
33
+ """Signal that the peer disconnected. Sets alive to False."""
34
+
35
+ @abstractmethod
36
+ def close(self) -> None:
37
+ """Release all resources. Idempotent."""
vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py
new
+233
@@ -0,0 +1,233 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ZMQ control-plane implementations for ECCPUConnector.
4
+
5
+ZmqClientConnection — one DEALER socket to a single producer peer.
6
+ZmqClientTransport — pool of ZmqClientConnection objects (consumer side).
7
+ZmqServerTransport — ROUTER socket (producer side).
8
+"""
9
+
10
+import contextlib
11
+from dataclasses import dataclass, field
12
+
13
+import zmq
14
+from zmq.utils.monitor import recv_monitor_message
15
+
16
+from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import (
17
+ ControlConnection,
18
+)
19
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr
20
+from vllm.logger import init_logger
21
+from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
22
+
23
+logger = init_logger(__name__)
24
+
25
+_HEARTBEAT_IVL_MS = 2000
26
+_HEARTBEAT_TIMEOUT_MS = 4000
27
+_HEARTBEAT_TTL_MS = 8000
28
+
29
+
30
+# ── Client side ───────────────────────────────────────────────────────────────
31
+
32
+
33
+@dataclass
34
+class ZmqClientConnection(ControlConnection):
35
+ """One DEALER socket to one producer peer.
36
+
37
+ Owns the socket and its optional monitor socket. NIXL state is NOT
38
+ stored here — it belongs in ECSession above this layer.
39
+ """
40
+
41
+ dealer: zmq.Socket
42
+ monitor: zmq.Socket | None
43
+ _alive: bool = field(default=True, init=False)
44
+
45
+ @property
46
+ def alive(self) -> bool:
47
+ return self._alive
48
+
49
+ def send(self, msg: bytes) -> None:
50
+ self.dealer.send_multipart([b"", msg])
51
+
52
+ def recv(self) -> list[bytes]:
53
+ msgs: list[bytes] = []
54
+ while True:
55
+ try:
56
+ frames = self.dealer.recv_multipart(flags=zmq.NOBLOCK)
57
+ except zmq.Again:
58
+ break
59
+ except Exception:
60
+ logger.exception("ec: DEALER recv failed")
61
+ break
62
+ if len(frames) != 2 or frames[0] != b"":
63
+ logger.warning(
64
+ "ec: dropped malformed frame envelope (%d frames)", len(frames)
65
+ )
66
+ continue
67
+ msgs.append(frames[1])
68
+ return msgs
69
+
70
+ def mark_dead(self) -> None:
71
+ self._alive = False
72
+
73
+ def close(self) -> None:
74
+ if self.monitor is not None:
75
+ with contextlib.suppress(Exception):
76
+ self.dealer.disable_monitor()
77
+ with contextlib.suppress(Exception):
78
+ self.monitor.close(linger=0)
79
+ with contextlib.suppress(Exception):
80
+ self.dealer.close(linger=0)
81
+ self._alive = False
82
+
83
+
84
+class ZmqClientTransport:
85
+ """Pool of ZmqClientConnection objects — one per producer peer.
86
+
87
+ Creates DEALER sockets lazily on first contact. Does not touch NIXL;
88
+ NIXL registration is handled by ECCPUConsumer once an XferAck arrives.
89
+ """
90
+
91
+ def __init__(self) -> None:
92
+ self._ctx = zmq.Context()
93
+ self._connections: dict[PeerAddr, ZmqClientConnection] = {}
94
+
95
+ def connect(self, addr: PeerAddr) -> ZmqClientConnection:
96
+ """Return the existing connection for addr, or create a new one."""
97
+ existing = self._connections.get(addr)
98
+ if existing is not None:
99
+ return existing
100
+
101
+ host, port = addr
102
+ dealer = make_zmq_socket(
103
+ ctx=self._ctx,
104
+ path=make_zmq_path(scheme="tcp", host=host, port=port),
105
+ socket_type=zmq.DEALER,
106
+ bind=False,
107
+ )
108
+ dealer.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS)
109
+ dealer.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS)
110
+ dealer.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS)
111
+
112
+ monitor_addr = f"inproc://ec-peer-mon-{host}-{port}"
113
+ dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED)
114
+ monitor = self._ctx.socket(zmq.PAIR)
115
+ monitor.connect(monitor_addr)
116
+
117
+ conn = ZmqClientConnection(dealer=dealer, monitor=monitor)
118
+ self._connections[addr] = conn
119
+ return conn
120
+
121
+ def poll(self) -> dict[PeerAddr, list[bytes]]:
122
+ """Non-blocking drain of messages from every peer's DEALER."""
123
+ msgs: dict[PeerAddr, list[bytes]] = {}
124
+ for addr, conn in list(self._connections.items()):
125
+ conn_msgs = conn.recv()
126
+ msgs[addr] = conn_msgs
127
+ return msgs
128
+
129
+ def poll_dead(self) -> list[PeerAddr]:
130
+ """Non-blocking sweep of peer monitors; return addresses that died."""
131
+ dead: list[PeerAddr] = []
132
+ for addr, conn in list(self._connections.items()):
133
+ if conn.monitor is None:
134
+ continue
135
+ try:
136
+ while True:
137
+ evt = recv_monitor_message(conn.monitor, flags=zmq.NOBLOCK)
138
+ if evt["event"] == zmq.EVENT_DISCONNECTED:
139
+ conn.mark_dead()
140
+ dead.append(addr)
141
+ break
142
+ except zmq.Again:
143
+ pass
144
+ except Exception:
145
+ logger.warning("ec: monitor poll failed for %s", addr, exc_info=True)
146
+ return dead
147
+
148
+ def remove(self, addr: PeerAddr) -> ZmqClientConnection | None:
149
+ """Remove and close the connection for addr, returning it (or None)."""
150
+ conn = self._connections.pop(addr, None)
151
+ if conn is not None:
152
+ conn.close()
153
+ return conn
154
+
155
+ def close(self) -> None:
156
+ for conn in list(self._connections.values()):
157
+ conn.close()
158
+ self._connections.clear()
159
+ try:
160
+ self._ctx.destroy(linger=0)
161
+ except Exception:
162
+ logger.debug("ec: zmq context destroy failed", exc_info=True)
163
+
164
+
165
+# ── Server side ───────────────────────────────────────────────────────────────
166
+
167
+
168
+class ZmqServerTransport:
169
+ """ROUTER socket for the producer side. Byte-level only, driven by
170
+ ProducerSession.
171
+
172
+ poll(timeout_ms) waits up to timeout_ms for messages, then drains all that
173
+ are available; timeout_ms=0 drains whatever has already arrived. send()
174
+ routes a reply back to a specific peer by ZMQ identity.
175
+ """
176
+
177
+ def __init__(self, host: str, port: int) -> None:
178
+ self._ctx = zmq.Context()
179
+ self._router = make_zmq_socket(
180
+ ctx=self._ctx,
181
+ path=make_zmq_path(scheme="tcp", host=host, port=port),
182
+ socket_type=zmq.ROUTER,
183
+ bind=True,
184
+ )
185
+ self._poller = zmq.Poller()
186
+ self._poller.register(self._router, zmq.POLLIN)
187
+
188
+ def poll(self, timeout_ms: int = 0) -> list[tuple[bytes, bytes]]:
189
+ """Wait up to timeout_ms for incoming messages, then drain all pending.
190
+
191
+ Returns a list of (identity, payload) pairs.
192
+ """
193
+ if timeout_ms > 0:
194
+ try:
195
+ events = dict(self._poller.poll(timeout=timeout_ms))
196
+ except zmq.ContextTerminated:
197
+ return []
198
+ except Exception:
199
+ logger.exception("ec: router poller failed")
200
+ return []
201
+ if self._router not in events:
202
+ return []
203
+
204
+ msgs: list[tuple[bytes, bytes]] = []
205
+ while True:
206
+ try:
207
+ identity, _, payload = self._router.recv_multipart(flags=zmq.NOBLOCK)
208
+ msgs.append((identity, payload))
209
+ except zmq.Again:
210
+ break
211
+ except zmq.ContextTerminated:
212
+ break
213
+ except Exception:
214
+ logger.exception("ec: router recv failed")
215
+ break
216
+ return msgs
217
+
218
+ def send(self, identity: bytes, payload: bytes) -> None:
219
+ """Send a reply payload to the peer identified by ZMQ identity."""
220
+ try:
221
+ self._router.send_multipart([identity, b"", payload])
222
+ except Exception:
223
+ logger.exception("ec: router send failed")
224
+
225
+ def close(self) -> None:
226
+ try:
227
+ self._router.close(linger=0)
228
+ except Exception:
229
+ logger.debug("ec: router close failed", exc_info=True)
230
+ try:
231
+ self._ctx.destroy(linger=0)
232
+ except Exception:
233
+ logger.debug("ec: zmq context destroy failed", exc_info=True)
vllm/distributed/ec_transfer/ec_connector/cpu/data/__init__.py
vllm/distributed/ec_transfer/ec_connector/cpu/data/base.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Abstract base class for the EC data-plane transport (NIXL/UCX)."""
4
+
5
+from abc import ABC, abstractmethod
6
+from typing import Any
7
+
8
+
9
+class DataTransport(ABC):
10
+ """Interface wrapping the NIXL data-plane for EC block transfers.
11
+
12
+ Peer lifecycle:
13
+ add_remote_peer() — register a peer and store its dlist handle internally
14
+ remove_remote_peer() — deregister and release the stored handle
15
+ post_read() — issue a READ using the stored handle for agent_name
16
+
17
+ The dlist handle (prepped xfer descriptor list) is an implementation detail
18
+ of the transport and is never exposed to callers.
19
+ """
20
+
21
+ @abstractmethod
22
+ def get_agent_metadata(self) -> bytes:
23
+ """Return this agent's serialized NIXL metadata blob."""
24
+
25
+ @abstractmethod
26
+ def get_mem_descriptor(self) -> bytes:
27
+ """Return the msgpack-encoded block descriptor list for this agent."""
28
+
29
+ @abstractmethod
30
+ def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str:
31
+ """Register a remote peer, prep its READ-source dlist, and return agent_name.
32
+
33
+ The dlist handle is stored internally and used by post_read().
34
+ """
35
+
36
+ @abstractmethod
37
+ def remove_remote_peer(self, agent_name: str) -> None:
38
+ """Deregister a previously registered remote peer and release its handle."""
39
+
40
+ @abstractmethod
41
+ def post_read(
42
+ self,
43
+ local_indices: list[int],
44
+ agent_name: str,
45
+ remote_indices: list[int],
46
+ notif_msg: bytes,
47
+ ) -> Any:
48
+ """Issue a consumer-initiated READ using the stored dlist for agent_name.
49
+
50
+ Returns a transfer handle for use with check_xfer_state / release_xfer_handle.
51
+ """
52
+
53
+ @abstractmethod
54
+ def get_new_notifs(self) -> dict[str, list[bytes]]:
55
+ """Drain NIXL completion notifications addressed to this agent."""
56
+
57
+ @abstractmethod
58
+ def check_xfer_state(self, handle: Any) -> str:
59
+ """Return the NIXL state string: 'DONE', 'PROC', or an error string."""
60
+
61
+ @abstractmethod
62
+ def release_xfer_handle(self, handle: Any) -> None:
63
+ """Release a transfer handle returned by post_read."""
64
+
65
+ @abstractmethod
66
+ def deregister(self) -> None:
67
+ """Deregister all memory and tear down the NIXL agent."""
vllm/distributed/ec_transfer/ec_connector/cpu/data/nixl.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NIXL-backed data-plane transport for ECCPUConnector."""
4
+
5
+from typing import Any
6
+
7
+from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
8
+ DataTransport,
9
+)
10
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
11
+ build_block_descs,
12
+ deserialize_mem_descriptor,
13
+ serialize_mem_descriptor,
14
+)
15
+from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
16
+from vllm.logger import init_logger
17
+
18
+logger = init_logger(__name__)
19
+_NIXL_DRAM = "DRAM"
20
+
21
+
22
+class NixlDataTransport(DataTransport):
23
+ """Adapts NixlWrapper to the DataTransport interface.
24
+
25
+ Dlist handles (prepped xfer descriptor lists for remote peers) are stored
26
+ internally keyed by agent_name. Callers never touch raw NIXL handles —
27
+ post_read() resolves the handle from agent_name internally.
28
+
29
+ NixlWrapper is not thread-safe, and neither is this class. Every call site
30
+ runs on the scheduler thread, which serializes the transport by
31
+ construction; transfers are async inside NIXL and are drained from that
32
+ same thread via get_new_notifs() / check_xfer_state().
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ agent_name: str,
38
+ base_ptr: int,
39
+ num_blocks: int,
40
+ block_size_bytes: int,
41
+ total_size_bytes: int,
42
+ num_threads: int = 1,
43
+ ) -> None:
44
+ if NixlWrapper is None or nixl_agent_config is None:
45
+ raise RuntimeError(
46
+ "ECCPUConnector requires NIXL; "
47
+ "install the `nixl` package or set a different ec_connector."
48
+ )
49
+ # agent_name → remote dlist handle (from add_remote_peer, used by post_read)
50
+ self._peer_handles: dict[str, int] = {}
51
+ # None until the local region is registered, so deregister() is safe on
52
+ # a transport whose registration failed part-way.
53
+ self._reg_descs: Any = None
54
+
55
+ self._nixl = NixlWrapper(
56
+ agent_name,
57
+ nixl_agent_config(num_threads=num_threads, capture_telemetry=True),
58
+ )
59
+ self._register_local_memory(
60
+ base_ptr, num_blocks, block_size_bytes, total_size_bytes
61
+ )
62
+
63
+ def _register_local_memory(
64
+ self,
65
+ base_ptr: int,
66
+ num_blocks: int,
67
+ block_size_bytes: int,
68
+ total_size_bytes: int,
69
+ ) -> None:
70
+ """Register the local region with NIXL and prep its transfer dlist.
71
+
72
+ Unwinds the registration if a later step fails: __init__ propagates the
73
+ error, so the caller never receives an object it could deregister with.
74
+ """
75
+ block_descs = build_block_descs(base_ptr, num_blocks, block_size_bytes)
76
+ self._reg_descs = self._nixl.get_reg_descs(
77
+ [(base_ptr, total_size_bytes, 0, "")], _NIXL_DRAM
78
+ )
79
+ self._nixl.register_memory(self._reg_descs, backends=["UCX"])
80
+ try:
81
+ xfer_descs = self._nixl.get_xfer_descs(block_descs, _NIXL_DRAM)
82
+ self._local_xfer_handle = self._nixl.prep_xfer_dlist(
83
+ "NIXL_INIT_AGENT", xfer_descs
84
+ )
85
+ self._agent_metadata = self._nixl.get_agent_metadata()
86
+ self._mem_descriptor_bytes = serialize_mem_descriptor(block_descs)
87
+ except Exception:
88
+ self.deregister()
89
+ raise
90
+
91
+ def get_agent_metadata(self) -> bytes:
92
+ return self._agent_metadata
93
+
94
+ def get_mem_descriptor(self) -> bytes:
95
+ return self._mem_descriptor_bytes
96
+
97
+ def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str:
98
+ """Register remote peer, prep its dlist, store handle internally.
99
+
100
+ Returns agent_name for use in post_read() and remove_remote_peer().
101
+ """
102
+ agent_name = self._nixl.add_remote_agent(metadata)
103
+ remote_blocks = deserialize_mem_descriptor(mem_descriptor)
104
+ remote_xfer_descs = self._nixl.get_xfer_descs(remote_blocks, _NIXL_DRAM)
105
+ self._peer_handles[agent_name] = self._nixl.prep_xfer_dlist(
106
+ agent_name, remote_xfer_descs
107
+ )
108
+ return agent_name
109
+
110
+ def remove_remote_peer(self, agent_name: str) -> None:
111
+ self._peer_handles.pop(agent_name, None)
112
+ try:
113
+ self._nixl.remove_remote_agent(agent_name)
114
+ except Exception:
115
+ logger.warning(
116
+ "EC: remove_remote_peer failed for %s", agent_name, exc_info=True
117
+ )
118
+
119
+ def post_read(
120
+ self,
121
+ local_indices: list[int],
122
+ agent_name: str,
123
+ remote_indices: list[int],
124
+ notif_msg: bytes,
125
+ ) -> Any:
126
+ if len(local_indices) != len(remote_indices):
127
+ raise ValueError(
128
+ f"EC: local/remote block count mismatch "
129
+ f"({len(local_indices)} vs {len(remote_indices)})"
130
+ )
131
+ remote_handle = self._peer_handles[agent_name]
132
+ handle = self._nixl.make_prepped_xfer(
133
+ "READ",
134
+ self._local_xfer_handle,
135
+ local_indices,
136
+ remote_handle,
137
+ remote_indices,
138
+ notif_msg=notif_msg,
139
+ )
140
+ self._nixl.transfer(handle)
141
+ return handle
142
+
143
+ def get_new_notifs(self) -> dict[str, list[bytes]]:
144
+ return self._nixl.get_new_notifs()
145
+
146
+ def check_xfer_state(self, handle: Any) -> str:
147
+ return self._nixl.check_xfer_state(handle)
148
+
149
+ def release_xfer_handle(self, handle: Any) -> None:
150
+ try:
151
+ self._nixl.release_xfer_handle(handle)
152
+ except Exception:
153
+ logger.warning("EC: release_xfer_handle failed", exc_info=True)
154
+
155
+ def deregister(self) -> None:
156
+ if self._reg_descs is None:
157
+ return
158
+ try:
159
+ self._nixl.deregister_memory(self._reg_descs)
160
+ except Exception:
161
+ logger.warning("EC: deregister failed", exc_info=True)
162
+ self._reg_descs = None
vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scheduler-side wire types for the ECCPUConnector.
4
+
5
+Two peer-to-peer message types exchanged between consumer and producer over
6
+ZMQ (XferReq, XferAck) plus a peer-compatibility fingerprint
7
+(compute_ec_compatibility_hash) sent on every XferReq so a producer refuses
8
+to serve a mismatched peer.
9
+"""
10
+
11
+import enum
12
+import hashlib
13
+import json
14
+
15
+import msgspec
16
+
17
+# Bump when the on-wire shape (XferReq/XferAck) changes in a
18
+# backward-incompatible way. Peers reject mismatched versions.
19
+EC_CONNECTOR_VERSION: int = 1
20
+
21
+
22
+class XferStatus(enum.IntEnum):
23
+ """Outcome the producer reports for an XferReq."""
24
+
25
+ OK = 0
26
+ NACK_MISSING = 1 # producer no longer holds the encoding
27
+ NACK_INCOMPAT = 2 # peer-compatibility hash mismatch
28
+ NACK_VERSION = 3 # connector wire-version mismatch
29
+ NACK_INTERNAL = 4 # producer hit an unexpected error
30
+ NACK_NOT_READY = 5 # announced, but its GPU->mmap save has not landed yet
31
+
32
+
33
+# The producer expects a retryable condition to clear on its own, so the
34
+# consumer re-requests the read rather than giving up on the encoding.
35
+RETRYABLE_NACKS = frozenset((XferStatus.NACK_NOT_READY,))
36
+# Ordinary outcomes rather than signs of a broken peer.
37
+EXPECTED_NACKS = RETRYABLE_NACKS | {XferStatus.NACK_MISSING}
38
+
39
+
40
+class XferReq( # type: ignore[call-arg]
41
+ msgspec.Struct,
42
+ tag="req",
43
+ omit_defaults=True,
44
+):
45
+ """Consumer → Producer: I want to READ mm_hash; pin it and tell me where."""
46
+
47
+ mm_hash: str
48
+ compatibility_hash: str
49
+ # Consumer's session identity. Together with mm_hash it uniquely identifies
50
+ # this transfer on both sides: the producer keys its active xfers by
51
+ # (session_id, mm_hash), and the NIXL notif_msg carries the same pair.
52
+ session_id: str = ""
53
+ connector_version: int = EC_CONNECTOR_VERSION
54
+
55
+
56
+class XferAck( # type: ignore[call-arg]
57
+ msgspec.Struct,
58
+ tag="ack",
59
+ omit_defaults=True,
60
+):
61
+ """Producer → Consumer: response to an XferReq.
62
+
63
+ On status == OK: producer has pinned source blocks and returns their
64
+ indices plus fresh NIXL metadata for the consumer to register and READ.
65
+ On any NACK_*: optional fields are empty; consumer falls back to local encode.
66
+ """
67
+
68
+ mm_hash: str
69
+ status: XferStatus
70
+ # Consumer's session_id echoed back from XferReq. Together with mm_hash it
71
+ # identifies the exact transfer; the consumer uses both to build notif_msg.
72
+ session_id: str = ""
73
+ src_block_indices: list[int] = []
74
+ agent_metadata: bytes = b""
75
+ mem_descriptor: bytes = b""
76
+
77
+
78
+def compute_ec_compatibility_hash(
79
+ vllm_version: str,
80
+ model: str,
81
+ dtype: str,
82
+ block_size_bytes: int,
83
+) -> str:
84
+ """Peer-compatibility fingerprint over the four factors that determine
85
+ byte-layout compatibility. Producer NACKs any XferReq whose hash differs.
86
+ """
87
+ factors = json.dumps(
88
+ [vllm_version, model, dtype, block_size_bytes],
89
+ separators=(",", ":"),
90
+ )
91
+ return hashlib.sha256(factors.encode("utf-8")).hexdigest()
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+416
-6
@@ -7,16 +7,23 @@ Owns the mmap region and the embedding cache, and handles the producer
7
for the ECCPUConnector.
8
"""
9
10
-from typing import TYPE_CHECKING
10
+from typing import TYPE_CHECKING, Any
11
+
12
+import torch
13
14
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
15
ECCPUConnectorMetadata,
16
ECCPUWorkerMetadata,
17
+ _get_encoder_cache_hidden_dim,
18
create_ec_shared_region,
19
)
20
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
21
EmbeddingCache,
22
)
23
+from vllm.distributed.ec_transfer.ec_connector.utils import (
24
+ PlaceholderMetadataResolver,
25
+ collect_ec_item_metadata,
26
+)
27
from vllm.logger import init_logger
28
29
if TYPE_CHECKING:
@@ -40,6 +47,7 @@ class ECCPUScheduler:
47
self._region = create_ec_shared_region(vllm_config)
48
# Block allocator + LRU eviction policy for the shared region.
49
self._cache = EmbeddingCache(self._region.num_blocks)
50
+ self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config)
51
52
# mm_hash → block IDs allocated this step for GPU→mmap saves.
53
self._pending_saves: dict[str, list[int]] = {}
@@ -63,6 +71,112 @@ class ECCPUScheduler:
71
else 1
72
)
73
74
+ self._ec_config = ec_config
75
+ # NIXL p2p is an option of this connector, not engine-wide behavior, so
76
+ # it lives in extra config. Extra config is not type coerced, so a value
77
+ # that arrived as a JSON/CLI string is parsed, not truth-tested.
78
+ raw_nixl = ec_config.get_from_extra_config("ec_enable_nixl", False)
79
+ self._nixl_enabled = (
80
+ raw_nixl
81
+ if isinstance(raw_nixl, bool)
82
+ else str(raw_nixl).strip().lower() in ("true", "1", "yes")
83
+ )
84
+ # NIXL fields default to None/empty so the gate-off path is untouched.
85
+ self._data: Any = None
86
+ self._compat_hash: str | None = None
87
+ self._first_in_batch = True
88
+ self._transport: Any = None
89
+ self._producer_session: Any = None
90
+ self._sessions: dict = {}
91
+ self._in_flight: set[str] = set()
92
+ self._tombstones: set[str] = set()
93
+ self._step_completed: set[str] = set()
94
+ self._peer_host: str | None = None
95
+ self._peer_port: int | None = None
96
+ # Model shape for size checks + compat hash; only set by
97
+ # _setup_nixl (the gate-off path never touches model_config).
98
+ self._dtype: torch.dtype | None = None
99
+ self._hidden_dim: int = 0
100
+ self._element_size: int = 0
101
+ self._ack_timeout_s: float = 0.0
102
+ if self._nixl_enabled:
103
+ self._setup_nixl(vllm_config)
104
+
105
+ def _setup_nixl(self, vllm_config: "VllmConfig") -> None:
106
+ # Lazy imports keep nixl/zmq off the gate-off path.
107
+ from vllm import envs
108
+ from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import (
109
+ ZmqClientTransport,
110
+ ZmqServerTransport,
111
+ )
112
+ from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
113
+ NixlDataTransport,
114
+ )
115
+ from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
116
+ compute_ec_compatibility_hash,
117
+ )
118
+ from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
119
+ ProducerSession,
120
+ )
121
+ from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
122
+ from vllm.version import __version__ as VLLM_VERSION
123
+
124
+ if NixlWrapper is None or nixl_agent_config is None:
125
+ raise RuntimeError(
126
+ "ec_enable_nixl requires NIXL; install the `nixl` package or "
127
+ "remove ec_enable_nixl from ec_connector_extra_config."
128
+ )
129
+ engine_id = self._ec_config.engine_id
130
+ assert engine_id is not None
131
+ self._dtype = vllm_config.model_config.dtype
132
+ self._hidden_dim = _get_encoder_cache_hidden_dim(vllm_config)
133
+ self._element_size = torch.empty(0, dtype=self._dtype).element_size()
134
+ # How long a consumer waits for an XferAck. The producer answers from
135
+ # its scheduler step, so its reply latency scales with the encoder's
136
+ # batch size: a deployment whose steps run longer must raise this.
137
+ self._ack_timeout_s = float(
138
+ self._ec_config.get_from_extra_config("consumer_ack_timeout_s", 2.0)
139
+ )
140
+ self._compat_hash = compute_ec_compatibility_hash(
141
+ vllm_version=VLLM_VERSION,
142
+ model=str(vllm_config.model_config.model),
143
+ dtype=str(self._dtype),
144
+ block_size_bytes=self._region.block_size_bytes,
145
+ )
146
+ if self._is_producer:
147
+ self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST
148
+ self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT
149
+
150
+ # Registering the region with NIXL and binding the control sockets are
151
+ # the first side effects here. __init__ propagates a failure, so the
152
+ # caller never receives a scheduler it could shut down: unwind through
153
+ # the same teardown shutdown() uses.
154
+ try:
155
+ self._data = NixlDataTransport(
156
+ agent_name=engine_id,
157
+ base_ptr=self._region.blocks.data_ptr(),
158
+ num_blocks=self._region.num_blocks,
159
+ block_size_bytes=self._region.block_size_bytes,
160
+ total_size_bytes=self._region.num_blocks
161
+ * self._region.block_size_bytes,
162
+ )
163
+ if self._is_producer:
164
+ assert self._peer_host is not None
165
+ assert self._peer_port is not None
166
+ self._producer_session = ProducerSession(
167
+ transport=ZmqServerTransport(
168
+ host=self._peer_host, port=self._peer_port
169
+ ),
170
+ data=self._data,
171
+ cache=self._cache,
172
+ compat_hash=self._compat_hash,
173
+ )
174
+ if self._is_consumer:
175
+ self._transport = ZmqClientTransport()
176
+ except Exception:
177
+ self._teardown_nixl()
178
+ raise
179
+
180
def has_cache_item(self, identifier: str) -> bool:
181
if not self._is_consumer:
182
return False
@@ -72,8 +186,202 @@ class ECCPUScheduler:
186
def ensure_cache_available(
187
self, request: "Request", num_computed_tokens: int
188
) -> bool:
189
+ if not self._nixl_enabled:
190
+ return True # CPU offload never blocks.
191
+ first = self._first_in_batch
192
+ self._first_in_batch = False
193
+ if not self._is_consumer:
194
+ return True
195
+ if first:
196
+ self._poll_step()
197
+ return self._nixl_consumer_admit(request, num_computed_tokens)
198
+
199
+ def _nixl_consumer_admit(
200
+ self, request: "Request", num_computed_tokens: int
201
+ ) -> bool:
202
+ """Admit a request only once all its remote encodings are cached.
203
+
204
+ Returns True when every encoding the request needs is either already
205
+ ready in the local cache or has no remote source. Any encoding whose
206
+ NIXL READ is still in flight (or was just started this step) defers
207
+ the request, which the scheduler re-presents on a later step.
208
+ """
209
+ params: dict[str, dict[str, Any]] = (
210
+ getattr(request, "ec_transfer_params", None) or {}
211
+ )
212
+ if not params:
213
+ return True
214
+ pending = False
215
+ for feature in request.mm_features:
216
+ pos = feature.mm_position
217
+ if pos.offset + pos.length <= num_computed_tokens:
218
+ continue
219
+ mm_hash = feature.identifier
220
+ entry = self._cache.get(mm_hash)
221
+ if entry is not None and entry.ready:
222
+ # Local hit: upstream's update_state_after_alloc pins and
223
+ # loads it through the same path as a natively cached entry.
224
+ continue
225
+ if mm_hash in self._in_flight:
226
+ pending = True
227
+ continue
228
+ if mm_hash in self._step_completed:
229
+ pending = True
230
+ continue
231
+ if mm_hash in self._tombstones:
232
+ self._tombstones.discard(mm_hash)
233
+ continue
234
+ if entry is not None:
235
+ # Present but not ready and not being fetched: its blocks are
236
+ # held by a quarantined/settling DMA and cannot be reused.
237
+ # Fall back to local compute this step; never re-alloc a
238
+ # mm_hash already in the cache.
239
+ continue
240
+ info = params.get(mm_hash)
241
+ # Absent entirely, or present with only placeholder metadata (no
242
+ # cache entry existed on the producer side at request_finished
243
+ # time): nothing to fetch, fall back to local compute.
244
+ if not info or "peer_host" not in info:
245
+ continue
246
+ expected = pos.length * self._hidden_dim * self._element_size
247
+ if int(info.get("size_bytes", -1)) != expected:
248
+ logger.warning(
249
+ "EC consumer: size mismatch mm_hash=%s; local encode", mm_hash
250
+ )
251
+ continue
252
+ try:
253
+ started = self._start_xfer(mm_hash, info, expected)
254
+ except Exception:
255
+ logger.exception(
256
+ "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash
257
+ )
258
+ continue
259
+ if not started:
260
+ continue
261
+ self._in_flight.add(mm_hash)
262
+ pending = True
263
+ return not pending
264
+
265
+ def _start_xfer(
266
+ self, mm_hash: str, info: "dict[str, Any]", size_bytes: int
267
+ ) -> bool:
268
+ """Allocate a not-ready cache entry and start a NIXL READ into it.
269
+
270
+ Returns True if the transfer was started. Returns False when the
271
+ cache cannot accommodate the encoding, in which case the request
272
+ falls back to local recomputation.
273
+ """
274
+ import time
275
+ from math import ceil
276
+
277
+ from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
278
+ ConsumerSession,
279
+ )
280
+
281
+ n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes))
282
+ entry = self._cache.alloc(mm_hash, n_blocks)
283
+ if entry is None:
284
+ logger.debug(
285
+ "EC consumer: cache full for mm_hash=%s (%d blocks); local encode",
286
+ mm_hash,
287
+ n_blocks,
288
+ )
289
+ return False
290
+ indices = list(entry.block_ids)
291
+ addr = (info["peer_host"], int(info["peer_port"]))
292
+ if addr not in self._sessions:
293
+ zmq_conn = self._transport.connect(addr)
294
+ assert self._compat_hash is not None
295
+ self._sessions[addr] = ConsumerSession(
296
+ addr=addr,
297
+ zmq_conn=zmq_conn,
298
+ transport=self._transport,
299
+ data=self._data,
300
+ compat_hash=self._compat_hash,
301
+ )
302
+ deadline = time.monotonic() + self._ack_timeout_s
303
+ try:
304
+ self._sessions[addr].start_xfer(mm_hash, indices, deadline)
305
+ except Exception:
306
+ self._cache.discard(mm_hash)
307
+ raise
308
+ logger.debug(
309
+ "EC consumer: starting NIXL xfer mm_hash=%s from %s:%d blocks=%d",
310
+ mm_hash,
311
+ addr[0],
312
+ addr[1],
313
+ n_blocks,
314
+ )
315
return True
316
317
+ def _poll_step(self) -> None:
318
+ import time
319
+
320
+ now = time.monotonic()
321
+ all_messages = self._transport.poll()
322
+ for addr, session in list(self._sessions.items()):
323
+ session.poll(all_messages.get(addr, []), now)
324
+ for addr in self._transport.poll_dead():
325
+ self._on_peer_down(addr)
326
+ for session in self._sessions.values():
327
+ self._process_session_results(session)
328
+
329
+ def _process_session_results(self, session: Any) -> None:
330
+ r = session.take_results()
331
+ for mm_hash in r.completed:
332
+ self._in_flight.discard(mm_hash)
333
+ self._cache.mark_ready(mm_hash)
334
+ self._step_completed.add(mm_hash)
335
+ logger.debug("EC consumer: NIXL xfer complete mm_hash=%s", mm_hash)
336
+ for mm_hash in r.tombstoned:
337
+ self._in_flight.discard(mm_hash)
338
+ self._cache.discard(mm_hash)
339
+ self._tombstones.add(mm_hash)
340
+ logger.debug("EC consumer: NIXL xfer failed mm_hash=%s", mm_hash)
341
+ for mm_hash in r.quarantined:
342
+ # DMA still running: keep the blocks reserved (entry stays
343
+ # not-ready, hence non-evictable) until the xfer settles.
344
+ self._in_flight.discard(mm_hash)
345
+ self._tombstones.add(mm_hash)
346
+ logger.debug(
347
+ "EC consumer: NIXL xfer mm_hash=%s quarantined; DMA still running",
348
+ mm_hash,
349
+ )
350
+ for mm_hash in r.retryable:
351
+ # No tombstone: dropping every trace of the attempt is what lets
352
+ # the next admit pass re-issue the read as if it were the first.
353
+ self._in_flight.discard(mm_hash)
354
+ self._cache.discard(mm_hash)
355
+ logger.debug(
356
+ "EC consumer: NIXL xfer mm_hash=%s retryable; will re-request",
357
+ mm_hash,
358
+ )
359
+ for mm_hash, _block_indices in r.settled:
360
+ self._cache.discard(mm_hash)
361
+ logger.debug(
362
+ "EC consumer: quarantined NIXL xfer mm_hash=%s settled", mm_hash
363
+ )
364
+
365
+ def _on_peer_down(self, addr: Any) -> None:
366
+ session = self._sessions.pop(addr, None)
367
+ if session is None:
368
+ return
369
+ session.on_peer_down()
370
+ self._process_session_results(session)
371
+ session.close()
372
+ logger.info("EC consumer: producer peer down addr=%s", addr)
373
+
374
+ def _promote_completed_reads(self) -> None:
375
+ """Drain the just-completed set built during ``_poll_step``.
376
+
377
+ Reads are marked ready in ``_process_session_results``, so a deferred
378
+ request is re-admitted by ``ensure_cache_available`` and loaded through
379
+ the same local path as a natively cached entry
380
+ (``update_state_after_alloc`` -> ``_pending_loads``). No explicit
381
+ pin/load bookkeeping is needed here beyond clearing the set.
382
+ """
383
+ self._step_completed.clear()
384
+
385
def update_state_after_alloc(self, request: "Request", index: int) -> None:
386
feature = request.mm_features[index]
387
mm_hash = feature.identifier
@@ -97,11 +405,17 @@ class ECCPUScheduler:
405
) -> ECCPUConnectorMetadata:
406
meta = ECCPUConnectorMetadata()
407
if self._is_producer:
408
+ if self._nixl_enabled and self._producer_session is not None:
409
+ self._producer_session.poll_step()
410
meta.saves = self._pending_saves
411
self._pending_saves = {}
412
if self._is_consumer:
413
+ if self._nixl_enabled:
414
+ self._promote_completed_reads()
415
meta.loads = self._pending_loads
416
self._pending_loads = {}
417
+ if self._nixl_enabled:
418
+ self._first_in_batch = True
419
return meta
420
421
def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:
@@ -117,8 +431,15 @@ class ECCPUScheduler:
431
return
432
for mm_hash in meta.completed_saves:
433
entry = self._cache.get(mm_hash)
120
- if entry is not None and not entry.ready:
434
+ if entry is None:
435
+ logger.debug(
436
+ "EC producer: worker reported completed save for unknown "
437
+ "mm_hash=%s (already discarded/evicted?)",
438
+ mm_hash,
439
+ )
440
+ elif not entry.ready:
441
self._cache.mark_ready(mm_hash)
442
+ logger.debug("EC producer: mm_hash=%s marked ready", mm_hash)
443
for transfer_id in meta.completed_loads:
444
pending = self._load_acks.get(transfer_id)
445
if pending is None:
@@ -131,14 +452,65 @@ class ECCPUScheduler:
452
# as stale rather than releasing the pin a second time.
453
del self._load_acks[transfer_id]
454
self._cache.unpin(mm_hash)
455
+ logger.debug("EC consumer: mm_hash=%s unpinned after load", mm_hash)
456
457
def has_pending_push_work(self) -> bool:
136
- """True while any dispatched save or load has not been confirmed done.
458
+ """Keep the engine stepping so this connector's polls keep running.
459
138
- Keeps the engine stepping so the worker's completion reports are polled
139
- even when no requests are otherwise runnable.
460
+ With NIXL enabled the engine tick is the only driver of
461
+ ``ProducerSession.poll_step()``, and a producer's work arrives from a
462
+ remote consumer, so nothing local would otherwise wake it. Without NIXL
463
+ this only has to outlive dispatched saves and loads, so the engine can
464
+ quiesce once the worker has confirmed them.
465
"""
141
- return self._cache.has_held_entries()
466
+ return self._nixl_enabled or self._cache.has_held_entries()
467
+
468
+ def request_finished(
469
+ self, request: "Request"
470
+ ) -> tuple[bool, "dict[str, Any] | None"]:
471
+ if not (self._nixl_enabled and self._is_producer):
472
+ return False, None
473
+ if not request.mm_features:
474
+ return False, None
475
+
476
+ items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver)
477
+
478
+ for feature in request.mm_features:
479
+ mm_hash = feature.identifier
480
+ entry = self._cache.get(mm_hash)
481
+ if entry is None:
482
+ # Never saved, or evicted since. Publishing placeholder
483
+ # metadata without an address invites an orchestrator to
484
+ # rewrite the media into a reference to an encoding no
485
+ # consumer can fetch, leaving the decoder nothing to embed.
486
+ # Publish neither, so the media stays on the request.
487
+ items[mm_hash]["metadata"] = {}
488
+ logger.debug(
489
+ "EC producer: mm_hash=%s absent at request_finished; "
490
+ "announcing no metadata so the media is not rewritten away",
491
+ mm_hash,
492
+ )
493
+ continue
494
+ # Announce even if the save's GPU->mmap copy hasn't been
495
+ # confirmed complete yet: a not-ready entry can't be evicted, so
496
+ # it will still be here by the time a consumer's XferReq arrives.
497
+ # A read arriving before the save lands is NACKed NACK_NOT_READY,
498
+ # which the consumer retries on a later step rather than treating
499
+ # as a miss.
500
+ size_bytes = (
501
+ feature.mm_position.length * self._hidden_dim * self._element_size
502
+ )
503
+ items[mm_hash].update(
504
+ peer_host=self._peer_host,
505
+ peer_port=self._peer_port,
506
+ size_bytes=size_bytes,
507
+ )
508
+ logger.debug(
509
+ "EC producer: announcing NIXL-readable encodings req_id=%s items=%s",
510
+ request.request_id,
511
+ items,
512
+ )
513
+ return False, items
514
515
def shutdown(self) -> None:
516
self._pending_saves.clear()
@@ -148,7 +520,45 @@ class ECCPUScheduler:
520
self._is_producer = False
521
self._is_consumer = False
522
523
+ if self._nixl_enabled:
524
+ self._teardown_nixl()
525
+
526
try:
527
self._region.cleanup()
528
except Exception:
529
logger.debug("ec: region cleanup failed", exc_info=True)
530
+
531
+ def _teardown_nixl(self) -> None:
532
+ """Close the control sockets and release the NIXL agent.
533
+
534
+ Guards each step and clears what it releases, so it is safe on a
535
+ scheduler whose NIXL setup failed part-way and safe to call twice.
536
+ The shared region is not touched — shutdown() owns that.
537
+ """
538
+ if self._producer_session is not None:
539
+ try:
540
+ self._producer_session.close()
541
+ except Exception:
542
+ logger.debug("ec: producer session close failed", exc_info=True)
543
+ self._producer_session = None
544
+
545
+ for session in list(self._sessions.values()):
546
+ try:
547
+ session.close()
548
+ except Exception:
549
+ logger.debug("ec: consumer session close failed", exc_info=True)
550
+ self._sessions.clear()
551
+
552
+ if self._transport is not None:
553
+ try:
554
+ self._transport.close()
555
+ except Exception:
556
+ logger.debug("ec: client transport close failed", exc_info=True)
557
+ self._transport = None
558
+
559
+ if self._data is not None:
560
+ try:
561
+ self._data.deregister()
562
+ except Exception:
563
+ logger.debug("ec: deregister failed", exc_info=True)
564
+ self._data = None
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py
+32
@@ -122,6 +122,23 @@ class EmbeddingCache:
122
self._evictable_block_count -= len(entry.block_ids)
123
entry.pin()
124
125
+ def pin_if_ready(self, key: str) -> CacheEntry | None:
126
+ """Atomically pin *key* if present and ready; return its entry.
127
+
128
+ Returns None when the key is absent. An entry returned with `ready`
129
+ False was not pinned: its save is still in flight, which the producer
130
+ reports differently from a miss because such a read can be retried.
131
+ """
132
+ with self._lock:
133
+ entry = self._entries.get(key)
134
+ if entry is None or not entry.ready:
135
+ return entry
136
+ if entry.evictable:
137
+ del self._entries_free_list[key]
138
+ self._evictable_block_count -= len(entry.block_ids)
139
+ entry.pin()
140
+ return entry
141
+
142
def unpin(self, key: str) -> None:
143
"""Unpin an entry. Asserts currently pinned."""
144
with self._lock:
@@ -134,6 +151,21 @@ class EmbeddingCache:
151
self._entries_free_list[key] = None
152
self._evictable_block_count += len(entry.block_ids)
153
154
+ def discard(self, key: str) -> None:
155
+ """Remove a not-ready in-flight entry, returning its blocks to the pool.
156
+
157
+ Used when an in-flight fill (e.g. a NIXL READ) fails before the entry
158
+ is marked ready. Asserts the entry is present and not ready — ready or
159
+ pinned entries are reclaimed through eviction/unpin, not discard.
160
+ """
161
+ with self._lock:
162
+ entry = self._entries[key]
163
+ assert entry._pin_count == -1, (
164
+ f"EmbeddingCache: discard of ready/pinned entry {key!r}"
165
+ )
166
+ del self._entries[key]
167
+ self._free_blocks.extend(entry.block_ids)
168
+
169
def _evict_until(self, n_blocks: int) -> None:
170
"""Evict ready+unpinned entries FIFO until enough space. Lock held."""
171
while len(self._free_blocks) < n_blocks:
vllm/distributed/ec_transfer/ec_connector/cpu/session.py
new
+633
@@ -0,0 +1,633 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Session objects for the ECCPUConnector.
4
+
5
+Layer hierarchy: connector → session → transport (connection)
6
+
7
+Two long-lived sessions:
8
+
9
+ ProducerSession — owns ZmqServerTransport. The connector calls poll_step()
10
+ once per engine step, which drains the transport without blocking and
11
+ feeds the raw bytes to poll(). poll() decodes XferReqs, grants or NACKs
12
+ them, pins source blocks, sends XferAck replies, drains NIXL completion
13
+ notifications, and sweeps expired xfers.
14
+
15
+ ConsumerSession — owns one ZmqClientConnection and the NIXL agent
16
+ registration for that peer. The connector calls poll(messages, now)
17
+ once per engine step with the raw bytes collected for this peer.
18
+ poll() decodes XferAcks, dispatches them to ConsumerXfer objects,
19
+ and advances all active xfer state machines.
20
+
21
+Internal data classes:
22
+
23
+ ProducerXfer — one pinned read grant, keyed by session_id.
24
+ ConsumerXfer — state machine for one NIXL READ, keyed by mm_hash.
25
+"""
26
+
27
+import enum
28
+import logging
29
+import time
30
+import uuid
31
+from dataclasses import dataclass, field
32
+from typing import Any
33
+
34
+import msgspec
35
+
36
+from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import (
37
+ ZmqClientConnection,
38
+ ZmqClientTransport,
39
+ ZmqServerTransport,
40
+)
41
+from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
42
+ DataTransport,
43
+)
44
+from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
45
+ EC_CONNECTOR_VERSION,
46
+ EXPECTED_NACKS,
47
+ RETRYABLE_NACKS,
48
+ XferAck,
49
+ XferReq,
50
+ XferStatus,
51
+)
52
+from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
53
+ EmbeddingCache,
54
+)
55
+from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr
56
+from vllm.logger import init_logger
57
+
58
+logger = init_logger(__name__)
59
+
60
+PRODUCER_PIN_LEASE_S = 30.0
61
+_CONSUMER_READ_TIMEOUT_S = 20.0
62
+_CONSUMER_QUARANTINE_TIMEOUT_S = 60.0
63
+
64
+
65
+# ── XferState ─────────────────────────────────────────────────────────────────
66
+
67
+
68
+class XferState(enum.Enum):
69
+ """Lifecycle states of a single ConsumerXfer."""
70
+
71
+ WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck
72
+ READING = "reading" # NIXL READ in flight
73
+ DONE = "done" # READ completed; promote to local cache
74
+ ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + tombstone
75
+ READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone
76
+ QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks
77
+ SETTLED = "settled" # Quarantined DMA done; safe to free blocks
78
+
79
+
80
+# ── ProducerXfer ──────────────────────────────────────────────────────────────
81
+
82
+
83
+@dataclass
84
+class ProducerXfer:
85
+ """One pinned read grant on the producer side.
86
+
87
+ Pure data — lifecycle operations (pin/unpin) are performed by ProducerSession.
88
+ Keyed by (consumer_session_id, mm_hash) in ProducerSession._active_xfers.
89
+ """
90
+
91
+ mm_hash: str
92
+ block_indices: list[int]
93
+ deadline: float
94
+
95
+ def is_expired(self) -> bool:
96
+ return time.monotonic() > self.deadline
97
+
98
+
99
+# ── ConsumerXfer ──────────────────────────────────────────────────────────────
100
+
101
+
102
+@dataclass
103
+class ConsumerXfer:
104
+ """State machine for one NIXL READ on the consumer side.
105
+
106
+ transfer_handle is the in-flight NIXL handle returned by post_read() —
107
+ used for check_xfer_state() / release_xfer_handle(). The dlist handle
108
+ for the remote peer lives inside DataTransport, not here.
109
+ """
110
+
111
+ mm_hash: str
112
+ block_indices: list[int]
113
+ addr: PeerAddr
114
+ deadline: float
115
+ data: DataTransport
116
+ consumer_session_id: str # session identity; combined with mm_hash for notif_msg
117
+
118
+ transfer_handle: Any | None = field(default=None)
119
+ _quarantined: bool = field(default=False, init=False)
120
+
121
+ def handle_ack(self, ack: XferAck, agent_name: str) -> bool:
122
+ """Issue the NIXL READ on an OK XferAck. Returns False on any NACK."""
123
+ if ack.status != XferStatus.OK:
124
+ return False
125
+ # notif_msg encodes both session_id and mm_hash so the producer can
126
+ # do a direct (session_id, mm_hash) → ProducerXfer lookup on completion.
127
+ notif_msg = f"{self.consumer_session_id}:{self.mm_hash}".encode()
128
+ self.transfer_handle = self.data.post_read(
129
+ self.block_indices,
130
+ agent_name,
131
+ ack.src_block_indices,
132
+ notif_msg=notif_msg,
133
+ )
134
+ self.deadline = time.monotonic() + _CONSUMER_READ_TIMEOUT_S
135
+ logger.debug(
136
+ "EC consumer: NIXL READ posted mm_hash=%s agent=%s local_blocks=%d "
137
+ "remote_blocks=%d",
138
+ self.mm_hash,
139
+ agent_name,
140
+ len(self.block_indices),
141
+ len(ack.src_block_indices),
142
+ )
143
+ return True
144
+
145
+ def poll(self, now: float) -> XferState:
146
+ if self._quarantined:
147
+ return self._poll_quarantined(now)
148
+ if self.transfer_handle is None:
149
+ return (
150
+ XferState.ACK_TIMEOUT if now > self.deadline else XferState.WAITING_ACK
151
+ )
152
+ return self._poll_reading(now)
153
+
154
+ def _poll_reading(self, now: float) -> XferState:
155
+ assert self.transfer_handle is not None
156
+ try:
157
+ state = self.data.check_xfer_state(self.transfer_handle)
158
+ except Exception:
159
+ logger.exception(
160
+ "EC Connector: check_xfer_state failed for mm_hash=%s", self.mm_hash
161
+ )
162
+ self.data.release_xfer_handle(self.transfer_handle)
163
+ self.transfer_handle = None
164
+ return XferState.READ_FAILED
165
+ if state == "DONE":
166
+ self.data.release_xfer_handle(self.transfer_handle)
167
+ self.transfer_handle = None
168
+ logger.debug("EC consumer: NIXL READ completed mm_hash=%s", self.mm_hash)
169
+ return XferState.DONE
170
+ if state == "PROC":
171
+ if now > self.deadline:
172
+ self._quarantined = True
173
+ self.deadline = now + _CONSUMER_QUARANTINE_TIMEOUT_S
174
+ logger.warning(
175
+ "EC Connector: READ for mm_hash=%s from %s:%d timed out; "
176
+ "quarantining",
177
+ self.mm_hash,
178
+ self.addr[0],
179
+ self.addr[1],
180
+ )
181
+ return XferState.QUARANTINED
182
+ return XferState.READING
183
+ logger.warning(
184
+ "EC Connector: READ for mm_hash=%s unexpected NIXL state %r",
185
+ self.mm_hash,
186
+ state,
187
+ )
188
+ self.data.release_xfer_handle(self.transfer_handle)
189
+ self.transfer_handle = None
190
+ return XferState.READ_FAILED
191
+
192
+ def _poll_quarantined(self, now: float) -> XferState:
193
+ if self.transfer_handle is None:
194
+ return XferState.SETTLED
195
+ # If quarantine has been too long, assume NIXL is stuck and evict blocks.
196
+ if now > self.deadline:
197
+ self.data.release_xfer_handle(self.transfer_handle)
198
+ self.transfer_handle = None
199
+ logger.warning(
200
+ "EC Connector: quarantined mm_hash=%s from %s:%d timed out; "
201
+ "evicting blocks",
202
+ self.mm_hash,
203
+ self.addr[0],
204
+ self.addr[1],
205
+ )
206
+ return XferState.READ_FAILED
207
+ try:
208
+ state = self.data.check_xfer_state(self.transfer_handle)
209
+ except Exception:
210
+ logger.exception(
211
+ "EC Connector: check_xfer_state failed for quarantined mm_hash=%s",
212
+ self.mm_hash,
213
+ )
214
+ state = None
215
+ if state == "PROC":
216
+ return XferState.QUARANTINED
217
+ self.data.release_xfer_handle(self.transfer_handle)
218
+ self.transfer_handle = None
219
+ logger.debug(
220
+ "EC Connector: quarantined mm_hash=%s settled (state=%s)",
221
+ self.mm_hash,
222
+ state,
223
+ )
224
+ return XferState.SETTLED
225
+
226
+ def cancel(self) -> None:
227
+ """Cancel a WAITING_ACK xfer (peer down). No tombstone — caller allows retry."""
228
+ assert self.transfer_handle is None, "cancel() requires WAITING_ACK state"
229
+
230
+ def release(self) -> None:
231
+ """Force-release the transfer handle. Shutdown only."""
232
+ if self.transfer_handle is not None:
233
+ self.data.release_xfer_handle(self.transfer_handle)
234
+ self.transfer_handle = None
235
+
236
+
237
+# ── ProducerSession ───────────────────────────────────────────────────────────
238
+
239
+
240
+class ProducerSession:
241
+ """Long-lived producer session. Owns ZmqServerTransport.
242
+
243
+ poll_step() runs once per engine step: drains the transport without
244
+ blocking, decodes XferReqs, grants or NACKs each one, drains NIXL completion
245
+ notifications, and sweeps expired xfers.
246
+ """
247
+
248
+ def __init__(
249
+ self,
250
+ transport: ZmqServerTransport,
251
+ data: DataTransport,
252
+ cache: EmbeddingCache,
253
+ compat_hash: str,
254
+ ) -> None:
255
+ self._transport = transport
256
+ self._data = data
257
+ self._cache = cache # EmbeddingCache — single block owner
258
+ self._compat_hash = compat_hash
259
+
260
+ # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly,
261
+ # so completion notification lookup is a direct dict pop with no parsing.
262
+ self._active_xfers: dict[str, ProducerXfer] = {}
263
+
264
+ self._req_decoder = msgspec.msgpack.Decoder(XferReq)
265
+ self._req_encoder = msgspec.msgpack.Encoder()
266
+
267
+ def close(self) -> None:
268
+ for xfer in self._active_xfers.values():
269
+ self._cache.unpin(xfer.mm_hash) # release any still-held pins
270
+ self._active_xfers.clear()
271
+ self._transport.close()
272
+
273
+ def poll_step(self) -> None:
274
+ """Drain the transport without blocking and process what arrived."""
275
+ self.poll(self._transport.poll(timeout_ms=0))
276
+
277
+ def poll(self, messages: list[tuple[bytes, bytes]]) -> None:
278
+ """Decode and process inbound XferReqs, drain NIXL notifs, sweep timeouts.
279
+
280
+ messages: (identity, payload) pairs from ZmqServerTransport.poll().
281
+ """
282
+ for identity, payload in messages:
283
+ try:
284
+ req = self._req_decoder.decode(payload)
285
+ except (msgspec.DecodeError, msgspec.ValidationError):
286
+ logger.warning("EC Connector: dropped malformed XferReq")
287
+ continue
288
+ ack = self._grant_or_nack(req)
289
+ self._transport.send(identity, self._req_encoder.encode(ack))
290
+
291
+ self._drain_notifs()
292
+ self._sweep_timeouts()
293
+
294
+ def _grant_or_nack(self, req: XferReq) -> XferAck:
295
+ if req.connector_version != EC_CONNECTOR_VERSION:
296
+ logger.warning(
297
+ "EC Connector: incompatible version for mm_hash=%s", req.mm_hash
298
+ )
299
+ return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_VERSION)
300
+ if req.compatibility_hash != self._compat_hash:
301
+ logger.warning(
302
+ "EC Connector: incompatible compat hash for mm_hash=%s", req.mm_hash
303
+ )
304
+ return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INCOMPAT)
305
+ if not req.session_id:
306
+ # Grants are keyed by (session_id, mm_hash), so an omitted id from
307
+ # one peer aliases another's: the first completion would unpin
308
+ # blocks the other peer is still reading.
309
+ logger.warning(
310
+ "EC Connector: XferReq without session_id for mm_hash=%s", req.mm_hash
311
+ )
312
+ return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INTERNAL)
313
+ entry = self._cache.pin_if_ready(req.mm_hash)
314
+ if entry is None or not entry.ready:
315
+ status = (
316
+ XferStatus.NACK_MISSING if entry is None else XferStatus.NACK_NOT_READY
317
+ )
318
+ logger.log(
319
+ logging.INFO if status in EXPECTED_NACKS else logging.WARNING,
320
+ "EC Connector: mm_hash=%s not serveable in local cache (%s); NACKing",
321
+ req.mm_hash,
322
+ status.name,
323
+ )
324
+ return XferAck(mm_hash=req.mm_hash, status=status)
325
+ block_indices = list(entry.block_ids)
326
+ key = f"{req.session_id}:{req.mm_hash}"
327
+ if key in self._active_xfers:
328
+ # This peer is re-requesting a grant it still holds, having timed
329
+ # out waiting. Overwriting it drops the pin taken for it, which
330
+ # neither a completion nor the sweep would ever release.
331
+ self._cache.unpin(req.mm_hash)
332
+ self._active_xfers[key] = ProducerXfer(
333
+ mm_hash=req.mm_hash,
334
+ block_indices=block_indices,
335
+ deadline=time.monotonic() + PRODUCER_PIN_LEASE_S,
336
+ )
337
+ logger.debug(
338
+ "EC producer: serving NIXL READ request mm_hash=%s key=%s blocks=%d",
339
+ req.mm_hash,
340
+ key,
341
+ len(block_indices),
342
+ )
343
+ return XferAck(
344
+ mm_hash=req.mm_hash,
345
+ status=XferStatus.OK,
346
+ session_id=req.session_id, # echo consumer's session_id back
347
+ src_block_indices=block_indices,
348
+ agent_metadata=self._data.get_agent_metadata(),
349
+ mem_descriptor=self._data.get_mem_descriptor(),
350
+ )
351
+
352
+ def _drain_notifs(self) -> None:
353
+ try:
354
+ notifs = self._data.get_new_notifs()
355
+ except Exception:
356
+ logger.exception("EC Connector: get_new_notifs failed")
357
+ return
358
+ for msgs in notifs.values():
359
+ for msg in msgs:
360
+ try:
361
+ key = msg.decode("utf-8") # "{session_id}:{mm_hash}"
362
+ except Exception:
363
+ logger.exception("EC Connector: Dropped malformed NIXL notif")
364
+ continue
365
+ xfer = self._active_xfers.pop(key, None)
366
+ if xfer is not None:
367
+ self._cache.unpin(xfer.mm_hash)
368
+ logger.debug(
369
+ "EC producer: NIXL READ completed mm_hash=%s key=%s",
370
+ xfer.mm_hash,
371
+ key,
372
+ )
373
+
374
+ def _sweep_timeouts(self) -> None:
375
+ for key, xfer in list(self._active_xfers.items()):
376
+ if xfer.is_expired():
377
+ logger.warning(
378
+ "EC producer: consumer did not complete NIXL READ "
379
+ "mm_hash=%s key=%s within the pin lease; releasing "
380
+ "pinned blocks",
381
+ xfer.mm_hash,
382
+ key,
383
+ )
384
+ self._cache.unpin(xfer.mm_hash)
385
+ del self._active_xfers[key]
386
+
387
+
388
+# ── ConsumerSession ───────────────────────────────────────────────────────────
389
+
390
+
391
+@dataclass
392
+class ConsumerSessionResults:
393
+ """Results from ConsumerSession.take_results()."""
394
+
395
+ completed: set[str] # promote to local cache
396
+ tombstoned: set[str] # free blocks + prevent retry this step
397
+ quarantined: set[str] # timed-out, DMA still running: prevent retry
398
+ # but DO NOT free blocks — session holds them
399
+ retryable: set[str] # free blocks, allow another attempt (not ready, peer down)
400
+ settled: list[tuple[str, list[int]]] # (mm_hash, block_indices) quarantine cleared
401
+
402
+
403
+class ConsumerSession:
404
+ """Long-lived per-producer session on the consumer side.
405
+
406
+ Owns the ZMQ DEALER connection and the NIXL agent registration for one
407
+ producer peer. The connector calls poll(messages, now) once per engine
408
+ step with the raw bytes collected from this peer's DEALER socket.
409
+
410
+ poll() decodes XferAcks, dispatches them to the right ConsumerXfer,
411
+ advances all active xfer state machines, and drains quarantined xfers.
412
+ Results accumulate in internal sets; call take_results() to collect them.
413
+ """
414
+
415
+ def __init__(
416
+ self,
417
+ addr: PeerAddr,
418
+ zmq_conn: ZmqClientConnection,
419
+ transport: ZmqClientTransport,
420
+ data: DataTransport,
421
+ compat_hash: str,
422
+ ) -> None:
423
+ self._addr = addr
424
+ self._zmq = zmq_conn
425
+ self._transport = transport
426
+ self._data = data
427
+ self._compat_hash = compat_hash
428
+ self._session_id = str(
429
+ uuid.uuid4()
430
+ ) # stable identity for (session_id:mm_hash) keys
431
+
432
+ self._nixl_agent_name: str | None = None
433
+ self._nixl_metadata_bytes: bytes | None = None
434
+
435
+ self._xfers: dict[str, ConsumerXfer] = {} # mm_hash → xfer
436
+ self._quarantined: list[ConsumerXfer] = [] # DMA pending, to poll
437
+ self._newly_quarantined: set[str] = set() # quarantined since last take_results
438
+
439
+ self._completed: set[str] = set()
440
+ self._tombstoned: set[str] = set()
441
+ self._retryable: set[str] = set()
442
+ self._settled: list[tuple[str, list[int]]] = []
443
+
444
+ self._encoder = msgspec.msgpack.Encoder()
445
+ self._decoder = msgspec.msgpack.Decoder(XferAck)
446
+
447
+ @property
448
+ def addr(self) -> PeerAddr:
449
+ return self._addr
450
+
451
+ def start_xfer(
452
+ self, mm_hash: str, block_indices: list[int], deadline: float
453
+ ) -> None:
454
+ """Create a ConsumerXfer and send the XferReq over the connection."""
455
+ xfer = ConsumerXfer(
456
+ mm_hash=mm_hash,
457
+ block_indices=block_indices,
458
+ addr=self._addr,
459
+ deadline=deadline,
460
+ data=self._data,
461
+ consumer_session_id=self._session_id,
462
+ )
463
+ self._zmq.send(
464
+ self._encoder.encode(
465
+ XferReq(
466
+ mm_hash=mm_hash,
467
+ compatibility_hash=self._compat_hash,
468
+ session_id=self._session_id,
469
+ )
470
+ )
471
+ )
472
+ self._xfers[mm_hash] = xfer
473
+ logger.debug(
474
+ "EC Connector: XferReq sent mm_hash=%s to %s:%d",
475
+ mm_hash,
476
+ self._addr[0],
477
+ self._addr[1],
478
+ )
479
+
480
+ def poll(self, messages: list[bytes], now: float) -> None:
481
+ """Decode inbound XferAcks, dispatch to xfers, and advance all state machines.
482
+
483
+ messages: raw bytes from ZmqClientTransport.poll() for this peer.
484
+ Results accumulate until take_results() is called.
485
+ """
486
+ for raw in messages:
487
+ try:
488
+ ack = self._decoder.decode(raw)
489
+ except (msgspec.DecodeError, msgspec.ValidationError):
490
+ logger.warning(
491
+ "EC Connector: dropped malformed XferAck from %s:%d",
492
+ self._addr[0],
493
+ self._addr[1],
494
+ )
495
+ continue
496
+ self._dispatch_ack(ack)
497
+
498
+ for mm_hash, xfer in list(self._xfers.items()):
499
+ state = xfer.poll(now)
500
+ if state == XferState.DONE:
501
+ self._completed.add(mm_hash)
502
+ del self._xfers[mm_hash]
503
+ elif state in (XferState.ACK_TIMEOUT, XferState.READ_FAILED):
504
+ self._tombstoned.add(mm_hash)
505
+ del self._xfers[mm_hash]
506
+ elif state == XferState.QUARANTINED:
507
+ # DMA may still be writing — caller must NOT free blocks.
508
+ self._quarantined.append(xfer)
509
+ self._newly_quarantined.add(mm_hash)
510
+ del self._xfers[mm_hash]
511
+
512
+ self._drain_quarantine(now)
513
+
514
+ def _dispatch_ack(self, ack: XferAck) -> None:
515
+ xfer = self._xfers.get(ack.mm_hash)
516
+ if xfer is None or xfer.transfer_handle is not None:
517
+ return # unknown mm_hash or duplicate ack
518
+
519
+ if ack.status != XferStatus.OK:
520
+ logger.log(
521
+ logging.INFO if ack.status in EXPECTED_NACKS else logging.WARNING,
522
+ "EC Connector: NACK %s from %s:%d for mm_hash=%s",
523
+ ack.status.name,
524
+ self._addr[0],
525
+ self._addr[1],
526
+ ack.mm_hash,
527
+ )
528
+ del self._xfers[ack.mm_hash]
529
+ # The producer announced this encoding and its save is still in
530
+ # flight, so another attempt on a later step can be granted. Every
531
+ # other NACK is terminal for this mm_hash.
532
+ if ack.status in RETRYABLE_NACKS:
533
+ self._retryable.add(ack.mm_hash)
534
+ else:
535
+ self._tombstoned.add(ack.mm_hash)
536
+ return
537
+
538
+ try:
539
+ agent_name = self._ensure_registered(ack.agent_metadata, ack.mem_descriptor)
540
+ except Exception:
541
+ logger.exception(
542
+ "EC Connector: failed to register peer %s:%d for mm_hash=%s",
543
+ self._addr[0],
544
+ self._addr[1],
545
+ ack.mm_hash,
546
+ )
547
+ del self._xfers[ack.mm_hash]
548
+ self._tombstoned.add(ack.mm_hash)
549
+ return
550
+
551
+ if not xfer.handle_ack(ack, agent_name):
552
+ del self._xfers[ack.mm_hash]
553
+ self._tombstoned.add(ack.mm_hash)
554
+
555
+ def _ensure_registered(self, metadata: bytes, mem_descriptor: bytes) -> str:
556
+ """Register or re-register the peer in DataTransport. Returns agent_name."""
557
+ if self._nixl_agent_name is not None and self._nixl_metadata_bytes == metadata:
558
+ return self._nixl_agent_name
559
+ if self._nixl_agent_name is not None:
560
+ logger.info(
561
+ "EC Connector: producer %s:%d restarted; re-registering",
562
+ self._addr[0],
563
+ self._addr[1],
564
+ )
565
+ self._data.remove_remote_peer(self._nixl_agent_name)
566
+ agent_name = self._data.add_remote_peer(metadata, mem_descriptor)
567
+ self._nixl_agent_name = agent_name
568
+ self._nixl_metadata_bytes = metadata
569
+ logger.debug(
570
+ "EC Connector: registered peer %s:%d agent=%s",
571
+ self._addr[0],
572
+ self._addr[1],
573
+ agent_name,
574
+ )
575
+ return agent_name
576
+
577
+ def _drain_quarantine(self, now: float) -> None:
578
+ still_pending: list[ConsumerXfer] = []
579
+ for xfer in list(self._quarantined):
580
+ state = xfer.poll(now)
581
+ if state == XferState.SETTLED:
582
+ self._settled.append((xfer.mm_hash, xfer.block_indices))
583
+ else:
584
+ still_pending.append(xfer)
585
+ self._quarantined = still_pending
586
+
587
+ def on_peer_down(self) -> None:
588
+ """WAITING_ACK xfers → cancel (retry allowed); READING xfers → quarantine."""
589
+ for mm_hash, xfer in list(self._xfers.items()):
590
+ if xfer.transfer_handle is None:
591
+ xfer.cancel()
592
+ self._retryable.add(mm_hash)
593
+ else:
594
+ self._quarantined.append(xfer)
595
+ self._newly_quarantined.add(mm_hash)
596
+ self._xfers.clear()
597
+
598
+ def take_results(self) -> ConsumerSessionResults:
599
+ """Return and clear all accumulated results since the last call."""
600
+ results = ConsumerSessionResults(
601
+ completed=self._completed,
602
+ tombstoned=self._tombstoned,
603
+ quarantined=self._newly_quarantined,
604
+ retryable=self._retryable,
605
+ settled=self._settled,
606
+ )
607
+ self._completed = set()
608
+ self._tombstoned = set()
609
+ self._newly_quarantined = set()
610
+ self._retryable = set()
611
+ self._settled = []
612
+ return results
613
+
614
+ def has_in_flight(self) -> bool:
615
+ return bool(self._xfers) or bool(self._quarantined)
616
+
617
+ def close(self) -> None:
618
+ """Release all transfer handles, deregister NIXL peer, close connection.
619
+
620
+ Connection teardown is routed through the transport so the pooled
621
+ connection is dropped in lock-step with the session — a session and its
622
+ DEALER connection have a 1:1 lifecycle, so a later connect() to the same
623
+ peer builds a fresh socket instead of handing back this dead one.
624
+ """
625
+ for xfer in self._xfers.values():
626
+ xfer.release()
627
+ self._xfers.clear()
628
+ for xfer in self._quarantined:
629
+ xfer.release()
630
+ self._quarantined.clear()
631
+ if self._nixl_agent_name is not None:
632
+ self._data.remove_remote_peer(self._nixl_agent_name)
633
+ self._transport.remove(self._addr)
vllm/distributed/ec_transfer/ec_connector/cpu/utils.py
new
+47
@@ -0,0 +1,47 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Supporting utilities for the ECCPUConnector scheduler."""
4
+
5
+import msgspec
6
+
7
+from vllm.logger import init_logger
8
+
9
+# (host, port) pair identifying a remote peer's ZMQ side-channel endpoint.
10
+PeerAddr = tuple[str, int]
11
+
12
+logger = init_logger(__name__)
13
+
14
+
15
+# Msgpack (de)serialization for the list-of-(addr, size, device_id)
16
+# descriptor form that NIXL's `get_xfer_descs` expects.
17
+_MemDescList = list[tuple[int, int, int]]
18
+
19
+_mem_desc_encoder = msgspec.msgpack.Encoder()
20
+_mem_desc_decoder = msgspec.msgpack.Decoder(_MemDescList)
21
+
22
+
23
+def serialize_mem_descriptor(descs: _MemDescList) -> bytes:
24
+ return _mem_desc_encoder.encode(descs)
25
+
26
+
27
+def deserialize_mem_descriptor(payload: bytes) -> _MemDescList:
28
+ return _mem_desc_decoder.decode(payload)
29
+
30
+
31
+def build_block_descs(
32
+ base_ptr: int,
33
+ num_blocks: int,
34
+ block_size_bytes: int,
35
+ device_id: int = 0,
36
+) -> _MemDescList:
37
+ """Per-block `(addr, size, device_id)` tuples for an `ECSharedRegion`.
38
+
39
+ Handed to `nixl.get_xfer_descs` to build a dlist where descriptor `i`
40
+ addresses block `i` in the mmap. This makes a WRITE to block indices
41
+ `[i1, i2, ...]` a simple `make_prepped_xfer(..., [i1, i2])` on both
42
+ ends — no per-request address math.
43
+ """
44
+ return [
45
+ (base_ptr + i * block_size_bytes, block_size_bytes, device_id)
46
+ for i in range(num_blocks)
47
+ ]
vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py
-1
@@ -55,7 +55,6 @@ def _coalesce_runs(
55
split. `EmbeddingCache` hands out ascending ids, which is what makes runs
56
findable at all.
57
"""
58
-
58
ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids))
59
breaks = np.flatnonzero(np.diff(ids) != 1) + 1
60
slots = np.concatenate(([0], breaks))
vllm/distributed/ec_transfer/ec_connector/example_connector.py
+7
-54
@@ -5,7 +5,6 @@ from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Any
6
7
import safetensors
8
-import torch
8
9
from vllm.config import VllmConfig
10
from vllm.distributed.ec_transfer.ec_connector.base import (
@@ -13,8 +12,11 @@ from vllm.distributed.ec_transfer.ec_connector.base import (
12
ECConnectorMetadata,
13
ECConnectorRole,
14
)
15
+from vllm.distributed.ec_transfer.ec_connector.utils import (
16
+ PlaceholderMetadataResolver,
17
+ collect_ec_item_metadata,
18
+)
19
from vllm.logger import init_logger
17
-from vllm.utils.collection_utils import is_list_of
20
from vllm.v1.core.sched.output import SchedulerOutput
21
22
if TYPE_CHECKING:
@@ -52,8 +54,7 @@ class ECExampleConnector(ECConnectorBase):
54
super().__init__(vllm_config=vllm_config, role=role)
55
# req_id -> index
56
self._mm_datas_need_loads: dict[str, int] = {}
55
- self._model_config = vllm_config.model_config
56
- self._metadata_fields_cache: dict[str, set[str]] = {}
57
+ self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config)
58
transfer_config = vllm_config.ec_transfer_config
59
if transfer_config is not None:
60
self._storage_path = transfer_config.get_from_extra_config(
@@ -169,36 +170,6 @@ class ECExampleConnector(ECConnectorBase):
170
self._mm_datas_need_loads.clear()
171
return meta
172
172
- def _placeholder_metadata_fields(self, modality: str) -> set[str]:
173
- """Which processed keys this model needs published for `modality`.
174
-
175
- Read from `MultiModalDataParser.embedding_fields`, the same declaration
176
- the consumer's parser requires, so the two cannot drift. An empty set
177
- means the modality cannot be delivered out of band, and the consumer
178
- will process the media itself.
179
- """
180
- if modality in self._metadata_fields_cache:
181
- return self._metadata_fields_cache[modality]
182
-
183
- fields: set[str] = set()
184
- try:
185
- from vllm.multimodal import MULTIMODAL_REGISTRY
186
-
187
- info = MULTIMODAL_REGISTRY.create_processor(self._model_config).info
188
- fields = info.data_parser.placeholder_metadata_fields(modality)
189
- except Exception:
190
- # Reporting nothing is a safe degradation: the consumer falls back to
191
- # processing the media itself.
192
- logger.warning(
193
- "Could not determine the placeholder metadata fields for "
194
- "modality %s; the consumer will preprocess the media itself.",
195
- modality,
196
- exc_info=True,
197
- )
198
-
199
- self._metadata_fields_cache[modality] = fields
200
- return fields
201
-
173
def request_finished(
174
self,
175
request: "Request",
@@ -214,28 +185,10 @@ class ECExampleConnector(ECConnectorBase):
185
if not self.is_producer:
186
return False, None
187
217
- items = []
218
- for feature in request.mm_features:
219
- metadata = {}
220
- # `data` is None for items served from the processor cache, in which
221
- # case the metadata is unavailable here and the consumer has to fall
222
- # back to processing the media itself.
223
- if feature.data is not None:
224
- wanted = self._placeholder_metadata_fields(feature.modality)
225
- for key, value in feature.data.get_data().items():
226
- if key not in wanted:
227
- continue
228
- if isinstance(value, torch.Tensor):
229
- metadata[key] = value.tolist()
230
- elif is_list_of(value, (int, float)):
231
- # Some metadata (e.g. Qwen3-VL video timestamps) is
232
- # produced as a plain list rather than a tensor.
233
- metadata[key] = value
234
- items.append({"mm_hash": feature.identifier, **metadata})
235
-
188
+ items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver)
189
if not items:
190
return False, None
238
- return False, {"ec_items": items}
191
+ return False, items
192
193
# ==============================
194
# Helper functions
vllm/distributed/ec_transfer/ec_connector/utils.py
+80
@@ -2,8 +2,88 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""EC connector helper utilities."""
4
5
+from typing import TYPE_CHECKING, Any
6
+
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
+from vllm.utils.collection_utils import is_list_of
11
from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
12
13
+if TYPE_CHECKING:
14
+ from vllm.config import ModelConfig
15
+ from vllm.multimodal.inputs import MultiModalFeatureSpec
16
+
17
+logger = init_logger(__name__)
18
+
19
+
20
+class PlaceholderMetadataResolver:
21
+ """Resolves which processed keys a model needs published per modality.
22
+
23
+ Reads `MultiModalDataParser.embedding_fields`, the same declaration the
24
+ consumer's parser requires, so the two cannot drift. An empty set means
25
+ the modality cannot be delivered out of band, and the consumer will
26
+ process the media itself.
27
+ """
28
+
29
+ def __init__(self, model_config: "ModelConfig") -> None:
30
+ self._model_config = model_config
31
+ self._cache: dict[str, set[str]] = {}
32
+
33
+ def fields_for(self, modality: str) -> set[str]:
34
+ if modality in self._cache:
35
+ return self._cache[modality]
36
+
37
+ fields: set[str] = set()
38
+ try:
39
+ from vllm.multimodal import MULTIMODAL_REGISTRY
40
+
41
+ info = MULTIMODAL_REGISTRY.create_processor(self._model_config).info
42
+ fields = info.data_parser.placeholder_metadata_fields(modality)
43
+ except Exception:
44
+ logger.warning(
45
+ "Could not determine the placeholder metadata fields for "
46
+ "modality %s; the consumer will preprocess the media itself.",
47
+ modality,
48
+ exc_info=True,
49
+ )
50
+
51
+ self._cache[modality] = fields
52
+ return fields
53
+
54
+
55
+def collect_ec_item_metadata(
56
+ mm_features: "list[MultiModalFeatureSpec]",
57
+ resolver: PlaceholderMetadataResolver,
58
+) -> dict[str, dict[str, Any]]:
59
+ """Build one `ec_transfer_params` entry per feature for `request_finished()`.
60
+
61
+ Keyed by mm_hash, each entry carries a `metadata` dict with whatever
62
+ placeholder fields `resolver` says this model needs published for its
63
+ modality, so a consumer can skip the image transform. `data` is None for
64
+ items served from the processor cache, in which case the metadata is
65
+ unavailable here and the consumer has to fall back to processing the
66
+ media itself. A connector that also has transfer coordinates to report
67
+ (e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
68
+ `metadata`, not into it.
69
+ """
70
+ items: dict[str, dict[str, Any]] = {}
71
+ for feature in mm_features:
72
+ metadata: dict[str, Any] = {}
73
+ if feature.data is not None:
74
+ wanted = resolver.fields_for(feature.modality)
75
+ for key, value in feature.data.get_data().items():
76
+ if key not in wanted:
77
+ continue
78
+ if isinstance(value, torch.Tensor):
79
+ metadata[key] = value.tolist()
80
+ elif is_list_of(value, (int, float)):
81
+ # Some metadata (e.g. Qwen3-VL video timestamps) is
82
+ # produced as a plain list rather than a tensor.
83
+ metadata[key] = value
84
+ items[feature.identifier] = {"metadata": metadata}
85
+ return items
86
+
87
88
class ECOutputAggregator:
89
"""Merge every worker's EC connector output onto the single