Skip to content

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

Open
jiaran-king wants to merge 1 commit into
kvcache-ai:add-mooncake-ec-connectorfrom
jiaran-king:codex/fix-mooncake-failure-paths
Open

[Bugfix][EPD] Fix Mooncake transfer lifecycle cleanup#8
jiaran-king wants to merge 1 commit into
kvcache-ai:add-mooncake-ec-connectorfrom
jiaran-king:codex/fix-mooncake-failure-paths

Conversation

@jiaran-king

@jiaran-king jiaran-king commented Aug 29, 2026

Copy link
Copy Markdown

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.

Wait for every sharded reservation and write before cleanup, attempt cancellation on every shard and retain tombstones after RPC failures, and unregister the producer staging pool on shutdown.

Assisted-by: OpenAI Codex
Signed-off-by: Zhou ziheng <jiaranran2@gmail.com>
@jiaran-king
jiaran-king force-pushed the codex/fix-mooncake-failure-paths branch from 78bb4e9 to 12d3476 Compare August 29, 2026 05:09
@jiaran-king jiaran-king changed the title [Bugfix][EPD] Make Mooncake shard failures lifecycle-safe [Bugfix][EPD] Fix Mooncake transfer lifecycle cleanup Aug 29, 2026
@jiaran-king
jiaran-king marked this pull request as draft August 29, 2026 06:15
@jiaran-king
jiaran-king marked this pull request as ready for review August 29, 2026 06:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant