568 lines
24 KiB
Python
568 lines
24 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from collections.abc import Iterator
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
|
|
from vllm.config import KVTransferConfig
|
|
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
|
|
from vllm.v1.outputs import KVConnectorOutput
|
|
|
|
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled
|
|
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
|
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
|
|
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def _output(*, submit=(), required=(), retire=(), metadata=None):
|
|
return DiffusionSchedulerOutput(
|
|
step_id=0,
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids=set(),
|
|
num_running_reqs=0,
|
|
num_waiting_reqs=0,
|
|
kv_connector_metadata=metadata,
|
|
kv_transfer_request_ids=set(submit),
|
|
kv_required_request_ids=set(required),
|
|
kv_finished_request_ids=set(retire),
|
|
kv_poll_only=metadata is None,
|
|
)
|
|
|
|
|
|
def _completed(*ids, invalid=()):
|
|
return KVConnectorOutput(finished_recving=set(ids), invalid_block_ids=set(invalid))
|
|
|
|
|
|
def test_prefetch_survives_post_forward_and_is_not_submitted_twice():
|
|
progress = KVReceiveProgress()
|
|
active = Mock()
|
|
# B finishes before A, in a different get_finished/post_forward call.
|
|
active.post_forward.side_effect = [_completed("B0"), _completed("A0"), _completed(), _completed()]
|
|
first = _output(submit=["A0", "B0"], required=["A0"], metadata=object())
|
|
assert progress.prepare(active, first, 1).finished_recving == {"A0", "B0"}
|
|
second = _output(required=["B0"], retire=["A0"])
|
|
assert progress.prepare(active, second, 1).finished_recving == {"B0"}
|
|
active.pre_forward.assert_called_once_with(first)
|
|
progress.prepare(active, _output(retire=["B0"]), 1)
|
|
assert not progress.submitted
|
|
assert not progress.received
|
|
assert not progress.deadlines
|
|
|
|
|
|
def test_submit_b_does_not_wait_for_b():
|
|
progress = KVReceiveProgress()
|
|
active = Mock()
|
|
active.post_forward.return_value = _completed("A0")
|
|
result = progress.prepare(active, _output(submit=["A0", "B0"], required=["A0"], metadata=object()), 10)
|
|
assert result.finished_recving == {"A0"}
|
|
assert progress.submitted == {"A0", "B0"}
|
|
active.post_forward.assert_called_once()
|
|
active.post_forward.side_effect = [_completed(), _completed("B0")]
|
|
progress.prepare(active, _output(required=["B0"]), 10)
|
|
assert progress.received == {"A0", "B0"}
|
|
assert active.pre_forward.call_count == 1
|
|
|
|
|
|
def test_duplicate_submission_and_unknown_required_fail_before_transfer():
|
|
progress = KVReceiveProgress(submitted={"B0"})
|
|
active = Mock()
|
|
with pytest.raises(RuntimeError, match="Duplicate"):
|
|
progress.prepare(active, _output(submit=["B0"], metadata=object()), 1)
|
|
with pytest.raises(RuntimeError, match="never submitted"):
|
|
progress.prepare(active, _output(required=["C0"]), 1)
|
|
active.pre_forward.assert_not_called()
|
|
|
|
|
|
def test_poll_checks_timeout_without_resetting_it(monkeypatch):
|
|
now = [1.0]
|
|
monkeypatch.setattr("vllm_omni.diffusion.diffusion_kv.kv_connector.time.monotonic", lambda: now[0])
|
|
progress = KVReceiveProgress()
|
|
active = Mock()
|
|
active.post_forward.side_effect = lambda *_: _completed()
|
|
progress.prepare(active, _output(submit=["B0"], metadata=object()), 5)
|
|
now[0] = 7
|
|
with pytest.raises(TimeoutError, match="B0"):
|
|
progress.prepare(active, _output(), 5)
|
|
assert progress.submitted == {"B0"}
|
|
|
|
|
|
def test_poll_checks_load_errors_and_refuses_early_retirement():
|
|
progress = KVReceiveProgress(submitted={"B0"})
|
|
active = Mock()
|
|
active.post_forward.return_value = _completed(invalid=[42])
|
|
with pytest.raises(RuntimeError, match="invalid remote pages"):
|
|
progress.prepare(active, _output(), 1)
|
|
active.post_forward.return_value = _completed()
|
|
with pytest.raises(RuntimeError, match="unfinished"):
|
|
progress.prepare(active, _output(retire=["B0"]), 1)
|
|
assert progress.submitted == {"B0"}
|
|
|
|
|
|
def test_late_unsubmitted_cancellation_ack_is_not_retained():
|
|
progress = KVReceiveProgress()
|
|
active = Mock()
|
|
active.post_forward.return_value = _completed("cancelled-before-admission")
|
|
result = progress.prepare(active, _output(), 1)
|
|
assert not result.finished_recving
|
|
assert not progress.received
|
|
|
|
|
|
def test_executor_intersects_rank_completions_and_waits_only_required():
|
|
executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2), collective_rpc=Mock())
|
|
output = _output(submit=["A0", "B0"], required=["A0"], metadata=object())
|
|
executor.collective_rpc.return_value = [_completed("A0", "B0"), _completed("A0")]
|
|
result = DiffusionExecutor.prepare_kv_for_forward(executor, output)
|
|
assert result.finished_recving == {"A0"}
|
|
assert executor.collective_rpc.call_args.kwargs["exec_all_ranks"] is True
|
|
executor.collective_rpc.return_value = [_completed("A0", "B0"), _completed("A0", "B0")]
|
|
assert DiffusionExecutor.prepare_kv_for_forward(executor, _output()).finished_recving == {"A0", "B0"}
|
|
executor.collective_rpc.return_value = [_completed("B0"), _completed()]
|
|
# Background polling can return early, but compute B requires every rank.
|
|
assert not DiffusionExecutor.prepare_kv_for_forward(executor, _output()).finished_recving
|
|
with pytest.raises(RuntimeError, match="every rank"):
|
|
DiffusionExecutor.prepare_kv_for_forward(executor, _output(required=["B0"]))
|
|
|
|
|
|
def test_executor_submits_prefetch_after_current_rank_barrier():
|
|
calls = []
|
|
|
|
def rpc(method, **kwargs):
|
|
output = kwargs["args"][0]
|
|
calls.append(output)
|
|
if len(calls) == 1:
|
|
assert output.kv_transfer_request_ids == output.kv_required_request_ids == {"A0"}
|
|
assert output.has_sync_kv_loads
|
|
else:
|
|
assert output.kv_transfer_request_ids == {"B0"}
|
|
assert not output.kv_required_request_ids
|
|
assert not output.has_sync_kv_loads
|
|
assert output.kv_connector_metadata == "B metadata"
|
|
# B is still in flight on both ranks. This must not block compute A.
|
|
return [_completed("A0"), _completed("A0")]
|
|
|
|
executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2), collective_rpc=rpc)
|
|
executor.prepare_kv_for_forward = lambda out: DiffusionExecutor.prepare_kv_for_forward(executor, out)
|
|
output = _output(submit=["A0", "B0"], required=["A0"], metadata="A metadata")
|
|
output.kv_prefetch_connector_metadata = "B metadata"
|
|
output.kv_prefetch_request_ids = {"B0"}
|
|
assert executor.prepare_kv_for_forward(output).finished_recving == {"A0"}
|
|
assert len(calls) == 2
|
|
|
|
|
|
def _config():
|
|
return SimpleNamespace(
|
|
kv_transfer_config=KVTransferConfig(
|
|
kv_connector="MooncakeConnector",
|
|
kv_role="kv_consumer",
|
|
engine_id="prefetch-test",
|
|
kv_connector_extra_config={"mooncake_protocol": "tcp", "enable_kv_async_prefetch": True},
|
|
),
|
|
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
|
max_num_seqs=1,
|
|
model_class_name="TestDiffusionPipeline",
|
|
)
|
|
|
|
|
|
def test_mock_config_does_not_enable_native_prefetch():
|
|
assert not native_prefetch_enabled(Mock())
|
|
|
|
|
|
@pytest.mark.parametrize("enabled", ["true", "false", 0, 1, None])
|
|
def test_native_prefetch_rejects_non_boolean_flag(enabled):
|
|
config = _config()
|
|
config.kv_transfer_config.kv_connector_extra_config["enable_kv_async_prefetch"] = enabled
|
|
with pytest.raises(ValueError, match="must be a boolean"):
|
|
native_prefetch_enabled(config)
|
|
|
|
|
|
def test_opt_in_and_platform_validation(monkeypatch):
|
|
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: True)
|
|
config = _config()
|
|
assert native_prefetch_enabled(config)
|
|
config.model_class_name = "AnotherDiffusionPipeline"
|
|
assert native_prefetch_enabled(config)
|
|
config.max_num_seqs = 2
|
|
with pytest.raises(ValueError, match="max_num_seqs=1"):
|
|
native_prefetch_enabled(config)
|
|
config.max_num_seqs = 1
|
|
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: False)
|
|
with pytest.raises(ValueError, match="CUDA"):
|
|
native_prefetch_enabled(config)
|
|
config.kv_transfer_config.kv_connector_extra_config["enable_kv_async_prefetch"] = False
|
|
assert not native_prefetch_enabled(config)
|
|
assert not native_prefetch_enabled(SimpleNamespace())
|
|
|
|
|
|
class _Connector:
|
|
def __init__(self):
|
|
self.commits = []
|
|
self.metas = []
|
|
|
|
def get_num_new_matched_tokens(self, request, computed):
|
|
return 4, True
|
|
|
|
def update_state_after_alloc(self, request, blocks, tokens):
|
|
self.commits.append(request.request_id)
|
|
self.metas.append(request.request_id)
|
|
# Upstream Mooncake mutates this flag after staging the receive.
|
|
request.kv_transfer_params["do_remote_prefill"] = False
|
|
|
|
def build_connector_meta(self, output):
|
|
meta, self.metas = tuple(self.metas), []
|
|
return meta
|
|
|
|
def update_connector_output(self, output):
|
|
pass
|
|
|
|
def request_finished(self, request, blocks):
|
|
return False, None
|
|
|
|
|
|
def _scheduler(monkeypatch, num_blocks=24):
|
|
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: True)
|
|
connector = _Connector()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.diffusion_kv.kv_connector.create_scheduler_kv_connector", lambda *args: connector
|
|
)
|
|
native_kv_managers.register_all_kvcache_specs(None)
|
|
spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
|
|
config = KVCacheConfig(
|
|
num_blocks=num_blocks,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(
|
|
size=spec.page_size_bytes * num_blocks,
|
|
layers=["layer0"],
|
|
layer_stride=spec.page_size_bytes,
|
|
block_stride=spec.page_size_bytes,
|
|
)
|
|
],
|
|
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
|
|
)
|
|
scheduler = RequestScheduler()
|
|
scheduler.initialize(
|
|
_config(),
|
|
kv_cache_config=config,
|
|
scheduler_block_size=4,
|
|
hash_block_size=4,
|
|
kv_vllm_config=SimpleNamespace(
|
|
model_config=SimpleNamespace(max_model_len=64),
|
|
max_in_flight_tokens=64,
|
|
cache_config=SimpleNamespace(enable_prefix_caching=False),
|
|
),
|
|
)
|
|
return scheduler, connector
|
|
|
|
|
|
def _request(rid):
|
|
return OmniDiffusionRequest(
|
|
prompt=rid,
|
|
request_id=rid,
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
|
diffusion_kv_requests=tuple(
|
|
DiffusionKVRequest(
|
|
f"{rid}/{row}", sequence_id=row, prefix_len=4, target_len=4, seq_len=8, prompt_token_ids=[1] * 4
|
|
)
|
|
for row in range(2)
|
|
),
|
|
kv_transfer_params={
|
|
"do_remote_prefill": True,
|
|
"transfer_id": f"xfer-{rid}",
|
|
"remote_engine_id": "ar",
|
|
"remote_bootstrap_addr": "localhost:8998",
|
|
"num_transfer_tokens": 4,
|
|
},
|
|
)
|
|
|
|
|
|
def test_scheduler_prefetches_one_waiting_request_and_reuses_allocation(monkeypatch):
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
for rid in ("A", "B", "C"):
|
|
scheduler.add_request(_request(rid))
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["A"]
|
|
assert output.kv_transfer_request_ids == {"A/0", "A/1", "B/0", "B/1"}
|
|
assert output.kv_required_request_ids == {"A/0", "A/1"}
|
|
assert output.kv_connector_metadata == ("A/0", "A/1")
|
|
assert output.kv_prefetch_connector_metadata == ("B/0", "B/1")
|
|
assert scheduler.get_request_state("B").status == DiffusionRequestStatus.WAITING
|
|
manager = scheduler._diffusion_kv_manager
|
|
b_meta = manager.get_metadata("B")
|
|
assert not manager.has_request("C")
|
|
a_blocks = {b for row in manager.get_metadata("A").sequences for group in row.block_ids for b in group}
|
|
b_blocks = {b for row in b_meta.sequences for group in row.block_ids for b in group}
|
|
assert a_blocks.isdisjoint(b_blocks)
|
|
scheduler.update_kv_connector_output(_completed("A/0", "A/1", "B/0", "B/1"))
|
|
scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED)
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["B"]
|
|
assert manager.get_metadata("B") is b_meta
|
|
assert output.kv_transfer_request_ids == {"C/0", "C/1"}
|
|
assert not output.kv_required_request_ids
|
|
assert output.kv_finished_request_ids == {"A/0", "A/1"}
|
|
assert connector.commits.count("B/0") == connector.commits.count("B/1") == 1
|
|
|
|
|
|
def test_prefetch_capacity_failure_rolls_back_all_cfg_rows(monkeypatch):
|
|
# One sentinel + A's four pages + only one of B's two CFG rows.
|
|
scheduler, connector = _scheduler(monkeypatch, num_blocks=7)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
output = scheduler.schedule()
|
|
assert output.kv_transfer_request_ids == {"A/0", "A/1"}
|
|
assert not scheduler._diffusion_kv_manager.has_request("B")
|
|
assert scheduler._native_prefetch_request_id is None
|
|
assert connector.commits == ["A/0", "A/1"]
|
|
|
|
|
|
def test_abort_drains_b_before_releasing_pages(monkeypatch):
|
|
scheduler, _ = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
scheduler.schedule()
|
|
scheduler.update_kv_connector_output(_completed("A/0", "A/1"))
|
|
scheduler.finish_requests("B", DiffusionRequestStatus.FINISHED_ABORTED)
|
|
assert "B" in scheduler._kv_draining_requests
|
|
assert scheduler._diffusion_kv_manager.has_request("B")
|
|
engine = object.__new__(DiffusionEngine)
|
|
engine.scheduler = scheduler
|
|
engine._remove_diffusion_kv_requests = Mock()
|
|
engine.abort_queue = SimpleNamespace(empty=lambda: True)
|
|
|
|
def drain(output):
|
|
assert output.kv_required_request_ids == {"B/0", "B/1"}
|
|
assert scheduler._diffusion_kv_manager.has_request("B")
|
|
engine._remove_diffusion_kv_requests.assert_not_called()
|
|
return _completed("A/0", "A/1", "B/0", "B/1")
|
|
|
|
engine.executor = SimpleNamespace(prepare_kv_for_forward=drain)
|
|
engine._abort_requests(["B"])
|
|
assert not scheduler._diffusion_kv_manager.has_request("B")
|
|
assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ABORTED
|
|
assert scheduler._native_prefetch_request_id is None
|
|
|
|
|
|
def test_pending_prefetch_reuses_transfer_when_admitted(monkeypatch):
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
scheduler.schedule()
|
|
scheduler.update_kv_connector_output(_completed("A/0", "A/1", "B/0"))
|
|
assert "B" in scheduler._kv_loading_request_ids
|
|
scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED)
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["B"]
|
|
assert output.kv_required_request_ids == {"B/0", "B/1"}
|
|
assert not output.kv_transfer_request_ids
|
|
assert len(connector.commits) == 4
|
|
|
|
|
|
def test_missing_source_skips_prefetch_without_allocating(monkeypatch):
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
scheduler.add_request(_request("A"))
|
|
b = _request("B")
|
|
del b.kv_transfer_params["remote_bootstrap_addr"]
|
|
scheduler.add_request(b)
|
|
scheduler.schedule()
|
|
assert not scheduler._diffusion_kv_manager.has_request("B")
|
|
assert connector.commits == ["A/0", "A/1"]
|
|
|
|
|
|
def test_prefetch_reservation_error_does_not_fail_current_request(monkeypatch):
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
manager = scheduler._diffusion_kv_manager
|
|
reserve = manager.reserve_request
|
|
|
|
def fail_b(request_id, requests):
|
|
if request_id == "B":
|
|
raise ValueError("invalid B layout")
|
|
return reserve(request_id, requests)
|
|
|
|
monkeypatch.setattr(manager, "reserve_request", fail_b)
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["A"]
|
|
assert connector.commits == ["A/0", "A/1"]
|
|
assert not manager.has_request("B")
|
|
scheduler.update_kv_connector_output(_completed("A/0", "A/1"))
|
|
scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED)
|
|
output = scheduler.schedule()
|
|
assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ERROR
|
|
assert "B" in output.finished_req_ids
|
|
|
|
|
|
def test_invalid_prefetch_boundary_is_rejected_before_reservation(monkeypatch):
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
b_rows = scheduler.get_request_state("B").diffusion_kv_requests
|
|
b_rows[1].kv_transfer_params["num_transfer_tokens"] = b_rows[1].num_tokens + 1
|
|
manager = scheduler._diffusion_kv_manager
|
|
reserve = manager.reserve_request
|
|
|
|
def fail_if_b_reserved(request_id, requests):
|
|
if request_id == "B":
|
|
raise AssertionError("invalid B must be rejected before reservation")
|
|
return reserve(request_id, requests)
|
|
|
|
monkeypatch.setattr(manager, "reserve_request", fail_if_b_reserved)
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["A"]
|
|
assert scheduler.get_request_state("B").status == DiffusionRequestStatus.WAITING
|
|
assert not manager.has_request("B")
|
|
|
|
|
|
def test_abort_timeout_keeps_target_pages(monkeypatch):
|
|
scheduler, _ = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
scheduler.schedule()
|
|
engine = object.__new__(DiffusionEngine)
|
|
engine.scheduler = scheduler
|
|
engine.executor = SimpleNamespace(prepare_kv_for_forward=Mock(side_effect=TimeoutError("receive")))
|
|
engine._fail_engine = Mock()
|
|
engine._remove_diffusion_kv_requests = Mock()
|
|
with pytest.raises(TimeoutError):
|
|
engine._abort_requests(["B"])
|
|
engine._fail_engine.assert_called_once()
|
|
engine._remove_diffusion_kv_requests.assert_not_called()
|
|
assert scheduler._diffusion_kv_manager.has_request("B")
|
|
|
|
|
|
def test_sleep_rejects_live_prefetch_before_touching_allocator():
|
|
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
|
|
|
worker = SimpleNamespace(model_runner=SimpleNamespace(_kv_receive_progress=KVReceiveProgress(submitted={"B/0"})))
|
|
with pytest.raises(RuntimeError, match="Cannot sleep"):
|
|
DiffusionWorker.sleep(worker)
|
|
|
|
|
|
def test_registration_failure_retires_b_and_flushes_cancellation(monkeypatch):
|
|
from vllm_omni.diffusion.diffusion_kv import kv_connector
|
|
|
|
scheduler, connector = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
commit = kv_connector.commit_kv_load
|
|
|
|
def fail_b(conn, manager, requests, matched):
|
|
if requests[0].request_id.startswith("B/"):
|
|
conn.metas.append("B cancellation")
|
|
raise kv_connector.KVTransferRegistrationError("registration rollback")
|
|
return commit(conn, manager, requests, matched)
|
|
|
|
monkeypatch.setattr("vllm_omni.diffusion.sched.base_scheduler.commit_kv_load", fail_b)
|
|
output = scheduler.schedule()
|
|
assert output.scheduled_request_ids == ["A"]
|
|
assert output.finished_req_ids == {"B"}
|
|
assert output.kv_finished_request_ids == {"B/0", "B/1"}
|
|
assert output.kv_prefetch_connector_metadata == ("B cancellation",)
|
|
assert not output.kv_prefetch_request_ids
|
|
assert output.num_waiting_reqs == 0
|
|
assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ERROR
|
|
assert not scheduler._diffusion_kv_manager.has_request("B")
|
|
assert scheduler._diffusion_kv_manager.has_request("A")
|
|
|
|
|
|
def test_partial_cfg_registration_rolls_back_every_row(monkeypatch):
|
|
from vllm_omni.diffusion.diffusion_kv import kv_connector
|
|
|
|
class FakeMooncakeConnector:
|
|
def __init__(self):
|
|
self.staged = []
|
|
self.cancelled = []
|
|
|
|
def update_state_after_alloc(self, request, blocks, tokens):
|
|
self.staged.append(request.request_id)
|
|
if request.request_id.endswith("/1"):
|
|
raise RuntimeError("second CFG row failed")
|
|
|
|
def request_finished(self, request, blocks):
|
|
self.cancelled.append(request.request_id)
|
|
return False, None
|
|
|
|
monkeypatch.setattr(
|
|
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.MooncakeConnector",
|
|
FakeMooncakeConnector,
|
|
)
|
|
scheduler, _ = _scheduler(monkeypatch)
|
|
scheduler.add_request(_request("B"))
|
|
state = scheduler.get_request_state("B")
|
|
rows = state.diffusion_kv_requests
|
|
for row in rows:
|
|
row.kv_transfer_params = dict(state.req.kv_transfer_params)
|
|
allocation = scheduler._diffusion_kv_manager.reserve_request("B", rows)
|
|
connector = FakeMooncakeConnector()
|
|
|
|
with pytest.raises(kv_connector.KVTransferRegistrationError, match="Could not register"):
|
|
kv_connector.commit_kv_load(
|
|
connector,
|
|
scheduler._diffusion_kv_manager.native_manager,
|
|
rows,
|
|
[4, 4],
|
|
)
|
|
assert allocation is not None
|
|
assert connector.staged == ["B/0", "B/1"]
|
|
assert connector.cancelled == ["B/0", "B/1"]
|
|
|
|
|
|
def test_drain_after_terminal_request_state_was_popped(monkeypatch):
|
|
scheduler, _ = _scheduler(monkeypatch)
|
|
for rid in ("A", "B"):
|
|
scheduler.add_request(_request(rid))
|
|
scheduler.schedule()
|
|
scheduler.finish_requests("B", DiffusionRequestStatus.FINISHED_ABORTED)
|
|
scheduler.pop_request_state("B")
|
|
output = scheduler.native_kv_poll_output(drain_request_ids=["B"])
|
|
assert output.kv_required_request_ids == {"B/0", "B/1"}
|
|
assert scheduler._diffusion_kv_manager.has_request("B")
|
|
|
|
|
|
@pytest.mark.parametrize("prefetch", [False, True])
|
|
def test_real_active_connector_starts_required_load_before_polling(prefetch, monkeypatch):
|
|
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
|
|
from vllm.v1.worker.gpu import kv_connector as native_kv_connector
|
|
|
|
from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load
|
|
|
|
connector = Mock(spec=KVConnectorBase_V1)
|
|
connector.get_block_ids_with_load_errors.return_value = set()
|
|
# Exercise upstream's compatibility adapter instead of returning an
|
|
# unconstrained Mock from the v0.30 get_transfer_results API.
|
|
connector.get_transfer_results.side_effect = lambda ids: KVConnectorBase_V1.get_transfer_results(connector, ids)
|
|
monkeypatch.setattr(native_kv_connector, "get_kv_transfer_group", lambda: connector)
|
|
context = object()
|
|
monkeypatch.setattr(native_kv_connector, "is_forward_context_available", lambda: True)
|
|
monkeypatch.setattr(native_kv_connector, "get_forward_context", lambda: context)
|
|
active = native_kv_connector.ActiveKVConnector(Mock(), {})
|
|
|
|
completions: Iterator[tuple[set[str], set[str]]] = iter([(set(), {"A0"}), (set(), set())])
|
|
|
|
def finished(_):
|
|
connector.start_load_kv.assert_called_once_with(context)
|
|
return next(completions)
|
|
|
|
connector.get_finished.side_effect = finished
|
|
output = _output(submit=["A0"], required=["A0"], metadata=object())
|
|
if prefetch:
|
|
result = KVReceiveProgress().prepare(active, output, 1)
|
|
else:
|
|
output.kv_required_request_ids = None
|
|
result = wait_for_kv_load(active, output, 1)
|
|
assert result.finished_recving == {"A0"}
|
|
assert result.finished_sending == set()
|
|
connector.get_transfer_results.assert_called_once_with(set())
|
|
connector.start_load_kv.assert_called_once_with(context)
|