kvcache-ai/vllm

Fork 0

mirror of https://github.com/kvcache-ai/vllm.git synced 2026-11-05 01:07:43 +00:00

[PR #8] [Bugfix][EPD] Fix Mooncake transfer lifecycle cleanup #8

New issue

Open

opened 2026-10-29 21:06:05 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:06:05 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/vllm/pull/8
Author: @jiaran-king
Created: 8/29/2026
Status: 🔄 Open

Base: add-mooncake-ec-connector ← Head: codex/fix-mooncake-failure-paths

📝 Commits (9)

5d5b523 [Bugfix][EPD] Fix DeepStack encoder cache reservation sizing (#4)
8246b86 Merge remote-tracking branch 'origin/main' into add-mooncake-ec-connector
4af268d [Bugfix] Prevent duplicate Mooncake push metadata across scheduler steps (#5)
623231e [Bugfix] Ignore late Mooncake ready events for cancelled transfers (#6)
134b7ec [Bugfix] Sweep Mooncake cancel records in deadline order
5be31a3 Merge branch 'main' into add-mooncake-ec-connector
76c582e Fix Mooncake EPD proxy flag in integration script (#7)
9a85544 [EPD] Make the proxy's per-request logging opt-in
12d3476 [Bugfix][EPD] Fix Mooncake transfer lifecycle cleanup

📊 Changes

17 files changed (+5943 additions, -261 deletions)

View changed files

📝 examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py (+255 -81)
📝 tests/v1/core/test_encoder_cache_manager.py (+32 -0)
📝 tests/v1/core/test_scheduler.py (+90 -3)
📝 tests/v1/ec_connector/integration/README.md (+170 -170)
➕ tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh (+226 -0)
➕ tests/v1/ec_connector/unit/test_ec_mooncake_connector.py (+2137 -0)
📝 tests/v1/ec_connector/unit/test_worker_ec_connector.py (+26 -0)
📝 vllm/config/ec_transfer.py (+5 -0)
📝 vllm/distributed/ec_transfer/ec_connector/base.py (+37 -1)
📝 vllm/distributed/ec_transfer/ec_connector/cpu/connector.py (+6 -2)
📝 vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py (+5 -1)
📝 vllm/distributed/ec_transfer/ec_connector/factory.py (+6 -0)
➕ vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py (+2907 -0)
📝 vllm/v1/core/sched/scheduler.py (+26 -3)
📝 vllm/v1/worker/ec_connector_model_runner_mixin.py (+7 -0)
📝 vllm/v1/worker/gpu/ec_connector.py (+3 -0)
📝 vllm/v1/worker/gpu_worker.py (+5 -0)

📄 Description

Summary

wait for every consumer-shard reservation future before rolling back every shard by transfer ID
wait for every started shard write before releasing producer staging memory or abandoning consumer allocations
attempt cancellation on every shard and keep cancelled-transfer tombstones when an RPC fails, so healthy shards are released and late ready events cannot recreate a cancelled transfer
explicitly unregister the producer staging pool during connector shutdown

Why

The Mooncake connector fans reservation and write work out across consumer TP shards. Previously, the result loop stopped at the first exception. Other shard operations could still be running while source staging memory or destination allocations were released, and successful reservations from the same fan-out could remain allocated until lease expiry.

A cancel RPC can also fail after some remote shards have already processed the cancellation. Removing the scheduler tombstone in that case lets a queued late ready event recreate a transfer whose remote state is already partially released.

Tests

.venv/bin/python -m pytest -q tests/v1/ec_connector/unit/test_ec_mooncake_connector.py (49 passed)
.venv/bin/python -m pytest -q tests/v1/ec_connector/unit/test_worker_ec_connector.py (6 passed)
pre-commit run ruff-check --files vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py tests/v1/ec_connector/unit/test_ec_mooncake_connector.py
pre-commit run ruff-format --files vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py tests/v1/ec_connector/unit/test_ec_mooncake_connector.py

The added fault-injection coverage verifies partial TP-shard reservation failure, partial TP-shard write failure, continued cancellation after a shard RPC failure, cancel-RPC failure with a late ready event, and producer-pool unregistration at shutdown.

Duplicate-work check

I searched the open PRs in vllm-project/vllm and kvcache-ai/vllm for the same Mooncake shard-failure and reservation-failure fixes and found no existing implementation.

AI assistance

AI assistance was used to trace the failure paths, prepare the implementation and tests, and run the checks above. The human submitter reviewed the change and its test evidence before submission.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/vllm/pull/8 **Author:** [@jiaran-king](https://github.com/jiaran-king) **Created:** 8/29/2026 **Status:** 🔄 Open **Base:** `add-mooncake-ec-connector` ← **Head:** `codex/fix-mooncake-failure-paths` --- ### 📝 Commits (9) - [`5d5b523`](https://github.com/kvcache-ai/vllm/commit/5d5b523be57d6493cae41fdfc9b741728810c1dd) [Bugfix][EPD] Fix DeepStack encoder cache reservation sizing (#4) - [`8246b86`](https://github.com/kvcache-ai/vllm/commit/8246b8690586ce5996f45ce3f19ed03702ede470) Merge remote-tracking branch 'origin/main' into add-mooncake-ec-connector - [`4af268d`](https://github.com/kvcache-ai/vllm/commit/4af268d2133051822c6649d0ebb99b5bf69b988a) [Bugfix] Prevent duplicate Mooncake push metadata across scheduler steps (#5) - [`623231e`](https://github.com/kvcache-ai/vllm/commit/623231e6d585cb1771c1edf497888311df2d2b23) [Bugfix] Ignore late Mooncake ready events for cancelled transfers (#6) - [`134b7ec`](https://github.com/kvcache-ai/vllm/commit/134b7ec4a9aea238c433db698800736b0e109994) [Bugfix] Sweep Mooncake cancel records in deadline order - [`5be31a3`](https://github.com/kvcache-ai/vllm/commit/5be31a37daf0a35526a73e38180ba5d03c6f63b4) Merge branch 'main' into add-mooncake-ec-connector - [`76c582e`](https://github.com/kvcache-ai/vllm/commit/76c582ebfa4716fafb78ed8ea17b93cb712a49d8) Fix Mooncake EPD proxy flag in integration script (#7) - [`9a85544`](https://github.com/kvcache-ai/vllm/commit/9a855442dbc6752b885e3560de5bac44fb16cb70) [EPD] Make the proxy's per-request logging opt-in - [`12d3476`](https://github.com/kvcache-ai/vllm/commit/12d3476443d150e09252ceba9ca3eb3f9579bced) [Bugfix][EPD] Fix Mooncake transfer lifecycle cleanup ### 📊 Changes **17 files changed** (+5943 additions, -261 deletions) <details> <summary>View changed files</summary> 📝 `examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py` (+255 -81) 📝 `tests/v1/core/test_encoder_cache_manager.py` (+32 -0) 📝 `tests/v1/core/test_scheduler.py` (+90 -3) 📝 `tests/v1/ec_connector/integration/README.md` (+170 -170) ➕ `tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh` (+226 -0) ➕ `tests/v1/ec_connector/unit/test_ec_mooncake_connector.py` (+2137 -0) 📝 `tests/v1/ec_connector/unit/test_worker_ec_connector.py` (+26 -0) 📝 `vllm/config/ec_transfer.py` (+5 -0) 📝 `vllm/distributed/ec_transfer/ec_connector/base.py` (+37 -1) 📝 `vllm/distributed/ec_transfer/ec_connector/cpu/connector.py` (+6 -2) 📝 `vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py` (+5 -1) 📝 `vllm/distributed/ec_transfer/ec_connector/factory.py` (+6 -0) ➕ `vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py` (+2907 -0) 📝 `vllm/v1/core/sched/scheduler.py` (+26 -3) 📝 `vllm/v1/worker/ec_connector_model_runner_mixin.py` (+7 -0) 📝 `vllm/v1/worker/gpu/ec_connector.py` (+3 -0) 📝 `vllm/v1/worker/gpu_worker.py` (+5 -0) </details> ### 📄 Description ## Summary - wait for every consumer-shard reservation future before rolling back every shard by transfer ID - wait for every started shard write before releasing producer staging memory or abandoning consumer allocations - attempt cancellation on every shard and keep cancelled-transfer tombstones when an RPC fails, so healthy shards are released and late ready events cannot recreate a cancelled transfer - explicitly unregister the producer staging pool during connector shutdown ## Why The Mooncake connector fans reservation and write work out across consumer TP shards. Previously, the result loop stopped at the first exception. Other shard operations could still be running while source staging memory or destination allocations were released, and successful reservations from the same fan-out could remain allocated until lease expiry. A cancel RPC can also fail after some remote shards have already processed the cancellation. Removing the scheduler tombstone in that case lets a queued late ready event recreate a transfer whose remote state is already partially released. ## Tests - `.venv/bin/python -m pytest -q tests/v1/ec_connector/unit/test_ec_mooncake_connector.py` (49 passed) - `.venv/bin/python -m pytest -q tests/v1/ec_connector/unit/test_worker_ec_connector.py` (6 passed) - `pre-commit run ruff-check --files vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py tests/v1/ec_connector/unit/test_ec_mooncake_connector.py` - `pre-commit run ruff-format --files vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py tests/v1/ec_connector/unit/test_ec_mooncake_connector.py` The added fault-injection coverage verifies partial TP-shard reservation failure, partial TP-shard write failure, continued cancellation after a shard RPC failure, cancel-RPC failure with a late ready event, and producer-pool unregistration at shutdown. ## Duplicate-work check I searched the open PRs in `vllm-project/vllm` and `kvcache-ai/vllm` for the same Mooncake shard-failure and reservation-failure fixes and found no existing implementation. ## AI assistance AI assistance was used to trace the failure paths, prepare the implementation and tests, and run the checks above. The human submitter reviewed the change and its test evidence before submission. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:06:05 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/vllm#8

No description provided.

Rows
Columns