Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,33 @@ def _add_unfinished_request(
)


def test_pending_load_for_non_chosen_connector_is_dropped():
"""A MultiConnector loser must not turn its proposed load into a save."""
scheduler = _make_bare_scheduler()
request = SimpleNamespace(
request_id="req-0",
block_hashes=[b"h0", b"h1", b"h2"],
)
blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
scheduler.load_specs["req-0"] = LoadSpec(
vllm_cached_tokens=0,
kvpool_cached_tokens=48,
can_load=False,
)

scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0)
meta = scheduler.build_connector_meta(_make_pending_load_scheduler_output())

# MultiConnector exposes the real allocation to every child, but only the
# chosen child receives external tokens. The losing store connector must
# neither retain those blocks for a pending load nor enqueue a save from
# the rejected speculative LoadSpec.
assert scheduler._unfinished_requests["req-0"][1] == ()
assert meta.requests == []
assert "req-0" not in scheduler.load_specs
assert "req-0" not in scheduler._request_trackers


def test_update_state_excludes_nontransfer_groups():
"""Store metadata must match the worker's registered cache groups."""
scheduler = _make_bare_scheduler()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -372,7 +372,12 @@ def build_connector_meta(
) in self._unfinished_requests.items():
if request_id not in request_ids and request_id not in cached_reqs.req_ids:
load_spec = self.load_specs.pop(request_id, None)
if not load_spec:
# A load spec may have been proposed by this connector's
# lookup but rejected by MultiConnector in favor of another
# connector. Only the chosen connector may issue the pending
# load; the normal store path gets its blocks later from
# SchedulerOutput once the request is actually scheduled.
if load_spec is None or not load_spec.can_load:
continue
num_tokens_to_compute = load_spec.kvpool_cached_tokens
request_tracker = RequestTracker(
Expand Down
Loading