[EPD] Add ECMooncakeConnector for encoder cache over Mooncake TransferEngine (#41567)
Signed-off-by: Teng Ma <sima.mt@alibaba-inc.com> Signed-off-by: Tianyu Guo <guoty@inferact.ai> Signed-off-by: Zhou ziheng <jiaranran2@gmail.com> Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: Zhou ziheng <jiaranran2@gmail.com> Co-authored-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Teng Ma committed
Sep 7, 2026 at 14:21 UTC
6fbb00b18874e27ba7d7adc0a3b8e93fee763ab1
30 files changed
+5777
-304
.buildkite/test_areas/disaggregated_mooncake.yaml
+54
@@ -19,3 +19,57 @@ steps:
19
commands:
20
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
21
- bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
22
+
23
+- label: ":nvidia: (L4) Mooncake EC TCP E2E"
24
+ key: mooncake-ec-tcp-e2e-2-gpus
25
+ timeout_in_minutes: 30
26
+ working_dir: "/vllm-workspace"
27
+ device: l4
28
+ num_devices: 2
29
+ source_file_dependencies:
30
+ - vllm/distributed/ec_transfer/
31
+ - vllm/config/ec_transfer.py
32
+ - vllm/config/multimodal.py
33
+ - vllm/config/vllm.py
34
+ - vllm/multimodal/
35
+ - vllm/v1/core/encoder_cache_manager.py
36
+ - vllm/v1/core/sched/
37
+ - vllm/v1/engine/
38
+ - vllm/v1/worker/ec_connector_model_runner_mixin.py
39
+ - vllm/v1/worker/gpu_model_runner.py
40
+ - vllm/v1/worker/gpu_worker.py
41
+ - vllm/v1/worker/gpu/ec_connector.py
42
+ - vllm/v1/worker/gpu/model_runner.py
43
+ - vllm/v1/worker/gpu/mm/
44
+ - examples/disaggregated/disaggregated_encoder/
45
+ - tests/v1/ec_connector/
46
+ - requirements/kv_connectors.txt
47
+ env:
48
+ PYTHON_BIN: "/vllm-workspace/.venv/bin/python"
49
+ PYTHONUNBUFFERED: "1"
50
+ VLLM_HOST_IP: "127.0.0.1"
51
+ VLLM_USE_V2_MODEL_RUNNER: "1"
52
+ MOONCAKE_EC_PROTOCOL: "tcp"
53
+ USE_MM_PROMPTS: "1"
54
+ SKIP_BASELINE: "0"
55
+ CONCURRENCY: "3"
56
+ REPEAT: "2"
57
+ LOG_PATH: "/tmp/mooncake-ec-e2e"
58
+ BASELINE_FILE: "/tmp/mooncake-ec-e2e/baseline.json"
59
+ commands:
60
+ - uv venv --system-site-packages --python 3.12 .venv
61
+ - |
62
+ uv pip install --python .venv/bin/python "$(.venv/bin/python - <<'PY'
63
+ from pathlib import Path
64
+ import torch
65
+
66
+ requirement = next(
67
+ line for line in Path("requirements/kv_connectors.txt").read_text().splitlines()
68
+ if line.startswith("mooncake-transfer-engine ")
69
+ )
70
+ if torch.version.cuda.split(".")[0] == "13":
71
+ requirement = requirement.replace("mooncake-transfer-engine", "mooncake-transfer-engine-cuda13")
72
+ print(requirement)
73
+ PY
74
+ )"
75
+ - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+336
-95
@@ -23,6 +23,7 @@ import argparse
23
import asyncio
24
import hashlib
25
import io
26
+import itertools
27
import json
28
import logging
29
import os
@@ -30,6 +31,7 @@ import random
31
import time
32
import uuid
33
from collections.abc import AsyncIterator
34
+from typing import Any
35
36
import aiohttp
37
import pybase64 as base64
@@ -41,9 +43,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
43
# FastAPI app & global state
44
###############################################################################
45
44
-logging.basicConfig(
45
- level=logging.DEBUG, format="%(asctime)s %(levelname)s: %(message)s"
46
-)
46
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
47
logger = logging.getLogger("proxy")
48
49
app = FastAPI()
@@ -87,9 +87,24 @@ def encoder_rr_assignment(
87
return urls, next_start
88
89
90
+def validate_ec_consumer_routing(
91
+ prefill_urls: list[str], consumer_addrs: list[str]
92
+) -> None:
93
+ """Reject the topology whose EC destination cannot be routed safely."""
94
+ if prefill_urls and consumer_addrs:
95
+ raise ValueError(
96
+ "Mooncake EC consumer routing supports E+PD only; disable independent "
97
+ "prefill or omit --ec-consumer-zmq-addrs."
98
+ )
99
+
100
+
101
# Diagnostic switch: forward the original request to the decoder so the
102
# only difference from the rewrite path is the rewrite itself.
103
NO_REWRITE = False
104
+# Decode-side retries for a retryable internal error (`finish_reason="error"`,
105
+# e.g. an encoder embedding the connector could not deliver). Re-issuing runs
106
+# the encode again, which produces a fresh transfer.
107
+DECODE_RETRIES = 1
108
109
110
# Grid metadata reported by the encoder instance, keyed by item index.
@@ -136,6 +151,7 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
151
here -- a second derivation could disagree with the encoder's.
152
"""
153
rewritten = 0
154
+ transfer_items = []
155
idx = 0
156
new_messages = []
157
for msg in req_data.get("messages", []):
@@ -151,26 +167,42 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
167
meta = dict(item_meta.get(idx) or {})
168
idx += 1
169
item_uuid = meta.pop("mm_hash", None)
170
+ ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid
171
+ transfer_id = meta.pop("transfer_id", None)
172
# Whatever keys the encoder reported are the metadata its model
173
# declared as needed to size the placeholder range; the proxy does
174
# not need to know their names.
175
metadata = {k: _b64_tensor(v) for k, v in meta.items()}
176
if not metadata or not item_uuid:
159
- # The encoder reported no metadata (e.g. the item came from its
160
- # processor cache); let the decoder process the media itself.
177
+ # Nothing to size the placeholder range with. A processor cache
178
+ # hit is not a cause on its own: with the default `lru` type the
179
+ # engine restores the item before the scheduler reports it. It
180
+ # goes missing when the encode request failed, or under
181
+ # `--mm-processor-cache-type shm`, where a hit replaces the item
182
+ # with its shared-memory address and only the worker restores
183
+ # it. Send the media so the decoder can derive the grid itself.
184
new_content.append(item)
185
continue
186
embeds_type = EMBEDS_TYPES[item["type"]]
187
new_content.append(
188
{"type": embeds_type, embeds_type: metadata, "uuid": item_uuid}
189
)
190
+ if transfer_id is not None:
191
+ transfer_items.append(
192
+ {"mm_hash": ec_mm_hash, "transfer_id": transfer_id}
193
+ )
194
rewritten += 1
195
new_messages.append({**msg, "content": new_content})
196
197
if not rewritten:
198
return req_data
199
logger.info("Rewrote %d media item(s) as metadata references", rewritten)
173
- return {**req_data, "messages": new_messages}
200
+ rewritten_request = {**req_data, "messages": new_messages}
201
+ if transfer_items:
202
+ ec_transfer_params = dict(req_data.get("ec_transfer_params") or {})
203
+ ec_transfer_params["ec_items"] = transfer_items
204
+ rewritten_request["ec_transfer_params"] = ec_transfer_params
205
+ return rewritten_request
206
207
208
def extract_mm_items(request_data: dict) -> list[dict]:
@@ -196,7 +228,8 @@ async def fanout_encoder_primer(
228
orig_request: dict,
229
e_urls: list[str],
230
req_id: str,
199
-) -> dict[int, dict]:
231
+ consumer_zmq: str | None = None,
232
+) -> tuple[dict[int, dict], dict[str, Any]]:
233
"""
234
1. Build one request *per MM item* with all text removed.
235
2. Send them concurrently to the encode cluster.
@@ -206,19 +239,25 @@ async def fanout_encoder_primer(
239
`ec_transfer_params`: its EC cache key and the grid its processor produced.
240
The proxy still supplies the uuid so both sides key the cache the same way;
241
the grid can only come from the encoder, which is the side that computed it.
242
+
243
+ Also returns the connector handles to put on the decode body, as a fresh
244
+ mapping. `orig_request` is left untouched so a retry re-encodes from the
245
+ original request instead of carrying the previous attempt's handles.
246
"""
247
logger.info("[%s] Processing multimodal items...", req_id)
248
249
mm_items = extract_mm_items(orig_request)
250
if not mm_items:
251
logger.info("[%s] No multimodal items, skipping encoder", req_id)
215
- return {} # nothing to do
252
+ return {}, {} # nothing to do
253
254
logger.info("[%s] got %d multimodal items...", req_id, len(mm_items))
255
256
tasks = []
257
item_uuids: dict[int, str] = {}
258
+ item_transfer_ids: dict[int, str] = {}
259
item_meta: dict[int, dict] = {}
260
+ ec_params: dict[str, Any] = {}
261
262
# Round-robin over encode servers to distribute load a bit. The cursor
263
# persists across requests so fan-out doesn't restart at e_urls[0] every
@@ -241,6 +280,8 @@ async def fanout_encoder_primer(
280
item_uuid = None if NO_REWRITE else content_uuid(item)
281
if item_uuid is not None:
282
item_uuids[idx] = item_uuid
283
+ transfer_id = uuid.uuid4().hex
284
+ item_transfer_ids[idx] = transfer_id
285
286
encoder_req = {
287
# You *may* need to keep additional fields
@@ -257,6 +298,19 @@ async def fanout_encoder_primer(
298
# once the prompt is encoded and its embeddings are published.
299
"stream": False,
300
}
301
+ if consumer_zmq is not None:
302
+ # No mm_hash here on purpose. The encoder's own
303
+ # `mm_features[i].identifier` is derived from the uuid *and* the
304
+ # engine's media_io_kwargs / mm_processor_kwargs, so this proxy
305
+ # cannot know it before the encoder runs. Sending the bare uuid
306
+ # would fail the connector's hash match and make the producer
307
+ # invent its own transfer id, which the consumer could then never
308
+ # cancel. Omitting it lets the connector match by position, which
309
+ # is exact: one encoder request carries exactly one item.
310
+ encoder_req["ec_transfer_params"] = {
311
+ "consumer_zmq": consumer_zmq,
312
+ "ec_items": [{"transfer_id": transfer_id}],
313
+ }
314
tasks.append(
315
encode_session.post(
316
f"{target_url}/v1/chat/completions",
@@ -304,24 +358,42 @@ async def fanout_encoder_primer(
358
except Exception:
359
logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
360
params = {}
307
- if idx in item_uuids:
308
- reported = params.get(item_uuids[idx])
361
+ if params:
362
+ # One encoder request carries exactly one item, so there is a
363
+ # single reported entry. Do not key it by this proxy's uuid: when
364
+ # media_io_kwargs or mm_processor_kwargs are set the engine
365
+ # re-hashes the uuid together with them, so the encoder's own
366
+ # `mm_features[i].identifier` is a derived value this proxy cannot
367
+ # predict. Fall back to the sole entry, and carry the key the
368
+ # encoder actually used through as `ec_mm_hash`.
369
+ ec_mm_hash = item_uuids.get(idx)
370
+ reported = params.get(ec_mm_hash)
371
+ if reported is None and len(params) == 1:
372
+ ((ec_mm_hash, reported),) = params.items()
373
if reported:
374
metadata = reported.get("metadata") or {}
375
if metadata:
312
- item_meta[idx] = {**metadata, "mm_hash": item_uuids[idx]}
376
+ item_meta[idx] = {
377
+ **metadata,
378
+ "mm_hash": item_uuids.get(idx, ec_mm_hash),
379
+ "ec_mm_hash": ec_mm_hash,
380
+ }
381
+ if idx in item_transfer_ids:
382
+ item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
383
# Whatever the encoder reported alongside `metadata` is the
384
# connector's own handle on the published embedding (for NIXL,
385
# peer_host/peer_port/size_bytes). The decoder's connector
386
# 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
- )
387
+ ec_params[item_uuids.get(idx, ec_mm_hash)] = reported
388
+ if NO_REWRITE and consumer_zmq is not None:
389
+ ec_params.setdefault("ec_items", []).append(
390
+ {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]}
391
+ )
392
393
logger.info(
394
"[%s] All %d encoder requests completed successfully", req_id, len(mm_items)
395
)
324
- return item_meta
396
+ return item_meta, ec_params
397
398
399
async def maybe_prefill(
@@ -331,8 +403,11 @@ async def maybe_prefill(
403
) -> dict:
404
"""
405
- Do prefill-only task if p_url exist;
334
- - Return modified request data with kv transfer params (for nixl connector)
406
+ - Return a new body carrying kv transfer params (for nixl connector)
407
- Else, skip and return the original request data for decode
408
+
409
+ `req_data` is never mutated: a decode retry re-enters this function with the
410
+ same body, and one attempt's `remote_block_ids` must not reach the next.
411
"""
412
if p_url:
413
logger.info("[%s] Processing through prefill: %s", req_id, p_url)
@@ -342,11 +417,9 @@ async def maybe_prefill(
417
prefill_response_json = await prefill_response.json()
418
kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
419
if kv_transfer_params:
345
- req_data["kv_transfer_params"] = kv_transfer_params
420
+ return {**req_data, "kv_transfer_params": kv_transfer_params}
421
347
- return req_data
348
- else:
349
- return req_data
422
+ return req_data
423
424
425
async def process_prefill_stage(
@@ -409,7 +482,6 @@ async def process_prefill_stage(
482
###############################################################################
483
484
412
-@app.middleware("http")
485
async def log_requests(request: Request, call_next):
486
"""Middleware to log all incoming requests and responses"""
487
req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
@@ -458,7 +530,20 @@ async def log_requests(request: Request, call_next):
530
async def on_startup() -> None:
531
global encode_session, prefill_session, decode_session
532
timeout = aiohttp.ClientTimeout(total=100_000)
461
- connector = aiohttp.TCPConnector(limit=0, force_close=False)
533
+ # vLLM closes an idle keep-alive connection after
534
+ # VLLM_HTTP_TIMEOUT_KEEP_ALIVE seconds (5 by default), while aiohttp keeps
535
+ # pooling it for 15. Reusing one it has already closed fails the request
536
+ # with ServerDisconnectedError, and the server logs nothing at all: it
537
+ # closed the socket before the request arrived. Retire ours first.
538
+ server_keep_alive = float(os.getenv("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5"))
539
+ connector = aiohttp.TCPConnector(
540
+ limit=0,
541
+ **(
542
+ {"keepalive_timeout": server_keep_alive / 2}
543
+ if server_keep_alive > 0
544
+ else {"force_close": True}
545
+ ),
546
+ )
547
encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
548
if app.state.p_urls:
549
# only setup if prefill instance(s) exist
@@ -482,40 +567,97 @@ async def on_shutdown() -> None:
567
###############################################################################
568
569
570
+async def prepare_for_decode(
571
+ req_data: dict,
572
+ req_id: str,
573
+ e_urls: list[str],
574
+ p_url: str,
575
+ consumer_zmq: str | None,
576
+) -> tuple[dict, float, float]:
577
+ """Encode, rewrite and prefill, returning the body to send to decode.
578
+
579
+ `req_data` is left untouched so a retry starts from the original media
580
+ rather than from a body whose images are already metadata references.
581
+ """
582
+ _t0 = time.perf_counter()
583
+ item_meta, ec_params = await fanout_encoder_primer(
584
+ req_data, e_urls, req_id, consumer_zmq
585
+ )
586
+ _t1 = time.perf_counter()
587
+ prepared = req_data if NO_REWRITE else rewrite_for_decode(req_data, item_meta)
588
+ if ec_params:
589
+ # A fresh body every time: `rewrite_for_decode` hands back `req_data`
590
+ # itself when it rewrote nothing, and this attempt's handles must not
591
+ # outlive it into a retry.
592
+ handles = dict(prepared.get("ec_transfer_params") or {})
593
+ handles.update(ec_params)
594
+ prepared = {**prepared, "ec_transfer_params": handles}
595
+ _t2 = time.perf_counter()
596
+ prepared = await maybe_prefill(prepared, p_url, req_id)
597
+ return prepared, _t1 - _t0, _t2 - _t1
598
+
599
+
600
async def forward_non_stream(
486
- req_data: dict, req_id: str, e_urls: list[str], p_url: str, d_url: str
601
+ req_data: dict,
602
+ req_id: str,
603
+ e_urls: list[str],
604
+ p_url: str,
605
+ d_url: str,
606
+ consumer_zmq: str | None,
607
+ dp_rank: int | None = None,
608
) -> dict:
609
try:
489
- # Step 1: Process through Encoder instance (if has MM input)
490
- _t0 = time.perf_counter()
491
- item_meta = await fanout_encoder_primer(req_data, e_urls, req_id)
492
- _t1 = time.perf_counter()
493
- req_data = req_data if NO_REWRITE else rewrite_for_decode(req_data, item_meta)
494
- _t2 = time.perf_counter()
495
-
496
- # Step 2: Process through Prefill instance
497
- req_data = await maybe_prefill(req_data, p_url, req_id)
498
-
499
- # Step 3: Process through Decode instance
500
- logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
501
- headers = {"x-request-id": req_id}
502
-
503
- # Non-streaming response
504
- async with decode_session.post(
505
- f"{d_url}/v1/chat/completions", json=req_data, headers=headers
506
- ) as resp:
507
- resp.raise_for_status()
508
- out = await resp.json()
509
- _t3 = time.perf_counter()
510
- logger.info(
511
- "STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f",
512
- "no-rewrite" if NO_REWRITE else "rewrite",
513
- (_t1 - _t0) * 1e3,
514
- (_t2 - _t1) * 1e3,
515
- (_t3 - _t2) * 1e3,
516
- (_t3 - _t0) * 1e3,
610
+ for attempt in range(DECODE_RETRIES + 1):
611
+ _t0 = time.perf_counter()
612
+ prepared, encode_s, rewrite_s = await prepare_for_decode(
613
+ req_data, req_id, e_urls, p_url, consumer_zmq
614
)
518
- return out
615
+ _t2 = time.perf_counter()
616
+
617
+ logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
618
+ headers = {"x-request-id": req_id}
619
+ if dp_rank is not None:
620
+ headers["X-data-parallel-rank"] = str(dp_rank)
621
+
622
+ async with decode_session.post(
623
+ f"{d_url}/v1/chat/completions", json=prepared, headers=headers
624
+ ) as resp:
625
+ if resp.status >= 400:
626
+ detail = await resp.text()
627
+ # 500 is the decoder's retryable internal error, which
628
+ # includes an encoder embedding it could not obtain. Redoing
629
+ # the encode publishes the item again.
630
+ if resp.status == 500 and attempt < DECODE_RETRIES:
631
+ logger.warning(
632
+ "[%s] Decode returned 500, re-encoding and retrying "
633
+ "(attempt %d/%d): %s",
634
+ req_id,
635
+ attempt + 1,
636
+ DECODE_RETRIES,
637
+ detail[:200],
638
+ )
639
+ continue
640
+ logger.error(
641
+ "[%s] Decode request returned status %s: %s",
642
+ req_id,
643
+ resp.status,
644
+ detail,
645
+ )
646
+ raise HTTPException(status_code=resp.status, detail=detail)
647
+ out = await resp.json()
648
+ _t3 = time.perf_counter()
649
+ logger.info(
650
+ "STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f "
651
+ "attempt=%d",
652
+ "no-rewrite" if NO_REWRITE else "rewrite",
653
+ encode_s * 1e3,
654
+ rewrite_s * 1e3,
655
+ (_t3 - _t2) * 1e3,
656
+ (_t3 - _t0) * 1e3,
657
+ attempt,
658
+ )
659
+ return out
660
+ raise HTTPException(status_code=500, detail="Decode failed after re-encoding")
661
662
except HTTPException:
663
raise
@@ -525,47 +667,66 @@ async def forward_non_stream(
667
668
669
async def forward_stream(
528
- req_data: dict, req_id: str, e_urls: list[str], p_url: str, d_url: str
670
+ req_data: dict,
671
+ req_id: str,
672
+ e_urls: list[str],
673
+ p_url: str,
674
+ d_url: str,
675
+ consumer_zmq: str | None,
676
+ dp_rank: int | None = None,
677
) -> AsyncIterator[str]:
678
try:
531
- # Step 1: Process through Encoder instance (if has MM input)
532
- _t0 = time.perf_counter()
533
- item_meta = await fanout_encoder_primer(req_data, e_urls, req_id)
534
- _t1 = time.perf_counter()
535
- req_data = req_data if NO_REWRITE else rewrite_for_decode(req_data, item_meta)
536
- _t2 = time.perf_counter()
537
-
538
- # Step 2: Process through Prefill instance
539
- req_data = await maybe_prefill(req_data, p_url, req_id)
540
-
541
- # Step 3: Process through Decode instance
542
- logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
543
- headers = {"x-request-id": req_id}
544
-
545
- # Streaming response
546
- _first = None
547
- async with decode_session.post(
548
- f"{d_url}/v1/chat/completions",
549
- json=req_data,
550
- headers=headers,
551
- ) as resp:
552
- resp.raise_for_status()
553
- async for chunk in resp.content.iter_chunked(1024):
554
- if chunk:
555
- if _first is None:
556
- _first = time.perf_counter()
557
- yield chunk.decode("utf-8", errors="ignore")
558
- _t3 = time.perf_counter()
679
+ for attempt in range(DECODE_RETRIES + 1):
680
+ _t0 = time.perf_counter()
681
+ prepared, encode_s, rewrite_s = await prepare_for_decode(
682
+ req_data, req_id, e_urls, p_url, consumer_zmq
683
+ )
684
+ _t2 = time.perf_counter()
685
560
- logger.info(
561
- "STAGE %s encode=%.1f rewrite=%.2f decode_ttfb=%.1f decode_total=%.1f",
562
- "no-rewrite" if NO_REWRITE else "rewrite",
563
- (_t1 - _t0) * 1e3,
564
- (_t2 - _t1) * 1e3,
565
- ((_first or _t3) - _t2) * 1e3,
566
- (_t3 - _t2) * 1e3,
567
- )
568
- logger.info("[%s] Streaming completed", req_id)
686
+ logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
687
+ headers = {"x-request-id": req_id}
688
+ if dp_rank is not None:
689
+ headers["X-data-parallel-rank"] = str(dp_rank)
690
+
691
+ _first = None
692
+ async with decode_session.post(
693
+ f"{d_url}/v1/chat/completions",
694
+ json=prepared,
695
+ headers=headers,
696
+ ) as resp:
697
+ # Retry only before the first chunk: once anything reached the
698
+ # client the response cannot be replaced.
699
+ if resp.status == 500 and attempt < DECODE_RETRIES:
700
+ detail = await resp.text()
701
+ logger.warning(
702
+ "[%s] Decode returned 500 before streaming, re-encoding "
703
+ "and retrying (attempt %d/%d): %s",
704
+ req_id,
705
+ attempt + 1,
706
+ DECODE_RETRIES,
707
+ detail[:200],
708
+ )
709
+ continue
710
+ resp.raise_for_status()
711
+ async for chunk in resp.content.iter_chunked(1024):
712
+ if chunk:
713
+ if _first is None:
714
+ _first = time.perf_counter()
715
+ yield chunk.decode("utf-8", errors="ignore")
716
+ _t3 = time.perf_counter()
717
+
718
+ logger.info(
719
+ "STAGE %s encode=%.1f rewrite=%.2f decode_ttfb=%.1f "
720
+ "decode_total=%.1f attempt=%d",
721
+ "no-rewrite" if NO_REWRITE else "rewrite",
722
+ encode_s * 1e3,
723
+ rewrite_s * 1e3,
724
+ ((_first or _t3) - _t2) * 1e3,
725
+ (_t3 - _t2) * 1e3,
726
+ attempt,
727
+ )
728
+ logger.info("[%s] Streaming completed", req_id)
729
+ return
730
731
except HTTPException:
732
logger.exception("[%s] HTTPException in forward_stream", req_id)
@@ -590,16 +751,29 @@ async def chat_completions(request: Request):
751
752
e_urls = app.state.e_urls # we want the full list for fan-out
753
p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
593
- d_url = random.choice(app.state.d_urls)
754
+ decode_index = random.randrange(len(app.state.d_urls))
755
+ d_url = app.state.d_urls[decode_index]
756
+ dp_size = app.state.ec_consumer_dp_size
757
+ # Round-robin the replica, then name it to both halves: the decoder
758
+ # honours the rank header instead of its own balancer, and the encoder
759
+ # pushes to that replica's control channel. Choosing once here means a
760
+ # decode retry re-encodes to the same replica.
761
+ dp_rank = next(app.state.replica_counter) % dp_size if dp_size > 1 else None
762
+ ec_index = decode_index * dp_size + (dp_rank or 0)
763
+ consumer_zmq = app.state.d_ec_urls[ec_index] if app.state.d_ec_urls else None
764
765
is_streaming = req_data.get("stream", False)
766
767
if is_streaming:
768
return StreamingResponse(
599
- forward_stream(req_data, req_id, e_urls, p_url, d_url),
769
+ forward_stream(
770
+ req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank
771
+ ),
772
media_type="text/event-stream",
773
)
602
- result = await forward_non_stream(req_data, req_id, e_urls, p_url, d_url)
774
+ result = await forward_non_stream(
775
+ req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank
776
+ )
777
return JSONResponse(content=result)
778
779
except HTTPException:
@@ -750,6 +924,14 @@ if __name__ == "__main__":
924
parser = argparse.ArgumentParser()
925
parser.add_argument("--host", default="0.0.0.0")
926
parser.add_argument("--port", type=int, default=8000)
927
+ parser.add_argument(
928
+ "--log-requests",
929
+ action="store_true",
930
+ help=(
931
+ "Log every request in and out, and raise the log level to DEBUG. "
932
+ "Off by default: the proxy is on the request path."
933
+ ),
934
+ )
935
parser.add_argument(
936
"--no-rewrite",
937
action="store_true",
@@ -764,8 +946,8 @@ if __name__ == "__main__":
946
"--prefill-servers-urls",
947
required=True,
948
help=(
767
- 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ',
768
- 'to enable E->P->D, set "disable" or "none" to enable E->PD',
949
+ 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") '
950
+ 'to enable E->P->D, set "disable" or "none" to enable E->PD'
951
),
952
)
953
parser.add_argument(
@@ -773,15 +955,64 @@ if __name__ == "__main__":
955
required=True,
956
help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
957
)
958
+ parser.add_argument(
959
+ "--decode-retries",
960
+ type=int,
961
+ default=1,
962
+ help=(
963
+ "Re-encode and re-send when decode returns 500, which is its "
964
+ "retryable internal error (an undeliverable encoder embedding "
965
+ "among them). 0 disables."
966
+ ),
967
+ )
968
+ parser.add_argument(
969
+ "--ec-consumer-zmq-addrs",
970
+ default="",
971
+ help=(
972
+ "Comma-separated Mooncake EC consumer control addresses, aligned "
973
+ "with --decode-servers-urls. Required for Mooncake EC consumers and "
974
+ "supported only in E+PD mode. With --ec-consumer-dp-size > 1, list "
975
+ "each server's replicas consecutively: s0r0,s0r1,s1r0,s1r1."
976
+ ),
977
+ )
978
+ parser.add_argument(
979
+ "--ec-consumer-dp-size",
980
+ type=int,
981
+ default=1,
982
+ help=(
983
+ "Data-parallel replicas per EC consumer. The proxy picks a replica "
984
+ "round-robin and names it to both halves of the request, because an "
985
+ "encoder push has to land where the request will run."
986
+ ),
987
+ )
988
989
args = parser.parse_args()
990
+ if args.log_requests:
991
+ logging.getLogger().setLevel(logging.DEBUG)
992
+ app.middleware("http")(log_requests)
993
NO_REWRITE = args.no_rewrite
994
+ DECODE_RETRIES = max(0, args.decode_retries)
995
app.state.e_urls = [
996
u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
997
]
998
app.state.d_urls = [
999
u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
1000
]
1001
+ app.state.d_ec_urls = [
1002
+ u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
1003
+ ]
1004
+ if args.ec_consumer_dp_size < 1:
1005
+ parser.error("--ec-consumer-dp-size must be at least 1")
1006
+ app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
1007
+ app.state.replica_counter = itertools.count()
1008
+ expected = len(app.state.d_urls) * args.ec_consumer_dp_size
1009
+ if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
1010
+ parser.error(
1011
+ "--ec-consumer-zmq-addrs must contain one address per consumer "
1012
+ f"replica: expected {expected} "
1013
+ f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), "
1014
+ f"got {len(app.state.d_ec_urls)}"
1015
+ )
1016
# handle prefill instances
1017
if args.prefill_servers_urls.lower() in ("disable", "none", ""):
1018
app.state.p_urls = []
@@ -793,11 +1024,21 @@ if __name__ == "__main__":
1024
u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
1025
]
1026
logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
1027
+ try:
1028
+ validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
1029
+ except ValueError as exc:
1030
+ parser.error(str(exc))
1031
1032
logger.info("Proxy listening on %s:%s", args.host, args.port)
1033
logger.info("Encode servers: %s", app.state.e_urls)
1034
logger.info("Prefill instances %s", app.state.p_urls)
1035
logger.info("Decode servers: %s", app.state.d_urls)
1036
+ if app.state.ec_consumer_dp_size > 1:
1037
+ logger.info(
1038
+ "EC consumer replicas per server: %d (control addresses: %s)",
1039
+ app.state.ec_consumer_dp_size,
1040
+ app.state.d_ec_urls,
1041
+ )
1042
1043
uvicorn.run(
1044
app,
tests/v1/core/test_encoder_cache_manager.py
+32
@@ -167,6 +167,38 @@ def test_get_freed_mm_hashes_clears_freed_list():
167
assert manager.get_freed_mm_hashes() == []
168
169
170
+def test_referencing_a_freeable_entry_protects_it_from_eviction():
171
+ """A waiting request can hold an entry it is not scheduled for yet.
172
+
173
+ Entries with no referent are evictable. A request that is deferred (e.g.
174
+ by an EC connector waiting on another item) still needs the items it
175
+ already has, and a connector that hands out one transfer per request
176
+ cannot always produce the same item a second time.
177
+ """
178
+ manager = EncoderCacheManager(cache_size=10)
179
+ owner = MockRequest("owner", ["a"], [5])
180
+ waiter = MockRequest("waiter", ["a"], [5])
181
+ newcomer = MockRequest("newcomer", ["b"], [6])
182
+
183
+ manager.allocate(owner, 0)
184
+ manager.free_encoder_input(owner, 0)
185
+ assert "a" in manager.freeable
186
+
187
+ # The waiting request references it; it is no longer reclaimable.
188
+ assert manager.check_and_update_cache(waiter, 0)
189
+ assert "a" not in manager.freeable
190
+
191
+ # 'b' no longer fits, and 'a' must not be taken away to make room.
192
+ assert not manager.can_allocate(newcomer, 0, int(1e9), 0)
193
+ assert "a" in manager.cached
194
+ assert manager.get_freed_mm_hashes() == []
195
+
196
+ # Once the waiter is done with it, the entry is reclaimable again.
197
+ manager.free_encoder_input(waiter, 0)
198
+ assert manager.can_allocate(newcomer, 0, int(1e9), 0)
199
+ assert manager.get_freed_mm_hashes() == ["a"]
200
+
201
+
202
def test_reallocated_hash_is_not_reported_as_freed():
203
manager = EncoderCacheManager(cache_size=8)
204
req_a = MockRequest("reqA", ["a"], [4])
tests/v1/core/test_scheduler.py
+89
-3
@@ -5643,6 +5643,68 @@ def test_free_encoder_inputs_respects_unconfirmed_placeholders():
5643
assert manager.get_cached_input_ids(request) == set()
5644
5645
5646
+def test_free_encoder_inputs_notifies_the_ec_connector():
5647
+ """The connector learns the item is consumed, not just that the request ended.
5648
+
5649
+ Connectors hold per-item transfer state (remote buffers, reservations);
5650
+ waiting for `request_finished` pins it for the whole generation.
5651
+ """
5652
+ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
5653
+ mm_start_pos, mm_length = 50, 100
5654
+ request = create_requests(
5655
+ num_requests=1,
5656
+ num_tokens=mm_start_pos + mm_length + 10,
5657
+ mm_positions=[[PlaceholderRange(offset=mm_start_pos, length=mm_length)]],
5658
+ )[0]
5659
+ scheduler.encoder_cache_manager.allocate(request, 0)
5660
+ scheduler.ec_connector = Mock()
5661
+
5662
+ request.num_computed_tokens = mm_start_pos + mm_length - 1
5663
+ scheduler._free_encoder_inputs(request)
5664
+ scheduler.ec_connector.update_state_after_free.assert_not_called()
5665
+
5666
+ request.num_computed_tokens = mm_start_pos + mm_length
5667
+ scheduler._free_encoder_inputs(request)
5668
+ scheduler.ec_connector.update_state_after_free.assert_called_once_with(request, 0)
5669
+
5670
+
5671
+def test_unavailable_encoder_input_fails_the_request_as_retryable():
5672
+ """An encoder input the connector gave up on must end the request.
5673
+
5674
+ `FinishReason.ERROR` is the retryable channel the KV connector already uses
5675
+ for load failures, so the caller can re-issue; deferring instead parked the
5676
+ request until the client timed out.
5677
+ """
5678
+ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
5679
+ request = create_requests(
5680
+ num_requests=1,
5681
+ num_tokens=160,
5682
+ mm_positions=[[PlaceholderRange(offset=50, length=100)]],
5683
+ )[0]
5684
+ scheduler.add_request(request)
5685
+ scheduler.ec_connector = Mock()
5686
+ scheduler.ec_connector.take_unavailable_requests.return_value = {request.request_id}
5687
+ # `request_finished` reports (delay_free, params) for the request teardown.
5688
+ scheduler.ec_connector.request_finished.return_value = (False, None)
5689
+
5690
+ scheduler_output = scheduler.schedule()
5691
+ outputs = scheduler.update_from_output(
5692
+ scheduler_output,
5693
+ ModelRunnerOutput(
5694
+ req_ids=[request.request_id],
5695
+ req_id_to_index={request.request_id: 0},
5696
+ sampled_token_ids=[[]],
5697
+ logprobs=None,
5698
+ prompt_logprobs_dict={},
5699
+ pooler_output=[],
5700
+ ),
5701
+ )
5702
+
5703
+ assert request.status == RequestStatus.FINISHED_ERROR
5704
+ engine_outputs = outputs[0].outputs
5705
+ assert [o.finish_reason for o in engine_outputs] == [FinishReason.ERROR]
5706
+
5707
+
5708
def test_free_encoder_inputs_defers_for_eagle_lookahead():
5709
"""With EAGLE speculative decoding, the encoder input is retained one extra
5710
position so the drafter's +1 look-ahead mm-embedding gather (which reads one
@@ -5840,9 +5902,8 @@ def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector):
5902
5903
# ensure_cache_available must have been called with (request, num_computed_tokens=0)
5904
# for a brand-new request that has no cached tokens yet.
5843
- scheduler.ec_connector.ensure_cache_available.assert_called_once_with(
5844
- request_deferred, 0
5845
- )
5905
+ ensure_call = scheduler.ec_connector.ensure_cache_available.call_args
5906
+ assert ensure_call.args == (request_deferred, 0)
5907
# Deferred request must NOT be scheduled
5908
assert request_deferred.request_id not in output.num_scheduled_tokens
5909
_assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0)
@@ -5874,6 +5935,31 @@ def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector):
5935
_assert_right_encoder_inputs(output, expected_total_reqs=0)
5936
5937
5938
+def test_ec_connector_defers_running_request_for_async_reload():
5939
+ scheduler = create_scheduler(
5940
+ model="llava-hf/llava-1.5-7b-hf",
5941
+ max_num_batched_tokens=32,
5942
+ use_ec_connector=True,
5943
+ ec_role="ec_consumer",
5944
+ )
5945
+ request = create_requests(
5946
+ num_requests=1,
5947
+ num_tokens=128,
5948
+ mm_positions=[[PlaceholderRange(offset=48, length=32)]],
5949
+ req_ids=["request"],
5950
+ )[0]
5951
+ scheduler.ec_connector.ensure_cache_available = Mock(side_effect=[True, False])
5952
+
5953
+ scheduler.add_request(request)
5954
+ first_output = scheduler.schedule()
5955
+ assert first_output.num_scheduled_tokens[request.request_id] == 32
5956
+
5957
+ second_output = scheduler.schedule()
5958
+ assert request.request_id not in second_output.num_scheduled_tokens
5959
+ ensure_call = scheduler.ec_connector.ensure_cache_available.call_args
5960
+ assert ensure_call.args[:2] == (request, 32)
5961
+
5962
+
5963
def test_ec_connector_pending_prefetch_only_checks_future_mm_features():
5964
"""Test that future mm feature filtering only yields features beyond
5965
the computed token frontier.
tests/v1/ec_connector/integration/README.md
+212
-170
@@ -1,171 +1,213 @@
1
-# EPD Correctness Test
2
-
3
-This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance.
4
-
5
-## What It Tests
6
-
7
-- **Baseline**: Single vLLM instance serving a multimodal model
8
-- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance
9
-- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance
10
-- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance
11
-
12
-The test ensures that disaggregated encoding produces **identical** outputs to the baseline.
13
-
14
-Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D
15
-
1
+# EPD Correctness Test
2
+
3
+This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance.
4
+
5
+## What It Tests
6
+
7
+- **Baseline**: Single vLLM instance serving a multimodal model
8
+- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance
9
+- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance
10
+- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance
11
+
12
+The test ensures that disaggregated encoding produces **identical** outputs to the baseline.
13
+
14
+Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D
15
+
16
Please refer to [Disaggregated Encoder Feature](../../../../docs/features/disagg_encoder.md) for the detailed explanation for the EPD features.
17
-
18
-## Files
19
-
20
-- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests)
21
-- `test_epd_correctness.py` - Python test script (compares outputs)
22
-
23
-## Usage
24
-
25
-### Multimodal Prompts (Default)
26
-
27
-```bash
28
-cd vllm
29
-./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
30
-```
31
-
32
-This runs the test with actual multimodal (image) prompts.
33
-
34
-### Text-Only Prompts
35
-
36
-```bash
37
-cd vllm
38
-USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
39
-```
40
-
41
-This runs a quick test with text-only prompts to verify the setup works.
42
-
43
-### Custom Configuration
44
-
45
-```bash
46
-# Use specific GPUs
47
-GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
48
-
49
-# Use specific ports
50
-ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
51
-
52
-# Use specific model
53
-MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
54
-
55
-# Use specific storage path
56
-EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
57
-```
58
-
59
-## How It Works
60
-
61
-### Step 1: Baseline
62
-
63
-1. Start single vLLM instance on GPU
64
-2. Run test prompts (multimodal or text-only)
65
-3. Save outputs to `.vllm_epd_baseline.txt`
66
-4. Shutdown instance
67
-
68
-### Step 2: EPD (1E + 1PD)
69
-
70
-1. Clear encoder cache storage
71
-2. Start instances and proxy
72
-3. Run same test prompts
73
-4. Assert outputs match baseline exactly
74
-5. Shutdown instances
75
-
76
-### Step 3: EPD (1E + 1P + 1D)
77
-
78
-1. Clear encoder cache storage
79
-2. Start instances and proxy
80
-3. Run same test prompts
81
-4. Assert outputs match baseline exactly
82
-5. Shutdown instances
83
-
84
-## Test Scenarios
85
-
86
-### Multimodal Prompts (--use_mm_prompts)
87
-
88
-Tests encoder cache transfer:
89
-
90
-- Single image query
91
-- Multiple images in one request
92
-- Mixed image and text
93
-- Image with detailed questions
94
-
95
-### Text-Only Prompts (default)
96
-
97
-Quick sanity check:
98
-
99
-- Simple text queries
100
-- Text-only explanations
101
-- Verifies proxy routing works
102
-
103
-## Expected Behavior
104
-
105
-### ✅ Test Passes When
106
-
107
-- All disagg outputs match baseline outputs exactly
108
-- No errors during instance startup
109
-- Encoder cache is properly saved and loaded
110
-- Proxy correctly routes requests
111
-
112
-### ❌ Test Fails When
113
-
114
-- Outputs differ between baseline and disagg
115
-- Server startup fails
116
-- Encoder cache not found (should fall back to local execution)
117
-- Proxy routing errors
118
-
119
-## Notes
120
-
121
-- The test uses deterministic generation (`temperature=0.0`, `seed=42`)
122
-- Encoder cache should enable exact output reproduction
123
-- Test cleans up all instances and cache files after completion
124
-- Safe to run multiple times (idempotent)
125
-- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/disaggregated/disaggregated_encoder/README.md`
126
-
127
-## Requirements
128
-
129
-- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline)
130
- - 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now.
131
-- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct)
132
-- Internet access (for accessing vllm test images)
133
-
134
-## Debugging
135
-
136
-### Check Logs
137
-
138
-Logs and baseline output are saved in `/tmp/` by default.
139
-Can be customized by changing the environment variables.
140
-
141
-### Check Encoder Cache
142
-
143
-```bash
144
-# Verify cache files are created
145
-ls -la $EC_SHARED_STORAGE_PATH/
146
-
147
-# Should see directories with mm_hash names
148
-# Each containing encoder_cache.safetensors
149
-```
150
-
151
-### Manual Testing
152
-
153
-Run individual components:
154
-
155
-```bash
156
-# Baseline only
157
-python test_epd_correctness.py \
158
- --service_url http://localhost:8000 \
159
- --model_name Qwen/Qwen2.5-VL-3B-Instruct \
160
- --mode baseline \
161
- --baseline_file test_output.txt \
162
- --use_mm_prompts
163
-
164
-# Disagg only (requires baseline output file!)
165
-python test_epd_correctness.py \
166
- --service_url http://localhost:8000 \
167
- --model_name Qwen/Qwen2.5-VL-3B-Instruct \
168
- --mode disagg \
169
- --baseline_file test_output.txt \
170
- --use_mm_prompts
171
-```
17
+
18
+## Files
19
+
20
+- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests)
21
+- `test_epd_correctness.py` - Python test script (compares outputs)
22
+
23
+## Usage
24
+
25
+### Multimodal Prompts (Default)
26
+
27
+```bash
28
+cd vllm
29
+./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
30
+```
31
+
32
+This runs the test with actual multimodal (image) prompts.
33
+
34
+### Text-Only Prompts
35
+
36
+```bash
37
+cd vllm
38
+USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
39
+```
40
+
41
+This runs a quick test with text-only prompts to verify the setup works.
42
+
43
+### Custom Configuration
44
+
45
+```bash
46
+# Use specific GPUs
47
+GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
48
+
49
+# Use specific ports
50
+ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
51
+
52
+# Use specific model
53
+MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
54
+
55
+# Use specific storage path
56
+EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
57
+```
58
+
59
+## How It Works
60
+
61
+### Mooncake EC (1E + 1PD)
62
+
63
+```bash
64
+PYTHON_BIN="$PWD/.venv/bin/python" \
65
+ bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
66
+```
67
+
68
+Requires two GPUs and Mooncake TransferEngine. TCP is the default transport;
69
+no RDMA-capable network hardware is required.
70
+The script sets `MC_FORCE_TCP=1` in TCP mode so Mooncake cannot auto-select RDMA.
71
+The baseline runs first on GPU 0, followed by E on GPU 0 and PD on GPU 1.
72
+`MOONCAKE_EC_PROTOCOL=rdma` selects RDMA instead; host-specific transport
73
+environment variables should be set by the caller.
74
+
75
+Three black-box cases check fixed short answers and compare with the baseline:
76
+
77
+- One image: read the STOP sign.
78
+- Two different images (including a local file): identify flowers and birds.
79
+- The same image twice: read both STOP signs.
80
+
81
+By default, all three requests run concurrently for two rounds, exercising
82
+shared hashes across requests and reuse after completion. Every response is
83
+compared, not just the final round. Set `CONCURRENCY` and `REPEAT` to override.
84
+Prefix caching is disabled and CUDA graphs remain enabled. This is a small
85
+correctness suite, not a performance benchmark or failure-injection suite.
86
+
87
+`LOG_PATH` and `BASELINE_FILE` select the log directory and reference output.
88
+`SKIP_BASELINE=1` reuses a reference generated with the same model and test
89
+configuration. `USE_MM_PROMPTS=0` only checks text routing, not EC transfer.
90
+
91
+Buildkite runs this script as `mooncake-ec-tcp-e2e-2-gpus` on two L4 GPUs.
92
+The job is defined in `.buildkite/test_areas/disaggregated_mooncake.yaml` and
93
+selected for changes to EC, multimodal processing, scheduler/model-runner
94
+integration, the proxy, or these tests (subject to normal PR CI approval).
95
+It reuses the CI image's Python packages through a system-site-packages venv
96
+and installs the CUDA-compatible Mooncake wheel. It uses loopback networking
97
+and TCP only, without RDMA devices or peer-memory setup.
98
+
99
+The job fails on startup errors, request errors, or answer mismatches. On
100
+failure, the script prints the last 100 lines of each server/proxy log to
101
+the CI job log; full files remain under `LOG_PATH` while the container exists.
102
+
103
+### Step 1: Baseline
104
+
105
+1. Start single vLLM instance on GPU
106
+2. Run test prompts (multimodal or text-only)
107
+3. Save outputs to `.vllm_epd_baseline.txt`
108
+4. Shutdown instance
109
+
110
+### Step 2: EPD (1E + 1PD)
111
+
112
+1. Clear encoder cache storage
113
+2. Start instances and proxy
114
+3. Run same test prompts
115
+4. Assert outputs match baseline exactly
116
+5. Shutdown instances
117
+
118
+### Step 3: EPD (1E + 1P + 1D)
119
+
120
+1. Clear encoder cache storage
121
+2. Start instances and proxy
122
+3. Run same test prompts
123
+4. Assert outputs match baseline exactly
124
+5. Shutdown instances
125
+
126
+## Test Scenarios
127
+
128
+### Multimodal Prompts (--use_mm_prompts)
129
+
130
+Tests encoder cache transfer:
131
+
132
+- Single image query
133
+- Multiple images in one request
134
+- Mixed image and text
135
+- Image with detailed questions
136
+
137
+### Text-Only Prompts (default)
138
+
139
+Quick sanity check:
140
+
141
+- Simple text queries
142
+- Text-only explanations
143
+- Verifies proxy routing works
144
+
145
+## Expected Behavior
146
+
147
+### ✅ Test Passes When
148
+
149
+- All disagg outputs match baseline outputs exactly
150
+- No errors during instance startup
151
+- Encoder cache is properly saved and loaded
152
+- Proxy correctly routes requests
153
+
154
+### ❌ Test Fails When
155
+
156
+- Outputs differ between baseline and disagg
157
+- Server startup fails
158
+- Encoder cache not found (should fall back to local execution)
159
+- Proxy routing errors
160
+
161
+## Notes
162
+
163
+- The test uses deterministic generation (`temperature=0.0`, `seed=42`)
164
+- Encoder cache should enable exact output reproduction
165
+- Test cleans up all instances and cache files after completion
166
+- Safe to run multiple times (idempotent)
167
+- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/disaggregated/disaggregated_encoder/README.md`
168
+
169
+## Requirements
170
+
171
+- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline)
172
+ - 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now.
173
+- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct)
174
+- Internet access (for accessing vllm test images)
175
+
176
+## Debugging
177
+
178
+### Check Logs
179
+
180
+Logs and baseline output are saved in `/tmp/` by default.
181
+Can be customized by changing the environment variables.
182
+
183
+### Check Encoder Cache
184
+
185
+```bash
186
+# Verify cache files are created
187
+ls -la $EC_SHARED_STORAGE_PATH/
188
+
189
+# Should see directories with mm_hash names
190
+# Each containing encoder_cache.safetensors
191
+```
192
+
193
+### Manual Testing
194
+
195
+Run individual components:
196
+
197
+```bash
198
+# Baseline only
199
+python test_epd_correctness.py \
200
+ --service_url http://localhost:8000 \
201
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
202
+ --mode baseline \
203
+ --baseline_file test_output.txt \
204
+ --use_mm_prompts
205
+
206
+# Disagg only (requires baseline output file!)
207
+python test_epd_correctness.py \
208
+ --service_url http://localhost:8000 \
209
+ --model_name Qwen/Qwen2.5-VL-3B-Instruct \
210
+ --mode disagg \
211
+ --baseline_file test_output.txt \
212
+ --use_mm_prompts
213
+```
tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
new
+289
@@ -0,0 +1,289 @@
1
+#!/usr/bin/env bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Full-stack EPD validation with ECMooncakeConnector (Mooncake TransferEngine):
6
+# 1) Single-GPU baseline (multimodal) -> saves baseline JSON
7
+# 2) 1 Encoder + 1 PD with Mooncake EC -> compare outputs to baseline
8
+#
9
+# Usage (from repo root):
10
+# ./tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
11
+#
12
+# Env:
13
+# MODEL HF model id (default: Qwen/Qwen2.5-VL-3B-Instruct)
14
+# GPU_SINGLE / GPU_E / GPU_PD GPU ids (defaults 0 / 0 / 1)
15
+# ENDPOINT_PORT, ENCODE_PORT, PREFILL_DECODE_PORT
16
+# EC_MOONCAKE_RESERVATION_HOST consumer control host (default 127.0.0.1)
17
+# MOONCAKE_EC_PROTOCOL tcp | rdma (default tcp)
18
+# USE_MM_PROMPTS 1 (default) or 0 for text-only quick sanity
19
+# TIMEOUT_SECONDS wait_for_server timeout (default 1200)
20
+# SKIP_BASELINE set to 1 to reuse existing BASELINE_FILE
21
+# CONCURRENCY / REPEAT concurrent requests / rounds (defaults 3 / 2)
22
+# MAX_MODEL_LEN context length (default 16384)
23
+
24
+set -euo pipefail
25
+
26
+GIT_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)
27
+cd "$GIT_ROOT" || exit 1
28
+export PYTHONPATH="${GIT_ROOT}:${PYTHONPATH:-}"
29
+PYTHON_BIN="${PYTHON_BIN:-${GIT_ROOT}/.venv/bin/python}"
30
+
31
+MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
32
+USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
33
+TEST_ARGS=(--concurrency "${CONCURRENCY:-3}" --repeat "${REPEAT:-2}")
34
+if [[ "$USE_MM_PROMPTS" == "1" ]]; then
35
+ TEST_ARGS+=(--use_mm_prompts --mm_smoke_test)
36
+fi
37
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
38
+
39
+GPU_SINGLE="${GPU_SINGLE:-0}"
40
+GPU_E="${GPU_E:-0}"
41
+GPU_PD="${GPU_PD:-1}"
42
+
43
+ENCODE_PORT="${ENCODE_PORT:-19534}"
44
+PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}"
45
+ENDPOINT_PORT="${ENDPOINT_PORT:-10002}"
46
+BASELINE_PORT="${BASELINE_PORT:-10003}"
47
+
48
+EC_MOONCAKE_RESERVATION_HOST="${EC_MOONCAKE_RESERVATION_HOST:-127.0.0.1}"
49
+EC_MOONCAKE_RESERVATION_PORT="${EC_MOONCAKE_RESERVATION_PORT:-19019}"
50
+MOONCAKE_EC_PROTOCOL="${MOONCAKE_EC_PROTOCOL:-tcp}"
51
+export EC_MOONCAKE_RESERVATION_HOST
52
+export EC_MOONCAKE_RESERVATION_PORT
53
+export MOONCAKE_EC_PROTOCOL
54
+if [[ "$MOONCAKE_EC_PROTOCOL" == "tcp" ]]; then
55
+ # TransferEngine may otherwise auto-select RDMA on hosts with an HCA.
56
+ export MC_FORCE_TCP=1
57
+else
58
+ unset MC_FORCE_TCP
59
+fi
60
+
61
+LOG_PATH="${LOG_PATH:-/tmp}"
62
+BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_epd_mooncake_baseline.txt}"
63
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-1200}"
64
+PIDS=()
65
+
66
+mkdir -p "$LOG_PATH"
67
+
68
+VLLM_SERVE=("$PYTHON_BIN" -m vllm.entrypoints.cli.main serve)
69
+
70
+ENC_EC_JSON=$("$PYTHON_BIN" <<PY
71
+import json, os
72
+print(json.dumps({
73
+ "ec_connector": "ECMooncakeConnector",
74
+ "ec_role": "ec_producer",
75
+ "ec_connector_extra_config": {
76
+ "mooncake_protocol": os.environ.get("MOONCAKE_EC_PROTOCOL", "tcp"),
77
+ },
78
+}, separators=(",", ":")))
79
+PY
80
+)
81
+
82
+PD_EC_JSON=$("$PYTHON_BIN" <<PY
83
+import json, os
84
+print(json.dumps({
85
+ "ec_connector": "ECMooncakeConnector",
86
+ "ec_role": "ec_consumer",
87
+ "ec_ip": os.environ["EC_MOONCAKE_RESERVATION_HOST"],
88
+ "ec_port": int(os.environ.get("EC_MOONCAKE_RESERVATION_PORT", "19019")),
89
+ "ec_connector_extra_config": {
90
+ "mooncake_protocol": os.environ.get("MOONCAKE_EC_PROTOCOL", "tcp"),
91
+ },
92
+}, separators=(",", ":")))
93
+PY
94
+)
95
+
96
+EC_MOONCAKE_RESERVATION_ADDR=$("$PYTHON_BIN" <<'PY'
97
+import os
98
+
99
+from vllm.utils.network_utils import make_zmq_path
100
+
101
+print(make_zmq_path(
102
+ "tcp", os.environ["EC_MOONCAKE_RESERVATION_HOST"],
103
+ int(os.environ["EC_MOONCAKE_RESERVATION_PORT"]),
104
+))
105
+PY
106
+)
107
+
108
+wait_for_server() {
109
+ local port=$1
110
+ local pid=$2
111
+ local deadline=$((SECONDS + TIMEOUT_SECONDS))
112
+ while ((SECONDS < deadline)); do
113
+ kill -0 "$pid" 2>/dev/null || return 1
114
+ if curl --max-time 2 -fsS "http://localhost:${port}/health" >/dev/null 2>&1; then
115
+ return 0
116
+ fi
117
+ sleep 2
118
+ done
119
+ return 1
120
+}
121
+
122
+cleanup_instances() {
123
+ if ((${#PIDS[@]} == 0)); then
124
+ return
125
+ fi
126
+ echo "Cleaning up tracked vLLM / proxy processes..."
127
+ for pid in "${PIDS[@]}"; do
128
+ kill "$pid" 2>/dev/null || true
129
+ done
130
+ for _ in {1..10}; do
131
+ local alive=0
132
+ for pid in "${PIDS[@]}"; do
133
+ if kill -0 "$pid" 2>/dev/null; then
134
+ alive=1
135
+ fi
136
+ done
137
+ ((alive == 0)) && break
138
+ sleep 1
139
+ done
140
+ for pid in "${PIDS[@]}"; do
141
+ if kill -0 "$pid" 2>/dev/null; then
142
+ kill -KILL "$pid" 2>/dev/null || true
143
+ fi
144
+ wait "$pid" 2>/dev/null || true
145
+ done
146
+ PIDS=()
147
+}
148
+
149
+finish() {
150
+ local status=$?
151
+ cleanup_instances
152
+ if ((status != 0)); then
153
+ for log in "${LOG_PATH}"/mooncake_epd_*.log; do
154
+ [[ -f "$log" ]] || continue
155
+ echo "=== $log (last 100 lines) ==="
156
+ tail -100 "$log"
157
+ done
158
+ fi
159
+ exit "$status"
160
+}
161
+
162
+trap finish EXIT
163
+trap 'exit 130' INT
164
+trap 'exit 143' TERM
165
+
166
+run_baseline() {
167
+ echo "================================"
168
+ echo "BASELINE (single vLLM, MM if enabled)"
169
+ echo "================================"
170
+ cleanup_instances
171
+ local PORT=$BASELINE_PORT
172
+ echo "Starting baseline on GPU $GPU_SINGLE port $PORT"
173
+ CUDA_VISIBLE_DEVICES="$GPU_SINGLE" "${VLLM_SERVE[@]}" "$MODEL" \
174
+ --port "$PORT" \
175
+ --gpu-memory-utilization 0.75 \
176
+ --max-num-seqs 32 \
177
+ --max-model-len "$MAX_MODEL_LEN" \
178
+ --no-enable-prefix-caching \
179
+ --limit-mm-per-prompt '{"image":2,"video":0}' \
180
+ --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
181
+ >"${LOG_PATH}/mooncake_epd_baseline.log" 2>&1 &
182
+ local BASELINE_PID=$!
183
+ PIDS+=("$BASELINE_PID")
184
+ echo "Waiting for baseline..."
185
+ wait_for_server "$PORT" "$BASELINE_PID" || { echo "Baseline failed to start"; return 1; }
186
+ curl -s "http://127.0.0.1:${PORT}/v1/models" | head -c 200 || true
187
+ echo ""
188
+ "$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
189
+ --service_url "http://localhost:$PORT" \
190
+ --model_name "$MODEL" \
191
+ --mode baseline \
192
+ --baseline_file "$BASELINE_FILE" \
193
+ "${TEST_ARGS[@]}"
194
+ cleanup_instances
195
+}
196
+
197
+run_epd_mooncake() {
198
+ echo "================================"
199
+ echo "EPD 1E + 1PD with ECMooncakeConnector"
200
+ echo "Reservation port on consumer: $EC_MOONCAKE_RESERVATION_PORT"
201
+ echo "Mooncake protocol: $MOONCAKE_EC_PROTOCOL"
202
+ echo "================================"
203
+ cleanup_instances
204
+
205
+ echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT"
206
+ CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
207
+ --port "$ENCODE_PORT" \
208
+ --gpu-memory-utilization 0.35 \
209
+ --mm-tensor-ipc torch_shm \
210
+ --enable-request-id-headers \
211
+ --no-enable-prefix-caching \
212
+ --max-num-batched-tokens "$MAX_MODEL_LEN" \
213
+ --max-num-seqs 32 \
214
+ --max-model-len "$MAX_MODEL_LEN" \
215
+ --limit-mm-per-prompt '{"image":2,"video":0}' \
216
+ --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
217
+ --ec-transfer-config "$ENC_EC_JSON" \
218
+ >"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 &
219
+ local ENCODER_PID=$!
220
+ PIDS+=("$ENCODER_PID")
221
+
222
+ echo "Starting PD on GPU $GPU_PD port $PREFILL_DECODE_PORT"
223
+ CUDA_VISIBLE_DEVICES="$GPU_PD" "${VLLM_SERVE[@]}" "$MODEL" \
224
+ --port "$PREFILL_DECODE_PORT" \
225
+ --gpu-memory-utilization 0.75 \
226
+ --mm-tensor-ipc torch_shm \
227
+ --enable-mm-embeds \
228
+ --enable-request-id-headers \
229
+ --max-num-seqs 32 \
230
+ --max-model-len "$MAX_MODEL_LEN" \
231
+ --no-enable-prefix-caching \
232
+ --limit-mm-per-prompt '{"image":2,"video":0}' \
233
+ --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
234
+ --ec-transfer-config "$PD_EC_JSON" \
235
+ >"${LOG_PATH}/mooncake_epd_pd.log" 2>&1 &
236
+ local PD_PID=$!
237
+ PIDS+=("$PD_PID")
238
+
239
+ echo "Waiting for encoder..."
240
+ wait_for_server "$ENCODE_PORT" "$ENCODER_PID" || { echo "Encoder failed to start"; return 1; }
241
+ echo "Waiting for PD..."
242
+ wait_for_server "$PREFILL_DECODE_PORT" "$PD_PID" || { echo "PD failed to start"; return 1; }
243
+
244
+ echo "Starting EPD proxy on $ENDPOINT_PORT"
245
+ "$PYTHON_BIN" "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \
246
+ --host "0.0.0.0" \
247
+ --port "$ENDPOINT_PORT" \
248
+ --encode-servers-urls "http://localhost:$ENCODE_PORT" \
249
+ --prefill-servers-urls "disable" \
250
+ --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
251
+ --ec-consumer-zmq-addrs \
252
+ "$EC_MOONCAKE_RESERVATION_ADDR" \
253
+ >"${LOG_PATH}/mooncake_epd_proxy.log" 2>&1 &
254
+ local PROXY_PID=$!
255
+ PIDS+=("$PROXY_PID")
256
+
257
+ echo "Waiting for proxy..."
258
+ wait_for_server "$ENDPOINT_PORT" "$PROXY_PID" || { echo "Proxy failed to start"; return 1; }
259
+ curl -s "http://127.0.0.1:${ENDPOINT_PORT}/health" || true
260
+ echo ""
261
+
262
+ "$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
263
+ --service_url "http://localhost:$ENDPOINT_PORT" \
264
+ --model_name "$MODEL" \
265
+ --mode disagg \
266
+ --baseline_file "$BASELINE_FILE" \
267
+ "${TEST_ARGS[@]}"
268
+
269
+ cleanup_instances
270
+}
271
+
272
+echo "================================"
273
+echo "1E + 1PD ECMooncake end-to-end correctness"
274
+echo "MODEL=$MODEL"
275
+echo "================================"
276
+
277
+if [[ "${SKIP_BASELINE:-0}" != "1" ]]; then
278
+ run_baseline
279
+else
280
+ echo "SKIP_BASELINE=1 -> using existing $BASELINE_FILE"
281
+ [[ -f "$BASELINE_FILE" ]] || { echo "Missing baseline file"; exit 1; }
282
+fi
283
+
284
+run_epd_mooncake
285
+
286
+echo "================================"
287
+echo "PASS: Mooncake EC EPD matches baseline"
288
+echo "Logs: ${LOG_PATH}/mooncake_epd_*.log"
289
+echo "================================"
tests/v1/ec_connector/integration/test_epd_correctness.py
+88
-23
@@ -26,6 +26,7 @@ import argparse
26
import json
27
import os
28
import time
29
+from concurrent.futures import ThreadPoolExecutor
30
31
import openai
32
import requests
@@ -138,17 +139,20 @@ def run_chat_completion(
139
Returns:
140
Generated text content
141
"""
141
- client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
142
-
143
- completion = client.chat.completions.create(
144
- model=model_name,
145
- messages=messages,
146
- max_tokens=max_tokens,
147
- temperature=0.0,
148
- seed=42,
149
- )
142
+ with openai.OpenAI(
143
+ api_key="EMPTY", base_url=base_url, timeout=120, max_retries=0
144
+ ) as client:
145
+ completion = client.chat.completions.create(
146
+ model=model_name,
147
+ messages=messages,
148
+ max_tokens=max_tokens,
149
+ temperature=0.0,
150
+ seed=42,
151
+ )
152
151
- return completion.choices[0].message.content
153
+ content = completion.choices[0].message.content
154
+ assert content, "Expected a nonempty completion"
155
+ return content
156
157
158
def main():
@@ -198,7 +202,19 @@ def main():
202
help="Skip the two-image multimodal prompt",
203
)
204
205
+ parser.add_argument(
206
+ "--mm_smoke_test",
207
+ action="store_true",
208
+ help="Use three short, fixed-answer image cases, including duplicate images",
209
+ )
210
+ parser.add_argument("--concurrency", type=int, default=1)
211
+ parser.add_argument("--repeat", type=int, default=1)
212
+
213
args = parser.parse_args()
214
+ if args.concurrency < 1 or args.repeat < 1:
215
+ parser.error("--concurrency and --repeat must be positive")
216
+ if args.mm_smoke_test and (not args.use_mm_prompts or args.skip_two_image_prompt):
217
+ parser.error("--mm_smoke_test requires all multimodal prompts")
218
219
print(f"Service URL: {args.service_url}")
220
print(f"Model: {args.model_name}")
@@ -235,27 +251,76 @@ def main():
251
test_prompts = SAMPLE_PROMPTS_TEXT
252
print("Using text-only prompts for quick testing")
253
254
+ if args.mm_smoke_test:
255
+ image = SAMPLE_PROMPTS_MM[0]["messages"][0]["content"][0]
256
+ image_pair = SAMPLE_PROMPTS_MM[1]["messages"][0]["content"][:2]
257
+ cases = [
258
+ (
259
+ "Single image",
260
+ [image],
261
+ "What word is on the red road sign? Reply with only that word "
262
+ "in uppercase.",
263
+ "STOP",
264
+ ),
265
+ (
266
+ "Two different images",
267
+ image_pair,
268
+ "Do these pictures contain both flowers and birds? "
269
+ "Reply with only YES or NO.",
270
+ "YES",
271
+ ),
272
+ (
273
+ "Same image twice",
274
+ [image, image],
275
+ "Read the red road sign in each image, in image order. Reply "
276
+ "with only the two uppercase words separated by a comma and a space.",
277
+ "STOP, STOP",
278
+ ),
279
+ ]
280
+ test_prompts = [
281
+ {
282
+ "description": description,
283
+ "expected": expected,
284
+ "messages": [
285
+ {
286
+ "role": "user",
287
+ "content": [*images, {"type": "text", "text": question}],
288
+ }
289
+ ],
290
+ }
291
+ for description, images, question, expected in cases
292
+ ]
293
+
294
# Run completions
295
service_url = f"{args.service_url}/v1"
296
output_strs = {}
297
242
- for i, prompt_data in enumerate(test_prompts):
243
- print(
244
- f"\nRunning prompt {i + 1}/{len(test_prompts)}: "
245
- f"{prompt_data['description']}"
246
- )
247
-
248
- output_str = run_chat_completion(
298
+ def complete(prompt_data):
299
+ output = run_chat_completion(
300
base_url=service_url,
301
model_name=args.model_name,
302
messages=prompt_data["messages"],
252
- max_tokens=MAX_OUTPUT_LEN,
303
+ max_tokens=16 if args.mm_smoke_test else MAX_OUTPUT_LEN,
304
)
254
-
255
- # Use description as key for comparison
256
- key = prompt_data["description"]
257
- output_strs[key] = output_str
258
- print(f"Output: {output_str}")
305
+ if args.mm_smoke_test:
306
+ output = output.strip()
307
+ assert output == prompt_data["expected"], (
308
+ f"{prompt_data['description']}: expected {prompt_data['expected']!r}, "
309
+ f"got {output!r}"
310
+ )
311
+ return output
312
+
313
+ # Each round includes concurrent requests sharing image hashes; later rounds
314
+ # exercise reuse after previous requests have finished.
315
+ with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
316
+ for repeat in range(args.repeat):
317
+ outputs = executor.map(complete, test_prompts)
318
+ for prompt_data, output_str in zip(test_prompts, outputs):
319
+ key = prompt_data["description"]
320
+ if args.repeat > 1:
321
+ key = f"{key} (round {repeat + 1})"
322
+ output_strs[key] = output_str
323
+ print(f"{key}: {output_str}")
324
325
if args.mode in ("baseline", "baseline_pd"):
326
# Baseline mode: Save outputs
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
new
+172
@@ -0,0 +1,172 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Body handling across the EPD proxy's decode retries.
4
+
5
+Exercises the REAL helpers loaded from the ``examples/`` proxy, so a future
6
+change to them is what these tests catch.
7
+"""
8
+
9
+import asyncio
10
+import importlib.util
11
+from pathlib import Path
12
+
13
+import pytest
14
+
15
+PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
16
+
17
+
18
+@pytest.fixture(scope="module")
19
+def proxy():
20
+ path = Path(__file__).parents[4] / PROXY_REL
21
+ spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path)
22
+ assert spec is not None and spec.loader is not None
23
+ module = importlib.util.module_from_spec(spec)
24
+ spec.loader.exec_module(module)
25
+ return module
26
+
27
+
28
+class _Response:
29
+ def __init__(self, params):
30
+ self._params = params
31
+
32
+ async def json(self):
33
+ return {"kv_transfer_params": self._params}
34
+
35
+
36
+def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
37
+ """A decode retry re-enters this function with the body it was given.
38
+
39
+ Mutating that body in place let one attempt's `remote_block_ids` survive
40
+ into the next, so a retry whose prefill returns nothing sent decode blocks
41
+ the prefiller may already have freed.
42
+ """
43
+ served = [{"remote_block_ids": [1, 2]}, {}]
44
+
45
+ async def _stage(req_data, p_url, req_id):
46
+ assert "kv_transfer_params" not in req_data
47
+ return _Response(served.pop(0))
48
+
49
+ monkeypatch.setattr(proxy, "process_prefill_stage", _stage)
50
+
51
+ body = {"messages": [], "stream": False}
52
+ first = asyncio.run(proxy.maybe_prefill(body, "http://prefill", "r1"))
53
+ assert first["kv_transfer_params"] == {"remote_block_ids": [1, 2]}
54
+ assert "kv_transfer_params" not in body
55
+
56
+ second = asyncio.run(proxy.maybe_prefill(body, "http://prefill", "r1"))
57
+ assert "kv_transfer_params" not in second
58
+
59
+
60
+class _EncoderResponse:
61
+ def __init__(self, params):
62
+ self.status = 200
63
+ self._params = params
64
+
65
+ async def json(self):
66
+ return {"ec_transfer_params": self._params}
67
+
68
+ async def text(self):
69
+ return ""
70
+
71
+
72
+class _EncoderSession:
73
+ """Serve one canned encoder reply per attempt."""
74
+
75
+ def __init__(self, replies):
76
+ self._replies = list(replies)
77
+
78
+ async def post(self, url, json=None, headers=None):
79
+ return _EncoderResponse(self._replies.pop(0))
80
+
81
+
82
+def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch):
83
+ """The retry loop re-enters `prepare_for_decode` with the same body.
84
+
85
+ Recording the encoder's connector handles on that body in place let
86
+ attempt 1's handle survive into attempt 2, so a second encode that
87
+ reported nothing still sent decode a handle on an embedding the encoder
88
+ no longer publishes -- the exact state the retry exists to leave behind.
89
+ """
90
+ handle = {"metadata": {"image_grid_thw": [1, 2, 2]}, "peer_port": 1234}
91
+ monkeypatch.setattr(
92
+ proxy,
93
+ "encode_session",
94
+ _EncoderSession([{"encoder-side-hash": handle}, {}]),
95
+ )
96
+
97
+ async def _no_prefill(req_data, p_url, req_id):
98
+ return req_data
99
+
100
+ monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill)
101
+
102
+ body = {
103
+ "messages": [
104
+ {
105
+ "role": "user",
106
+ "content": [{"type": "image_url", "image_url": {"url": "image"}}],
107
+ }
108
+ ],
109
+ "stream": False,
110
+ }
111
+ args = ("r1", ["http://encoder"], "http://prefill", None)
112
+
113
+ first, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args))
114
+ reported = first["ec_transfer_params"]
115
+ assert [handle] == [value for key, value in reported.items() if key != "ec_items"]
116
+ assert "ec_transfer_params" not in body
117
+
118
+ # Attempt 2's encode reports nothing: decode must be told nothing.
119
+ second, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args))
120
+ assert "ec_transfer_params" not in second
121
+
122
+
123
+def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
124
+ handle = {"metadata": {"image_grid_thw": [1, 2, 2]}}
125
+ monkeypatch.setattr(proxy, "NO_REWRITE", True)
126
+ monkeypatch.setattr(
127
+ proxy, "encode_session", _EncoderSession([{"encoded-hash": handle}])
128
+ )
129
+ body = {
130
+ "messages": [
131
+ {
132
+ "role": "user",
133
+ "content": [{"type": "image_url", "image_url": {"url": "image"}}],
134
+ }
135
+ ]
136
+ }
137
+ prepared, _, _ = asyncio.run(
138
+ proxy.prepare_for_decode(
139
+ body, "request", ["http://encoder"], "", "tcp://consumer:1"
140
+ )
141
+ )
142
+ assert prepared["messages"] == body["messages"]
143
+ assert "ec_transfer_params" not in body
144
+ params = prepared["ec_transfer_params"]
145
+ assert params["encoded-hash"] == handle
146
+ assert params["ec_items"][0]["mm_hash"] == "encoded-hash"
147
+ assert params["ec_items"][0]["transfer_id"]
148
+
149
+
150
+@pytest.mark.parametrize("server_keep_alive", ["0.1", "1", "5", "2", "30"])
151
+def test_pooled_connections_are_retired_before_the_server_closes_them(
152
+ proxy, monkeypatch, server_keep_alive
153
+):
154
+ """The proxy must not hand a request a connection the server has dropped.
155
+
156
+ vLLM closes idle keep-alive connections at `VLLM_HTTP_TIMEOUT_KEEP_ALIVE`
157
+ seconds while aiohttp pools them for 15, so a slow hop leaves a dead
158
+ connection in the pool. The next request fails with
159
+ ServerDisconnectedError and the server logs nothing, because it closed the
160
+ socket before the request arrived.
161
+ """
162
+ monkeypatch.setenv("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", server_keep_alive)
163
+ monkeypatch.setattr(proxy.app.state, "p_urls", [], raising=False)
164
+
165
+ asyncio.run(proxy.on_startup())
166
+ try:
167
+ # No public accessor for the pool's idle timeout.
168
+ pooled_for = proxy.encode_session.connector._keepalive_timeout
169
+ assert pooled_for < float(server_keep_alive)
170
+ assert proxy.decode_session.connector._keepalive_timeout == pooled_for
171
+ finally:
172
+ asyncio.run(proxy.on_shutdown())
tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
+48
-2
@@ -32,8 +32,13 @@ def _load_proxy_module():
32
33
34
@pytest.fixture(scope="module")
35
-def assign():
36
- return _load_proxy_module().encoder_rr_assignment
35
+def proxy():
36
+ return _load_proxy_module()
37
+
38
+
39
+@pytest.fixture(scope="module")
40
+def assign(proxy):
41
+ return proxy.encoder_rr_assignment
42
43
44
def _drive(assign, e_urls, counts):
@@ -80,3 +85,44 @@ def test_single_encoder_always_resolves_to_it(assign):
85
urls, next_cursor = assign(["E0"], 0, 4)
86
assert urls == ["E0"] * 4
87
assert next_cursor == 0
88
+
89
+
90
+def test_generic_e_p_d_routing_without_mooncake_is_supported(proxy):
91
+ proxy.validate_ec_consumer_routing(["http://prefill"], [])
92
+
93
+
94
+def test_mooncake_e_pd_routing_is_supported(proxy):
95
+ proxy.validate_ec_consumer_routing([], ["tcp://decode:19019"])
96
+
97
+
98
+def test_mooncake_independent_prefill_routing_fails_fast(proxy):
99
+ with pytest.raises(ValueError, match=r"supports E\+PD only"):
100
+ proxy.validate_ec_consumer_routing(["http://prefill"], ["tcp://decode:19019"])
101
+
102
+
103
+def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy):
104
+ request = {
105
+ "messages": [
106
+ {
107
+ "role": "user",
108
+ "content": [{"type": "image_url", "image_url": {"url": "image"}}],
109
+ }
110
+ ]
111
+ }
112
+
113
+ rewritten = proxy.rewrite_for_decode(
114
+ request,
115
+ {
116
+ 0: {
117
+ "mm_hash": "proxy-uuid",
118
+ "ec_mm_hash": "engine-derived-hash",
119
+ "transfer_id": "transfer",
120
+ "image_grid_thw": [1, 2, 3],
121
+ }
122
+ },
123
+ )
124
+
125
+ assert rewritten["messages"][0]["content"][0]["uuid"] == "proxy-uuid"
126
+ assert rewritten["ec_transfer_params"]["ec_items"] == [
127
+ {"mm_hash": "engine-derived-hash", "transfer_id": "transfer"}
128
+ ]
tests/v1/ec_connector/unit/test_worker_ec_connector.py
+26
@@ -12,6 +12,9 @@ from vllm.distributed.ec_transfer.ec_connector.base import (
12
ECConnectorMetadata,
13
)
14
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
15
+from vllm.v1.worker.ec_connector_model_runner_mixin import (
16
+ ECConnectorModelRunnerMixin,
17
+)
18
from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector
19
20
pytestmark = pytest.mark.cpu_test
@@ -52,6 +55,9 @@ def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer):
55
56
saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list]
57
assert saved == (["mm_new"] if is_producer else [])
58
+ assert fake.start_save_caches.called == is_producer
59
+ if is_producer:
60
+ assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache
61
assert fake.start_load_caches.called == is_consumer
62
63
@@ -66,6 +72,26 @@ def test_worker_meta_is_reported_on_context_exit():
72
assert fake.clear_connector_metadata.called
73
74
75
+def test_v1_producer_receives_existing_encoder_cache():
76
+ encoder_cache = {"mm_cached": None}
77
+ fake = MagicMock(spec=ECConnectorBase)
78
+ fake.is_producer = True
79
+ fake.is_consumer = False
80
+ fake.get_finished.return_value = (None, None)
81
+
82
+ module = "vllm.v1.worker.ec_connector_model_runner_mixin"
83
+ with (
84
+ patch(f"{module}.has_ec_transfer", return_value=True),
85
+ patch(f"{module}.get_ec_transfer", return_value=fake),
86
+ ECConnectorModelRunnerMixin.maybe_get_ec_connector_output(
87
+ _scheduler_output(), encoder_cache
88
+ ),
89
+ ):
90
+ pass
91
+
92
+ assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache
93
+
94
+
95
def test_no_forward_reports_without_running_the_model():
96
connector, _ = _connector()
97
vllm/config/ec_transfer.py
+5
@@ -18,6 +18,11 @@ class ECTransferConfig:
18
19
ec_connector: str | None = None
20
"""The EC connector for vLLM to transmit EC caches between vLLM instances.
21
+
22
+ Built-in options include ``ECExampleConnector`` (shared filesystem via
23
+ safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA;
24
+ requires ``mooncake-transfer-engine`` and matching producer/consumer
25
+ ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring).
26
"""
27
28
engine_id: str | None = None
vllm/distributed/ec_transfer/ec_connector/base.py
+19
@@ -261,6 +261,25 @@ class ECConnectorBase(ABC):
261
"""
262
return True
263
264
+ def start_save_caches(self, **kwargs: Any) -> None:
265
+ """Prepare this step's outbound pushes before the model runs."""
266
+ return
267
+
268
+ def start_worker_services(self) -> None:
269
+ """Start Worker-side services once the model is resident."""
270
+ return
271
+
272
+ def take_unavailable_requests(self) -> set[str]:
273
+ """Request IDs whose encoder inputs the connector can no longer obtain.
274
+
275
+ The Scheduler fails these; re-issuing the request re-runs the encode.
276
+ """
277
+ return set()
278
+
279
+ def update_state_after_free(self, request: "Request", index: int):
280
+ """Notify the connector that an encoder cache entry was released."""
281
+ return
282
+
283
@abstractmethod
284
def update_state_after_alloc(self, request: "Request", index: int):
285
"""
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+5
-1
@@ -7,6 +7,7 @@ Owns the mmap region and the embedding cache, and handles the producer
7
for the ECCPUConnector.
8
"""
9
10
+from collections.abc import Collection
11
from typing import TYPE_CHECKING, Any
12
13
import torch
@@ -184,7 +185,10 @@ class ECCPUScheduler:
185
return entry is not None and entry.ready
186
187
def ensure_cache_available(
187
- self, request: "Request", num_computed_tokens: int
188
+ self,
189
+ request: "Request",
190
+ num_computed_tokens: int,
191
+ local_cache_hashes: Collection[str] | None = None,
192
) -> bool:
193
if not self._nixl_enabled:
194
return True # CPU offload never blocks.
vllm/distributed/ec_transfer/ec_connector/factory.py
+6
@@ -89,3 +89,9 @@ ECConnectorFactory.register_connector(
89
"vllm.distributed.ec_transfer.ec_connector.cpu.connector",
90
"ECCPUConnector",
91
)
92
+
93
+ECConnectorFactory.register_connector(
94
+ "ECMooncakeConnector",
95
+ "vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector",
96
+ "ECMooncakeConnector",
97
+)
vllm/distributed/ec_transfer/ec_connector/mooncake/__init__.py
new
+9
@@ -0,0 +1,9 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Internal building blocks for the Mooncake encoder-cache connector.
4
+
5
+The public connector remains in :mod:`mooncake_ec_connector`. This package
6
+separates configuration, control-plane messaging, transfer state, registered
7
+memory, and worker/scheduler orchestration so that each component has one
8
+owner and can be tested independently.
9
+"""
vllm/distributed/ec_transfer/ec_connector/mooncake/config.py
new
+123
@@ -0,0 +1,123 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from __future__ import annotations
4
+
5
+import math
6
+from dataclasses import dataclass
7
+from typing import TYPE_CHECKING
8
+
9
+from vllm.utils.network_utils import make_zmq_path
10
+
11
+if TYPE_CHECKING:
12
+ from vllm.config import VllmConfig
13
+
14
+# Scheduler and Worker must agree on when a remote reservation becomes stale.
15
+_RESERVATION_TTL_SECONDS = 300
16
+
17
+
18
+def _positive_int(name: str, value: object) -> int:
19
+ message = f"ECMooncakeConnector requires {name} to be a positive integer."
20
+ if isinstance(value, bool):
21
+ raise ValueError(message)
22
+ if isinstance(value, int):
23
+ result = value
24
+ elif isinstance(value, float) and math.isfinite(value) and value.is_integer():
25
+ result = int(value)
26
+ elif isinstance(value, str):
27
+ try:
28
+ result = int(value)
29
+ except ValueError as error:
30
+ raise ValueError(message) from error
31
+ else:
32
+ raise ValueError(message)
33
+ if result <= 0:
34
+ raise ValueError(f"ECMooncakeConnector requires {name} > 0.")
35
+ return result
36
+
37
+
38
+def _positive_float(name: str, value: object) -> float:
39
+ message = f"ECMooncakeConnector requires {name} > 0."
40
+ if isinstance(value, bool):
41
+ raise ValueError(message)
42
+ try:
43
+ result = float(value) # type: ignore[arg-type]
44
+ except (TypeError, ValueError) as error:
45
+ raise ValueError(message) from error
46
+ if not math.isfinite(result) or result <= 0:
47
+ raise ValueError(message)
48
+ return result
49
+
50
+
51
+@dataclass(frozen=True)
52
+class MooncakeECConfig:
53
+ """Validated settings shared by the Scheduler and Worker roles.
54
+
55
+ ``control_port`` is the first TP-shard port after the DP offset;
56
+ ``control_addr`` targets that shard, which advertises the full topology.
57
+ ``control_host`` is the address this instance both advertises and binds,
58
+ so reaching a Consumer means being routed to it rather than finding it on
59
+ every interface.
60
+ """
61
+
62
+ is_producer: bool
63
+ is_consumer: bool
64
+ protocol: str
65
+ buffer_device: str
66
+ control_host: str
67
+ control_port: int
68
+ control_addr: str
69
+ control_timeout_ms: int
70
+ push_wait_timeout_s: float
71
+ pool_size: int
72
+
73
+ @classmethod
74
+ def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig:
75
+ parallel_config = vllm_config.parallel_config
76
+ ec_config = vllm_config.ec_transfer_config
77
+ assert ec_config is not None
78
+
79
+ if ec_config.is_ec_producer:
80
+ if parallel_config.tensor_parallel_size > 1:
81
+ raise ValueError(
82
+ "ECMooncakeConnector producers require tensor_parallel_size=1."
83
+ )
84
+ if parallel_config.pipeline_parallel_size > 1:
85
+ raise ValueError(
86
+ "ECMooncakeConnector producers do not support pipeline parallelism."
87
+ )
88
+ if parallel_config.data_parallel_size > 1:
89
+ raise ValueError(
90
+ "ECMooncakeConnector producers require data_parallel_size=1."
91
+ )
92
+
93
+ registered_buffer_size = _positive_int(
94
+ "ec_buffer_size", ec_config.ec_buffer_size
95
+ )
96
+ get = ec_config.get_from_extra_config
97
+ control_port = int(ec_config.ec_port) + (
98
+ parallel_config.data_parallel_index * parallel_config.tensor_parallel_size
99
+ )
100
+ highest_port = control_port + parallel_config.tensor_parallel_size - 1
101
+ if not 1 <= control_port <= highest_port <= 65535:
102
+ raise ValueError("ECMooncakeConnector ec_port must be in 1..65535.")
103
+
104
+ return cls(
105
+ is_producer=ec_config.is_ec_producer,
106
+ is_consumer=ec_config.is_ec_consumer,
107
+ protocol=str(get("mooncake_protocol", "rdma")),
108
+ buffer_device=str(ec_config.ec_buffer_device or "cuda").lower(),
109
+ control_host=str(ec_config.ec_ip),
110
+ control_port=control_port,
111
+ control_addr=make_zmq_path("tcp", ec_config.ec_ip, control_port),
112
+ control_timeout_ms=max(
113
+ 1,
114
+ math.ceil(
115
+ _positive_float("control_timeout_s", get("control_timeout_s", 30))
116
+ * 1000
117
+ ),
118
+ ),
119
+ push_wait_timeout_s=_positive_float(
120
+ "push_wait_timeout_s", get("push_wait_timeout_s", 60)
121
+ ),
122
+ pool_size=registered_buffer_size,
123
+ )
vllm/distributed/ec_transfer/ec_connector/mooncake/control.py
new
+460
@@ -0,0 +1,460 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ZMQ control plane for Mooncake encoder-cache reservations and events.
4
+
5
+Tensor bytes never travel through this module. It coordinates destination
6
+reservations, completion/cancellation, TP-shard discovery, and readiness
7
+notifications while :mod:`transfer` owns the Mooncake data plane.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import threading
13
+import time
14
+from collections import deque
15
+from collections.abc import Callable
16
+from concurrent.futures import Executor, Future
17
+from typing import Any
18
+
19
+import torch
20
+import zmq
21
+
22
+from vllm.logger import init_logger
23
+from vllm.utils.network_utils import make_zmq_path
24
+
25
+logger = init_logger(__name__)
26
+
27
+_MAX_PENDING_EVENTS = 4096
28
+_RESERVATION_REAP_INTERVAL_SECONDS = 1
29
+
30
+
31
+ControlRequest = dict[str, Any]
32
+ControlResponse = dict[str, Any]
33
+
34
+
35
+class ControlClient:
36
+ """Send control requests through reusable, thread-local REQ sockets."""
37
+
38
+ def __init__(self, timeout_ms: int) -> None:
39
+ self._context = zmq.Context()
40
+ self._timeout_ms = timeout_ms
41
+ self._local = threading.local()
42
+ self._topologies: dict[str, list[str]] = {}
43
+ self._closed = False
44
+
45
+ def _sockets(self) -> dict[str, zmq.Socket]:
46
+ sockets = getattr(self._local, "sockets", None)
47
+ if sockets is None:
48
+ sockets = {}
49
+ self._local.sockets = sockets
50
+ return sockets
51
+
52
+ def _discard(self, addr: str) -> None:
53
+ socket = self._sockets().pop(addr, None)
54
+ if socket is not None:
55
+ socket.close(linger=0)
56
+
57
+ def _exchange(self, addr: str, payload: ControlRequest) -> ControlResponse:
58
+ sockets = self._sockets()
59
+ socket = sockets.get(addr)
60
+ if socket is None:
61
+ socket = self._context.socket(zmq.REQ)
62
+ socket.setsockopt(zmq.IPV6, 1)
63
+ socket.setsockopt(zmq.RCVTIMEO, self._timeout_ms)
64
+ socket.setsockopt(zmq.SNDTIMEO, self._timeout_ms)
65
+ socket.setsockopt(zmq.LINGER, 0)
66
+ socket.connect(addr)
67
+ sockets[addr] = socket
68
+ try:
69
+ socket.send_json(payload)
70
+ response = socket.recv_json()
71
+ except Exception:
72
+ self._discard(addr)
73
+ raise
74
+ assert isinstance(response, dict)
75
+ return response
76
+
77
+ def request(self, addr: str, payload: ControlRequest) -> Any:
78
+ response = self._exchange(addr, payload)
79
+ if not response.get("ok"):
80
+ raise RuntimeError(response.get("error", "EC control request failed"))
81
+ return response.get("result")
82
+
83
+ def discover_shards(self, base_addr: str) -> list[str] | None:
84
+ if base_addr in self._topologies:
85
+ return self._topologies[base_addr]
86
+ try:
87
+ reply = self.request(base_addr, {"op": "peers"})
88
+ ports = reply.get("ports") if isinstance(reply, dict) else None
89
+ if not isinstance(ports, list) or not ports:
90
+ raise ValueError("invalid or empty peer list")
91
+ prefix = base_addr.rsplit(":", 1)[0]
92
+ shards = [f"{prefix}:{int(port)}" for port in ports]
93
+ except Exception:
94
+ logger.warning(
95
+ "EC Mooncake consumer at %s did not report its shards",
96
+ base_addr,
97
+ exc_info=True,
98
+ )
99
+ return None
100
+ self._topologies[base_addr] = shards
101
+ return shards
102
+
103
+ def close(self) -> None:
104
+ if self._closed:
105
+ return
106
+ self._closed = True
107
+ self._context.destroy(linger=0)
108
+
109
+
110
+def make_cancel_request(
111
+ transfer_id: str,
112
+ reservation_id: str = "",
113
+ *,
114
+ abandon: bool = False,
115
+ refresh: bool = False,
116
+) -> ControlRequest:
117
+ request: ControlRequest = {
118
+ "op": "cancel",
119
+ "transfer_id": transfer_id,
120
+ "reservation_id": reservation_id,
121
+ }
122
+ if abandon:
123
+ request["abandon"] = True
124
+ if refresh:
125
+ request["refresh"] = True
126
+ return request
127
+
128
+
129
+class EventInbox:
130
+ """Receive Consumer readiness events without blocking the Scheduler.
131
+
132
+ Subscribing needs one blocking control request per shard, so it runs on
133
+ `executor` when one is given: an unreachable Consumer must cost the
134
+ Scheduler nothing, not one control timeout per drain.
135
+ """
136
+
137
+ def __init__(self, client: ControlClient, executor: Executor | None = None) -> None:
138
+ self._client = client
139
+ self._executor = executor
140
+ self._discovery: Future[list[str] | None] | None = None
141
+ self._context: zmq.Context | None = None
142
+ self._socket: zmq.Socket | None = None
143
+ self._closed = False
144
+ self.shard_count = 1
145
+
146
+ def _endpoints(self, base_addr: str) -> list[str] | None:
147
+ """Ask every shard where it publishes readiness events.
148
+
149
+ Blocking: called on `self._executor` unless there is none.
150
+ """
151
+ shards = self._client.discover_shards(base_addr)
152
+ if shards is None:
153
+ return None
154
+ endpoints = []
155
+ for addr in shards:
156
+ try:
157
+ event_port = self._client.request(addr, {"op": "event_port"})
158
+ address, _ = addr.rsplit(":", 1)
159
+ endpoints.append(f"{address}:{int(event_port)}")
160
+ except Exception:
161
+ logger.warning(
162
+ "EC Mooncake could not subscribe to the event channel of "
163
+ "consumer shard %s; retrying the complete topology later.",
164
+ addr,
165
+ )
166
+ return None
167
+ return endpoints
168
+
169
+ def _connect(self, base_addr: str) -> None:
170
+ if self._socket is not None or self._closed:
171
+ return
172
+ if self._executor is None:
173
+ self._install(self._endpoints(base_addr))
174
+ return
175
+ if self._discovery is None:
176
+ self._discovery = self._executor.submit(self._endpoints, base_addr)
177
+ return
178
+ if not self._discovery.done():
179
+ return
180
+ discovery, self._discovery = self._discovery, None
181
+ try:
182
+ self._install(discovery.result())
183
+ except Exception:
184
+ logger.warning(
185
+ "EC Mooncake event-channel discovery for %s failed; retrying.",
186
+ base_addr,
187
+ exc_info=True,
188
+ )
189
+
190
+ def _install(self, endpoints: list[str] | None) -> None:
191
+ """Adopt a discovered topology.
192
+
193
+ Runs on the caller's thread so the socket is only ever touched there.
194
+ """
195
+ if not endpoints or self._closed:
196
+ return
197
+ context = zmq.Context()
198
+ socket = context.socket(zmq.PULL)
199
+ socket.setsockopt(zmq.IPV6, 1)
200
+ for endpoint in endpoints:
201
+ socket.connect(endpoint)
202
+ self._context = context
203
+ self._socket = socket
204
+ self.shard_count = len(endpoints)
205
+
206
+ def drain(self, base_addr: str) -> list[dict[str, Any]]:
207
+ self._connect(base_addr)
208
+ if self._socket is None:
209
+ return []
210
+ events = []
211
+ while True:
212
+ try:
213
+ events.append(self._socket.recv_json(flags=zmq.DONTWAIT))
214
+ except zmq.Again:
215
+ return events
216
+ except zmq.ZMQError:
217
+ logger.warning("EC Mooncake event channel failed", exc_info=True)
218
+ return events
219
+ except (ValueError, UnicodeError):
220
+ # The event channel is a plain PULL socket: an undecodable
221
+ # frame must cost one frame, not the engine.
222
+ logger.warning(
223
+ "Discarding an undecodable EC Mooncake event.", exc_info=True
224
+ )
225
+
226
+ def close(self) -> None:
227
+ if self._closed:
228
+ return
229
+ self._closed = True
230
+ if self._discovery is not None:
231
+ self._discovery.cancel()
232
+ self._discovery = None
233
+ if self._socket is not None:
234
+ self._socket.close(linger=0)
235
+ self._socket = None
236
+ if self._context is not None:
237
+ self._context.term()
238
+
239
+
240
+class ConsumerControlServer:
241
+ """Expose Consumer reservations and readiness events over ZMQ.
242
+
243
+ One server runs on every receiving TP rank. The REP channel handles
244
+ reservation operations, while a PUSH channel publishes newly ready items
245
+ to the Scheduler.
246
+ """
247
+
248
+ def __init__(
249
+ self,
250
+ host: str,
251
+ port: int,
252
+ reserve: Callable[[dict[str, Any]], dict[str, Any]],
253
+ status: Callable[[str], dict[str, Any] | None],
254
+ complete: Callable[[str, str], tuple[bool, bool]],
255
+ cancel: Callable[[str, str, bool, bool], bool],
256
+ reap: Callable[[], int],
257
+ peer_ports: list[int] | None = None,
258
+ device: torch.device | None = None,
259
+ drain_ready: Callable[[], list[str]] = lambda: [],
260
+ ) -> None:
261
+ self.host = host
262
+ self.port = port
263
+ self.peer_ports = peer_ports or [port]
264
+ self._device = device
265
+ self.event_port: int | None = None
266
+ self._reserve = reserve
267
+ self._status = status
268
+ self._complete = complete
269
+ self._cancel = cancel
270
+ self._reap = reap
271
+ self._drain_ready = drain_ready
272
+ self._stop = threading.Event()
273
+ self._started = threading.Event()
274
+ self._thread: threading.Thread | None = None
275
+ self._startup_error: Exception | None = None
276
+
277
+ def start(self) -> None:
278
+ def loop() -> None:
279
+ if self._device is not None and self._device.type == "cuda":
280
+ torch.accelerator.set_device_index(self._device.index or 0)
281
+ context = zmq.Context()
282
+ socket = context.socket(zmq.REP)
283
+ event_socket = context.socket(zmq.PUSH)
284
+ socket.setsockopt(zmq.IPV6, 1)
285
+ event_socket.setsockopt(zmq.IPV6, 1)
286
+ pending_events: deque[dict[str, Any]] = deque()
287
+
288
+ def queue_event(event: dict[str, Any]) -> None:
289
+ event["shard"] = self.port
290
+ if len(pending_events) >= _MAX_PENDING_EVENTS:
291
+ dropped = pending_events.popleft()
292
+ logger.warning(
293
+ "EC Mooncake event backlog full on port %d; dropping "
294
+ "readiness for transfer_id=%s",
295
+ self.port,
296
+ dropped.get("transfer_id"),
297
+ )
298
+ pending_events.append(event)
299
+
300
+ def queue_ready(transfer_id: str) -> None:
301
+ status = self._status(transfer_id)
302
+ if status is not None:
303
+ queue_event({"transfer_id": transfer_id, **status})
304
+
305
+ last_reap_at = time.monotonic()
306
+ socket.setsockopt(zmq.RCVTIMEO, 100)
307
+ try:
308
+ socket.bind(make_zmq_path("tcp", self.host, self.port))
309
+ event_socket.bind(make_zmq_path("tcp", self.host, 0))
310
+ self.event_port = int(
311
+ event_socket.getsockopt(zmq.LAST_ENDPOINT)
312
+ .decode()
313
+ .rsplit(":", 1)[1]
314
+ )
315
+ except Exception as e:
316
+ self._startup_error = e
317
+ self._started.set()
318
+ socket.close(linger=0)
319
+ event_socket.close(linger=0)
320
+ context.term()
321
+ return
322
+ self._started.set()
323
+ try:
324
+ while not self._stop.is_set():
325
+ while pending_events:
326
+ try:
327
+ event_socket.send_json(
328
+ pending_events[0], flags=zmq.DONTWAIT
329
+ )
330
+ except zmq.Again:
331
+ break
332
+ pending_events.popleft()
333
+ now = time.monotonic()
334
+ if now - last_reap_at >= _RESERVATION_REAP_INTERVAL_SECONDS:
335
+ self._reap()
336
+ last_reap_at = now
337
+ try:
338
+ request = socket.recv_json()
339
+ except zmq.Again:
340
+ continue
341
+ except Exception:
342
+ # The frame arrived but did not decode. REP still owes a
343
+ # reply, so answer before returning to the loop.
344
+ logger.exception(
345
+ "EC Mooncake control channel on port %d received an "
346
+ "undecodable request",
347
+ self.port,
348
+ )
349
+ socket.send_json(
350
+ {"ok": False, "error": "malformed control request"}
351
+ )
352
+ continue
353
+ try:
354
+ op = request.get("op")
355
+ result: Any = None
356
+ if op in ("reserve", "reserve_batch"):
357
+ items = (
358
+ request["items"] if op == "reserve_batch" else [request]
359
+ )
360
+ results = []
361
+ for item in items:
362
+ try:
363
+ reserved = self._reserve(item)
364
+ results.append({"ok": True, "result": reserved})
365
+ if reserved.get("ready"):
366
+ queue_ready(str(item["transfer_id"]))
367
+ except Exception as exc:
368
+ results.append({"ok": False, "error": str(exc)})
369
+ if op == "reserve_batch":
370
+ result = {"items": results}
371
+ elif not results[0]["ok"]:
372
+ raise RuntimeError(results[0]["error"])
373
+ else:
374
+ result = results[0]["result"]
375
+ elif op == "status":
376
+ result = self._status(str(request["transfer_id"]))
377
+ elif op == "event_port":
378
+ result = self.event_port
379
+ elif op == "peers":
380
+ result = {"ports": self.peer_ports}
381
+ elif op in ("complete", "complete_batch"):
382
+ items = (
383
+ request["items"]
384
+ if op == "complete_batch"
385
+ else [request]
386
+ )
387
+ completions = []
388
+ for item in items:
389
+ transfer_id = str(item["transfer_id"])
390
+ accepted, became_ready = self._complete(
391
+ transfer_id, str(item["reservation_id"])
392
+ )
393
+ for ready_id in self._drain_ready():
394
+ queue_ready(ready_id)
395
+ completions.append(
396
+ {
397
+ "completed": accepted,
398
+ "became_ready": became_ready,
399
+ }
400
+ )
401
+ if not became_ready:
402
+ continue
403
+ queue_ready(transfer_id)
404
+ result = (
405
+ {"items": completions}
406
+ if op == "complete_batch"
407
+ else completions[0]
408
+ )
409
+ elif op == "cancel":
410
+ result = {
411
+ "cancelled": self._cancel(
412
+ str(request["transfer_id"]),
413
+ str(request.get("reservation_id", "")),
414
+ bool(request.get("abandon", False)),
415
+ bool(request.get("refresh", False)),
416
+ )
417
+ }
418
+ else:
419
+ raise ValueError(f"unknown control op: {op!r}")
420
+ socket.send_json({"ok": True, "result": result})
421
+ except Exception as e:
422
+ socket.send_json({"ok": False, "error": str(e)})
423
+ except Exception:
424
+ # `finally` closes the sockets and the thread ends, so without
425
+ # this every later reserve against this shard would surface
426
+ # only as a control timeout with nothing to attribute it to.
427
+ logger.exception(
428
+ "EC Mooncake control channel on port %d stopped serving",
429
+ self.port,
430
+ )
431
+ raise
432
+ finally:
433
+ socket.close(linger=0)
434
+ event_socket.close(linger=0)
435
+ context.term()
436
+
437
+ self._thread = threading.Thread(
438
+ target=loop, name="ec-mooncake-control", daemon=True
439
+ )
440
+ self._thread.start()
441
+ if not self._started.wait(timeout=5):
442
+ raise RuntimeError("EC Mooncake control channel failed to start")
443
+ if self._startup_error is not None:
444
+ raise RuntimeError("EC Mooncake control channel failed to bind") from (
445
+ self._startup_error
446
+ )
447
+ logger.info(
448
+ "EC Mooncake control channel listening on tcp://%s:%d (events tcp://%s:%d)",
449
+ self.host,
450
+ self.port,
451
+ self.host,
452
+ self.event_port,
453
+ )
454
+
455
+ def close(self) -> None:
456
+ if self._thread is None:
457
+ return
458
+ self._stop.set()
459
+ self._thread.join()
460
+ self._thread = None
vllm/distributed/ec_transfer/ec_connector/mooncake/memory.py
new
+504
@@ -0,0 +1,504 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Registered-memory allocation and residency for Mooncake transfers.
4
+
5
+The Producer pool stages source tensors in a registered slab. The Consumer
6
+pool owns destination allocations from reservation through publication,
7
+resident reuse, CUDA-safe retirement, and pressure-driven reclamation.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import bisect
13
+import threading
14
+from collections import OrderedDict
15
+from collections.abc import Callable
16
+from dataclasses import dataclass
17
+from typing import Generic, TypeVar
18
+
19
+import torch
20
+
21
+from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
22
+ MooncakeTransfer,
23
+)
24
+from vllm.logger import init_logger
25
+
26
+logger = init_logger(__name__)
27
+
28
+_T = TypeVar("_T")
29
+
30
+
31
+@dataclass
32
+class MemoryAllocation:
33
+ """Describe a tensor view carved from the Consumer receive slab."""
34
+
35
+ offset: int
36
+ size: int
37
+ tensor: torch.Tensor
38
+
39
+
40
+@dataclass
41
+class _ResidentEntry(Generic[_T]):
42
+ """Store one resident value and its ownership accounting."""
43
+
44
+ value: _T
45
+ pinned: bool = True
46
+ leases: int = 0
47
+
48
+
49
+@dataclass
50
+class ResidentLease(Generic[_T]):
51
+ """Represent a borrow of one resident entry."""
52
+
53
+ key: str
54
+ _entry: _ResidentEntry[_T]
55
+ _active: bool = True
56
+
57
+ @property
58
+ def value(self) -> _T:
59
+ return self._entry.value
60
+
61
+
62
+class ContiguousAllocator:
63
+ """Own a registered slab and its aligned regions; callers serialize access."""
64
+
65
+ def __init__(self, capacity: int, alignment: int = 256):
66
+ self.alignment = alignment
67
+ self._capacity = capacity
68
+ self._free = [(0, capacity)]
69
+ self.tensor: torch.Tensor | None = None
70
+ self._disabled = False
71
+
72
+ def prepare(self, device: torch.device, transfer: MooncakeTransfer) -> None:
73
+ if self.tensor is not None or self._disabled:
74
+ return
75
+ try:
76
+ tensor = torch.empty(self._capacity, dtype=torch.uint8, device=device)
77
+ ret = transfer.register_memory(tensor)
78
+ if ret != 0:
79
+ raise RuntimeError(f"Mooncake returned {ret}")
80
+ except (RuntimeError, torch.OutOfMemoryError) as error:
81
+ self._disabled = True
82
+ logger.warning("Could not initialize the EC registered buffer: %s", error)
83
+ return
84
+ self.tensor = tensor
85
+ self._free = [(0, tensor.nbytes)]
86
+ logger.info("Registered %d-byte buffer for Mooncake EC", tensor.nbytes)
87
+
88
+ def view(
89
+ self, offset: int, nbytes: int, shape: tuple[int, ...], dtype: torch.dtype
90
+ ) -> torch.Tensor:
91
+ assert self.tensor is not None
92
+ return self.tensor.narrow(0, offset, nbytes).view(dtype).view(shape)
93
+
94
+ def close(self, transfer: MooncakeTransfer) -> bool:
95
+ if self.tensor is not None and not transfer.unregister_memory(self.tensor):
96
+ return False
97
+ self.tensor = None
98
+ self._free.clear()
99
+ return True
100
+
101
+ def allocate(self, nbytes: int) -> tuple[int, int] | None:
102
+ size = (nbytes + self.alignment - 1) // self.alignment * self.alignment
103
+ for index, (offset, available) in enumerate(self._free):
104
+ if size > available:
105
+ continue
106
+ if size == available:
107
+ self._free.pop(index)
108
+ else:
109
+ self._free[index] = (offset + size, available - size)
110
+ return offset, size
111
+ return None
112
+
113
+ def free(self, offset: int, size: int) -> None:
114
+ index = bisect.bisect_left(self._free, (offset, size))
115
+ self._free.insert(index, (offset, size))
116
+ if index + 1 < len(self._free):
117
+ next_offset, next_size = self._free[index + 1]
118
+ if offset + size == next_offset:
119
+ self._free[index] = (offset, size + next_size)
120
+ self._free.pop(index + 1)
121
+ if index > 0:
122
+ previous_offset, previous_size = self._free[index - 1]
123
+ current_offset, current_size = self._free[index]
124
+ if previous_offset + previous_size == current_offset:
125
+ self._free[index - 1] = (
126
+ previous_offset,
127
+ previous_size + current_size,
128
+ )
129
+ self._free.pop(index)
130
+
131
+
132
+class ResidentPool(Generic[_T]):
133
+ """Track resident values, active leases, and LRU eviction eligibility."""
134
+
135
+ def __init__(self):
136
+ self._entries: dict[str, _ResidentEntry[_T]] = {}
137
+ self._evictable: OrderedDict[str, None] = OrderedDict()
138
+
139
+ def referenced(self) -> list[str]:
140
+ return [key for key, entry in self._entries.items() if entry.pinned]
141
+
142
+ def get(self, key: str) -> _T | None:
143
+ entry = self._entries.get(key)
144
+ return entry.value if entry is not None else None
145
+
146
+ def insert(
147
+ self,
148
+ key: str,
149
+ value: _T,
150
+ ) -> _T | None:
151
+ """Pin an entry and return a displaced value that has no owners."""
152
+ previous = self._entries.get(key)
153
+ entry = _ResidentEntry(value)
154
+ self._entries[key] = entry
155
+ self._evictable.pop(key, None)
156
+ if previous is not None and previous.leases == 0:
157
+ return previous.value
158
+ return None
159
+
160
+ def pin(self, key: str) -> _T | None:
161
+ entry = self._entries.get(key)
162
+ if entry is None:
163
+ return None
164
+ self._evictable.pop(key, None)
165
+ entry.pinned = True
166
+ return entry.value
167
+
168
+ def acquire(self, key: str) -> ResidentLease[_T] | None:
169
+ entry = self._entries.get(key)
170
+ if entry is None:
171
+ return None
172
+ entry.leases += 1
173
+ self._evictable.pop(key, None)
174
+ return ResidentLease(key, entry)
175
+
176
+ def release(self, lease: ResidentLease[_T]) -> _T | None:
177
+ if not lease._active:
178
+ return None
179
+ lease._active = False
180
+ entry = lease._entry
181
+ entry.leases -= 1
182
+ current = self._entries.get(lease.key)
183
+ if current is not entry:
184
+ if entry.leases == 0:
185
+ return entry.value
186
+ return None
187
+ if not entry.pinned and entry.leases == 0:
188
+ self._evictable[lease.key] = None
189
+ return None
190
+
191
+ def consume(self, lease: ResidentLease[_T]) -> tuple[_T, _T | None]:
192
+ current = self._entries[lease.key]
193
+ current.pinned = True
194
+ self._evictable.pop(lease.key, None)
195
+ released = self.release(lease)
196
+ return current.value, released
197
+
198
+ def retire(self, key: str) -> None:
199
+ if key not in self._entries:
200
+ return
201
+ entry = self._entries[key]
202
+ entry.pinned = False
203
+ if entry.leases == 0:
204
+ self._evictable[key] = None
205
+
206
+ def evict_lru(self, evict: Callable[[str, _T], bool]) -> str | None:
207
+ for key in list(self._evictable):
208
+ entry = self._entries[key]
209
+ if not evict(key, entry.value):
210
+ continue
211
+ self._evictable.pop(key, None)
212
+ del self._entries[key]
213
+ return key
214
+ return None
215
+
216
+ def clear(self) -> None:
217
+ self._entries.clear()
218
+ self._evictable.clear()
219
+
220
+
221
+@dataclass
222
+class StagedSources:
223
+ """Own Producer tensor views and their staging-slab regions."""
224
+
225
+ tensors: list[torch.Tensor]
226
+ regions: list[tuple[int, int]]
227
+
228
+
229
+class ProducerMemoryPool:
230
+ """Own the Producer staging slab and regions carved from it."""
231
+
232
+ def __init__(self, capacity: int, transfer: MooncakeTransfer) -> None:
233
+ self._transfer = transfer
234
+ self._allocator = ContiguousAllocator(capacity)
235
+ self._lock = threading.Lock()
236
+ self._local = threading.local()
237
+
238
+ @property
239
+ def tensor(self) -> torch.Tensor | None:
240
+ return self._allocator.tensor
241
+
242
+ def _free_regions(self, regions: list[tuple[int, int]]) -> None:
243
+ for offset, size in regions:
244
+ self._allocator.free(offset, size)
245
+
246
+ def stage(self, tensors: list[torch.Tensor]) -> StagedSources | None:
247
+ """Copy tensors into one registered slab, or return None for fallback."""
248
+ if not tensors:
249
+ return StagedSources([], [])
250
+ allocator = self._allocator
251
+ staged: list[torch.Tensor] = []
252
+ regions: list[tuple[int, int]] = []
253
+ with self._lock:
254
+ allocator.prepare(tensors[0].device, self._transfer)
255
+ pool = allocator.tensor
256
+ if pool is None:
257
+ return None
258
+ for tensor in tensors:
259
+ region = allocator.allocate(tensor.nbytes)
260
+ if region is None:
261
+ self._free_regions(regions)
262
+ return None
263
+ regions.append(region)
264
+ staged.append(
265
+ allocator.view(
266
+ region[0], tensor.nbytes, tuple(tensor.shape), tensor.dtype
267
+ )
268
+ )
269
+ if pool.device.type == "cuda":
270
+ stream = getattr(self._local, "stream", None)
271
+ if stream is None:
272
+ stream = torch.cuda.Stream(device=pool.device)
273
+ self._local.stream = stream
274
+ with torch.cuda.stream(stream):
275
+ for destination, source in zip(staged, tensors):
276
+ destination.copy_(source, non_blocking=True)
277
+ stream.synchronize()
278
+ else:
279
+ for destination, source in zip(staged, tensors):
280
+ destination.copy_(source)
281
+ return StagedSources(staged, regions)
282
+
283
+ def release(self, staged: StagedSources) -> None:
284
+ if not staged.regions:
285
+ return
286
+ with self._lock:
287
+ self._free_regions(staged.regions)
288
+
289
+ def close(self) -> None:
290
+ with self._lock:
291
+ self._allocator.close(self._transfer)
292
+
293
+
294
+class ConsumerMemoryPool:
295
+ """Own the registered receive slab and resident allocation lifecycle."""
296
+
297
+ def __init__(
298
+ self,
299
+ capacity: int,
300
+ transfer: MooncakeTransfer,
301
+ ) -> None:
302
+ self._transfer = transfer
303
+ self._allocator = ContiguousAllocator(capacity)
304
+ self._residents: ResidentPool[MemoryAllocation] = ResidentPool()
305
+ self._retire_events: dict[str, torch.Event] = {}
306
+ self._pending_frees: list[tuple[torch.Event, MemoryAllocation]] = []
307
+ self._reclaimed: set[str] = set()
308
+ self.lock = threading.RLock()
309
+
310
+ @property
311
+ def tensor(self) -> torch.Tensor | None:
312
+ return self._allocator.tensor
313
+
314
+ def prepare(
315
+ self,
316
+ device: torch.device,
317
+ ) -> None:
318
+ with self.lock:
319
+ self._allocator.prepare(device, self._transfer)
320
+
321
+ def _free(self, allocation: MemoryAllocation) -> None:
322
+ self._allocator.free(allocation.offset, allocation.size)
323
+
324
+ def free(self, allocation: MemoryAllocation) -> None:
325
+ with self.lock:
326
+ self._free(allocation)
327
+
328
+ def _defer_or_free(
329
+ self, allocation: MemoryAllocation, event: torch.Event | None
330
+ ) -> None:
331
+ if event is None or event.query():
332
+ self._free(allocation)
333
+ else:
334
+ self._pending_frees.append((event, allocation))
335
+
336
+ def _poll_frees_locked(self) -> None:
337
+ pending = []
338
+ for event, allocation in self._pending_frees:
339
+ if event.query():
340
+ self._free(allocation)
341
+ else:
342
+ pending.append((event, allocation))
343
+ self._pending_frees = pending
344
+
345
+ def _reclaim_locked(self, nbytes: int) -> tuple[int, int] | None:
346
+ def evict(mm_hash: str, allocation: MemoryAllocation) -> bool:
347
+ event = self._retire_events.pop(mm_hash, None)
348
+ self._defer_or_free(allocation, event)
349
+ self._reclaimed.add(mm_hash)
350
+ return True
351
+
352
+ while self._residents.evict_lru(evict) is not None:
353
+ # Eviction defers any free whose CUDA event is still pending, so
354
+ # those bytes reach the allocator only once the event is polled.
355
+ self._poll_frees_locked()
356
+ region = self._allocator.allocate(nbytes)
357
+ if region is not None:
358
+ return region
359
+ self._poll_frees_locked()
360
+ return self._allocator.allocate(nbytes)
361
+
362
+ def _make_allocation(
363
+ self,
364
+ region: tuple[int, int],
365
+ nbytes: int,
366
+ shape: tuple[int, ...],
367
+ dtype: torch.dtype,
368
+ ) -> MemoryAllocation:
369
+ offset, size = region
370
+ tensor = self._allocator.view(offset, nbytes, shape, dtype)
371
+ return MemoryAllocation(offset, size, tensor)
372
+
373
+ def try_allocate(
374
+ self, nbytes: int, shape: tuple[int, ...], dtype: torch.dtype
375
+ ) -> MemoryAllocation | None:
376
+ with self.lock:
377
+ allocator = self._allocator
378
+ assert allocator.tensor is not None
379
+ self._poll_frees_locked()
380
+ region = allocator.allocate(nbytes)
381
+ if region is None:
382
+ return None
383
+ return self._make_allocation(region, nbytes, shape, dtype)
384
+
385
+ def reclaim_and_allocate(
386
+ self, nbytes: int, shape: tuple[int, ...], dtype: torch.dtype
387
+ ) -> MemoryAllocation | None:
388
+ with self.lock:
389
+ region = self._reclaim_locked(nbytes)
390
+ if region is None:
391
+ return None
392
+ return self._make_allocation(region, nbytes, shape, dtype)
393
+
394
+ def acquire_cached(
395
+ self, mm_hash: str, shape: tuple[int, ...], dtype: torch.dtype
396
+ ) -> ResidentLease[MemoryAllocation] | None:
397
+ with self.lock:
398
+ allocation = self._residents.get(mm_hash)
399
+ if allocation is None:
400
+ return None
401
+ if (
402
+ tuple(allocation.tensor.shape) != shape
403
+ or allocation.tensor.dtype != dtype
404
+ ):
405
+ raise ValueError("conflicting cached tensor for mm_hash")
406
+ return self._residents.acquire(mm_hash)
407
+
408
+ def take_resident(
409
+ self, mm_hash: str, shape: tuple[int, ...], dtype_name: str
410
+ ) -> torch.Tensor | None:
411
+ with self.lock:
412
+ allocation = self._residents.get(mm_hash)
413
+ if allocation is None:
414
+ return None
415
+ tensor = allocation.tensor
416
+ if (
417
+ tuple(tensor.shape) != shape
418
+ or str(tensor.dtype).split(".")[-1] != dtype_name
419
+ ):
420
+ return None
421
+ self._residents.pin(mm_hash)
422
+ self._retire_events.pop(mm_hash, None)
423
+ return tensor
424
+
425
+ def _record_release_event(self) -> torch.Event | None:
426
+ pool = self.tensor
427
+ if pool is None or pool.device.type != "cuda":
428
+ return None
429
+ event = torch.Event()
430
+ event.record(torch.accelerator.current_stream(pool.device))
431
+ return event
432
+
433
+ def release_cached(self, lease: ResidentLease[MemoryAllocation]) -> None:
434
+ with self.lock:
435
+ released = self._residents.release(lease)
436
+ if released is not None:
437
+ self._defer_or_free(released, self._record_release_event())
438
+
439
+ def publish(
440
+ self,
441
+ mm_hash: str,
442
+ allocation: MemoryAllocation,
443
+ lease: ResidentLease[MemoryAllocation] | None = None,
444
+ *,
445
+ pin: bool = True,
446
+ ) -> MemoryAllocation:
447
+ with self.lock:
448
+ if lease is not None:
449
+ canonical, released = self._residents.consume(lease)
450
+ self._retire_events.pop(mm_hash, None)
451
+ if released is not None:
452
+ self._defer_or_free(released, self._record_release_event())
453
+ return canonical
454
+ previous = self._residents.get(mm_hash)
455
+ displaced = self._residents.insert(mm_hash, allocation)
456
+ if not pin:
457
+ self._residents.retire(mm_hash)
458
+ event = None
459
+ if previous is not None and previous is not allocation:
460
+ event = self._retire_events.pop(mm_hash, None)
461
+ if displaced is None or displaced is allocation:
462
+ return allocation
463
+ if event is None:
464
+ event = self._record_release_event()
465
+ self._defer_or_free(displaced, event)
466
+ return allocation
467
+
468
+ def retire_stale(
469
+ self,
470
+ encoder_cache: dict[str, torch.Tensor],
471
+ reserved_hashes: set[str] | None = None,
472
+ *,
473
+ freed: list[str] | None = None,
474
+ ) -> None:
475
+ if self.tensor is None:
476
+ return
477
+ with self.lock:
478
+ for mm_hash in self._residents.referenced() if freed is None else freed:
479
+ allocation = self._residents.get(mm_hash)
480
+ if allocation is None:
481
+ continue
482
+ if encoder_cache.get(mm_hash) is allocation.tensor:
483
+ continue
484
+ if reserved_hashes and mm_hash in reserved_hashes:
485
+ continue
486
+ event = self._record_release_event()
487
+ if event is not None:
488
+ self._retire_events[mm_hash] = event
489
+ self._residents.retire(mm_hash)
490
+ self._poll_frees_locked()
491
+
492
+ def drain_reclaimed(self) -> set[str]:
493
+ with self.lock:
494
+ reclaimed = self._reclaimed
495
+ self._reclaimed = set()
496
+ return reclaimed
497
+
498
+ def close(self) -> None:
499
+ with self.lock:
500
+ if not self._allocator.close(self._transfer):
501
+ return
502
+ self._residents.clear()
503
+ self._retire_events.clear()
504
+ self._pending_frees.clear()
vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Metadata exchanged by the Mooncake encoder-cache connector."""
4
+
5
+from __future__ import annotations
6
+
7
+from dataclasses import dataclass, field
8
+
9
+from vllm.distributed.ec_transfer.ec_connector.base import (
10
+ ECConnectorMetadata,
11
+ ECConnectorWorkerMetadata,
12
+)
13
+
14
+
15
+@dataclass
16
+class ECMooncakeLoadSpec:
17
+ """Describe a remote reservation or resident allocation to load."""
18
+
19
+ mm_hash: str
20
+ nbytes: int
21
+ shape: tuple[int, ...]
22
+ dtype: str
23
+ transfer_id: str
24
+ # The consumer pool still holds this item, so the load is a local handoff:
25
+ # no transfer, no producer.
26
+ local: bool = False
27
+
28
+
29
+@dataclass
30
+class ECMooncakePushSpec:
31
+ """Describe a destination reservation prepared before a tensor is ready."""
32
+
33
+ mm_hash: str
34
+ nbytes: int
35
+ shape: tuple[int, ...]
36
+ dtype: str
37
+ consumer_zmq: str
38
+ transfer_id: str
39
+ request_id: str = ""
40
+
41
+
42
+@dataclass
43
+class ECMooncakeConnectorMetadata(ECConnectorMetadata):
44
+ """Worker operations emitted for one Scheduler step."""
45
+
46
+ loads: list[ECMooncakeLoadSpec] = field(default_factory=list)
47
+ pushes: list[ECMooncakePushSpec] = field(default_factory=list)
48
+ freed: list[str] | None = None
49
+
50
+
51
+@dataclass
52
+class ECMooncakeWorkerMetadata(ECConnectorWorkerMetadata):
53
+ """Completion state reported from Workers to the Scheduler."""
54
+
55
+ loaded: set[str] = field(default_factory=set)
56
+ failed_loads: set[str] = field(default_factory=set)
57
+ # Items the receive pool dropped under pressure. The scheduler assumes an
58
+ # evicted item stays resident until told otherwise.
59
+ reclaimed: set[str] = field(default_factory=set)
60
+ pending_saves: bool = False
61
+ failed_saves: set[str] = field(default_factory=set)
62
+
63
+ def aggregate(self, other: ECConnectorWorkerMetadata) -> ECMooncakeWorkerMetadata:
64
+ assert isinstance(other, ECMooncakeWorkerMetadata)
65
+ return ECMooncakeWorkerMetadata(
66
+ # Every tensor-parallel rank gathers the embedding from its own
67
+ # cache, so an item counts as loaded only where all of them have
68
+ # it; one rank falling short must fail the load rather than leave
69
+ # the scheduler believing it is ready.
70
+ loaded=self.loaded & other.loaded,
71
+ failed_loads=self.failed_loads | other.failed_loads,
72
+ reclaimed=self.reclaimed | other.reclaimed,
73
+ pending_saves=self.pending_saves or other.pending_saves,
74
+ failed_saves=self.failed_saves | other.failed_saves,
75
+ )
vllm/distributed/ec_transfer/ec_connector/mooncake/producer.py
new
+418
@@ -0,0 +1,418 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Producer-side push lifecycle, futures, and source-tensor ownership.
4
+
5
+This module records when a push may advance or release its source. Worker
6
+orchestration performs the actual control exchanges and Mooncake writes.
7
+"""
8
+
9
+from __future__ import annotations
10
+
11
+import threading
12
+from collections import OrderedDict
13
+from collections.abc import Callable
14
+from concurrent.futures import Future, ThreadPoolExecutor
15
+from contextlib import suppress
16
+from dataclasses import dataclass, field
17
+from enum import Enum, auto
18
+from typing import Any
19
+
20
+import torch
21
+
22
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
23
+ ECMooncakePushSpec,
24
+)
25
+
26
+
27
+def _same_destination(
28
+ existing: ECMooncakePushSpec, incoming: ECMooncakePushSpec
29
+) -> bool:
30
+ """Whether two specs describe one push, ignoring which request asked.
31
+
32
+ The same encoding requested twice shares a transfer id legitimately; only
33
+ a different payload or destination is a collision.
34
+ """
35
+ return (
36
+ existing.mm_hash == incoming.mm_hash
37
+ and existing.nbytes == incoming.nbytes
38
+ and existing.shape == incoming.shape
39
+ and existing.dtype == incoming.dtype
40
+ and existing.consumer_zmq == incoming.consumer_zmq
41
+ )
42
+
43
+
44
+class ProducerPushState(Enum):
45
+ """Lifecycle of one Producer push from reservation to terminal state."""
46
+
47
+ # The reservation reply and source tensor may arrive in either order.
48
+ WAITING_INPUTS = auto()
49
+ WRITING = auto()
50
+ NOTIFYING = auto()
51
+ DONE = auto()
52
+ CANCEL_PENDING = auto()
53
+ CANCELLED = auto()
54
+ FAILED = auto()
55
+
56
+
57
+@dataclass
58
+class ProducerPushRecord:
59
+ """Collect all asynchronous state for one Producer push."""
60
+
61
+ spec: ECMooncakePushSpec
62
+ state: ProducerPushState
63
+ reservation_future: Future[list[dict[str, Any]]]
64
+ reservations: list[dict[str, Any]] = field(default_factory=list)
65
+ shard_futures: list[Future[Any]] = field(default_factory=list)
66
+ source_tensor: torch.Tensor | None = None
67
+ source_event: torch.Event | None = None
68
+ batch_future: Future[None] | None = None
69
+ error: str | None = None
70
+
71
+
72
+_ALLOWED_TRANSITIONS = {
73
+ ProducerPushState.WAITING_INPUTS: {
74
+ ProducerPushState.WRITING,
75
+ ProducerPushState.CANCEL_PENDING,
76
+ ProducerPushState.FAILED,
77
+ },
78
+ ProducerPushState.WRITING: {
79
+ ProducerPushState.NOTIFYING,
80
+ ProducerPushState.FAILED,
81
+ },
82
+ ProducerPushState.NOTIFYING: {
83
+ ProducerPushState.DONE,
84
+ ProducerPushState.FAILED,
85
+ },
86
+ ProducerPushState.CANCEL_PENDING: {
87
+ ProducerPushState.CANCELLED,
88
+ ProducerPushState.FAILED,
89
+ },
90
+ ProducerPushState.DONE: set(),
91
+ ProducerPushState.CANCELLED: set(),
92
+ ProducerPushState.FAILED: set(),
93
+}
94
+
95
+_TERMINAL_STATES = {
96
+ ProducerPushState.DONE,
97
+ ProducerPushState.CANCELLED,
98
+ ProducerPushState.FAILED,
99
+}
100
+_TERMINAL_LIMIT = 1 << 16
101
+
102
+
103
+class ProducerPushManager:
104
+ """Own Producer push records, transitions, and source tensor leases.
105
+
106
+ Terminal records are retained for duplicate-metadata idempotency. Separate
107
+ ordered indexes keep hot polling paths from scanning those tombstones.
108
+ """
109
+
110
+ def __init__(self, wake: Callable[[], None] = lambda: None) -> None:
111
+ self._records: OrderedDict[str, ProducerPushRecord] = OrderedDict()
112
+ self._active_ids: OrderedDict[str, None] = OrderedDict()
113
+ self._reapable_terminal_ids: OrderedDict[str, None] = OrderedDict()
114
+ self._unreported_ids: OrderedDict[str, None] = OrderedDict()
115
+ self._batch_ids: OrderedDict[str, None] = OrderedDict()
116
+ self._source_waiters: dict[str, OrderedDict[str, None]] = {}
117
+ self._lock = threading.RLock()
118
+ self._wake = wake
119
+
120
+ def reserve(
121
+ self,
122
+ spec: ECMooncakePushSpec,
123
+ submit: Callable[[], Future[list[dict[str, Any]]]],
124
+ ) -> tuple[ProducerPushRecord, bool]:
125
+ with self._lock:
126
+ existing = self._records.get(spec.transfer_id)
127
+ if existing is not None:
128
+ if not _same_destination(existing.spec, spec):
129
+ raise ValueError(
130
+ f"Conflicting EC destination for transfer_id={spec.transfer_id}"
131
+ )
132
+ return existing, False
133
+ record = ProducerPushRecord(
134
+ spec=spec,
135
+ state=ProducerPushState.WAITING_INPUTS,
136
+ reservation_future=submit(),
137
+ )
138
+ self._records[spec.transfer_id] = record
139
+ self._active_ids[spec.transfer_id] = None
140
+ self._source_waiters.setdefault(spec.mm_hash, OrderedDict())[
141
+ spec.transfer_id
142
+ ] = None
143
+ record.reservation_future.add_done_callback(
144
+ lambda future: self._reservation_done(record, future)
145
+ )
146
+ return record, True
147
+
148
+ def bind_source(
149
+ self,
150
+ mm_hash: str,
151
+ tensor: torch.Tensor,
152
+ ready_event: torch.Event | None,
153
+ ) -> None:
154
+ with self._lock:
155
+ waiters = self._source_waiters.pop(mm_hash, OrderedDict())
156
+ for transfer_id in waiters:
157
+ record = self._records[transfer_id]
158
+ if (
159
+ record.source_tensor is not None
160
+ or record.state is not ProducerPushState.WAITING_INPUTS
161
+ ):
162
+ continue
163
+ record.source_tensor = tensor
164
+ record.source_event = ready_event
165
+ self._wake()
166
+
167
+ def submit_batches(
168
+ self,
169
+ executor: ThreadPoolExecutor,
170
+ run_batch: Callable[[list[ProducerPushRecord]], None],
171
+ *,
172
+ wait: bool = False,
173
+ ) -> bool:
174
+ pending_event = False
175
+ with self._lock:
176
+ grouped: dict[str, list[ProducerPushRecord]] = {}
177
+ for transfer_id in list(self._active_ids):
178
+ record = self._records[transfer_id]
179
+ if (
180
+ record.source_tensor is not None
181
+ and record.batch_future is None
182
+ and record.state is ProducerPushState.WAITING_INPUTS
183
+ ):
184
+ if not record.reservation_future.done():
185
+ continue
186
+ if record.source_event is not None:
187
+ if wait:
188
+ record.source_event.synchronize()
189
+ elif not record.source_event.query():
190
+ pending_event = True
191
+ continue
192
+ grouped.setdefault(record.spec.consumer_zmq, []).append(record)
193
+ batches = list(grouped.values())
194
+ for records in batches:
195
+ future = executor.submit(run_batch, records)
196
+ for record in records:
197
+ record.batch_future = future
198
+ self._batch_ids[record.spec.transfer_id] = None
199
+ return pending_event
200
+
201
+ def resolve_reservations(self, record: ProducerPushRecord) -> list[dict[str, Any]]:
202
+ results = record.reservation_future.result()
203
+ with self._lock:
204
+ record.reservations = results
205
+ return list(results)
206
+
207
+ def finish_reservations(
208
+ self,
209
+ outcomes: list[tuple[ProducerPushRecord, list[dict[str, Any]] | Exception]],
210
+ ) -> None:
211
+ with self._lock:
212
+ for record, result in outcomes:
213
+ if isinstance(result, Exception):
214
+ record.reservation_future.set_exception(result)
215
+ else:
216
+ record.reservation_future.set_result(result)
217
+
218
+ def _reservation_done(
219
+ self,
220
+ record: ProducerPushRecord,
221
+ future: Future[list[dict[str, Any]]],
222
+ ) -> None:
223
+ try:
224
+ future.result()
225
+ except Exception as exc:
226
+ with self._lock:
227
+ self._set_error(record, exc)
228
+ if (
229
+ record.state is ProducerPushState.WAITING_INPUTS
230
+ and record.source_tensor is None
231
+ ):
232
+ self._transition(record, ProducerPushState.FAILED)
233
+ finally:
234
+ self._wake()
235
+
236
+ def settle_all(self, records: list[ProducerPushRecord]) -> None:
237
+ for record in records:
238
+ with suppress(Exception):
239
+ record.reservation_future.result()
240
+
241
+ def replace_reservations(
242
+ self,
243
+ record: ProducerPushRecord,
244
+ reservations: list[dict[str, Any]],
245
+ ) -> None:
246
+ with self._lock:
247
+ record.reservations = reservations
248
+
249
+ def track_shard_futures(
250
+ self,
251
+ records: list[ProducerPushRecord],
252
+ futures: list[Future[Any]],
253
+ ) -> None:
254
+ with self._lock:
255
+ for record in records:
256
+ record.shard_futures.extend(futures)
257
+
258
+ def begin_writing(self, record: ProducerPushRecord) -> None:
259
+ with self._lock:
260
+ if record.source_tensor is None:
261
+ raise RuntimeError(
262
+ f"Producer push {record.spec.transfer_id!r} has no source tensor"
263
+ )
264
+ self._transition(record, ProducerPushState.WRITING)
265
+
266
+ def begin_notifying(self, records: list[ProducerPushRecord]) -> None:
267
+ with self._lock:
268
+ for record in records:
269
+ self._transition(record, ProducerPushState.NOTIFYING)
270
+
271
+ def complete(self, records: list[ProducerPushRecord]) -> None:
272
+ with self._lock:
273
+ self._check_sources_releasable(records)
274
+ for record in records:
275
+ self._transition(record, ProducerPushState.DONE)
276
+ self._release_source(record)
277
+
278
+ def fail(self, records: list[ProducerPushRecord], error: Exception) -> None:
279
+ with self._lock:
280
+ self._check_sources_releasable(records)
281
+ for record in records:
282
+ if record.state not in _TERMINAL_STATES:
283
+ self._set_error(record, error)
284
+ self._transition(record, ProducerPushState.FAILED)
285
+ self._release_source(record)
286
+
287
+ def cancel_requests(self, request_ids: set[str] | None) -> list[ProducerPushRecord]:
288
+ """Cancel source-less waiters for requests, or all waiters at shutdown."""
289
+ cancelled = []
290
+ with self._lock:
291
+ for transfer_id in list(self._active_ids):
292
+ record = self._records[transfer_id]
293
+ if (
294
+ (
295
+ request_ids is not None
296
+ and record.spec.request_id not in request_ids
297
+ )
298
+ or record.source_tensor is not None
299
+ or record.batch_future is not None
300
+ ):
301
+ continue
302
+ if record.state is not ProducerPushState.WAITING_INPUTS:
303
+ continue
304
+ self._transition(record, ProducerPushState.CANCEL_PENDING)
305
+ cancelled.append(record)
306
+ return cancelled
307
+
308
+ def finish_cancel(self, record: ProducerPushRecord) -> None:
309
+ with self._lock:
310
+ if record.state is ProducerPushState.CANCEL_PENDING:
311
+ self._transition(record, ProducerPushState.CANCELLED)
312
+
313
+ def submit_cancel(
314
+ self,
315
+ record: ProducerPushRecord,
316
+ executor: ThreadPoolExecutor,
317
+ run_cancel: Callable[[ProducerPushRecord], None],
318
+ ) -> None:
319
+ with self._lock:
320
+ future = executor.submit(run_cancel, record)
321
+ record.batch_future = future
322
+ transfer_id = record.spec.transfer_id
323
+ self._batch_ids[transfer_id] = None
324
+ if record.state in _TERMINAL_STATES:
325
+ self._reapable_terminal_ids.pop(transfer_id, None)
326
+
327
+ def poll(self) -> list[tuple[str, str]]:
328
+ failures = []
329
+ with self._lock:
330
+ for transfer_id in list(self._batch_ids):
331
+ record = self._records[transfer_id]
332
+ future = record.batch_future
333
+ assert future is not None
334
+ if not future.done():
335
+ continue
336
+ try:
337
+ future.result()
338
+ except Exception as exc:
339
+ self._set_error(record, exc)
340
+ self._batch_ids.pop(transfer_id)
341
+ self._mark_reapable(record)
342
+ for transfer_id in list(self._unreported_ids):
343
+ if transfer_id in self._batch_ids:
344
+ continue
345
+ record = self._records[transfer_id]
346
+ assert record.error is not None
347
+ failures.append((record.spec.mm_hash, record.error))
348
+ self._unreported_ids.pop(transfer_id)
349
+ self._mark_reapable(record)
350
+ self._reap_terminals()
351
+ return failures
352
+
353
+ @property
354
+ def pending(self) -> bool:
355
+ with self._lock:
356
+ return bool(self._active_ids or self._batch_ids)
357
+
358
+ def _release_source(self, record: ProducerPushRecord) -> None:
359
+ record.source_tensor = None
360
+ record.source_event = None
361
+
362
+ def _set_error(self, record: ProducerPushRecord, error: BaseException) -> None:
363
+ if record.error is None:
364
+ record.error = str(error)
365
+ if record.state in _TERMINAL_STATES:
366
+ transfer_id = record.spec.transfer_id
367
+ self._unreported_ids[transfer_id] = None
368
+ self._reapable_terminal_ids.pop(transfer_id, None)
369
+
370
+ def _reap_terminals(self) -> None:
371
+ while len(self._reapable_terminal_ids) > _TERMINAL_LIMIT:
372
+ transfer_id, _ = self._reapable_terminal_ids.popitem(last=False)
373
+ self._records.pop(transfer_id)
374
+
375
+ def _mark_reapable(self, record: ProducerPushRecord) -> None:
376
+ transfer_id = record.spec.transfer_id
377
+ if (
378
+ record.state in _TERMINAL_STATES
379
+ and transfer_id not in self._batch_ids
380
+ and transfer_id not in self._unreported_ids
381
+ ):
382
+ self._reapable_terminal_ids[transfer_id] = None
383
+
384
+ def _drop_source_waiter(self, record: ProducerPushRecord) -> None:
385
+ waiters = self._source_waiters.get(record.spec.mm_hash)
386
+ if waiters is None:
387
+ return
388
+ waiters.pop(record.spec.transfer_id, None)
389
+ if not waiters:
390
+ self._source_waiters.pop(record.spec.mm_hash)
391
+
392
+ @staticmethod
393
+ def _check_sources_releasable(records: list[ProducerPushRecord]) -> None:
394
+ for record in records:
395
+ futures = [record.reservation_future, *record.shard_futures]
396
+ if record.source_tensor is not None and not all(
397
+ future.done() for future in futures
398
+ ):
399
+ raise RuntimeError(
400
+ f"Producer push {record.spec.transfer_id!r} released its "
401
+ "source too early"
402
+ )
403
+
404
+ def _transition(self, record: ProducerPushRecord, state: ProducerPushState) -> None:
405
+ if state not in _ALLOWED_TRANSITIONS[record.state]:
406
+ raise RuntimeError(
407
+ f"Cannot transition producer push {record.spec.transfer_id!r} from "
408
+ f"{record.state.name} to {state.name}"
409
+ )
410
+ record.state = state
411
+ if state is not ProducerPushState.WAITING_INPUTS:
412
+ self._drop_source_waiter(record)
413
+ if state in _TERMINAL_STATES:
414
+ transfer_id = record.spec.transfer_id
415
+ self._active_ids.pop(transfer_id, None)
416
+ if record.error is not None:
417
+ self._unreported_ids[transfer_id] = None
418
+ self._mark_reapable(record)
vllm/distributed/ec_transfer/ec_connector/mooncake/reservation.py
new
+487
@@ -0,0 +1,487 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Consumer-side reservation lifecycle and destination-memory ownership.
4
+
5
+Reservations make remote writes idempotent and ensure cancellation or expiry
6
+cannot free a destination while Mooncake may still be writing into it.
7
+"""
8
+
9
+from __future__ import annotations
10
+
11
+import threading
12
+import time
13
+import uuid
14
+from collections import OrderedDict
15
+from dataclasses import dataclass
16
+from enum import Enum, auto
17
+
18
+import torch
19
+
20
+from vllm.distributed.ec_transfer.ec_connector.mooncake.memory import (
21
+ ConsumerMemoryPool,
22
+ MemoryAllocation,
23
+ ResidentLease,
24
+)
25
+
26
+
27
+class ConsumerReservationState(Enum):
28
+ """Lifecycle of one Consumer destination reservation."""
29
+
30
+ WRITING = auto()
31
+ READY = auto()
32
+ CANCEL_PENDING = auto()
33
+ EXPIRE_PENDING = auto()
34
+ CANCELLED = auto()
35
+ EXPIRED = auto()
36
+
37
+
38
+_ALLOWED_TRANSITIONS = {
39
+ ConsumerReservationState.WRITING: {
40
+ ConsumerReservationState.READY,
41
+ ConsumerReservationState.CANCEL_PENDING,
42
+ ConsumerReservationState.EXPIRE_PENDING,
43
+ ConsumerReservationState.CANCELLED,
44
+ },
45
+ ConsumerReservationState.READY: {
46
+ ConsumerReservationState.CANCELLED,
47
+ ConsumerReservationState.EXPIRED,
48
+ },
49
+ ConsumerReservationState.CANCEL_PENDING: {ConsumerReservationState.CANCELLED},
50
+ ConsumerReservationState.EXPIRE_PENDING: {
51
+ ConsumerReservationState.CANCELLED,
52
+ ConsumerReservationState.EXPIRED,
53
+ },
54
+ ConsumerReservationState.CANCELLED: set(),
55
+ ConsumerReservationState.EXPIRED: set(),
56
+}
57
+
58
+_ACTIVE_STATES = {
59
+ ConsumerReservationState.WRITING,
60
+ ConsumerReservationState.READY,
61
+ ConsumerReservationState.CANCEL_PENDING,
62
+ ConsumerReservationState.EXPIRE_PENDING,
63
+}
64
+_DEFERRED_STATES = {
65
+ ConsumerReservationState.CANCEL_PENDING,
66
+ ConsumerReservationState.EXPIRE_PENDING,
67
+}
68
+_WRITER_OWNED_STATES = {
69
+ ConsumerReservationState.WRITING,
70
+ ConsumerReservationState.CANCEL_PENDING,
71
+ ConsumerReservationState.EXPIRE_PENDING,
72
+}
73
+
74
+
75
+@dataclass
76
+class ConsumerReservation:
77
+ """Own the identity and destination allocation of one remote push."""
78
+
79
+ transfer_id: str
80
+ mm_hash: str
81
+ reservation_id: str
82
+ state: ConsumerReservationState
83
+ shape: tuple[int, ...] = ()
84
+ dtype: str = ""
85
+ allocation: MemoryAllocation | None = None
86
+ lease: ResidentLease[MemoryAllocation] | None = None
87
+ expires_at: float = 0
88
+ writer_id: str = ""
89
+
90
+
91
+class ConsumerReservationManager:
92
+ """Own reservation transitions and destination allocation releases.
93
+
94
+ A remote writer owns a ``WRITING`` allocation. Cancellation and expiry are
95
+ therefore deferred until completion, unless refresh explicitly abandons
96
+ that writer before replacing its reservation.
97
+ """
98
+
99
+ def __init__(
100
+ self,
101
+ memory: ConsumerMemoryPool,
102
+ lease_ttl: float,
103
+ tombstone_limit: int,
104
+ ) -> None:
105
+ self._memory = memory
106
+ self._lease_ttl = lease_ttl
107
+ self._tombstone_limit = tombstone_limit
108
+ self._records: dict[str, ConsumerReservation] = {}
109
+ self._active_ids: dict[str, None] = {}
110
+ self._tombstones: OrderedDict[str, None] = OrderedDict()
111
+ self._writers: dict[str, ConsumerReservation] = {}
112
+ self._followers: dict[str, dict[str, None]] = {}
113
+ self._ready: list[str] = []
114
+ self._condition = threading.Condition(memory.lock)
115
+ self._shutting_down = False
116
+
117
+ def reserve(
118
+ self,
119
+ transfer_id: str,
120
+ mm_hash: str,
121
+ nbytes: int,
122
+ shape: tuple[int, ...],
123
+ dtype_name: str,
124
+ dtype: torch.dtype,
125
+ ) -> tuple[ConsumerReservation | None, bool]:
126
+ with self._condition:
127
+ if self._shutting_down:
128
+ raise RuntimeError("Consumer reservation manager is shutting down")
129
+ existing = self._records.get(transfer_id)
130
+ if (
131
+ existing is not None
132
+ and existing.state is ConsumerReservationState.CANCELLED
133
+ ):
134
+ return existing, False
135
+ if (
136
+ existing is not None
137
+ and existing.state is ConsumerReservationState.EXPIRED
138
+ ):
139
+ self._remove(transfer_id)
140
+ existing = None
141
+ if (
142
+ existing is not None
143
+ and existing.state is ConsumerReservationState.EXPIRE_PENDING
144
+ ):
145
+ raise RuntimeError(
146
+ f"Transfer {transfer_id!r} still has an active writer"
147
+ )
148
+ if existing is not None:
149
+ if (
150
+ existing.mm_hash != mm_hash
151
+ or existing.shape != shape
152
+ or existing.dtype != dtype_name
153
+ ):
154
+ raise ValueError("conflicting reservation for transfer_id")
155
+ if existing.state not in _ACTIVE_STATES:
156
+ raise RuntimeError(
157
+ f"Cannot reserve transfer {transfer_id!r} in "
158
+ f"{existing.state.name}"
159
+ )
160
+ if existing.state is ConsumerReservationState.WRITING:
161
+ existing.expires_at = time.monotonic() + self._lease_ttl
162
+ return existing, False
163
+
164
+ lease = self._memory.acquire_cached(mm_hash, shape, dtype)
165
+ now = time.monotonic()
166
+ if lease is not None:
167
+ record = ConsumerReservation(
168
+ transfer_id,
169
+ mm_hash,
170
+ uuid.uuid4().hex,
171
+ ConsumerReservationState.READY,
172
+ shape,
173
+ dtype_name,
174
+ lease.value,
175
+ lease,
176
+ now + self._lease_ttl,
177
+ )
178
+ self._insert(record)
179
+ return record, False
180
+ writer = self._writers.get(mm_hash)
181
+ if writer is not None and writer.state is ConsumerReservationState.WRITING:
182
+ if writer.shape != shape or writer.dtype != dtype_name:
183
+ raise ValueError("conflicting in-flight tensor for mm_hash")
184
+ record = ConsumerReservation(
185
+ transfer_id,
186
+ mm_hash,
187
+ uuid.uuid4().hex,
188
+ ConsumerReservationState.WRITING,
189
+ shape,
190
+ dtype_name,
191
+ writer.allocation,
192
+ expires_at=now + self._lease_ttl,
193
+ writer_id=writer.transfer_id,
194
+ )
195
+ self._insert(record)
196
+ self._followers.setdefault(writer.transfer_id, {})[transfer_id] = None
197
+ return record, False
198
+ allocation = self._memory.try_allocate(nbytes, shape, dtype)
199
+ if allocation is None:
200
+ self._expire_locked(time.monotonic())
201
+ allocation = self._memory.try_allocate(nbytes, shape, dtype)
202
+ if allocation is None:
203
+ allocation = self._memory.reclaim_and_allocate(nbytes, shape, dtype)
204
+ if allocation is None:
205
+ return None, False
206
+ record = ConsumerReservation(
207
+ transfer_id,
208
+ mm_hash,
209
+ uuid.uuid4().hex,
210
+ ConsumerReservationState.WRITING,
211
+ shape,
212
+ dtype_name,
213
+ allocation,
214
+ None,
215
+ now + self._lease_ttl,
216
+ )
217
+ self._insert(record)
218
+ self._writers[mm_hash] = record
219
+ return record, True
220
+
221
+ def status(self, transfer_id: str) -> ConsumerReservation | None:
222
+ with self._memory.lock:
223
+ record = self._records.get(transfer_id)
224
+ if record is None or record.state not in _ACTIVE_STATES:
225
+ return None
226
+ return record
227
+
228
+ def complete(self, transfer_id: str, reservation_id: str) -> tuple[bool, bool]:
229
+ with self._condition:
230
+ try:
231
+ record = self._records.get(transfer_id)
232
+ if record is None or record.reservation_id != reservation_id:
233
+ return False, False
234
+ if record.state is ConsumerReservationState.READY:
235
+ return True, False
236
+ if record.writer_id:
237
+ return False, False
238
+ if record.state in _WRITER_OWNED_STATES:
239
+ self._publish_followers(record)
240
+ if record.state in _DEFERRED_STATES:
241
+ terminal = (
242
+ ConsumerReservationState.CANCELLED
243
+ if record.state is ConsumerReservationState.CANCEL_PENDING
244
+ else ConsumerReservationState.EXPIRED
245
+ )
246
+ self._terminate(record, terminal)
247
+ return True, False
248
+ if record.state is not ConsumerReservationState.WRITING:
249
+ return False, False
250
+ self._transition(record, ConsumerReservationState.READY)
251
+ record.expires_at = time.monotonic() + self._lease_ttl
252
+ return True, True
253
+ finally:
254
+ self._condition.notify_all()
255
+
256
+ def _publish_followers(self, writer: ConsumerReservation) -> None:
257
+ if self._writers.get(writer.mm_hash) is writer:
258
+ self._writers.pop(writer.mm_hash)
259
+ followers = self._followers.pop(writer.transfer_id, {})
260
+ if not followers:
261
+ return
262
+ assert writer.allocation is not None
263
+ allocation = self._memory.publish(writer.mm_hash, writer.allocation, pin=False)
264
+ for record in [writer, *(self._records[key] for key in followers)]:
265
+ record.allocation = allocation
266
+ record.lease = self._memory.acquire_cached(
267
+ record.mm_hash, record.shape, allocation.tensor.dtype
268
+ )
269
+ assert record.lease is not None
270
+ if record is not writer:
271
+ record.writer_id = ""
272
+ self._transition(record, ConsumerReservationState.READY)
273
+ record.expires_at = time.monotonic() + self._lease_ttl
274
+ self._ready.append(record.transfer_id)
275
+
276
+ def drain_ready(self) -> list[str]:
277
+ with self._memory.lock:
278
+ ready, self._ready = self._ready, []
279
+ return ready
280
+
281
+ def begin_shutdown(self) -> None:
282
+ """Stop new reservations and cancel everything without a remote writer."""
283
+ with self._condition:
284
+ if self._shutting_down:
285
+ return
286
+ self._shutting_down = True
287
+ for transfer_id in list(self._active_ids):
288
+ record = self._records[transfer_id]
289
+ if record.writer_id or record.state is ConsumerReservationState.READY:
290
+ self._terminate(record, ConsumerReservationState.CANCELLED)
291
+ elif record.state is ConsumerReservationState.WRITING:
292
+ self._defer(record, ConsumerReservationState.CANCEL_PENDING)
293
+ self._condition.notify_all()
294
+
295
+ def wait_for_writers(self, timeout: float) -> bool:
296
+ """Wait until no reservation is still owned by a remote writer."""
297
+
298
+ def writers_finished() -> bool:
299
+ return not any(
300
+ self._records[transfer_id].state in _WRITER_OWNED_STATES
301
+ for transfer_id in self._active_ids
302
+ )
303
+
304
+ with self._condition:
305
+ return self._condition.wait_for(writers_finished, timeout=max(0, timeout))
306
+
307
+ def cancel(
308
+ self,
309
+ transfer_id: str,
310
+ reservation_id: str,
311
+ abandon: bool = False,
312
+ refresh: bool = False,
313
+ ) -> bool:
314
+ with self._memory.lock:
315
+ record = self._records.get(transfer_id)
316
+ if (
317
+ record is not None
318
+ and reservation_id
319
+ and record.reservation_id != reservation_id
320
+ ):
321
+ return False
322
+ if record is None:
323
+ record = ConsumerReservation(
324
+ transfer_id,
325
+ "",
326
+ "",
327
+ ConsumerReservationState.CANCELLED,
328
+ )
329
+ self._insert(record)
330
+ self._set_tombstone_deadline(record)
331
+ self._reap_tombstones(time.monotonic())
332
+ return True
333
+ if record.state is ConsumerReservationState.CANCELLED:
334
+ self._set_tombstone_deadline(record)
335
+ self._reap_tombstones(time.monotonic())
336
+ return True
337
+ if refresh:
338
+ if not abandon or record.state not in {
339
+ ConsumerReservationState.WRITING,
340
+ ConsumerReservationState.EXPIRE_PENDING,
341
+ }:
342
+ return False
343
+ if record.state is ConsumerReservationState.WRITING:
344
+ self._transition(record, ConsumerReservationState.EXPIRE_PENDING)
345
+ self._terminate(record, ConsumerReservationState.EXPIRED)
346
+ self._condition.notify_all()
347
+ self._reap_tombstones(time.monotonic())
348
+ return True
349
+ if record.writer_id:
350
+ self._terminate(record, ConsumerReservationState.CANCELLED)
351
+ return True
352
+ if record.state in _DEFERRED_STATES and not abandon:
353
+ return True
354
+ if record.state is ConsumerReservationState.WRITING and not abandon:
355
+ self._defer(record, ConsumerReservationState.CANCEL_PENDING)
356
+ return True
357
+ if record.state not in _ACTIVE_STATES:
358
+ return False
359
+ self._terminate(record, ConsumerReservationState.CANCELLED)
360
+ if self._shutting_down:
361
+ self._condition.notify_all()
362
+ self._reap_tombstones(time.monotonic())
363
+ return True
364
+
365
+ def take(self, transfer_id: str, mm_hash: str) -> MemoryAllocation:
366
+ with self._memory.lock:
367
+ record = self._records.get(transfer_id)
368
+ if (
369
+ record is None
370
+ or record.state is not ConsumerReservationState.READY
371
+ or record.mm_hash != mm_hash
372
+ or record.allocation is None
373
+ ):
374
+ raise RuntimeError(
375
+ f"Pushed EC tensor is not ready for mm_hash={mm_hash}"
376
+ )
377
+ allocation = self._memory.publish(mm_hash, record.allocation, record.lease)
378
+ record.allocation = None
379
+ record.lease = None
380
+ self._remove(transfer_id)
381
+ return allocation
382
+
383
+ def expire(self) -> int:
384
+ with self._memory.lock:
385
+ return self._expire_locked(time.monotonic())
386
+
387
+ def retire_stale(
388
+ self, encoder_cache: dict[str, torch.Tensor], freed: list[str] | None = None
389
+ ) -> None:
390
+ with self._memory.lock:
391
+ self._memory.retire_stale(encoder_cache, freed=freed)
392
+
393
+ def _expire_locked(self, now: float) -> int:
394
+ expired = 0
395
+ for transfer_id in list(self._active_ids):
396
+ record = self._records[transfer_id]
397
+ if record.expires_at > now:
398
+ continue
399
+ if record.state is ConsumerReservationState.READY:
400
+ self._terminate(record, ConsumerReservationState.EXPIRED)
401
+ expired += 1
402
+ elif record.state is ConsumerReservationState.WRITING:
403
+ if record.writer_id:
404
+ self._terminate(record, ConsumerReservationState.CANCELLED)
405
+ else:
406
+ self._defer(record, ConsumerReservationState.EXPIRE_PENDING)
407
+ self._reap_tombstones(now)
408
+ return expired
409
+
410
+ def _defer(
411
+ self, record: ConsumerReservation, state: ConsumerReservationState
412
+ ) -> None:
413
+ """Keep the destination until the writer completes or abandons it."""
414
+ self._transition(record, state)
415
+ record.expires_at = float("inf")
416
+
417
+ def _terminate(
418
+ self, record: ConsumerReservation, state: ConsumerReservationState
419
+ ) -> None:
420
+ self._transition(record, state)
421
+ self._release(record)
422
+ self._set_tombstone_deadline(record)
423
+
424
+ def _release(self, record: ConsumerReservation) -> None:
425
+ if record.writer_id:
426
+ followers = self._followers.get(record.writer_id, {})
427
+ followers.pop(record.transfer_id, None)
428
+ record.allocation = None
429
+ return
430
+ if self._writers.get(record.mm_hash) is record:
431
+ self._writers.pop(record.mm_hash)
432
+ for transfer_id in list(self._followers.pop(record.transfer_id, {})):
433
+ self._terminate(
434
+ self._records[transfer_id], ConsumerReservationState.CANCELLED
435
+ )
436
+ allocation = record.allocation
437
+ if allocation is None:
438
+ return
439
+ if record.lease is not None:
440
+ self._memory.release_cached(record.lease)
441
+ else:
442
+ self._memory.free(allocation)
443
+ record.allocation = None
444
+ record.lease = None
445
+
446
+ def _set_tombstone_deadline(self, record: ConsumerReservation) -> None:
447
+ record.expires_at = time.monotonic() + self._lease_ttl
448
+ self._tombstones[record.transfer_id] = None
449
+ self._tombstones.move_to_end(record.transfer_id)
450
+
451
+ def _reap_tombstones(self, now: float) -> int:
452
+ dropped = 0
453
+ while self._tombstones:
454
+ transfer_id = next(iter(self._tombstones))
455
+ record = self._records[transfer_id]
456
+ if (
457
+ record.expires_at > now
458
+ and len(self._tombstones) <= self._tombstone_limit
459
+ ):
460
+ break
461
+ self._remove(transfer_id)
462
+ dropped += 1
463
+ return dropped
464
+
465
+ def _insert(self, record: ConsumerReservation) -> None:
466
+ self._records[record.transfer_id] = record
467
+ if record.state in _ACTIVE_STATES:
468
+ self._active_ids[record.transfer_id] = None
469
+
470
+ def _remove(self, transfer_id: str) -> None:
471
+ self._records.pop(transfer_id, None)
472
+ self._active_ids.pop(transfer_id, None)
473
+ self._tombstones.pop(transfer_id, None)
474
+
475
+ def _transition(
476
+ self, record: ConsumerReservation, state: ConsumerReservationState
477
+ ) -> None:
478
+ if state not in _ALLOWED_TRANSITIONS[record.state]:
479
+ raise RuntimeError(
480
+ f"Cannot transition {record.transfer_id!r} from "
481
+ f"{record.state.name} to {state.name}"
482
+ )
483
+ record.state = state
484
+ if state in _ACTIVE_STATES:
485
+ self._active_ids[record.transfer_id] = None
486
+ else:
487
+ self._active_ids.pop(record.transfer_id, None)
vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py
new
+537
@@ -0,0 +1,537 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scheduler-side planning and observation for Mooncake cache transfers.
4
+
5
+The Scheduler prepares Producer reservations, consumes Consumer readiness
6
+events, emits per-step Worker metadata, and converts Worker results back into
7
+request availability without owning tensor memory or running data transfers.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import math
13
+import time
14
+from collections import OrderedDict
15
+from concurrent.futures import Future, ThreadPoolExecutor
16
+from typing import TYPE_CHECKING, Any
17
+
18
+import torch
19
+
20
+from vllm.distributed.ec_transfer.ec_connector.base import (
21
+ ECConnectorMetadata,
22
+)
23
+from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
24
+ _get_encoder_cache_hidden_dim,
25
+)
26
+from vllm.distributed.ec_transfer.ec_connector.mooncake.config import (
27
+ _RESERVATION_TTL_SECONDS,
28
+ MooncakeECConfig,
29
+)
30
+from vllm.distributed.ec_transfer.ec_connector.mooncake.control import (
31
+ ControlClient,
32
+ EventInbox,
33
+ make_cancel_request,
34
+)
35
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
36
+ ECMooncakeConnectorMetadata,
37
+ ECMooncakeLoadSpec,
38
+ ECMooncakePushSpec,
39
+ ECMooncakeWorkerMetadata,
40
+)
41
+from vllm.distributed.ec_transfer.ec_connector.mooncake.state import (
42
+ SchedulerTransferState,
43
+ SchedulerTransferTable,
44
+)
45
+from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
46
+ ensure_mooncake_available,
47
+)
48
+from vllm.distributed.ec_transfer.ec_connector.utils import (
49
+ PlaceholderMetadataResolver,
50
+ collect_ec_item_metadata,
51
+)
52
+from vllm.logger import init_logger
53
+from vllm.v1.core.sched.output import SchedulerOutput
54
+from vllm.v1.outputs import ECConnectorOutput
55
+
56
+if TYPE_CHECKING:
57
+ from vllm.config import VllmConfig
58
+
59
+logger = init_logger(__name__)
60
+
61
+# Unmatched `ec_items` entries are a protocol desync, not a per-request
62
+# event: warn on the first few, then only once in a while.
63
+_MAX_UNRESOLVED_TRANSFER_ID_WARNINGS = 5
64
+
65
+_DRAIN_MIN_INTERVAL = 0.005
66
+_MAX_PENDING_EVENTS = 4096
67
+_MAX_TERMINAL_TRANSFER_RECORDS = 1 << 16
68
+_CANCEL_ATTEMPTS = 2
69
+_CONTROL_WORKERS = 8
70
+
71
+
72
+class ECMooncakeScheduler:
73
+ """Coordinate Mooncake transfers from the vLLM Scheduler process."""
74
+
75
+ def __init__(self, vllm_config: VllmConfig) -> None:
76
+ ensure_mooncake_available()
77
+ config = MooncakeECConfig.from_vllm_config(vllm_config)
78
+
79
+ self._is_producer = config.is_producer
80
+ self._is_consumer = config.is_consumer
81
+ self._control_addr = config.control_addr
82
+ self._push_wait_timeout = config.push_wait_timeout_s
83
+ self._encoder_cache_hidden_dim = (
84
+ _get_encoder_cache_hidden_dim(vllm_config) if config.is_producer else None
85
+ )
86
+ self._model_config = vllm_config.model_config
87
+ self._control_client = ControlClient(config.control_timeout_ms)
88
+ self._control_executor = ThreadPoolExecutor(
89
+ max_workers=_CONTROL_WORKERS,
90
+ thread_name_prefix="ec-mooncake-control",
91
+ )
92
+ self._event_inbox = EventInbox(self._control_client, self._control_executor)
93
+
94
+ self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config)
95
+ self._unresolved_transfer_ids = 0
96
+ self._drain_pending = True
97
+ self._drained_at = 0.0
98
+ self._pending_cancels: dict[str, Future[Any]] = {}
99
+ self._transfers = SchedulerTransferTable(
100
+ config.pool_size, _RESERVATION_TTL_SECONDS
101
+ )
102
+ self._scheduler_pending_work = False
103
+ self._pushes_to_prepare: dict[str, ECMooncakePushSpec] = {}
104
+ self._prepared_push_transfer_ids: set[str] = set()
105
+ self._event_ready_shards: OrderedDict[str, set[int]] = OrderedDict()
106
+ # Mirror of the engine's encoder cache, maintained from the alloc and
107
+ # free notifications the Scheduler already sends. Tracking it here
108
+ # keeps `ensure_cache_available` on the upstream two-argument shape.
109
+ self._local_cache: set[str] = set()
110
+ self._failed_saves: set[str] = set()
111
+
112
+ def _cancel_remote(self, consumer_zmq: str, transfer_id: str) -> bool:
113
+ pending = None
114
+ for _ in range(_CANCEL_ATTEMPTS):
115
+ pending = self._control_client.discover_shards(consumer_zmq)
116
+ if pending is not None:
117
+ break
118
+ if pending is None:
119
+ raise RuntimeError(
120
+ f"Could not discover every EC consumer shard at {consumer_zmq}"
121
+ )
122
+
123
+ cancelled = False
124
+ error: BaseException | None = None
125
+ for _ in range(_CANCEL_ATTEMPTS):
126
+ failed = []
127
+ for addr in pending:
128
+ try:
129
+ result = self._control_client.request(
130
+ addr,
131
+ make_cancel_request(transfer_id),
132
+ )
133
+ except Exception as exc:
134
+ if error is None:
135
+ error = exc
136
+ failed.append(addr)
137
+ continue
138
+ cancelled |= isinstance(result, dict) and bool(result.get("cancelled"))
139
+ if not failed:
140
+ return cancelled
141
+ pending = failed
142
+ if error is not None:
143
+ raise error
144
+ return cancelled
145
+
146
+ def _note_awaiting_push(
147
+ self,
148
+ mm_hash: str,
149
+ transfer_id: str,
150
+ request_id: str,
151
+ now: float,
152
+ ) -> None:
153
+ record = self._transfers.wait_for_event(
154
+ transfer_id,
155
+ request_id,
156
+ mm_hash,
157
+ now + self._push_wait_timeout,
158
+ )
159
+ if record is None:
160
+ # The id names another encoding; the request is already failed.
161
+ return
162
+ if record.state is not SchedulerTransferState.WAITING_EVENT:
163
+ return
164
+ assert record.deadline is not None
165
+ if now < record.deadline:
166
+ return
167
+ elapsed = now - record.deadline + self._push_wait_timeout
168
+ self._transfers.mark_unavailable(transfer_id, now)
169
+ logger.warning(
170
+ "EC Mooncake waited %.1fs for a push of mm_hash=%s "
171
+ "(transfer_id=%s) that never arrived; "
172
+ "requests needing it fail with a retryable error.",
173
+ elapsed,
174
+ mm_hash,
175
+ transfer_id,
176
+ )
177
+
178
+ def take_unavailable_requests(self) -> set[str]:
179
+ failed, self._failed_saves = self._failed_saves, set()
180
+ return failed | self._transfers.take_unavailable_requests()
181
+
182
+ def _poll_pending_cancels(self) -> None:
183
+ pending = {}
184
+ for transfer_id, future in self._pending_cancels.items():
185
+ if not future.done():
186
+ pending[transfer_id] = future
187
+ continue
188
+ try:
189
+ future.result()
190
+ except Exception:
191
+ logger.warning(
192
+ "EC Mooncake reservation cancellation failed", exc_info=True
193
+ )
194
+ self._pending_cancels = pending
195
+
196
+ def _note_shard_ready(self, data: dict[str, Any]) -> bool:
197
+ shard_count = self._event_inbox.shard_count
198
+ if shard_count <= 1:
199
+ return True
200
+ transfer_id = str(data["transfer_id"])
201
+ record = self._transfers.get(transfer_id)
202
+ if (
203
+ record is not None
204
+ and record.state is not SchedulerTransferState.WAITING_EVENT
205
+ ):
206
+ return False
207
+ shard = data.get("shard")
208
+ shards = self._event_ready_shards.setdefault(transfer_id, set())
209
+ self._event_ready_shards.move_to_end(transfer_id)
210
+ shards.add(int(shard) if shard is not None else len(shards))
211
+ if len(shards) < shard_count:
212
+ while len(self._event_ready_shards) > _MAX_PENDING_EVENTS:
213
+ self._event_ready_shards.popitem(last=False)
214
+ return False
215
+ self._event_ready_shards.pop(transfer_id, None)
216
+ return True
217
+
218
+ def _forget_shard_readiness(self, transfer_id: str) -> None:
219
+ self._event_ready_shards.pop(transfer_id, None)
220
+
221
+ def _store_pushed_spec(self, data: dict[str, Any]) -> bool:
222
+ transfer_id = str(data["transfer_id"])
223
+ identifier = str(data["mm_hash"])
224
+ _, accepted = self._transfers.observe_ready(
225
+ ECMooncakeLoadSpec(
226
+ mm_hash=identifier,
227
+ nbytes=int(data["nbytes"]),
228
+ shape=tuple(int(value) for value in data["shape"]),
229
+ dtype=str(data["dtype"]),
230
+ transfer_id=transfer_id,
231
+ ),
232
+ time.monotonic() + _RESERVATION_TTL_SECONDS,
233
+ )
234
+ return accepted
235
+
236
+ def _queue_cancel(
237
+ self,
238
+ transfer_id: str,
239
+ mm_hash: str = "",
240
+ request_id: str = "",
241
+ ) -> None:
242
+ if not self._transfers.cancel(
243
+ transfer_id,
244
+ time.monotonic(),
245
+ mm_hash=mm_hash,
246
+ request_id=request_id,
247
+ ):
248
+ return
249
+ self._forget_shard_readiness(transfer_id)
250
+ self._pending_cancels[transfer_id] = self._control_executor.submit(
251
+ self._cancel_remote,
252
+ self._control_addr,
253
+ transfer_id,
254
+ )
255
+
256
+ def _expire_transfers(self) -> None:
257
+ now = time.monotonic()
258
+ expired = self._transfers.expire(now, _MAX_TERMINAL_TRANSFER_RECORDS)
259
+ for record in expired:
260
+ self._queue_cancel(record.transfer_id)
261
+
262
+ def _drain_push_notifications(self) -> None:
263
+ now = time.monotonic()
264
+ if not self._drain_pending and now - self._drained_at < _DRAIN_MIN_INTERVAL:
265
+ return
266
+ self._drain_pending = False
267
+ self._drained_at = now
268
+ self._poll_pending_cancels()
269
+ self._expire_transfers()
270
+ events = self._event_inbox.drain(self._control_addr)
271
+ for data in events:
272
+ try:
273
+ if not isinstance(data, dict):
274
+ raise TypeError("EC readiness event must be an object")
275
+ if not data.get("ready"):
276
+ continue
277
+ self._accept_ready_event(data)
278
+ except (KeyError, TypeError, ValueError):
279
+ # Readiness events cross a plain PULL socket, so a malformed
280
+ # one costs that event and not the engine.
281
+ logger.warning(
282
+ "Discarding a malformed EC readiness event.", exc_info=True
283
+ )
284
+
285
+ def _accept_ready_event(self, data: dict[str, Any]) -> None:
286
+ transfer_id = str(data["transfer_id"])
287
+ record = self._transfers.get(transfer_id)
288
+ if record is not None and record.state in {
289
+ SchedulerTransferState.CANCELLED,
290
+ SchedulerTransferState.UNAVAILABLE,
291
+ SchedulerTransferState.EXPIRED,
292
+ SchedulerTransferState.FAILED,
293
+ }:
294
+ return
295
+ if not self._note_shard_ready(data):
296
+ return
297
+ self._store_pushed_spec(data)
298
+
299
+ def has_cache_item(self, identifier: str) -> bool:
300
+ if not self._is_consumer:
301
+ return False
302
+ self._drain_push_notifications()
303
+ return self._transfers.has_state(
304
+ identifier,
305
+ (
306
+ SchedulerTransferState.READY,
307
+ SchedulerTransferState.RESIDENT,
308
+ SchedulerTransferState.AVAILABLE,
309
+ ),
310
+ )
311
+
312
+ def _warn_unresolved_transfer_id(
313
+ self, request: Any, index: int, where: str
314
+ ) -> None:
315
+ """Report an `ec_items` entry that cannot be matched to a feature.
316
+
317
+ A caller that cannot resolve a transfer id has to fall back to a
318
+ locally invented one or skip its bookkeeping entirely, and either way
319
+ the two sides of a transfer stop agreeing on its name. That desyncs
320
+ silently -- reservations are never released and only surface minutes
321
+ later as a full buffer pool -- so say it out loud the first time.
322
+
323
+ `mm_hash` in `ec_items` must be a value some engine reported, never one
324
+ the caller derived itself: `mm_features[i].identifier` folds in the
325
+ engine's media_io_kwargs and mm_processor_kwargs, so a media uuid alone
326
+ no longer equals it. Omit the field to match by position instead.
327
+ """
328
+ self._unresolved_transfer_ids += 1
329
+ count = self._unresolved_transfer_ids
330
+ if count > _MAX_UNRESOLVED_TRANSFER_ID_WARNINGS and count % 1000:
331
+ return
332
+ params = getattr(request, "ec_transfer_params", None)
333
+ items = (params or {}).get("ec_items") or []
334
+ logger.warning(
335
+ "EC Mooncake could not resolve a transfer id at %s for req=%s "
336
+ "index=%d: feature identifier=%s does not match any of the %d "
337
+ "ec_items %s (ec_transfer_params present=%s). Occurrence %d. "
338
+ "ec_items[].mm_hash must echo an engine-reported identifier, or "
339
+ "be omitted so the entry is matched by position.",
340
+ where,
341
+ request.request_id,
342
+ index,
343
+ request.mm_features[index].identifier[:16],
344
+ len(items),
345
+ [str(item.get("mm_hash"))[:16] for item in items[:4]],
346
+ params is not None,
347
+ count,
348
+ )
349
+
350
+ @staticmethod
351
+ def _request_transfer_id(request: Any, index: int) -> str | None:
352
+ params = getattr(request, "ec_transfer_params", None) or {}
353
+ items = params.get("ec_items") or []
354
+ mm_hash = request.mm_features[index].identifier
355
+ if index < len(items):
356
+ item = items[index]
357
+ if item.get("mm_hash") in (None, mm_hash) and item.get("transfer_id"):
358
+ return str(item["transfer_id"])
359
+ for item in items:
360
+ if item.get("mm_hash") == mm_hash and item.get("transfer_id"):
361
+ return str(item["transfer_id"])
362
+ return None
363
+
364
+ def ensure_cache_available(self, request: Any, num_computed_tokens: int) -> bool:
365
+ if self._is_producer:
366
+ for index, feature in enumerate(request.mm_features):
367
+ if (
368
+ feature.mm_position.offset + feature.mm_position.length
369
+ > num_computed_tokens
370
+ ):
371
+ self._prepare_push_spec(request, index)
372
+ if not self._is_consumer:
373
+ return True
374
+
375
+ self._drain_push_notifications()
376
+ # One timestamp for the whole decision, so the deadlines it hands out
377
+ # cannot disagree between features.
378
+ now = time.monotonic()
379
+ all_ready = True
380
+ for index, feature in enumerate(request.mm_features):
381
+ if (
382
+ feature.mm_position.offset + feature.mm_position.length
383
+ <= num_computed_tokens
384
+ ):
385
+ continue
386
+ mm_hash = feature.identifier
387
+ transfer_id = self._request_transfer_id(request, index)
388
+ if transfer_id is not None:
389
+ self._transfers.touch_available(
390
+ transfer_id, now + _RESERVATION_TTL_SECONDS
391
+ )
392
+ if mm_hash in self._local_cache:
393
+ continue
394
+ if self._transfers.has_state(mm_hash, (SchedulerTransferState.READY,)):
395
+ continue
396
+ if self._transfers.has_state(mm_hash, (SchedulerTransferState.LOADING,)):
397
+ all_ready = False
398
+ continue
399
+ record = self._transfers.begin_load(
400
+ mm_hash,
401
+ transfer_id,
402
+ request.request_id,
403
+ now + self._push_wait_timeout,
404
+ )
405
+ if record is not None:
406
+ self._scheduler_pending_work = True
407
+ all_ready = False
408
+ else:
409
+ waiting_id = transfer_id or f"{request.request_id}:{index}"
410
+ self._note_awaiting_push(mm_hash, waiting_id, request.request_id, now)
411
+ all_ready = False
412
+ return all_ready
413
+
414
+ def _prepare_push_spec(self, request: Any, index: int) -> None:
415
+ params = getattr(request, "ec_transfer_params", None) or {}
416
+ consumer_zmq = params.get("consumer_zmq")
417
+ mm_hash = request.mm_features[index].identifier
418
+ transfer_id = self._request_transfer_id(request, index)
419
+ if transfer_id is None:
420
+ # Still push: after the proxy has rewritten the item to embeds the
421
+ # consumer has no media left to fall back on, so a nameless push
422
+ # beats none. But the consumer knows this transfer by the id it
423
+ # sent, not by the one invented here, so its cancel will never
424
+ # reach the reservation this push is about to take.
425
+ if consumer_zmq:
426
+ self._warn_unresolved_transfer_id(request, index, "push prepare")
427
+ transfer_id = f"{request.request_id}:{index}"
428
+ if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids:
429
+ return
430
+ num_tokens = request.get_num_encoder_embeds(index)
431
+ dtype = self._model_config.dtype
432
+ assert isinstance(dtype, torch.dtype)
433
+ assert self._encoder_cache_hidden_dim is not None
434
+ dtype_name = str(dtype).split(".")[-1]
435
+ shape = (num_tokens, self._encoder_cache_hidden_dim)
436
+ nbytes = math.prod(shape) * dtype.itemsize
437
+ self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec(
438
+ mm_hash=mm_hash,
439
+ nbytes=nbytes,
440
+ shape=shape,
441
+ dtype=dtype_name,
442
+ consumer_zmq=str(consumer_zmq),
443
+ transfer_id=transfer_id,
444
+ request_id=request.request_id,
445
+ )
446
+ self._prepared_push_transfer_ids.add(transfer_id)
447
+
448
+ def update_state_after_alloc(self, request: Any, index: int) -> None:
449
+ self._local_cache.add(request.mm_features[index].identifier)
450
+ if self._is_producer:
451
+ self._prepare_push_spec(request, index)
452
+
453
+ def update_state_after_free(self, request: Any, index: int) -> None:
454
+ if not self._is_consumer:
455
+ return
456
+ transfer_id = self._request_transfer_id(request, index)
457
+ if transfer_id is None:
458
+ self._warn_unresolved_transfer_id(request, index, "encoder-cache free")
459
+ return
460
+ self._queue_cancel(
461
+ transfer_id,
462
+ mm_hash=request.mm_features[index].identifier,
463
+ request_id=request.request_id,
464
+ )
465
+
466
+ def build_connector_meta(
467
+ self, scheduler_output: SchedulerOutput
468
+ ) -> ECConnectorMetadata:
469
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
470
+ self._local_cache.discard(mm_hash)
471
+ self._transfers.release_ready(mm_hash, time.monotonic())
472
+ for transfer_id in self._transfers.drain_orphaned():
473
+ self._queue_cancel(transfer_id)
474
+ meta = ECMooncakeConnectorMetadata(
475
+ freed=scheduler_output.free_encoder_mm_hashes
476
+ )
477
+ for push_spec in self._pushes_to_prepare.values():
478
+ meta.pushes.append(push_spec)
479
+ self._pushes_to_prepare.clear()
480
+ for record in self._transfers.take_loads_to_dispatch():
481
+ assert record.spec is not None
482
+ meta.loads.append(record.spec)
483
+ self._poll_pending_cancels()
484
+ self._drain_pending = True
485
+ return meta
486
+
487
+ def update_connector_output(self, connector_output: ECConnectorOutput) -> None:
488
+ meta = connector_output.ec_connector_worker_meta
489
+ if not isinstance(meta, ECMooncakeWorkerMetadata):
490
+ return
491
+ for mm_hash in meta.loaded:
492
+ self._transfers.complete_load(mm_hash)
493
+ for mm_hash in meta.failed_loads:
494
+ self._transfers.fail_load(mm_hash, time.monotonic())
495
+ for mm_hash in meta.reclaimed:
496
+ self._transfers.reclaim(mm_hash, time.monotonic())
497
+ self._scheduler_pending_work = meta.pending_saves
498
+ self._failed_saves.update(meta.failed_saves)
499
+
500
+ def has_pending_push_work(self) -> bool:
501
+ return self._scheduler_pending_work
502
+
503
+ def request_finished(self, request: Any) -> tuple[bool, dict[str, Any] | None]:
504
+ if self._is_consumer:
505
+ for index in range(len(request.mm_features)):
506
+ transfer_id = self._request_transfer_id(request, index)
507
+ if transfer_id is None:
508
+ self._warn_unresolved_transfer_id(request, index, "request finish")
509
+ continue
510
+ self._queue_cancel(
511
+ transfer_id,
512
+ mm_hash=request.mm_features[index].identifier,
513
+ request_id=request.request_id,
514
+ )
515
+ if self._is_producer and self._prepared_push_transfer_ids:
516
+ for index in range(len(request.mm_features)):
517
+ transfer_id = self._request_transfer_id(request, index)
518
+ if transfer_id is None:
519
+ transfer_id = f"{request.request_id}:{index}"
520
+ self._prepared_push_transfer_ids.discard(transfer_id)
521
+ if not self._is_producer:
522
+ return False, None
523
+
524
+ items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver)
525
+ for index, feature in enumerate(request.mm_features):
526
+ transfer_id = self._request_transfer_id(request, index)
527
+ if transfer_id is not None:
528
+ items[feature.identifier]["transfer_id"] = transfer_id
529
+
530
+ if not items:
531
+ return False, None
532
+ return False, items
533
+
534
+ def close(self) -> None:
535
+ self._control_executor.shutdown(wait=True, cancel_futures=True)
536
+ self._control_client.close()
537
+ self._event_inbox.close()
vllm/distributed/ec_transfer/ec_connector/mooncake/state.py
new
+448
@@ -0,0 +1,448 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scheduler-owned lifecycle and indexes for Consumer-bound transfers."""
4
+
5
+from __future__ import annotations
6
+
7
+from collections import OrderedDict, deque
8
+from collections.abc import Iterable
9
+from dataclasses import dataclass, field, replace
10
+from enum import Enum, auto
11
+
12
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
13
+ ECMooncakeLoadSpec,
14
+)
15
+from vllm.logger import init_logger
16
+
17
+logger = init_logger(__name__)
18
+
19
+
20
+class SchedulerTransferState(Enum):
21
+ """Lifecycle states visible to the Scheduler.
22
+
23
+ The states cover waiting for a Consumer event, dispatching a Worker load,
24
+ retaining a locally reusable tensor, and terminal failure conditions.
25
+ """
26
+
27
+ WAITING_EVENT = auto()
28
+ AVAILABLE = auto()
29
+ LOADING = auto()
30
+ READY = auto()
31
+ RESIDENT = auto()
32
+ UNAVAILABLE = auto()
33
+ EXPIRED = auto()
34
+ FAILED = auto()
35
+ CANCELLED = auto()
36
+
37
+
38
+@dataclass
39
+class SchedulerTransfer:
40
+ """Track one transfer as observed by the Scheduler."""
41
+
42
+ transfer_id: str
43
+ request_id: str
44
+ mm_hash: str
45
+ state: SchedulerTransferState
46
+ spec: ECMooncakeLoadSpec | None
47
+ deadline: float | None
48
+ notified_requests: set[str] = field(default_factory=set, repr=False)
49
+
50
+
51
+_TERMINAL_STATES = {
52
+ SchedulerTransferState.UNAVAILABLE,
53
+ SchedulerTransferState.EXPIRED,
54
+ SchedulerTransferState.FAILED,
55
+ SchedulerTransferState.CANCELLED,
56
+}
57
+
58
+
59
+class SchedulerTransferTable:
60
+ """Own Scheduler transfer state, lookup indexes, and dispatch queues.
61
+
62
+ The Scheduler is single-threaded, so direct transitions are sufficient;
63
+ the Producer and Consumer managers retain stricter transition matrices
64
+ because callbacks and control requests can race there.
65
+ """
66
+
67
+ def __init__(self, resident_capacity: int, tombstone_ttl: float) -> None:
68
+ self._resident_capacity = resident_capacity
69
+ self._tombstone_ttl = tombstone_ttl
70
+ self._records: OrderedDict[str, SchedulerTransfer] = OrderedDict()
71
+ self._active_ids: dict[str, None] = {}
72
+ self._terminal_ids: OrderedDict[str, None] = OrderedDict()
73
+ self._resident_ids: OrderedDict[str, None] = OrderedDict()
74
+ self._resident_bytes = 0
75
+ self._hash_index: dict[str, deque[str]] = {}
76
+ self._loads_to_dispatch: OrderedDict[str, None] = OrderedDict()
77
+ self._unavailable_requests: set[str] = set()
78
+ # Records expired straight out of READY/RESIDENT: the consumer worker
79
+ # still holds their push reservation, and nothing else tells it to let
80
+ # go, so the destination buffer would sit pinned until the lease TTL.
81
+ self._orphaned: list[str] = []
82
+
83
+ def get(self, transfer_id: str) -> SchedulerTransfer | None:
84
+ return self._records.get(transfer_id)
85
+
86
+ def drain_orphaned(self) -> list[str]:
87
+ """Transfer IDs expired out of READY/RESIDENT since the last drain.
88
+
89
+ The caller must release the consumer worker's reservation for each:
90
+ `cancel()` still accepts an EXPIRED record, so the usual cancel path
91
+ applies.
92
+ """
93
+ orphaned = self._orphaned
94
+ self._orphaned = []
95
+ return orphaned
96
+
97
+ def records_for_hash(
98
+ self,
99
+ mm_hash: str,
100
+ states: Iterable[SchedulerTransferState],
101
+ ) -> list[SchedulerTransfer]:
102
+ wanted = set(states)
103
+ return [
104
+ record
105
+ for transfer_id in self._hash_index.get(mm_hash, ())
106
+ if (record := self._records.get(transfer_id)) is not None
107
+ and record.state in wanted
108
+ ]
109
+
110
+ def first_for_hash(
111
+ self,
112
+ mm_hash: str,
113
+ states: Iterable[SchedulerTransferState],
114
+ ) -> SchedulerTransfer | None:
115
+ states = tuple(states)
116
+ if len(states) == 1:
117
+ wanted_state = states[0]
118
+ for transfer_id in self._hash_index.get(mm_hash, ()):
119
+ record = self._records.get(transfer_id)
120
+ if record is not None and record.state is wanted_state:
121
+ return record
122
+ return None
123
+ wanted_states = set(states)
124
+ for transfer_id in self._hash_index.get(mm_hash, ()):
125
+ record = self._records.get(transfer_id)
126
+ if record is not None and record.state in wanted_states:
127
+ return record
128
+ return None
129
+
130
+ def has_state(self, mm_hash: str, states: Iterable[SchedulerTransferState]) -> bool:
131
+ return self.first_for_hash(mm_hash, states) is not None
132
+
133
+ def wait_for_event(
134
+ self,
135
+ transfer_id: str,
136
+ request_id: str,
137
+ mm_hash: str,
138
+ deadline: float,
139
+ ) -> SchedulerTransfer | None:
140
+ """Track a request waiting for a push.
141
+
142
+ Returns None when `transfer_id` already names a different encoding.
143
+ The id comes from the request, so a collision has to fail that one
144
+ request; re-issuing it re-runs the encode under a fresh id.
145
+ """
146
+ record = self._records.get(transfer_id)
147
+ if record is None:
148
+ record = SchedulerTransfer(
149
+ transfer_id=transfer_id,
150
+ request_id=request_id,
151
+ mm_hash=mm_hash,
152
+ state=SchedulerTransferState.WAITING_EVENT,
153
+ spec=None,
154
+ deadline=deadline,
155
+ )
156
+ self._insert(record)
157
+ elif not self._identity_matches(record, mm_hash):
158
+ self._refuse_colliding_id(record, mm_hash, request_id)
159
+ return None
160
+ else:
161
+ if not record.request_id:
162
+ record.request_id = request_id
163
+ if record.state in _TERMINAL_STATES and request_id:
164
+ self._notify_unavailable(record, request_id)
165
+ return record
166
+
167
+ def observe_ready(
168
+ self, spec: ECMooncakeLoadSpec, deadline: float
169
+ ) -> tuple[SchedulerTransfer, bool]:
170
+ transfer_id = spec.transfer_id
171
+ record = self._records.get(transfer_id)
172
+ if record is None:
173
+ record = SchedulerTransfer(
174
+ transfer_id=transfer_id,
175
+ request_id="",
176
+ mm_hash=spec.mm_hash,
177
+ state=SchedulerTransferState.WAITING_EVENT,
178
+ spec=None,
179
+ deadline=None,
180
+ )
181
+ self._insert(record)
182
+ elif not self._identity_matches(record, spec.mm_hash):
183
+ self._refuse_colliding_id(record, spec.mm_hash, record.request_id)
184
+ return record, False
185
+ if record.state is not SchedulerTransferState.WAITING_EVENT:
186
+ return record, False
187
+ record.spec = spec
188
+ record.deadline = deadline
189
+ self._transition(record, SchedulerTransferState.AVAILABLE)
190
+ return record, True
191
+
192
+ def touch_available(self, transfer_id: str, deadline: float) -> None:
193
+ record = self._records.get(transfer_id)
194
+ if record is not None and record.state is SchedulerTransferState.AVAILABLE:
195
+ record.deadline = deadline
196
+
197
+ def begin_load(
198
+ self,
199
+ mm_hash: str,
200
+ transfer_id: str | None = None,
201
+ request_id: str = "",
202
+ deadline: float | None = None,
203
+ ) -> SchedulerTransfer | None:
204
+ record = self._records.get(transfer_id) if transfer_id else None
205
+ if record is not None and (
206
+ record.mm_hash != mm_hash
207
+ or record.state is not SchedulerTransferState.AVAILABLE
208
+ ):
209
+ record = None
210
+ if record is None:
211
+ record = self.first_for_hash(
212
+ mm_hash,
213
+ (
214
+ SchedulerTransferState.AVAILABLE,
215
+ SchedulerTransferState.RESIDENT,
216
+ ),
217
+ )
218
+ if record is None or record.spec is None:
219
+ return None
220
+ if not record.request_id:
221
+ record.request_id = request_id
222
+ # A dispatched load that never reports back would otherwise hold this
223
+ # hash in LOADING for good, deferring every later request for it.
224
+ record.deadline = deadline
225
+ self._transition(record, SchedulerTransferState.LOADING)
226
+ self._loads_to_dispatch[record.transfer_id] = None
227
+ return record
228
+
229
+ def take_loads_to_dispatch(self) -> list[SchedulerTransfer]:
230
+ records = [
231
+ record
232
+ for transfer_id in self._loads_to_dispatch
233
+ if (record := self._records.get(transfer_id)) is not None
234
+ and record.state is SchedulerTransferState.LOADING
235
+ ]
236
+ self._loads_to_dispatch.clear()
237
+ return records
238
+
239
+ def complete_load(self, mm_hash: str) -> bool:
240
+ record = self.first_for_hash(mm_hash, (SchedulerTransferState.LOADING,))
241
+ if record is None:
242
+ return self.has_state(mm_hash, (SchedulerTransferState.READY,))
243
+ self._transition(record, SchedulerTransferState.READY)
244
+ return True
245
+
246
+ def fail_load(self, mm_hash: str, now: float) -> bool:
247
+ record = self.first_for_hash(mm_hash, (SchedulerTransferState.LOADING,))
248
+ if record is None:
249
+ return self.has_state(mm_hash, (SchedulerTransferState.FAILED,))
250
+ self._transition(record, SchedulerTransferState.FAILED, now=now)
251
+ return True
252
+
253
+ def release_ready(self, mm_hash: str, now: float) -> None:
254
+ ready = self.records_for_hash(mm_hash, (SchedulerTransferState.READY,))
255
+ if ready:
256
+ canonical = ready[-1]
257
+ for record in self.records_for_hash(
258
+ mm_hash,
259
+ (SchedulerTransferState.READY, SchedulerTransferState.RESIDENT),
260
+ ):
261
+ if record is not canonical:
262
+ self._transition(record, SchedulerTransferState.EXPIRED, now=now)
263
+ if canonical.spec is None:
264
+ self._transition(canonical, SchedulerTransferState.EXPIRED, now=now)
265
+ else:
266
+ canonical.spec = replace(canonical.spec, local=True)
267
+ self._transition(canonical, SchedulerTransferState.RESIDENT)
268
+ self._evict_residents(now)
269
+
270
+ def reclaim(self, mm_hash: str, now: float) -> None:
271
+ for record in self.records_for_hash(
272
+ mm_hash,
273
+ (SchedulerTransferState.READY, SchedulerTransferState.RESIDENT),
274
+ ):
275
+ if record.state is SchedulerTransferState.READY:
276
+ record.spec = None
277
+ else:
278
+ self._transition(record, SchedulerTransferState.EXPIRED, now=now)
279
+
280
+ def mark_unavailable(self, transfer_id: str, now: float) -> None:
281
+ record = self._records[transfer_id]
282
+ self._transition(record, SchedulerTransferState.UNAVAILABLE, now=now)
283
+ if record.request_id:
284
+ self._notify_unavailable(record, record.request_id)
285
+
286
+ def cancel(
287
+ self,
288
+ transfer_id: str,
289
+ now: float,
290
+ mm_hash: str = "",
291
+ request_id: str = "",
292
+ ) -> bool:
293
+ record = self._records.get(transfer_id)
294
+ if (
295
+ record is not None
296
+ and mm_hash
297
+ and not self._identity_matches(record, mm_hash)
298
+ ):
299
+ return False
300
+ if record is None:
301
+ record = SchedulerTransfer(
302
+ transfer_id=transfer_id,
303
+ request_id=request_id,
304
+ mm_hash=mm_hash,
305
+ state=SchedulerTransferState.WAITING_EVENT,
306
+ spec=None,
307
+ deadline=None,
308
+ )
309
+ self._insert(record)
310
+ if record.state in {
311
+ SchedulerTransferState.CANCELLED,
312
+ SchedulerTransferState.READY,
313
+ SchedulerTransferState.RESIDENT,
314
+ SchedulerTransferState.FAILED,
315
+ }:
316
+ return False
317
+ self._transition(record, SchedulerTransferState.CANCELLED, now=now)
318
+ self._loads_to_dispatch.pop(transfer_id, None)
319
+ return True
320
+
321
+ def expire(self, now: float, terminal_limit: int) -> list[SchedulerTransfer]:
322
+ if terminal_limit < 0:
323
+ raise ValueError("terminal_limit must be non-negative")
324
+ expired = []
325
+ for transfer_id in list(self._active_ids):
326
+ record = self._records[transfer_id]
327
+ if record.deadline is None or record.deadline > now:
328
+ continue
329
+ if record.state is SchedulerTransferState.AVAILABLE:
330
+ self._transition(
331
+ record,
332
+ SchedulerTransferState.EXPIRED,
333
+ now=now,
334
+ )
335
+ expired.append(record)
336
+ elif record.state is SchedulerTransferState.LOADING:
337
+ self._transition(
338
+ record,
339
+ SchedulerTransferState.UNAVAILABLE,
340
+ now=now,
341
+ )
342
+ if record.request_id:
343
+ self._notify_unavailable(record, record.request_id)
344
+ expired.append(record)
345
+ while self._terminal_ids:
346
+ transfer_id = next(iter(self._terminal_ids))
347
+ record = self._records[transfer_id]
348
+ if (
349
+ record.deadline is not None
350
+ and record.deadline > now
351
+ and len(self._terminal_ids) <= terminal_limit
352
+ ):
353
+ break
354
+ self._remove(transfer_id)
355
+ return expired
356
+
357
+ def take_unavailable_requests(self) -> set[str]:
358
+ unavailable = self._unavailable_requests
359
+ self._unavailable_requests = set()
360
+ return unavailable
361
+
362
+ def _notify_unavailable(self, record: SchedulerTransfer, request_id: str) -> None:
363
+ if request_id not in record.notified_requests:
364
+ record.notified_requests.add(request_id)
365
+ self._unavailable_requests.add(request_id)
366
+
367
+ def _insert(self, record: SchedulerTransfer) -> None:
368
+ self._records[record.transfer_id] = record
369
+ self._active_ids[record.transfer_id] = None
370
+ if record.mm_hash:
371
+ self._hash_index.setdefault(record.mm_hash, deque()).append(
372
+ record.transfer_id
373
+ )
374
+
375
+ @staticmethod
376
+ def _identity_matches(record: SchedulerTransfer, mm_hash: str) -> bool:
377
+ return not record.mm_hash or record.mm_hash == mm_hash
378
+
379
+ def _refuse_colliding_id(
380
+ self, record: SchedulerTransfer, mm_hash: str, request_id: str
381
+ ) -> None:
382
+ """Report a transfer id that already names another encoding.
383
+
384
+ Transfer ids arrive on the request, so a collision is reachable from
385
+ outside and must not reach the engine as an exception.
386
+ """
387
+ logger.warning(
388
+ "EC Mooncake transfer_id=%s already names mm_hash=%s; refusing "
389
+ "mm_hash=%s for request %s",
390
+ record.transfer_id,
391
+ record.mm_hash[:16],
392
+ mm_hash[:16],
393
+ request_id or "<unknown>",
394
+ )
395
+ if request_id:
396
+ self._unavailable_requests.add(request_id)
397
+
398
+ def _transition(
399
+ self,
400
+ record: SchedulerTransfer,
401
+ state: SchedulerTransferState,
402
+ now: float | None = None,
403
+ ) -> None:
404
+ if record.state is SchedulerTransferState.RESIDENT:
405
+ self._resident_ids.pop(record.transfer_id, None)
406
+ if record.spec is not None:
407
+ self._resident_bytes -= record.spec.nbytes
408
+ if state is SchedulerTransferState.RESIDENT:
409
+ self._resident_ids[record.transfer_id] = None
410
+ if record.spec is not None:
411
+ self._resident_bytes += record.spec.nbytes
412
+ if state in _TERMINAL_STATES:
413
+ if now is None:
414
+ raise ValueError("Terminal transition requires a timestamp")
415
+ # Keep a bounded tombstone so late events and repeat request IDs
416
+ # remain idempotent instead of reviving a finished transfer.
417
+ record.deadline = now + self._tombstone_ttl
418
+ self._active_ids.pop(record.transfer_id, None)
419
+ self._terminal_ids[record.transfer_id] = None
420
+ self._terminal_ids.move_to_end(record.transfer_id)
421
+ if state is SchedulerTransferState.EXPIRED and record.state in (
422
+ SchedulerTransferState.READY,
423
+ SchedulerTransferState.RESIDENT,
424
+ ):
425
+ self._orphaned.append(record.transfer_id)
426
+ record.state = state
427
+ self._records.move_to_end(record.transfer_id)
428
+
429
+ def _evict_residents(self, now: float) -> None:
430
+ while self._resident_bytes > self._resident_capacity and self._resident_ids:
431
+ record = self._records[next(iter(self._resident_ids))]
432
+ self._transition(record, SchedulerTransferState.EXPIRED, now=now)
433
+
434
+ def _remove(self, transfer_id: str) -> None:
435
+ record = self._records.pop(transfer_id, None)
436
+ self._active_ids.pop(transfer_id, None)
437
+ self._terminal_ids.pop(transfer_id, None)
438
+ if transfer_id in self._resident_ids:
439
+ self._resident_ids.pop(transfer_id)
440
+ if record is not None and record.spec is not None:
441
+ self._resident_bytes -= record.spec.nbytes
442
+ self._loads_to_dispatch.pop(transfer_id, None)
443
+ if record is None or not record.mm_hash:
444
+ return
445
+ transfer_ids = self._hash_index[record.mm_hash]
446
+ transfer_ids.remove(transfer_id)
447
+ if not transfer_ids:
448
+ self._hash_index.pop(record.mm_hash)
vllm/distributed/ec_transfer/ec_connector/mooncake/transfer.py
new
+209
@@ -0,0 +1,209 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Mooncake data-plane engine and memory-registration ownership.
4
+
5
+The wrapper isolates optional-engine initialization, session addressing,
6
+synchronous batched writes, and the lifetime of transient source registrations.
7
+"""
8
+
9
+from __future__ import annotations
10
+
11
+import threading
12
+from dataclasses import dataclass
13
+
14
+import torch
15
+
16
+from vllm.logger import init_logger
17
+
18
+logger = init_logger(__name__)
19
+
20
+_MOONCAKE_IMPORT_ERROR: ImportError | None = None
21
+try:
22
+ from mooncake.engine import TransferEngine
23
+except ImportError as error:
24
+ TransferEngine = None # type: ignore[misc, assignment]
25
+ _MOONCAKE_IMPORT_ERROR = error
26
+
27
+
28
+def ensure_mooncake_available() -> None:
29
+ if _MOONCAKE_IMPORT_ERROR is not None:
30
+ raise ImportError(
31
+ "Install mooncake-transfer-engine to use ECMooncakeConnector."
32
+ ) from _MOONCAKE_IMPORT_ERROR
33
+
34
+
35
+@dataclass
36
+class _SourceRegistration:
37
+ """Retain one transient source range while batches reference it."""
38
+
39
+ tensor: torch.Tensor
40
+ nbytes: int
41
+ users: int = 1
42
+
43
+
44
+class MooncakeTransfer:
45
+ """Own a lazy Mooncake engine and transient memory registrations."""
46
+
47
+ def __init__(self, hostname: str, protocol: str) -> None:
48
+ self._hostname = hostname
49
+ self._protocol = protocol
50
+ self._engine: TransferEngine | None = None
51
+ self._engine_lock = threading.Lock()
52
+ self._source_registrations: dict[int, _SourceRegistration] = {}
53
+ self._pending_unregister: dict[int, torch.Tensor] = {}
54
+ self._registration_lock = threading.Lock()
55
+ self._closed = False
56
+
57
+ def _ensure_engine(self) -> TransferEngine:
58
+ if self._engine is not None:
59
+ return self._engine
60
+ with self._engine_lock:
61
+ if self._engine is not None:
62
+ return self._engine
63
+ engine = TransferEngine()
64
+ ret = engine.initialize(self._hostname, "P2PHANDSHAKE", self._protocol, "")
65
+ if ret != 0:
66
+ raise RuntimeError("Mooncake TransferEngine initialization failed.")
67
+ self._engine = engine
68
+ logger.info(
69
+ "ECMooncakeConnector TransferEngine ready at %s:%d",
70
+ self._hostname,
71
+ engine.get_rpc_port(),
72
+ )
73
+ return self._engine
74
+
75
+ def ensure_ready(self) -> None:
76
+ self._ensure_engine()
77
+
78
+ def local_session(self) -> str:
79
+ engine = self._ensure_engine()
80
+ return f"{self._hostname}:{engine.get_rpc_port()}"
81
+
82
+ def register_memory(self, tensor: torch.Tensor) -> int:
83
+ return self._ensure_engine().batch_register_memory(
84
+ [tensor.data_ptr()], [tensor.nbytes]
85
+ )
86
+
87
+ def unregister_memory(self, tensor: torch.Tensor) -> bool:
88
+ engine = self._ensure_engine()
89
+ address = tensor.data_ptr()
90
+ ret = engine.unregister_memory(address)
91
+ with self._registration_lock:
92
+ if ret != 0:
93
+ logger.error(
94
+ "Mooncake EC memory unregistration failed for address %d: %d",
95
+ address,
96
+ ret,
97
+ )
98
+ self._pending_unregister[address] = tensor
99
+ return False
100
+ self._pending_unregister.pop(address, None)
101
+ return True
102
+
103
+ @staticmethod
104
+ def _source_range(tensor: torch.Tensor) -> tuple[int, int]:
105
+ # Encoder batches commonly split one storage into sibling tensor views.
106
+ # Registering each view's exact bytes avoids overlapping memory regions.
107
+ return tensor.data_ptr(), tensor.nbytes
108
+
109
+ def acquire_sources(self, tensors: list[torch.Tensor]) -> list[int]:
110
+ ranges: dict[int, tuple[int, torch.Tensor]] = {}
111
+ for tensor in tensors:
112
+ address, nbytes = self._source_range(tensor)
113
+ ranges.setdefault(address, (nbytes, tensor))
114
+
115
+ engine = self._ensure_engine()
116
+ acquired: list[int] = []
117
+ new_addresses: list[int] = []
118
+ new_lengths: list[int] = []
119
+ with self._registration_lock:
120
+ for address, (nbytes, tensor) in ranges.items():
121
+ entry = self._source_registrations.get(address)
122
+ if entry is not None:
123
+ if entry.nbytes != nbytes:
124
+ raise RuntimeError(
125
+ "Mooncake EC source storage changed size while registered"
126
+ )
127
+ entry.users += 1
128
+ acquired.append(address)
129
+ continue
130
+ new_addresses.append(address)
131
+ new_lengths.append(nbytes)
132
+ self._source_registrations[address] = _SourceRegistration(
133
+ tensor=tensor,
134
+ nbytes=nbytes,
135
+ )
136
+ acquired.append(address)
137
+
138
+ if new_addresses:
139
+ ret = engine.batch_register_memory(new_addresses, new_lengths)
140
+ if ret != 0:
141
+ for address in acquired:
142
+ entry = self._source_registrations[address]
143
+ entry.users -= 1
144
+ if entry.users == 0:
145
+ del self._source_registrations[address]
146
+ raise RuntimeError("Mooncake EC source registration failed")
147
+ return acquired
148
+
149
+ def release_sources(self, addresses: list[int]) -> bool:
150
+ if not addresses:
151
+ return True
152
+ with self._registration_lock:
153
+ unused = []
154
+ for address in addresses:
155
+ entry = self._source_registrations.get(address)
156
+ if entry is None:
157
+ continue
158
+ entry.users -= 1
159
+ if entry.users == 0:
160
+ unused.append(address)
161
+ if not unused:
162
+ return True
163
+ ret = self._ensure_engine().batch_unregister_memory(unused)
164
+ if ret != 0:
165
+ logger.warning(
166
+ "Keeping %d EC source tensors registered after Mooncake "
167
+ "unregistration failure",
168
+ len(unused),
169
+ )
170
+ return False
171
+ for address in unused:
172
+ del self._source_registrations[address]
173
+ self._pending_unregister.pop(address, None)
174
+ return True
175
+
176
+ def write(
177
+ self,
178
+ session: str,
179
+ sources: list[int],
180
+ destinations: list[int],
181
+ lengths: list[int],
182
+ ) -> None:
183
+ """Write one synchronous batch, returning only at terminal status."""
184
+ ret = self._ensure_engine().batch_transfer_sync_write(
185
+ session, sources, destinations, lengths
186
+ )
187
+ if ret != 0:
188
+ raise RuntimeError(
189
+ f"Mooncake EC push to {session} failed with status {ret}"
190
+ )
191
+
192
+ def close(self) -> None:
193
+ if self._closed:
194
+ return
195
+ self._closed = True
196
+ engine = self._engine
197
+ if engine is None:
198
+ return
199
+ with self._registration_lock:
200
+ addresses = list(self._source_registrations)
201
+ addresses.extend(self._pending_unregister)
202
+ if not addresses:
203
+ return
204
+ ret = engine.batch_unregister_memory(list(dict.fromkeys(addresses)))
205
+ if ret != 0:
206
+ logger.error("Mooncake EC batch memory unregistration failed: %d", ret)
207
+ return
208
+ self._source_registrations.clear()
209
+ self._pending_unregister.clear()
vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py
new
+940
@@ -0,0 +1,940 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Worker-side orchestration of Mooncake control, memory, and data planes.
4
+
5
+Consumer Workers expose rank-local reservations and publish received tensors.
6
+Producer Workers reserve every destination shard, bind computed sources, run
7
+batched Mooncake writes, and report asynchronous completion to the Scheduler.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import math
13
+import threading
14
+import time
15
+from collections.abc import Callable
16
+from concurrent.futures import Future, ThreadPoolExecutor
17
+from functools import partial
18
+from typing import TYPE_CHECKING, Any, TypeVar, cast
19
+
20
+import torch
21
+
22
+from vllm.distributed.ec_transfer.ec_connector.mooncake.config import (
23
+ _RESERVATION_TTL_SECONDS,
24
+ MooncakeECConfig,
25
+)
26
+from vllm.distributed.ec_transfer.ec_connector.mooncake.control import (
27
+ ConsumerControlServer,
28
+ ControlClient,
29
+ make_cancel_request,
30
+)
31
+from vllm.distributed.ec_transfer.ec_connector.mooncake.memory import (
32
+ ConsumerMemoryPool,
33
+ ProducerMemoryPool,
34
+)
35
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
36
+ ECMooncakeConnectorMetadata,
37
+ ECMooncakePushSpec,
38
+ ECMooncakeWorkerMetadata,
39
+)
40
+from vllm.distributed.ec_transfer.ec_connector.mooncake.producer import (
41
+ ProducerPushManager,
42
+ ProducerPushRecord,
43
+)
44
+from vllm.distributed.ec_transfer.ec_connector.mooncake.reservation import (
45
+ ConsumerReservationManager,
46
+ ConsumerReservationState,
47
+)
48
+from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
49
+ MooncakeTransfer,
50
+ ensure_mooncake_available,
51
+)
52
+from vllm.logger import init_logger
53
+from vllm.utils.network_utils import get_ip
54
+
55
+logger = init_logger(__name__)
56
+
57
+_T = TypeVar("_T")
58
+
59
+if TYPE_CHECKING:
60
+ from vllm.config import VllmConfig
61
+
62
+_RESERVATION_REFRESH_SECONDS = _RESERVATION_TTL_SECONDS / 2
63
+_MAX_CANCELLED_TRANSFER_IDS = 1 << 16
64
+_CANCEL_ATTEMPTS = 2
65
+_TRANSFER_WORKERS = 4
66
+_CONTROL_WORKERS = 8
67
+
68
+
69
+class _FanoutError(RuntimeError):
70
+ """Retain shard outcomes after every started task settles."""
71
+
72
+ def __init__(self, error: BaseException, results: list[Any | None]) -> None:
73
+ self.results = results
74
+ super().__init__(str(error))
75
+
76
+
77
+class ECMooncakeWorker:
78
+ """Orchestrate consumer reservations and producer push batches.
79
+
80
+ Consumers may use TP, PP, and DP. Only the first PP stage receives encoder
81
+ outputs, and each TP rank exposes a consecutive control port and receives
82
+ the same source concurrently. Producers remain unsharded and unreplicated.
83
+ With DP, the caller must route both halves of a request to the same replica
84
+ and pass that replica's control address to the producer.
85
+ """
86
+
87
+ def __init__(self, vllm_config: VllmConfig) -> None:
88
+ ensure_mooncake_available()
89
+ config = MooncakeECConfig.from_vllm_config(vllm_config)
90
+ self.is_producer = config.is_producer
91
+ self.is_consumer = config.is_consumer
92
+ self._buffer_device = config.buffer_device
93
+ self._control_host = config.control_host
94
+ self._control_port = config.control_port
95
+ self._transfer = MooncakeTransfer(get_ip(), config.protocol)
96
+ self._consumer_memory = ConsumerMemoryPool(
97
+ config.pool_size,
98
+ self._transfer,
99
+ )
100
+ self._reservations = ConsumerReservationManager(
101
+ self._consumer_memory,
102
+ _RESERVATION_TTL_SECONDS,
103
+ _MAX_CANCELLED_TRANSFER_IDS,
104
+ )
105
+ self._consumer_rank_resolved = False
106
+ self._is_receiving_rank = True
107
+ self._tp_rank = 0
108
+ self._tp_size = 1
109
+ self._control_server: ConsumerControlServer | None = None
110
+ # Worker producer
111
+ self._producer_memory = ProducerMemoryPool(
112
+ config.pool_size,
113
+ self._transfer,
114
+ )
115
+ self._control_client = ControlClient(config.control_timeout_ms)
116
+ self._shutdown_drain_timeout_s = config.control_timeout_ms / 1000
117
+ self._io_executor = ThreadPoolExecutor(
118
+ max_workers=_TRANSFER_WORKERS,
119
+ thread_name_prefix="ec-mooncake-transfer",
120
+ )
121
+ self._control_executor = ThreadPoolExecutor(
122
+ max_workers=_CONTROL_WORKERS,
123
+ thread_name_prefix="ec-mooncake-control",
124
+ )
125
+ self._shard_pool: ThreadPoolExecutor | None = None
126
+ self._shard_pool_lock = threading.Lock()
127
+ self._push_ready = threading.Event()
128
+ self._producer_pushes = ProducerPushManager(self._push_ready.set)
129
+ self._dispatch_stop = threading.Event()
130
+ self._dispatcher: threading.Thread | None = None
131
+ self._failed_saves: set[str] = set()
132
+ self._collecting_sources = False
133
+ self._completed_loads: set[str] = set()
134
+ self._failed_loads: set[str] = set()
135
+ self._shutdown = False
136
+ if self.is_producer:
137
+ self._dispatcher = threading.Thread(
138
+ target=self._dispatch_pushes, name="ec-mooncake-ready", daemon=True
139
+ )
140
+ self._dispatcher.start()
141
+
142
+ def _dispatch_pushes(self) -> None:
143
+ pending_event = False
144
+ while not self._dispatch_stop.is_set():
145
+ self._push_ready.wait(timeout=0.001 if pending_event else None)
146
+ self._push_ready.clear()
147
+ if self._dispatch_stop.is_set():
148
+ break
149
+ if self._collecting_sources:
150
+ pending_event = False
151
+ continue
152
+ pending_event = self._producer_pushes.submit_batches(
153
+ self._io_executor, self._push_batch
154
+ )
155
+
156
+ def _resolve_consumer_rank(self) -> None:
157
+ """Place this worker in the consumer receive topology."""
158
+ if self._consumer_rank_resolved:
159
+ return
160
+ self._consumer_rank_resolved = True
161
+ try:
162
+ from vllm.distributed.parallel_state import get_pp_group, get_tp_group
163
+
164
+ tp_group = get_tp_group()
165
+ self._tp_rank = tp_group.rank_in_group
166
+ self._tp_size = tp_group.world_size
167
+ self._is_receiving_rank = get_pp_group().is_first_rank
168
+ except AssertionError:
169
+ # Groups are only absent outside a distributed run, where this
170
+ # worker is the whole consumer.
171
+ self._tp_rank = 0
172
+ self._tp_size = 1
173
+ self._is_receiving_rank = True
174
+
175
+ def start_services(self) -> None:
176
+ if not self.is_consumer or self._control_server is not None:
177
+ return
178
+ self._resolve_consumer_rank()
179
+ if not self._is_receiving_rank:
180
+ # Later pipeline stages hold no encoder outputs, so they need
181
+ # neither a receive pool nor a control channel.
182
+ return
183
+ self._consumer_memory.prepare(torch.device(self._buffer_device))
184
+ consumer_pool = self._consumer_memory.tensor
185
+ if consumer_pool is None:
186
+ raise RuntimeError(
187
+ "Mooncake push mode requires a registered consumer buffer pool."
188
+ )
189
+ base_port = self._control_port
190
+ self._control_server = ConsumerControlServer(
191
+ # The reservation channel hands out registered-memory addresses
192
+ # and takes cancellations, so it listens where the Consumer
193
+ # advertises itself (`ec_ip`), not on every interface.
194
+ self._control_host,
195
+ base_port + self._tp_rank,
196
+ self._reserve_push_destination,
197
+ self._push_status,
198
+ self._reservations.complete,
199
+ self._reservations.cancel,
200
+ self._reservations.expire,
201
+ peer_ports=[base_port + rank for rank in range(self._tp_size)],
202
+ device=consumer_pool.device,
203
+ drain_ready=self._reservations.drain_ready,
204
+ )
205
+ try:
206
+ self._control_server.start()
207
+ except Exception:
208
+ self._control_server.close()
209
+ self._control_server = None
210
+ raise
211
+
212
+ def _reserve_push_destination(self, payload: dict[str, Any]) -> dict[str, Any]:
213
+ transfer_id = str(payload["transfer_id"])
214
+ mm_hash = str(payload["mm_hash"])
215
+ nbytes = int(payload["nbytes"])
216
+ shape = tuple(int(value) for value in payload["shape"])
217
+ dtype_name = str(payload["dtype"])
218
+ dtype = getattr(torch, dtype_name, None)
219
+ if not isinstance(dtype, torch.dtype):
220
+ raise ValueError(f"Unsupported torch dtype string: {dtype_name!r}")
221
+ expected_nbytes = math.prod(shape) * dtype.itemsize
222
+ if expected_nbytes != nbytes:
223
+ raise ValueError("shape and dtype do not match nbytes")
224
+
225
+ self._reservations.expire()
226
+ reservation, should_write = self._reservations.reserve(
227
+ transfer_id, mm_hash, nbytes, shape, dtype_name, dtype
228
+ )
229
+ if reservation is None:
230
+ raise RuntimeError("EC consumer buffer pool is full")
231
+ if reservation.state in {
232
+ ConsumerReservationState.CANCEL_PENDING,
233
+ ConsumerReservationState.CANCELLED,
234
+ }:
235
+ return {
236
+ "reservation_id": "",
237
+ "dst_session": "",
238
+ "dst_ptr": 0,
239
+ "nbytes": nbytes,
240
+ "write": False,
241
+ "ready": False,
242
+ "cancelled": True,
243
+ }
244
+ assert reservation.allocation is not None
245
+
246
+ return {
247
+ "reservation_id": reservation.reservation_id,
248
+ "dst_session": self._transfer.local_session(),
249
+ "dst_ptr": reservation.allocation.tensor.data_ptr(),
250
+ "nbytes": reservation.allocation.tensor.nbytes,
251
+ "write": should_write,
252
+ "ready": reservation.state is ConsumerReservationState.READY,
253
+ "cached": reservation.lease is not None,
254
+ }
255
+
256
+ def _push_status(self, transfer_id: str) -> dict[str, Any] | None:
257
+ reservation = self._reservations.status(transfer_id)
258
+ if reservation is None:
259
+ return None
260
+ assert reservation.allocation is not None
261
+ return {
262
+ "mm_hash": reservation.mm_hash,
263
+ "ready": reservation.state is ConsumerReservationState.READY,
264
+ "reservation_id": reservation.reservation_id,
265
+ "nbytes": reservation.allocation.tensor.nbytes,
266
+ "shape": list(reservation.shape),
267
+ "dtype": reservation.dtype,
268
+ }
269
+
270
+ def _shard_executor(self) -> ThreadPoolExecutor:
271
+ """Use a separate pool so nested shard fan-out cannot deadlock."""
272
+ with self._shard_pool_lock:
273
+ if self._shard_pool is None:
274
+ self._shard_pool = ThreadPoolExecutor(
275
+ max_workers=32, thread_name_prefix="ec-mooncake-shard"
276
+ )
277
+ return self._shard_pool
278
+
279
+ def _reserve_one(self, addr: str, spec: ECMooncakePushSpec) -> dict[str, Any]:
280
+ result = self._control_client.request(
281
+ addr,
282
+ {
283
+ "op": "reserve",
284
+ "transfer_id": spec.transfer_id,
285
+ "mm_hash": spec.mm_hash,
286
+ "nbytes": spec.nbytes,
287
+ "shape": list(spec.shape),
288
+ "dtype": spec.dtype,
289
+ },
290
+ )
291
+ if not isinstance(result, dict):
292
+ raise RuntimeError("Invalid EC reservation response")
293
+ result["_received_at"] = time.monotonic()
294
+ result["addr"] = addr
295
+ return result
296
+
297
+ def _run_fanout(
298
+ self,
299
+ tasks: list[Callable[[], _T]],
300
+ on_submit: Callable[[int, Future[_T]], None] | None = None,
301
+ ) -> list[_T]:
302
+ """Run every started shard task and retain partial results on failure.
303
+
304
+ Waiting for all submitted tasks is what makes source-memory release and
305
+ partial-reservation cleanup safe after one shard fails.
306
+ """
307
+ if not tasks:
308
+ return []
309
+ futures: list[tuple[int, Future[_T]]] = []
310
+ results: list[_T | None] = [None] * len(tasks)
311
+ error: BaseException | None = None
312
+ for index, task in enumerate(tasks[1:], 1):
313
+ try:
314
+ future = self._shard_executor().submit(task)
315
+ except Exception as exc:
316
+ error = exc
317
+ break
318
+ futures.append((index, future))
319
+ if on_submit is not None:
320
+ on_submit(index, future)
321
+ if error is None:
322
+ try:
323
+ results[0] = tasks[0]()
324
+ except Exception as exc:
325
+ error = exc
326
+ for index, future in futures:
327
+ try:
328
+ results[index] = future.result()
329
+ except Exception as exc:
330
+ if error is None:
331
+ error = exc
332
+ if error is not None:
333
+ raise _FanoutError(error, results)
334
+ return cast(list[_T], results)
335
+
336
+ def _cancel_reservations(
337
+ self,
338
+ spec: ECMooncakePushSpec,
339
+ reservations: list[dict[str, Any]],
340
+ *,
341
+ refresh: bool = False,
342
+ record: ProducerPushRecord | None = None,
343
+ ) -> None:
344
+ reservations = [
345
+ shard for shard in reservations if not shard.get("cancelled", False)
346
+ ]
347
+ if not reservations:
348
+ return
349
+
350
+ def cancel(shard: dict[str, Any]) -> dict[str, Any]:
351
+ addr = shard.get("addr")
352
+ if not isinstance(addr, str) or not addr:
353
+ raise RuntimeError(
354
+ "EC reservation is missing a confirmed shard address"
355
+ )
356
+ result = self._control_client.request(
357
+ addr,
358
+ make_cancel_request(
359
+ spec.transfer_id,
360
+ str(shard.get("reservation_id", "")),
361
+ abandon=True,
362
+ refresh=refresh,
363
+ ),
364
+ )
365
+ if not isinstance(result, dict) or not result.get("cancelled"):
366
+ raise RuntimeError(
367
+ f"Could not cancel EC reservation for mm_hash={spec.mm_hash}"
368
+ )
369
+ return shard
370
+
371
+ def track(_index: int, future: Future[dict[str, Any]]) -> None:
372
+ if record is not None:
373
+ self._producer_pushes.track_shard_futures([record], [future])
374
+
375
+ self._run_fanout([partial(cancel, shard) for shard in reservations], track)
376
+
377
+ def _retry_cancel_reservations(
378
+ self,
379
+ spec: ECMooncakePushSpec,
380
+ reservations: list[dict[str, Any]],
381
+ *,
382
+ record: ProducerPushRecord | None = None,
383
+ ) -> None:
384
+ pending = [shard for shard in reservations if not shard.get("cancelled", False)]
385
+ error: _FanoutError | None = None
386
+ for _ in range(_CANCEL_ATTEMPTS):
387
+ try:
388
+ self._cancel_reservations(spec, pending, record=record)
389
+ except _FanoutError as exc:
390
+ error = exc
391
+ pending = [
392
+ shard
393
+ for index, shard in enumerate(pending)
394
+ if exc.results[index] is None
395
+ ]
396
+ continue
397
+ return
398
+ assert error is not None
399
+ raise error
400
+
401
+ def _reserve_remote(self, spec: ECMooncakePushSpec) -> list[dict[str, Any]]:
402
+ """Reserve a destination on every shard of the consumer."""
403
+ result = self._reserve_remote_many([spec])[0]
404
+ if isinstance(result, Exception):
405
+ raise result
406
+ return result
407
+
408
+ def _refresh_remote_reservations(
409
+ self,
410
+ spec: ECMooncakePushSpec,
411
+ reservations: list[dict[str, Any]],
412
+ record: ProducerPushRecord | None = None,
413
+ ) -> list[dict[str, Any]]:
414
+ stale = [
415
+ shard
416
+ for shard in reservations
417
+ if not shard.get("ready", False)
418
+ and not shard.get("cached", False)
419
+ and not shard.get("cancelled", False)
420
+ ]
421
+ try:
422
+ self._cancel_reservations(spec, stale, refresh=True, record=record)
423
+ except _FanoutError as exc:
424
+ pending = [
425
+ shard for index, shard in enumerate(stale) if exc.results[index] is None
426
+ ]
427
+ try:
428
+ self._retry_cancel_reservations(spec, pending, record=record)
429
+ except _FanoutError as cleanup_error:
430
+ raise exc from cleanup_error
431
+ raise
432
+ return self._reserve_remote(spec)
433
+
434
+ @staticmethod
435
+ def _validate_push_source(push: ProducerPushRecord) -> None:
436
+ tensor = push.source_tensor
437
+ assert tensor is not None
438
+ spec = push.spec
439
+ if tuple(tensor.shape) != tuple(spec.shape):
440
+ raise ValueError(f"EC source shape mismatch for mm_hash={spec.mm_hash}")
441
+ if str(tensor.dtype).split(".")[-1] != spec.dtype:
442
+ raise ValueError(f"EC source dtype mismatch for mm_hash={spec.mm_hash}")
443
+ if not tensor.is_contiguous():
444
+ raise ValueError(f"EC source must be contiguous for mm_hash={spec.mm_hash}")
445
+ if tensor.nbytes != spec.nbytes:
446
+ raise ValueError(f"EC source size mismatch for mm_hash={spec.mm_hash}")
447
+
448
+ def start_save_caches(
449
+ self,
450
+ metadata: ECMooncakeConnectorMetadata,
451
+ encoder_cache: dict[str, torch.Tensor] | None = None,
452
+ **kwargs: Any,
453
+ ) -> None:
454
+ self._collecting_sources = True
455
+ new: dict[str, list[ProducerPushRecord]] = {}
456
+ for spec in metadata.pushes:
457
+ try:
458
+ record, created = self._producer_pushes.reserve(spec, Future)
459
+ if created:
460
+ new.setdefault(spec.consumer_zmq, []).append(record)
461
+ except ValueError:
462
+ logger.warning("Rejected conflicting EC push", exc_info=True)
463
+ self._failed_saves.add(spec.request_id)
464
+ for records in new.values():
465
+ future = self._control_executor.submit(self._reserve_batch, records)
466
+ future.add_done_callback(partial(self._reservation_batch_done, records))
467
+ if not isinstance(encoder_cache, dict):
468
+ return
469
+ for mm_hash in dict.fromkeys(spec.mm_hash for spec in metadata.pushes):
470
+ tensor = encoder_cache.get(mm_hash)
471
+ if tensor is not None:
472
+ self._bind_push_source(tensor, mm_hash)
473
+
474
+ def _reserve_batch(self, records: list[ProducerPushRecord]) -> None:
475
+ """Resolve independent item futures with one reserve RPC per shard."""
476
+ if len(records) == 1:
477
+ record = records[0]
478
+ try:
479
+ record.reservation_future.set_result(self._reserve_remote(record.spec))
480
+ except Exception as exc:
481
+ record.reservation_future.set_exception(exc)
482
+ return
483
+ outcomes = self._reserve_remote_many([record.spec for record in records])
484
+ self._producer_pushes.finish_reservations(list(zip(records, outcomes)))
485
+
486
+ def _reserve_remote_many(
487
+ self, specs: list[ECMooncakePushSpec]
488
+ ) -> list[list[dict[str, Any]] | Exception]:
489
+ shards = None
490
+ for _ in range(_CANCEL_ATTEMPTS):
491
+ shards = self._control_client.discover_shards(specs[0].consumer_zmq)
492
+ if shards is not None:
493
+ break
494
+ if shards is None:
495
+ raise RuntimeError(
496
+ f"Could not discover every EC consumer shard at {specs[0].consumer_zmq}"
497
+ )
498
+
499
+ def reserve(addr: str) -> list[dict[str, Any]]:
500
+ if len(specs) == 1:
501
+ return [{"ok": True, "result": self._reserve_one(addr, specs[0])}]
502
+ response = self._control_client.request(
503
+ addr,
504
+ {
505
+ "op": "reserve_batch",
506
+ "items": [
507
+ {
508
+ "transfer_id": spec.transfer_id,
509
+ "mm_hash": spec.mm_hash,
510
+ "nbytes": spec.nbytes,
511
+ "shape": list(spec.shape),
512
+ "dtype": spec.dtype,
513
+ }
514
+ for spec in specs
515
+ ],
516
+ },
517
+ )
518
+ items = response["items"]
519
+ if len(items) != len(specs):
520
+ raise RuntimeError("Malformed EC reservation batch response")
521
+ for item in items:
522
+ if item.get("ok"):
523
+ item["result"]["_received_at"] = time.monotonic()
524
+ return items
525
+
526
+ fanout_error = None
527
+ results: list[Any]
528
+ try:
529
+ results = self._run_fanout([partial(reserve, addr) for addr in shards])
530
+ except _FanoutError as exc:
531
+ results = exc.results
532
+ fanout_error = exc
533
+ outcomes: list[list[dict[str, Any]] | Exception] = []
534
+ for index, spec in enumerate(specs):
535
+ reservations: list[dict[str, Any]] = []
536
+ error = None
537
+ for addr, items in zip(shards, results):
538
+ item = items[index] if items is not None else {}
539
+ if item.get("ok"):
540
+ reservations.append({**item["result"], "addr": addr})
541
+ else:
542
+ error = fanout_error or RuntimeError(item["error"])
543
+ reservations.append({"addr": addr, "reservation_id": ""})
544
+ if error is None:
545
+ outcomes.append(reservations)
546
+ else:
547
+ failure = _FanoutError(error, list(reservations))
548
+ try:
549
+ self._retry_cancel_reservations(spec, reservations)
550
+ except Exception as cleanup_error:
551
+ failure.__cause__ = cleanup_error
552
+ logger.exception("Failed to release partial EC reservation batch")
553
+ outcomes.append(failure)
554
+ return outcomes
555
+
556
+ @staticmethod
557
+ def _reservation_batch_done(
558
+ records: list[ProducerPushRecord], future: Future[None]
559
+ ) -> None:
560
+ error = future.exception()
561
+ if error is not None:
562
+ for record in records:
563
+ if not record.reservation_future.done():
564
+ record.reservation_future.set_exception(error)
565
+
566
+ def start_load_caches(
567
+ self,
568
+ metadata: ECMooncakeConnectorMetadata,
569
+ encoder_cache: dict[str, torch.Tensor],
570
+ **kwargs: Any,
571
+ ) -> None:
572
+ self._resolve_consumer_rank()
573
+ if not self._is_receiving_rank:
574
+ # Later pipeline stages never gather multimodal embeddings.
575
+ return
576
+ self._transfer.ensure_ready()
577
+ if self._buffer_device == "cuda" and not torch.accelerator.is_available():
578
+ raise RuntimeError(
579
+ "ECMooncakeConnector requires CUDA for ec_buffer_device=cuda"
580
+ )
581
+ self._reservations.retire_stale(encoder_cache, metadata.freed)
582
+
583
+ for spec in metadata.loads:
584
+ if spec.mm_hash in encoder_cache:
585
+ if not spec.local:
586
+ # The spec's id is one shard's; cancel by transfer.
587
+ self._reservations.cancel(spec.transfer_id, "")
588
+ self._completed_loads.add(spec.mm_hash)
589
+ continue
590
+ if spec.local:
591
+ tensor = self._consumer_memory.take_resident(
592
+ spec.mm_hash, tuple(spec.shape), spec.dtype
593
+ )
594
+ else:
595
+ try:
596
+ allocation = self._reservations.take(spec.transfer_id, spec.mm_hash)
597
+ tensor = allocation.tensor
598
+ except RuntimeError as e:
599
+ logger.warning("EC Mooncake pushed load failed: %s", e)
600
+ tensor = None
601
+ if tensor is None:
602
+ self._failed_loads.add(spec.mm_hash)
603
+ else:
604
+ encoder_cache[spec.mm_hash] = tensor
605
+ self._completed_loads.add(spec.mm_hash)
606
+
607
+ def _push_batch(self, pushes: list[ProducerPushRecord]) -> None:
608
+ started_at = time.monotonic()
609
+ ready: list[tuple[ProducerPushRecord, dict[str, Any]]] = []
610
+ written_pushes: dict[str, ProducerPushRecord] = {}
611
+ failure: Exception | None = None
612
+ try:
613
+ for push in pushes:
614
+ self._validate_push_source(push)
615
+ reservations = self._producer_pushes.resolve_reservations(push)
616
+ stale = [
617
+ index
618
+ for index, shard in enumerate(reservations)
619
+ if not shard.get("ready", False)
620
+ and not shard.get("cancelled", False)
621
+ and time.monotonic() - float(shard.get("_received_at", started_at))
622
+ >= _RESERVATION_REFRESH_SECONDS
623
+ ]
624
+ if stale:
625
+ reservations = self._refresh_remote_reservations(
626
+ push.spec, reservations, push
627
+ )
628
+ self._producer_pushes.replace_reservations(push, reservations)
629
+ self._producer_pushes.begin_writing(push)
630
+ writable = [
631
+ shard
632
+ for shard in reservations
633
+ if not shard.get("cached", False)
634
+ and not shard.get("cancelled", False)
635
+ and shard.get("write", True)
636
+ ]
637
+ source = push.source_tensor
638
+ assert source is not None
639
+ for shard in writable:
640
+ if int(shard["nbytes"]) != source.nbytes:
641
+ raise RuntimeError(
642
+ "Reserved EC size does not match tensor for "
643
+ f"mm_hash={push.spec.mm_hash}"
644
+ )
645
+ ready.append((push, shard))
646
+ written_pushes.setdefault(push.spec.transfer_id, push)
647
+ if ready:
648
+ # Stage each source once, then write it to every destination.
649
+ source_index = {
650
+ push.spec.transfer_id: index
651
+ for index, push in enumerate(written_pushes.values())
652
+ }
653
+ tensors = [
654
+ cast(torch.Tensor, push.source_tensor)
655
+ for push in written_pushes.values()
656
+ if push.source_tensor is not None
657
+ ]
658
+ lengths = [tensor.nbytes for tensor in tensors]
659
+ staged = self._producer_memory.stage(tensors)
660
+ registered_sources: list[int] = []
661
+ if staged is not None:
662
+ sources = staged.tensors
663
+ else:
664
+ sources = tensors
665
+ registered_sources = self._transfer.acquire_sources(tensors)
666
+ addresses = [tensor.data_ptr() for tensor in sources]
667
+ try:
668
+ by_session: dict[str, list[tuple[int, int]]] = {}
669
+ session_records: dict[str, dict[str, ProducerPushRecord]] = {}
670
+ for push, shard in ready:
671
+ session = str(shard["dst_session"])
672
+ by_session.setdefault(session, []).append(
673
+ (source_index[push.spec.transfer_id], int(shard["dst_ptr"]))
674
+ )
675
+ session_records.setdefault(session, {})[
676
+ push.spec.transfer_id
677
+ ] = push
678
+
679
+ def write(session: str, items: list[tuple[int, int]]) -> None:
680
+ self._transfer.write(
681
+ session,
682
+ [addresses[index] for index, _ in items],
683
+ [dst for _, dst in items],
684
+ [lengths[index] for index, _ in items],
685
+ )
686
+
687
+ sessions = list(by_session.items())
688
+
689
+ # Write shards concurrently to avoid serial TP latency.
690
+ def track_write(index: int, future: Future[None]) -> None:
691
+ session = sessions[index][0]
692
+ self._producer_pushes.track_shard_futures(
693
+ list(session_records[session].values()), [future]
694
+ )
695
+
696
+ writes: list[Callable[[], None]] = [
697
+ partial(write, *session) for session in sessions
698
+ ]
699
+ self._run_fanout(writes, track_write)
700
+ finally:
701
+ if staged is not None:
702
+ self._producer_memory.release(staged)
703
+ self._transfer.release_sources(registered_sources)
704
+
705
+ self._producer_pushes.begin_notifying(pushes)
706
+ self._notify_completions(ready)
707
+ self._producer_pushes.complete(pushes)
708
+ except Exception as exc:
709
+ # Report asynchronously; raising here would fail EngineCore.
710
+ failure = exc
711
+ logger.exception(
712
+ "EC Mooncake push batch failed for mm_hashes=%s",
713
+ [push.spec.mm_hash for push in pushes],
714
+ )
715
+ self._producer_pushes.settle_all(pushes)
716
+ self._abandon_pushes(pushes)
717
+ finally:
718
+ if failure is not None:
719
+ self._producer_pushes.fail(pushes, failure)
720
+
721
+ def _notify_completions(
722
+ self, notifications: list[tuple[ProducerPushRecord, dict[str, Any]]]
723
+ ) -> None:
724
+ """Tell the consumer, in one message per destination, what landed."""
725
+ if not notifications:
726
+ return
727
+ by_destination: dict[str, list[tuple[ProducerPushRecord, dict[str, Any]]]] = {}
728
+ for push, reservation in notifications:
729
+ by_destination.setdefault(
730
+ str(reservation.get("addr", push.spec.consumer_zmq)), []
731
+ ).append((push, reservation))
732
+ destinations = list(by_destination.items())
733
+
734
+ def notify(
735
+ consumer_zmq: str,
736
+ items: list[tuple[ProducerPushRecord, dict[str, Any]]],
737
+ ) -> None:
738
+ result = self._control_client.request(
739
+ consumer_zmq,
740
+ {
741
+ "op": "complete_batch",
742
+ "items": [
743
+ {
744
+ "transfer_id": push.spec.transfer_id,
745
+ "reservation_id": reservation["reservation_id"],
746
+ }
747
+ for push, reservation in items
748
+ ],
749
+ },
750
+ )
751
+ completions = result.get("items", []) if isinstance(result, dict) else []
752
+ if len(completions) != len(items):
753
+ raise RuntimeError("Malformed EC completion response")
754
+ for (push, _), completion in zip(items, completions):
755
+ if not completion.get("completed"):
756
+ raise RuntimeError(
757
+ f"Unknown EC reservation for mm_hash={push.spec.mm_hash}"
758
+ )
759
+
760
+ def track(index: int, future: Future[None]) -> None:
761
+ records = {
762
+ push.spec.transfer_id: push for push, _ in destinations[index][1]
763
+ }
764
+ self._producer_pushes.track_shard_futures(list(records.values()), [future])
765
+
766
+ self._run_fanout(
767
+ [
768
+ partial(notify, destination, items)
769
+ for destination, items in destinations
770
+ ],
771
+ track,
772
+ )
773
+
774
+ @staticmethod
775
+ def _known_reservations(record: ProducerPushRecord) -> list[dict[str, Any]]:
776
+ if record.reservations:
777
+ return list(record.reservations)
778
+ try:
779
+ return list(record.reservation_future.result())
780
+ except _FanoutError as exc:
781
+ return [result for result in exc.results if isinstance(result, dict)]
782
+ except Exception:
783
+ return []
784
+
785
+ def _abandon_pushes(self, pushes: list[ProducerPushRecord]) -> None:
786
+ """Release the consumer-side reservations of a batch that failed."""
787
+ for push in pushes:
788
+ shards = self._known_reservations(push)
789
+ if not shards:
790
+ # No confirmed topology means there is no safe address to
791
+ # cancel. The original push failure remains observable via
792
+ # ProducerPushManager.fail below.
793
+ continue
794
+ try:
795
+ self._retry_cancel_reservations(push.spec, shards, record=push)
796
+ except _FanoutError:
797
+ logger.exception(
798
+ "Failed to abandon EC reservations for transfer_id=%s",
799
+ push.spec.transfer_id,
800
+ )
801
+
802
+ def _flush_pending_pushes(self) -> None:
803
+ self._producer_pushes.submit_batches(
804
+ self._io_executor,
805
+ self._push_batch,
806
+ )
807
+
808
+ def _bind_push_source(self, tensor: torch.Tensor, mm_hash: str) -> None:
809
+ ready_event = None
810
+ if tensor.device.type == "cuda":
811
+ ready_event = torch.Event()
812
+ ready_event.record(torch.accelerator.current_stream(tensor.device))
813
+ self._producer_pushes.bind_source(mm_hash, tensor, ready_event)
814
+
815
+ def _cancel_orphaned_reservation(self, record: ProducerPushRecord) -> None:
816
+ resolution_error: Exception | None = None
817
+ try:
818
+ reservations = self._producer_pushes.resolve_reservations(record)
819
+ except Exception as exc:
820
+ resolution_error = exc
821
+ reservations = self._known_reservations(record)
822
+ reservations = [
823
+ shard for shard in reservations if not shard.get("cancelled", False)
824
+ ]
825
+ if not reservations:
826
+ if resolution_error is not None:
827
+ self._producer_pushes.fail([record], resolution_error)
828
+ else:
829
+ self._producer_pushes.finish_cancel(record)
830
+ return
831
+ try:
832
+ self._retry_cancel_reservations(record.spec, reservations, record=record)
833
+ except Exception as cleanup_error:
834
+ if resolution_error is not None:
835
+ combined_error = RuntimeError(
836
+ f"EC reservation resolution failed ({resolution_error}); "
837
+ f"cleanup also failed ({cleanup_error})"
838
+ )
839
+ combined_error.__cause__ = resolution_error
840
+ cleanup_error = combined_error
841
+ self._producer_pushes.fail([record], cleanup_error)
842
+ return
843
+ if resolution_error is not None:
844
+ self._producer_pushes.fail([record], resolution_error)
845
+ else:
846
+ self._producer_pushes.finish_cancel(record)
847
+
848
+ def get_finished(
849
+ self, finished_req_ids: set[str]
850
+ ) -> tuple[set[str] | None, set[str] | None]:
851
+ if not self.is_producer:
852
+ return None, None
853
+
854
+ for record in self._producer_pushes.cancel_requests(finished_req_ids):
855
+ self._producer_pushes.submit_cancel(
856
+ record,
857
+ self._io_executor,
858
+ self._cancel_orphaned_reservation,
859
+ )
860
+ return None, None
861
+
862
+ def save_caches(
863
+ self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs: Any
864
+ ) -> None:
865
+ if not self.is_producer:
866
+ return
867
+ tensor = encoder_cache[mm_hash]
868
+ self._bind_push_source(tensor, mm_hash)
869
+
870
+ def build_connector_worker_meta(self) -> ECMooncakeWorkerMetadata | None:
871
+ if self.is_consumer and not self._is_receiving_rank:
872
+ # `loaded` is intersected across reporting ranks, so a stage that
873
+ # never loads must not report at all rather than report nothing.
874
+ return None
875
+
876
+ self._collecting_sources = False
877
+ self._push_ready.set()
878
+ self._flush_pending_pushes()
879
+ failures = self._producer_pushes.poll()
880
+ for mm_hash, error in failures:
881
+ logger.error(
882
+ "EC Mooncake async save failed for mm_hash=%s: %s",
883
+ mm_hash,
884
+ error,
885
+ )
886
+ reclaimed = self._consumer_memory.drain_reclaimed()
887
+ meta = ECMooncakeWorkerMetadata(
888
+ loaded=self._completed_loads,
889
+ failed_loads=self._failed_loads,
890
+ reclaimed=reclaimed,
891
+ pending_saves=self._producer_pushes.pending,
892
+ failed_saves=self._failed_saves,
893
+ )
894
+ self._failed_saves = set()
895
+ self._completed_loads = set()
896
+ self._failed_loads = set()
897
+ return meta
898
+
899
+ def close(self) -> None:
900
+ if self._shutdown:
901
+ return
902
+ self._shutdown = True
903
+ dispatcher = getattr(self, "_dispatcher", None)
904
+ if dispatcher is not None:
905
+ self._dispatch_stop.set()
906
+ self._push_ready.set()
907
+ dispatcher.join()
908
+ if self._control_server is not None:
909
+ self._reservations.begin_shutdown()
910
+ self._flush_pending_pushes()
911
+ for record in self._producer_pushes.cancel_requests(None):
912
+ self._producer_pushes.submit_cancel(
913
+ record,
914
+ self._io_executor,
915
+ self._cancel_orphaned_reservation,
916
+ )
917
+ self._control_executor.shutdown(wait=True)
918
+ self._producer_pushes.submit_batches(
919
+ self._io_executor, self._push_batch, wait=True
920
+ )
921
+ self._io_executor.shutdown(wait=True)
922
+ if self._shard_pool is not None:
923
+ self._shard_pool.shutdown(wait=True)
924
+ # Every producer-side thread that could hold a control socket is stopped.
925
+ self._control_client.close()
926
+ drained = True
927
+ if self._control_server is not None:
928
+ drained = self._reservations.wait_for_writers(
929
+ self._shutdown_drain_timeout_s
930
+ )
931
+ self._control_server.close()
932
+ if drained:
933
+ self._consumer_memory.close()
934
+ else:
935
+ logger.error(
936
+ "Timed out waiting for Mooncake EC writers; keeping the consumer "
937
+ "receive pool registered"
938
+ )
939
+ self._producer_memory.close()
940
+ self._transfer.close()
vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Encoder-cache connector backed by Mooncake TransferEngine."""
4
+
5
+from __future__ import annotations
6
+
7
+from typing import TYPE_CHECKING, Any
8
+
9
+from vllm.distributed.ec_transfer.ec_connector.base import (
10
+ ECConnectorBase,
11
+ ECConnectorMetadata,
12
+ ECConnectorRole,
13
+ ECConnectorWorkerMetadata,
14
+)
15
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
16
+ ECMooncakeConnectorMetadata,
17
+)
18
+from vllm.distributed.ec_transfer.ec_connector.mooncake.scheduler import (
19
+ ECMooncakeScheduler,
20
+)
21
+from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import ECMooncakeWorker
22
+
23
+if TYPE_CHECKING:
24
+ import torch
25
+
26
+ from vllm.config import VllmConfig
27
+ from vllm.v1.core.sched.output import SchedulerOutput
28
+ from vllm.v1.outputs import ECConnectorOutput
29
+ from vllm.v1.request import Request
30
+
31
+
32
+class ECMooncakeConnector(ECConnectorBase):
33
+ """Preserve the public API while delegating to one process-role component."""
34
+
35
+ def __init__(self, vllm_config: VllmConfig, role: ECConnectorRole):
36
+ super().__init__(vllm_config=vllm_config, role=role)
37
+
38
+ self._scheduler: ECMooncakeScheduler | None = None
39
+ self._worker: ECMooncakeWorker | None = None
40
+ self._closed = False
41
+
42
+ if role == ECConnectorRole.SCHEDULER:
43
+ self._scheduler = ECMooncakeScheduler(vllm_config)
44
+ elif role == ECConnectorRole.WORKER:
45
+ self._worker = ECMooncakeWorker(vllm_config)
46
+ else:
47
+ raise ValueError(f"Unknown EC connector role: {role}")
48
+
49
+ def start_worker_services(self) -> None:
50
+ assert self._worker is not None
51
+ self._worker.start_services()
52
+
53
+ def start_save_caches(self, **kwargs: Any) -> None:
54
+ assert self._worker is not None
55
+ metadata = self._get_connector_metadata()
56
+ assert isinstance(metadata, ECMooncakeConnectorMetadata)
57
+ self._worker.start_save_caches(metadata, **kwargs)
58
+
59
+ def start_load_caches(
60
+ self, encoder_cache: dict[str, torch.Tensor], **kwargs: Any
61
+ ) -> None:
62
+ assert self._worker is not None
63
+ metadata = self._get_connector_metadata()
64
+ assert isinstance(metadata, ECMooncakeConnectorMetadata)
65
+ self._worker.start_load_caches(metadata, encoder_cache, **kwargs)
66
+
67
+ def save_caches(
68
+ self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs: Any
69
+ ) -> None:
70
+ assert self._worker is not None
71
+ self._worker.save_caches(encoder_cache, mm_hash, **kwargs)
72
+
73
+ def get_finished(
74
+ self, finished_req_ids: set[str]
75
+ ) -> tuple[set[str] | None, set[str] | None]:
76
+ assert self._worker is not None
77
+ return self._worker.get_finished(finished_req_ids)
78
+
79
+ def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
80
+ assert self._worker is not None
81
+ return self._worker.build_connector_worker_meta()
82
+
83
+ def take_unavailable_requests(self) -> set[str]:
84
+ assert self._scheduler is not None
85
+ return self._scheduler.take_unavailable_requests()
86
+
87
+ def has_cache_item(self, identifier: str) -> bool:
88
+ assert self._scheduler is not None
89
+ return self._scheduler.has_cache_item(identifier)
90
+
91
+ def ensure_cache_available(
92
+ self, request: Request, num_computed_tokens: int
93
+ ) -> bool:
94
+ assert self._scheduler is not None
95
+ return self._scheduler.ensure_cache_available(request, num_computed_tokens)
96
+
97
+ def update_state_after_alloc(self, request: Request, index: int) -> None:
98
+ assert self._scheduler is not None
99
+ self._scheduler.update_state_after_alloc(request, index)
100
+
101
+ def update_state_after_free(self, request: Request, index: int) -> None:
102
+ assert self._scheduler is not None
103
+ self._scheduler.update_state_after_free(request, index)
104
+
105
+ def build_connector_meta(
106
+ self, scheduler_output: SchedulerOutput
107
+ ) -> ECConnectorMetadata:
108
+ assert self._scheduler is not None
109
+ return self._scheduler.build_connector_meta(scheduler_output)
110
+
111
+ def update_connector_output(self, connector_output: ECConnectorOutput) -> None:
112
+ assert self._scheduler is not None
113
+ self._scheduler.update_connector_output(connector_output)
114
+
115
+ def has_pending_push_work(self) -> bool:
116
+ assert self._scheduler is not None
117
+ return self._scheduler.has_pending_push_work()
118
+
119
+ def request_finished(self, request: Request) -> tuple[bool, dict[str, Any] | None]:
120
+ assert self._scheduler is not None
121
+ return self._scheduler.request_finished(request)
122
+
123
+ def shutdown(self) -> None:
124
+ if self._closed:
125
+ return
126
+ self._closed = True
127
+ if self._scheduler is not None:
128
+ self._scheduler.close()
129
+ if self._worker is not None:
130
+ self._worker.close()
vllm/v1/core/sched/scheduler.py
+41
-10
@@ -616,6 +616,17 @@ class Scheduler(SchedulerInterface):
616
req_index += 1
617
continue
618
619
+ if (
620
+ self.ec_connector is not None
621
+ and request.mm_features
622
+ and not self.ec_connector.ensure_cache_available(
623
+ request,
624
+ request.num_computed_tokens - request.num_output_placeholders,
625
+ )
626
+ ):
627
+ req_index += 1
628
+ continue
629
+
630
num_new_tokens = (
631
request.num_tokens_with_spec
632
+ request.num_output_placeholders
@@ -945,13 +956,7 @@ class Scheduler(SchedulerInterface):
956
assert num_computed_tokens <= request.num_tokens
957
958
# Skip request with pending mm encoding prefetches
948
- if (
949
- self.ec_connector is not None
950
- and request.mm_features
951
- and not self.ec_connector.ensure_cache_available(
952
- request, num_computed_tokens
953
- )
954
- ):
959
+ if self._ec_transfer_pending(request, num_computed_tokens):
960
request_queue.pop_request()
961
step_skipped_waiting.prepend_request(request)
962
continue
@@ -966,11 +971,18 @@ class Scheduler(SchedulerInterface):
971
)
972
else:
973
# KVTransfer: WAITING reqs have num_computed_tokens > 0
969
- # after async KV recvs are completed.
974
+ # after async KV recvs are completed. A streaming-input
975
+ # session resumes here too, carrying whatever media its
976
+ # latest chunk added, so this branch needs the same gate.
977
new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks
978
num_new_local_computed_tokens = 0
979
num_computed_tokens = request.num_computed_tokens
980
981
+ if self._ec_transfer_pending(request, num_computed_tokens):
982
+ request_queue.pop_request()
983
+ step_skipped_waiting.prepend_request(request)
984
+ continue
985
+
986
encoder_inputs_to_schedule = None
987
external_load_encoder_input = []
988
new_encoder_compute_budget = encoder_compute_budget
@@ -2162,6 +2174,10 @@ class Scheduler(SchedulerInterface):
2174
self.grammar_compile_error_reqs.clear()
2175
if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
2176
error_req_ids.update(failed_kv_load_req_ids)
2177
+ if self.ec_connector is not None:
2178
+ # An encoder input the connector can no longer obtain. Failing is
2179
+ # retryable: re-issuing the request re-runs the encode.
2180
+ error_req_ids.update(self.ec_connector.take_unavailable_requests())
2181
2182
if error_req_ids:
2183
error_reqs = self.finish_requests(
@@ -2258,6 +2274,16 @@ class Scheduler(SchedulerInterface):
2274
2275
return engine_core_outputs
2276
2277
+ def _ec_transfer_pending(self, request: Request, num_computed_tokens: int) -> bool:
2278
+ """Whether an encoder input this request needs is still in transit."""
2279
+ return (
2280
+ self.ec_connector is not None
2281
+ and bool(request.mm_features)
2282
+ and not self.ec_connector.ensure_cache_available(
2283
+ request, num_computed_tokens
2284
+ )
2285
+ )
2286
+
2287
@staticmethod
2288
def _is_blocked_waiting_status(status: RequestStatus) -> bool:
2289
return status in (
@@ -2344,7 +2370,7 @@ class Scheduler(SchedulerInterface):
2370
# With Whisper, as soon as we've generated a single token,
2371
# we know we're done with the encoder input. Cross Attention
2372
# KVs have been calculated and cached already.
2347
- self.encoder_cache_manager.free_encoder_input(request, input_id)
2373
+ self._free_encoder_input(request, input_id)
2374
elif (
2375
start_pos + num_tokens + spec_lookahead
2376
<= request.num_computed_tokens - request.num_output_placeholders
@@ -2352,7 +2378,12 @@ class Scheduler(SchedulerInterface):
2378
# Processed, stored in the decoder KV cache, and far enough past
2379
# the placeholder range (plus the drafter's look-ahead) that no
2380
# rejection or drafter gather can reference it.
2355
- self.encoder_cache_manager.free_encoder_input(request, input_id)
2381
+ self._free_encoder_input(request, input_id)
2382
+
2383
+ def _free_encoder_input(self, request: Request, input_id: int) -> None:
2384
+ self.encoder_cache_manager.free_encoder_input(request, input_id)
2385
+ if self.ec_connector is not None:
2386
+ self.ec_connector.update_state_after_free(request, input_id)
2387
2388
def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
2389
for req_id, spec_token_ids in zip(
vllm/v1/worker/ec_connector_model_runner_mixin.py
+7
@@ -64,6 +64,12 @@ class ECConnectorModelRunnerMixin:
64
assert scheduler_output.ec_connector_metadata is not None
65
ec_connector.bind_connector_metadata(scheduler_output.ec_connector_metadata)
66
67
+ if ec_connector.is_producer:
68
+ # Pass the cache explicitly: a producer needs it to publish an item
69
+ # that is already cached from an earlier step, which this step will
70
+ # not re-encode -- `save_caches` never fires for those.
71
+ ec_connector.start_save_caches(encoder_cache=encoder_cache, **kwargs)
72
+
73
# Load caches for consumer or both roles
74
if ec_connector.is_consumer:
75
ec_connector.start_load_caches(encoder_cache, **kwargs)
@@ -74,5 +80,6 @@ class ECConnectorModelRunnerMixin:
80
output.finished_sending, output.finished_recving = (
81
ec_connector.get_finished(scheduler_output.finished_req_ids)
82
)
83
+ output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
84
85
ec_connector.clear_connector_metadata()
vllm/v1/worker/gpu/ec_connector.py
+3
@@ -62,6 +62,9 @@ class ActiveECConnector(ECConnector):
62
assert scheduler_output.ec_connector_metadata is not None
63
ec_connector.bind_connector_metadata(scheduler_output.ec_connector_metadata)
64
65
+ if ec_connector.is_producer:
66
+ ec_connector.start_save_caches(encoder_cache=self.encoder_cache)
67
+
68
if ec_connector.is_consumer:
69
ec_connector.start_load_caches(self.encoder_cache)
70
vllm/v1/worker/gpu_worker.py
+5
@@ -28,6 +28,8 @@ from vllm.distributed import (
28
from vllm.distributed.ec_transfer import (
29
ensure_ec_transfer_initialized,
30
ensure_ec_transfer_shutdown,
31
+ get_ec_transfer,
32
+ has_ec_transfer,
33
)
34
from vllm.distributed.eplb.eplb_utils import override_envs_for_eplb
35
from vllm.distributed.kv_transfer import (
@@ -501,6 +503,9 @@ class Worker(WorkerBase):
503
):
504
self.model_runner.load_model(load_dummy_weights=load_dummy_weights)
505
506
+ if has_ec_transfer():
507
+ get_ec_transfer().start_worker_services()
508
+
509
if self.vllm_config.weight_transfer_config is not None:
510
self.weight_transfer_engine = WeightTransferEngineFactory.create_engine(
511
self.vllm_config.weight_transfer_config,