mooncake-store: cascade Scheduler.reset_connector_cache to remove_all(force=True)

[aoshen02]

Summary

When verl (or any other RL framework) calls engine.reset_prefix_cache(reset_connector=True) after a weight update, vLLM core already routes the request through Scheduler.reset_connector_cache() -> self.connector.reset_cache(). This patch wires MooncakeStoreConnector.reset_cache() so the cascade actually clears the Mooncake master, dropping all external KV blocks computed against the previous weights before any new rollout request can read them.

This is the RL-correct hard-reset path; without it, the external Mooncake store would silently serve KV computed against stale weights on the next rollout step.

Design

• SCHEDULER role MooncakeStoreConnector.reset_cache() -> MooncakeStoreScheduler.reset_store() -> LookupKeyClient.reset() (ZMQ send RESET_MAGIC frame).
• WORKER rank-0 LookupKeyServer recognizes the RESET_MAGIC (0xFFFFFFFF) discriminator frame and calls self.store_worker.store.remove_all(force=True) on the master.
• The reset path reuses the existing scheduler<->worker rank 0 ZMQ admin channel originally built for prefix lookups, so no new sockets / handshake. A reserved sentinel value (unreachable for any real prompt length) discriminates RESET from a normal lookup so the wire format stays backward compatible.
• WORKER role reset_cache() is a no-op: reset is driven from the scheduler side via the ZMQ admin channel.
• Soft failure semantics: MooncakeStoreScheduler.reset_store() returns False (rather than raising) when the RPC errors out, matching the soft-dependency contract verl uses to keep training alive if Mooncake hiccups.

Comparison vs SGLang

	SGLang	this patch
trigger	per-RPC opt-in flush_cache=True flag — forget it once -> silent stale-cache corruption	automatic data-plane cascade via reset_prefix_cache(reset_connector=True) (no flag)
multi-rank	each scheduler instance calls remove_all -> idempotent but N redundant RPCs	one ZMQ RPC to rank-0 worker -> one RemoveAll per reset
guard consistency	is_fully_idle() silently false on failure, no doc/test	cascades only when the existing scheduler guard passes; covered by the new test_scheduler_reset_prefix_cache_cascades_only_on_success analog
RL-correctness	client must remember flag	hard-reset is the default path

Tests

6 new mock-based unit tests in tests/v1/kv_connector/unit/test_mooncake_store_connector.py:

• test_reset_cache_scheduler_role_delegates_to_reset_store
• test_reset_cache_scheduler_role_propagates_failure
• test_reset_cache_worker_role_returns_none
• test_scheduler_reset_store_returns_client_reset_result
• test_scheduler_reset_store_handles_rpc_exception
• test_lookup_key_client_reset_uses_magic_protocol

Run:

.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_mooncake_store_connector.py -v -k "reset or lookup_key_client_reset"

Result: 6 new + 13 existing connector tests + 18 existing worker tests all pass (37/37).

End-to-end cascade verified inside the verl-bench-aoshen container against a per-run Mooncake master. Master Admin Metrics confirm the path traffic (smoke wrote 12 puts across 2 rounds + 3 exist queries; both DelAll RemoveAll calls landed):

PutStart:(Req=12/0/12, Item=12/12)
PutEnd:(Req=12/0/12, Item=12/12)
ExistKey:(Req=3/0/3, Item=18/18)
DelAll=2/2

Test plan

• 6 new unit tests pass (mock MooncakeDistributedStore)
• 31 existing mooncake store unit tests still pass (no regressions)
• End-to-end cascade against real Mooncake master (in container) — connector.reset_cache() -> remove_all(force=True) confirmed via master Admin Metrics
• Full verl GRPO 5-cycle smoke (Qwen3-0.6B / 1 tray) — recipe authored at recipe_aoshen/phase_b_mooncake.sh in the paired verl PR; deferred for a dedicated bench run

AI assistance

This PR was authored with the help of Claude Opus 4.7 (1M context). Each line was reviewed by me end-to-end; tests were verified locally on a GB200 host with the container verl-bench-aoshen. The cascade-design (re-using the existing LookupKey ZMQ channel with a RESET_MAGIC sentinel discriminator instead of opening a new admin socket) was chosen to minimize new surface area.

Duplicate-work check

No existing PR addresses MooncakeStoreConnector.reset_cache for the RL hard-reset path. KVConnectorBase_V1.reset_cache is the framework-level hook already shipped in vLLM 0.20.1; this PR is the connector-specific implementation.

[gemini-code-assist]

Code Review

This pull request implements a reset_cache mechanism for the Mooncake store connector, enabling global cache resets via a ZMQ-based administrative command. This is particularly important for RL workflows to prevent stale KV data from persisting after weight updates. The review feedback identifies several critical improvements: clearing pending KV events during a reset to prevent stale state, addressing a potential race condition where background threads might upload stale data post-reset, and adding a timeout to the blocking ZMQ receive call to prevent the engine from hanging if the worker becomes unresponsive.

[gemini-code-assist]

When the cache is reset, any previously collected but not yet taken KV events (self._kv_cache_events) should be cleared. These events represent blocks that were successfully stored in Mooncake, but since the store is being cleared, these events are no longer valid and could lead to incorrect prefix matching or stale state in the consumer if they are processed after the reset.

Suggested change

	if self.role == KVConnectorRole.SCHEDULER:
	assert self.connector_scheduler is not None
	return self.connector_scheduler.reset_store()
	if self.role == KVConnectorRole.SCHEDULER:
	assert self.connector_scheduler is not None
	self._kv_cache_events = None
	return self.connector_scheduler.reset_store()

[gemini-code-assist]

There is a potential race condition between the global remove_all(force=True) call and the background KVCacheStoreSendingThread on all worker ranks. Since reset_cache is currently a no-op on the worker side and only triggers a reset via rank 0, other ranks (or even rank 0's background thread) might still be processing their request_queue and performing batch_put operations for stale KV data (computed with old weights) that land in the store after the reset has completed. This would lead to the silent stale-cache corruption that this PR aims to prevent. Consider implementing a mechanism to flush or clear the background sending queues on all workers before performing the global store reset.

[gemini-code-assist]

The self.socket.recv() call is blocking and lacks a timeout. If the worker rank 0 becomes unresponsive or the ZMQ connection is interrupted, the scheduler (and thus the entire vLLM engine) will hang indefinitely during a prefix cache reset. Given that this is an administrative/control-plane operation, it should have a reasonable timeout and handle the failure gracefully. Note that setting RCVTIMEO on this shared socket will also affect the lookup method, which is likely desirable for robustness.

Suggested change

	def reset(self) -> bool:
	"""Trigger `store.remove_all(force=True)` on worker rank 0.
	Returns True on success, False on failure.
	"""
	self.socket.send_multipart([RESET_MAGIC.to_bytes(4, "big")], copy=False)
	resp = self.socket.recv()
	return int.from_bytes(resp, "big") == RESET_MAGIC
	def reset(self) -> bool:
	"""Trigger `store.remove_all(force=True)` on worker rank 0.
	Returns True on success, False on failure.
	"""
	self.socket.setsockopt(zmq.RCVTIMEO, 10000)
	try:
	self.socket.send_multipart([RESET_MAGIC.to_bytes(4, "big")], copy=False)
	resp = self.socket.recv()
	return int.from_bytes(resp, "big") == RESET_MAGIC
	except zmq.Again:
	return False

[ivanium]

left a minor comment mainly for splitting core changes into a separate PR to oss vllm. Also worthwhile to check gemini's review

[ivanium]

I feel these core changes can be a separate PR to oss vllm

[ivanium]

should we expose this rest_connector as part of the interface?