diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index c3ca30e1b..0bd7ac4d7 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -175,7 +175,7 @@ steps: grade: Blocking commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t" agent_pool: mi300_1 diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index dde5a17e7..7131041c0 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -122,6 +122,7 @@ steps: - label: "Omni · Qwen3-Omni Test" source_file_dependencies: - tests/e2e/offline_inference/test_qwen3_omni.py + - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py - tests/e2e/online_serving/test_qwen3_omni.py - tests/helpers/runtime.py - vllm_omni/entrypoints/openai/serving_chat.py @@ -132,9 +133,10 @@ steps: - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml + - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml timeout_in_minutes: 50 commands: - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Test" diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index 2e176b35d..f3f07a43f 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -37,6 +37,7 @@ validation_paths: - tests/engine/test_async_omni_engine_input.py - tests/engine/test_async_omni_engine_outputs.py - tests/engine/test_async_omni_engine_abort.py + - tests/e2e/offline_inference/test_qwen3_omni.py - tests/engine/test_async_omni_engine_stage_init.py - tests/engine/test_orchestrator.py - tests/engine/test_orchestrator_error_handling.py @@ -113,6 +114,15 @@ Numbers become append-only after normative promotion. Test routing, output ordering, queue correlation, cancellation, failure propagation, shutdown ordering, and representative multi-stage execution. +AR abort of a final-stage LLM request must deliver a terminal output with +`finish_reason="abort"` and the generated prefix so collocated training can +resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling +child must drop `parent_requests` once no children remain. EngineCore control +RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must +propagate worker exceptions rather than returning `{"supported": False}`. +When a request has multiple final output stages, only the last abort message +is request-terminal. Output-processor abort state is committed only after the +physical EngineCore abort succeeds. Changes that cross into stage runtime or public error behavior require review from that contract's owners. diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index 160f12514..fcdfe23ea 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -43,6 +43,7 @@ depends_on: validation_paths: - tests/entrypoints/test_omni_entrypoints.py - tests/entrypoints/test_async_omni.py + - tests/entrypoints/test_async_omni_pause_sleep_routing.py - tests/entrypoints/test_async_omni_duplex.py - tests/entrypoints/test_serve.py - tests/entrypoints/test_stream_finish_reason.py @@ -119,7 +120,13 @@ promotion. Test request validation, protocol conversion, model-adapter routing, streaming, session identity, disconnect/cancellation, and error mapping for -each affected offline or serving entrypoint. +each affected offline or serving entrypoint. Sleep must wait for in-flight +`generate()` admission before EngineCore offload; `wake_up` does not resume +admission — callers must `resume_generation()`. Sleeping tags are tracked per +stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`. +Streaming input pumps take an admission slot immediately before each EngineCore +ADD or update, not while waiting for the next client chunk. Frontend abort +keeps `request_states` until `generate()` consumes the terminal output. ## Promotion gate diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 52a0b3556..32c345194 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -3,6 +3,9 @@ """ E2E offline tests for Omni model with video input and audio output. + +Abort / sleep-admission lives in ``test_qwen3_omni_colocate_async.py`` so it +does not overlap this module's OmniRunners. """ import os diff --git a/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py new file mode 100644 index 000000000..54e9b8f62 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Isolated AsyncOmni abort / sleep-admission regression for Qwen3-Omni. + +This lives in its own module so the L3 merge job can run it *after* +``test_qwen3_omni.py`` (offline) and the online serving file. Those modules keep +module-scoped OmniRunner / OmniServer instances until teardown; overlapping a +CuMem sleep-mode engine with them kills StageEngineCoreProc at handshake +(``Failed core proc(s): {}``). +""" + +import asyncio +import os +from contextlib import ExitStack + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +from vllm import SamplingParams +from vllm.inputs import TokensPrompt + +from tests.helpers.clean import wait_for_gpu_memory_to_clear +from tests.helpers.mark import hardware_test +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.platforms import current_omni_platform + +MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + + +def _output_token_ids(output) -> list[int]: + if not getattr(output, "outputs", None): + return [] + completion = output.outputs[0] + token_ids = list(completion.token_ids or []) + if token_ids: + return token_ids + return list(getattr(completion, "cumulative_token_ids", None) or []) + + +def _finish_reason(output) -> str | None: + if not getattr(output, "outputs", None): + return None + return getattr(output.outputs[0], "finish_reason", None) + + +def _colocate_async_deploy() -> str: + """Thinker-only Instruct deploy with sleep mode on a single GPU. + + Abort / sleep admission is a thinker control-plane contract. The 3-stage + CI yaml plus CuMem on every stage fails engine-core startup in the L3 + merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible. + """ + return modify_stage_config( + get_deploy_config_path("qwen3_omni_moe_thinking.yaml"), + updates={ + "stages": { + 0: { + "enable_sleep_mode": True, + "enforce_eager": True, + "trust_remote_code": True, + "tensor_parallel_size": 1, + "devices": "0", + "gpu_memory_utilization": 0.9, + "max_num_seqs": 1, + "max_model_len": 8192, + "max_num_batched_tokens": 8192, + "enable_prefix_caching": False, + "skip_mm_profiling": True, + }, + }, + }, + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) +@pytest.mark.usefixtures("clean_gpu_memory_between_tests") +@pytest.mark.asyncio +async def test_colocate_async_abort_tokens_and_sleep_admission() -> None: + """**required** regression. Control-plane APIs (``abort`` / ``sleep`` / + ``wake_up`` / ``resume_generation``) are not in ``send_omni_request``. + + On ``origin/main`` this test fails: + - abort: ``generate()`` never yields ``finish_reason="abort"`` with the + tokens produced so far, so resume has an empty prefix (or the stream + hangs after frontend state is dropped). + - sleep: ``generate()`` is admitted into EngineCore while weights/KV are + offloaded, so the task errors (asleep / corrupted ADD frame) instead of + waiting for ``resume_generation()``. + + On this PR both legs pass: abort returns a resumeable prefix, and sleep + holds ``generate()`` until wake + resume. + """ + prompt = "What color is the sky? Write a long, detailed explanation." + device_count = current_omni_platform.device_count() + if device_count > 0: + # Fail closed if a prior module-scoped runner is still on the GPU. + # ``clean_gpu_memory_between_tests`` only logs a note on timeout. + wait_for_gpu_memory_to_clear( + devices=list(range(device_count)), + threshold_ratio=0.15, + timeout_s=120, + ) + + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + deploy_config=_colocate_async_deploy(), + enable_sleep_mode=True, + ) + after.callback(engine.shutdown) + + request_id = "qwen3-abort-partial" + outputs: list = [] + + async def _generate(max_tokens: int, req_id: str, gen_prompt) -> None: + async for output in engine.generate( + prompt=gen_prompt, + request_id=req_id, + sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens), + output_modalities=["text"], + ): + outputs.append(output) + + abort_task = asyncio.create_task(_generate(256, request_id, prompt)) + prefix: list[int] = [] + prompt_token_ids: list[int] = [] + for _ in range(600): + if abort_task.done(): + break + if outputs: + latest = outputs[-1] + prefix = _output_token_ids(latest) + prompt_token_ids = list(getattr(latest, "prompt_token_ids", None) or []) + if prefix: + break + await asyncio.sleep(0.1) + + assert prefix, "generate produced no tokens before abort" + assert not abort_task.done(), "generate finished before abort; raise max_tokens" + await engine.abort(request_id) + await asyncio.wait_for(abort_task, timeout=90) + + final = outputs[-1] + assert final.finished + assert _finish_reason(final) == "abort" + abort_prefix = _output_token_ids(final) + assert abort_prefix, "abort dropped the generated prefix (main-branch behavior)" + + resume_prompt: str | TokensPrompt = ( + TokensPrompt(prompt_token_ids=prompt_token_ids + abort_prefix) if prompt_token_ids else prompt + ) + outputs.clear() + async for output in engine.generate( + prompt=resume_prompt, + request_id="qwen3-abort-resume", + sampling_params=SamplingParams(temperature=0.0, max_tokens=8), + output_modalities=["text"], + ): + outputs.append(output) + assert outputs + assert any(_output_token_ids(out) for out in outputs) + + # Trainer order is pause → abort → sleep. Pause here after the + # resume generate so EngineCore is idle before CuMem offload. + await engine.pause_generation(mode="abort", clear_cache=True) + await engine.sleep(level=1) + outputs.clear() + sleep_task = asyncio.create_task(_generate(8, "qwen3-sleep-admission", prompt)) + await asyncio.sleep(1.0) + assert not sleep_task.done(), "generate() ran while EngineCore was sleeping (main-branch admission race)" + await engine.wake_up() + await asyncio.sleep(0.5) + assert not sleep_task.done(), "generate() resumed before resume_generation()" + await engine.resume_generation() + await asyncio.wait_for(sleep_task, timeout=180) + assert outputs + assert _finish_reason(outputs[-1]) != "abort" + await asyncio.sleep(5) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index 1923113d5..86276ec66 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -67,7 +67,8 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture): assert not task.done() rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True)) - await task + result = await task + assert result == [] try: asyncio.run(_run()) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 61a32fb86..aa7d1c4b2 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -256,8 +256,32 @@ class FakeOutputProcessor: ) def abort_requests(self, request_ids, internal: bool = False): - self.abort_calls.append(request_ids) - return request_ids + aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal) + return aborted_ids + + def abort_requests_collecting_outputs(self, request_ids, *, internal: bool = False, commit_state: bool = True): + """Mirror MultimodalOutputProcessor collecting for AR abort-prefix tests.""" + del internal, commit_state + ids = list(request_ids) + self.abort_calls.append(ids) + outputs: list[RequestOutput] = [] + for rid in ids: + seeded = next( + (ro for ro in self.request_outputs if getattr(ro, "request_id", None) == rid), + None, + ) + token_ids = list(seeded.outputs[0].token_ids) if seeded is not None and seeded.outputs else [1, 2] + # Intentionally stamp an internal-looking id on the RequestOutput so + # StagePool must re-key by the orchestrator id it passed in. + outputs.append( + _build_request_output( + f"engine-internal-{rid}", + token_ids=token_ids, + finished=True, + finish_reason="abort", + ) + ) + return ids, outputs def update_scheduler_stats(self, _scheduler_stats) -> None: return None @@ -320,6 +344,7 @@ def _build_request_output( prompt_token_ids: list[int] | None = None, finished: bool = True, text: str = "test", + finish_reason: str | None = None, ) -> RequestOutput: completion = CompletionOutput( index=0, @@ -327,7 +352,7 @@ def _build_request_output( token_ids=list(token_ids or [1, 2]), cumulative_logprob=0.0, logprobs=None, - finish_reason="stop" if finished else None, + finish_reason=(finish_reason if finish_reason is not None else ("stop" if finished else None)), stop_reason=None, ) return RequestOutput( @@ -983,7 +1008,11 @@ async def test_run_abort(orchestrator_factory) -> None: async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None: stages = [ FakeStageClient(stage_type="llm", final_output=False), - FakeStageClient(stage_type="llm", final_output=True), + FakeStageClient( + stage_type="llm", + final_output=True, + next_inputs=[{"prompt_token_ids": [7, 8, 9]}], + ), ] processors = [ FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]), @@ -1002,6 +1031,11 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N final_stage_id=1, ) await _wait_for(lambda: len(stages[0].add_request_calls) == 1) + # Abort outputs come from the final AR stage's output processor, and + # StagePool.abort_requests only collects for requests with a live + # replica binding. Forward off stage-0 first so stage-1 is bound. + stages[0].push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0)) + await _wait_for(lambda: len(stages[1].add_request_calls) == 1) orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1")) result = await _get_rpc_message(orchestrator_fixture) @@ -1011,11 +1045,53 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N assert result.error is None assert result.rpc_correlation_key == ("abort", "abort-1") assert stages[0].abort_calls == [["req-ack"]] + assert stages[1].abort_calls == [["req-ack"]] assert "req-ack" not in orchestrator_fixture.orchestrator.request_states + assert result.abort_outputs is not None + assert len(result.abort_outputs) == 1 + abort_msg = result.abort_outputs[0] + assert abort_msg.request_id == "req-ack" + assert abort_msg.finished is True + assert abort_msg.stage_id == 1 + assert list(abort_msg.engine_outputs.outputs[0].token_ids) == [2] + assert abort_msg.engine_outputs.outputs[0].finish_reason == "abort" finally: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_abort_marks_only_last_final_stage_output_finished() -> None: + """A request with two final AR stages must keep earlier abort outputs unfinished.""" + stages = [ + FakeStageClient(stage_type="llm", final_output=True), + FakeStageClient(stage_type="llm", final_output=True), + ] + processors = [ + FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[1], finished=True)]), + FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[2], finished=True)]), + ] + pools = _build_stage_pools([[stages[0]], [stages[1]]], output_processors=processors) + pools[0]._request_bindings["req-two"] = 0 + pools[1]._request_bindings["req-two"] = 0 + + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=pools, + ) + orchestrator.request_states["req-two"] = OrchestratorRequestState( + request_id="req-two", + final_stage_id=1, + final_output_stage_ids={0, 1}, + ) + + outputs = await orchestrator._abort_request_ids(["req-two"]) + assert [(msg.stage_id, msg.finished) for msg in outputs] == [(0, False), (1, True)] + assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [1] + assert list(outputs[1].engine_outputs.outputs[0].token_ids) == [2] + + @pytest.mark.asyncio async def test_handle_abort_emits_error_result_on_failure() -> None: rpc_queue: asyncio.Queue = asyncio.Queue() diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 076e8ad91..62ed8334b 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -700,3 +700,62 @@ def test_num_cache_creation_tokens_reaches_direct_request_output(): assert direct.num_cached_tokens == 3 assert direct.num_cache_creation_tokens == 2 + + +def _abort_processor_with_parent(child_ids: list[str]): + processor = object.__new__(MultimodalOutputProcessor) + processor.request_states = {} + processor.external_req_ids = {} + processor.parent_requests = {} + processor._native_text_metrics_by_request = {} + processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None) + parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids)) + processor.parent_requests["parent"] = parent + for child_id in child_ids: + req_state = SimpleNamespace( + parent_req=parent, + lora_name=None, + output_kind=RequestOutputKind.CUMULATIVE, + detokenizer=SimpleNamespace(output_token_ids=[7, 8]), + queue=None, + external_req_id="parent", + make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)), + ) + processor.request_states[child_id] = req_state + return processor, parent + + +def test_abort_last_parallel_child_drops_parent_request(): + processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"]) + + aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True) + assert aborted == ["0_parent"] + assert "parent" in processor.parent_requests + assert parent.child_requests == {"1_parent"} + + aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True) + assert aborted == ["1_parent"] + assert "parent" not in processor.parent_requests + assert not parent.child_requests + + +def test_abort_snapshot_leaves_state_until_commit(): + processor, parent = _abort_processor_with_parent(["0_parent"]) + processor.external_req_ids["parent"] = ["0_parent"] + + aborted, outputs = processor.abort_requests_collecting_outputs( + ["parent"], + internal=False, + commit_state=False, + ) + assert aborted == ["0_parent"] + assert outputs + assert "0_parent" in processor.request_states + assert processor.external_req_ids["parent"] == ["0_parent"] + assert "parent" in processor.parent_requests + assert parent.child_requests == {"0_parent"} + + processor.commit_aborted_request_state(["parent"], internal=False) + assert "0_parent" not in processor.request_states + assert "parent" not in processor.parent_requests + assert "parent" not in processor.external_req_ids diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index b252cf0b8..e8a228827 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -64,8 +64,61 @@ def test_collective_rpc_control_helper_honors_timeout(): return "slept" pool, _client = _make_pool(sleep_async=slow_sleep) - result = await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,)) - assert result != "slept" + with pytest.raises(asyncio.TimeoutError): + await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,)) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_control_method_reraises_worker_error(): + async def run() -> None: + pool, _client = _make_pool(wake_up_async=AsyncMock(side_effect=RuntimeError("worker died"))) + with pytest.raises(RuntimeError, match="worker died"): + await pool.collective_rpc(0, "wake_up") + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_non_control_method_still_returns_error_dict(): + async def run() -> None: + pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed"))) + result = await pool.collective_rpc(0, "reset_prefix_cache") assert result["supported"] is False + assert "probe failed" in result["error"] + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails(): + async def run() -> None: + class RecordingOutputProcessor: + def __init__(self) -> None: + self.collected = False + self.committed = False + + def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True): + del internal + self.collected = True + assert commit_state is False + return list(request_ids), [SimpleNamespace(request_id=request_ids[0])] + + def commit_aborted_request_state(self, request_ids, *, internal=False): + del request_ids, internal + self.committed = True + + abort = AsyncMock(side_effect=RuntimeError("engine abort failed")) + output_processor = RecordingOutputProcessor() + client = SimpleNamespace(stage_type="llm", abort_requests_async=abort) + pool = StagePool(0, [client], output_processor=output_processor) # type: ignore[arg-type] + pool._request_bindings["req-1"] = 0 + + with pytest.raises(RuntimeError, match="engine abort failed"): + await pool.abort_requests(["req-1"]) + assert output_processor.collected is True + assert output_processor.committed is False + abort.assert_awaited_once() asyncio.run(run()) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 88b835048..2549889b3 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -180,7 +180,6 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid) assert len(aborted_ids) == expected_cancel_count assert len(set(aborted_ids)) == expected_cancel_count - assert len(omni.request_states) == len(req_ids) - expected_cancel_count for t in tasks: t.cancel() @@ -190,7 +189,9 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi @pytest.mark.cpu -def test_abort_pops_request_states_only_after_ack(): +def test_abort_keeps_request_states_until_generate_cleanup(): + """Frontend abort must not drop request_states; generate() owns cleanup.""" + async def run(): release = asyncio.Event() seen_during_wait: list[int] = [] @@ -215,7 +216,7 @@ def test_abort_pops_request_states_only_after_ack(): release.set() await task - assert "req-1-aaaa" not in omni.request_states + assert "req-1-aaaa" in omni.request_states asyncio.run(run()) @@ -241,6 +242,93 @@ def test_abort_propagates_engine_errors_without_popping_state(): asyncio.run(run()) +@pytest.mark.cpu +def test_abort_enqueues_prefix_tokens_from_engine(): + """AR abort must deliver the generated prefix on the request queue (CPU).""" + + async def run(): + from vllm.outputs import CompletionOutput + + from vllm_omni.engine.messages import OutputMessage + + queue: asyncio.Queue = asyncio.Queue() + + async def abort_with_prefix(request_ids): + rid = request_ids[0] + engine_output = OmniRequestOutput( + request_id=rid, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="partial", + token_ids=[7, 8, 9], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return [ + OutputMessage( + request_id=rid, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + ] + + omni = get_async_omni_instance(fake_abort_request=abort_with_prefix) + omni.request_states["req-1-cccc"] = SimpleNamespace( + request_id="req-1-cccc", + external_request_id="req-1", + input_stream_task=None, + queue=queue, + ) + await omni._abort(["req-1-cccc"]) + msg = queue.get_nowait() + assert msg.finished is True + assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9] + assert msg.engine_outputs.outputs[0].finish_reason == "abort" + assert "req-1-cccc" in omni.request_states + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): + """Empty abort_async must still unblock generate() with finish_reason=abort.""" + + async def run(): + queue: asyncio.Queue = asyncio.Queue() + + async def empty_abort_async(request_ids): + del request_ids + return [] + + omni = get_async_omni_instance(fake_abort_request=empty_abort_async) + omni.request_states["req-1-dddd"] = SimpleNamespace( + request_id="req-1-dddd", + external_request_id="req-1", + input_stream_task=None, + queue=queue, + ) + await omni._abort(["req-1-dddd"]) + msg = queue.get_nowait() + assert msg.finished is True + assert msg.request_id == "req-1-dddd" + assert msg.engine_outputs.outputs[0].finish_reason == "abort" + assert "req-1-dddd" in omni.request_states + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 5518d0b8b..9600b990d 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -18,11 +18,14 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: omni = object.__new__(AsyncOmni) omni._pause_cond = asyncio.Condition() omni._paused = False + omni._admitting = 0 omni._hold_admission_until_resume = False omni._sleeping_tags = set() + omni._stage_sleeping_tags = {} omni._level2_sleeping = False omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock()) omni._final_output_handler = lambda: None + omni.reset_mm_cache = AsyncMock() stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] omni.engine = SimpleNamespace( @@ -171,6 +174,33 @@ def test_sleep_blocks_admission_before_engine_core_rpc(): asyncio.run(run()) +@pytest.mark.cpu +def test_sleep_waits_for_in_flight_generate_admission(): + """Sleep must not offload while generate() is still in add_request.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + omni._admitting = 1 + rpc_started = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not sleep_task.done() + assert not rpc_started.is_set() + await omni._release_generate_admission() + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + assert omni._admitting == 0 + + asyncio.run(run()) + + @pytest.mark.cpu def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): async def run() -> None: @@ -260,6 +290,80 @@ def test_sleep_diffusion_only_skips_engine_core_collective_rpc(): asyncio.run(run()) +@pytest.mark.cpu +def test_streaming_wait_for_first_chunk_does_not_block_sleep(): + """Waiting for the next client chunk must not hold an admission slot.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + rpc_started = asyncio.Event() + first_chunk = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + async def wait_then_submit() -> None: + await first_chunk.wait() + await omni._submit_with_admission(asyncio.sleep(0)) + + wait_task = asyncio.create_task(wait_then_submit()) + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not wait_task.done() + assert omni._admitting == 0 + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + + wait_task.cancel() + with pytest.raises(asyncio.CancelledError): + await wait_task + assert omni._admitting == 0 + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_streaming_in_flight_add_blocks_sleep(): + """Sleep must wait while a streaming ADD/update holds an admission slot.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + add_started = asyncio.Event() + add_release = asyncio.Event() + rpc_started = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + async def slow_add() -> None: + add_started.set() + await add_release.wait() + + add_task = asyncio.create_task(omni._submit_with_admission(slow_add())) + await add_started.wait() + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not sleep_task.done() + assert not rpc_started.is_set() + assert omni._admitting == 1 + + add_release.set() + await asyncio.wait_for(add_task, timeout=1) + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + assert omni._admitting == 0 + + asyncio.run(run()) + + @pytest.mark.cpu def test_wake_up_restores_admission_for_diffusion_only(): """Pure diffusion must keep sleep → wake → generate (no resume).""" @@ -308,3 +412,30 @@ def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion(): assert omni._paused is False asyncio.run(run()) + + +@pytest.mark.cpu +def test_partial_wake_does_not_skip_remaining_sleeping_stage(): + """sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1]).""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm", "llm"]) + await omni.sleep(stage_ids=[0, 1], level=1, mode="abort") + assert omni._stage_sleeping_tags.keys() == {0, 1} + + await omni.wake_up(stage_ids=[0]) + assert 0 not in omni._stage_sleeping_tags + assert 1 in omni._stage_sleeping_tags + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + + omni.collective_rpc.reset_mock() + acks = await omni.wake_up(stage_ids=[1]) + assert acks + omni.collective_rpc.assert_awaited_once() + wake_kwargs = omni.collective_rpc.await_args.kwargs + assert wake_kwargs["method"] == "wake_up" + assert wake_kwargs["stage_ids"] == [1] + assert not omni._sleeping_tags + assert not omni._stage_sleeping_tags + + asyncio.run(run()) diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index 821e5ac80..9567cbc96 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -109,7 +109,7 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str: async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: - """Best-effort full wake so the next shared-engine test starts active.""" + """Best-effort full wake + resume so the next shared-engine test starts active.""" try: await engine.wake_up(stage_ids=stage_ids) except Exception as e: @@ -532,6 +532,7 @@ async def test_pure_diffusion_sleep_wake(): try: await engine.sleep(level=1) await engine.wake_up() + await engine.resume_generation() async for _ in engine.generate( "test", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index e2890797f..73702df3a 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -216,6 +216,20 @@ def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytes assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4" +@pytest.mark.core_model +@pytest.mark.cpu +def test_resolve_stage_physical_devices_captured_unset_ignores_live_env(monkeypatch: pytest.MonkeyPatch): + """Parallel init captures an unset CUDA_VISIBLE_DEVICES as None. + + A sibling thread may then narrow the live env for spawn. The captured + None must not fall back to that live value, or logical device 0 would + map onto the other stage's GPU. + """ + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7") + assert resolve_stage_physical_devices(0, "0", visible_baseline=None) == "0" + assert resolve_stage_physical_devices(1, "1", visible_baseline=None) == "1" + + @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_raises_on_non_numeric(): diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index 96e1ea5d3..305b0778f 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -199,7 +199,14 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags): worker = object.__new__(OmniGPUWorkerBase) worker.rank = 0 worker.device = "cuda:0" - monkeypatch.setattr(base, "current_omni_platform", _fake_platform()) + order: list[str] = [] + platform = SimpleNamespace( + get_current_memory_usage=lambda device: 0, + empty_cache=lambda: order.append("empty_cache"), + synchronize=lambda: order.append("sync"), + collect=lambda: None, + ) + monkeypatch.setattr(base, "current_omni_platform", platform) calls = {} @@ -209,12 +216,17 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags): return cls() def sleep(self, offload_tags): + order.append("allocator_sleep") calls["offload_tags"] = offload_tags monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator) assert OmniGPUWorkerBase.sleep(worker, level=level) is True assert calls["offload_tags"] == expected_tags + assert order[0] == "sync" + assert "allocator_sleep" in order + assert order.index("sync") < order.index("allocator_sleep") + assert "empty_cache" not in order def test_wake_up_forwards_tags_to_allocator(monkeypatch): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 27b564529..4abccb6fe 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -67,6 +67,7 @@ from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, InteractionMessage, + OutputMessage, StageSubmissionMessage, ) from vllm_omni.engine.orchestrator import Orchestrator @@ -1920,12 +1921,16 @@ class AsyncOmniEngine: self, request_ids: list[str], timeout: float | None = None, - ) -> None: + ) -> list[OutputMessage]: """Abort requests and wait for orchestrator acknowledgment. Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and raises if the orchestrator reports failure or times out. + + Returns: + Final-stage AR abort ``OutputMessage`` list carrying partial + tokens generated before abort (empty for diffusion / no OP state). """ if not request_ids or getattr(self, "_shutdown_called", False): return @@ -1959,6 +1964,7 @@ class AsyncOmniEngine: raise if not result_msg.success: raise RuntimeError(result_msg.error or "abort failed") + return list(result_msg.abort_outputs or []) def submit_interaction( self, diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 0b6630ba7..f6a91dbcc 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -48,17 +48,6 @@ class AbortRequestMessage(EngineQueueMessage, kw_only=True): rpc_id: str | None = None -class AbortResultMessage(EngineQueueMessage, kw_only=True): - type: Literal["abort_result"] = "abort_result" - rpc_id: str - success: bool - error: str | None = None - - @property - def rpc_correlation_key(self) -> tuple[str, str]: - return ("abort", self.rpc_id) - - class InteractionMessage(EngineQueueMessage, kw_only=True): type: Literal["interaction"] = "interaction" request_id: str @@ -113,6 +102,20 @@ class OutputMessage(EngineQueueMessage, kw_only=True): stage_submit_ts: float | None = None +class AbortResultMessage(EngineQueueMessage, kw_only=True): + type: Literal["abort_result"] = "abort_result" + rpc_id: str + success: bool + error: str | None = None + # Final-stage AR abort outputs (partial tokens) for frontend delivery. + # Empty for diffusion / requests with no output-processor state. + abort_outputs: list[OutputMessage] | None = None + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return ("abort", self.rpc_id) + + class StageMetricsMessage(EngineQueueMessage, kw_only=True): type: Literal["stage_metrics"] = "stage_metrics" request_id: str diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 3fee36477..353ca7627 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -871,14 +871,19 @@ class Orchestrator: When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after stage aborts, binding release, and request cleanup complete so the caller can await acknowledgment. Fire-and-forget aborts omit rpc_id. + + AR final-stage abort outputs (partial tokens) are attached to the + result message when ``rpc_id`` is set, or enqueued on + ``output_async_queue`` for fire-and-forget aborts. """ request_ids = msg.request_ids error: str | None = None + abort_outputs: list[OutputMessage] = [] try: # _cleanup_request_ids is CFG-aware: it expands aborted parents to # their companions and fails a deferred parent whose companion is # aborted before its output arrived. - await self._cleanup_request_ids(list(request_ids), abort=True) + abort_outputs = await self._cleanup_request_ids(list(request_ids), abort=True) logger.info("[Orchestrator] Aborted request(s) %s", request_ids) except Exception as exc: error = str(exc) @@ -893,8 +898,12 @@ class Orchestrator: rpc_id=msg.rpc_id, success=error is None, error=error, + abort_outputs=abort_outputs or None, ) ) + elif abort_outputs: + for output_msg in abort_outputs: + await self.output_async_queue.put(output_msg) async def _handle_interaction(self, msg: InteractionMessage) -> None: """Handle a midway interaction for an active streaming diffusion request.""" @@ -934,13 +943,53 @@ class Orchestrator: ) ) - async def _abort_request_ids(self, request_ids: list[str]) -> None: - """Forward abort requests to all stage pools.""" + async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: + """Forward abort requests to all stage pools. + + Collects final-stage AR abort outputs (partial tokens) while request + state is still available for stage-id filtering. Diffusion pools do + not contribute abort outputs. + """ if not request_ids: - return + return [] + abort_outputs: list[OutputMessage] = [] for pool in self.stage_pools: - await pool.abort_requests(request_ids) + stage_outputs = await pool.abort_requests(request_ids) or [] + if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": + final_output_type = getattr(pool.stage_client, "final_output_type", None) or "text" + for orch_req_id, request_output in stage_outputs: + req_state = self.request_states.get(orch_req_id) + if req_state is None: + continue + final_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} + if pool.stage_id not in final_stage_ids: + continue + engine_output = OmniRequestOutput.from_stage_output( + request_output, + request_id=orch_req_id, + finished=True, + stage_id=pool.stage_id, + replica_id=pool.get_bound_replica_id(orch_req_id), + final_output_type=final_output_type, + ) + abort_outputs.append( + OutputMessage( + request_id=orch_req_id, + stage_id=pool.stage_id, + replica_id=pool.get_bound_replica_id(orch_req_id), + engine_outputs=engine_output, + metrics=None, + finished=False, + stage_submit_ts=req_state.stage_submit_ts.get(pool.stage_id), + ) + ) pool.release_bindings(request_ids) + last_index_by_req: dict[str, int] = {} + for index, output_msg in enumerate(abort_outputs): + last_index_by_req[output_msg.request_id] = index + for index, output_msg in enumerate(abort_outputs): + output_msg.finished = index == last_index_by_req[output_msg.request_id] + return abort_outputs def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" @@ -1688,7 +1737,7 @@ class Orchestrator: *, abort: bool = False, close_duplex_sessions: bool = False, - ) -> None: + ) -> list[OutputMessage]: """Release pool bindings and logical request state for the given ids. CFG-aware: cleaning a parent releases its tracker state and pulls its @@ -1697,9 +1746,13 @@ class Orchestrator: can never complete). Every teardown path (stage error, abort, replica loss, membership unregister) funnels through here, so tracker state cannot outlive its requests. + + Returns: + Final-stage AR abort ``OutputMessage`` list when ``abort=True``; + otherwise an empty list. """ if not request_ids: - return + return [] cleanup_ids = list(dict.fromkeys(request_ids)) batch = set(cleanup_ids) @@ -1741,9 +1794,10 @@ class Orchestrator: cleanup_ids.extend(stale_request_ids) cleanup_ids = list(dict.fromkeys(cleanup_ids)) + abort_outputs: list[OutputMessage] = [] try: if abort: - await self._abort_request_ids(cleanup_ids) + abort_outputs = await self._abort_request_ids(cleanup_ids) self._release_request_bindings(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) @@ -1759,6 +1813,7 @@ class Orchestrator: raise if closing_session_ids and self.duplex_control_plane is not None: self.duplex_control_plane.finalize_closed_sessions(closing_session_ids) + return abort_outputs async def _apply_raw_terminal_stage_finish( self, diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 5a8305803..9c8b3792f 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1189,14 +1189,19 @@ class StagePool: # ---- Stage-local control plane ---- - async def abort_requests(self, request_ids: list[str]) -> None: + async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: """Abort the given requests in this stage pool. Request-bound abort routing stays inside the pool because route affinity (``request_id -> replica_id``) is pool-owned. + + For AR stages, collect terminal abort ``RequestOutput`` objects from + output-processor state *before* EngineCore abort so partial tokens + generated so far are preserved. Diffusion stages return no abort + outputs (whole-sample retry remains the recovery path). """ if not request_ids: - return + return [] request_ids_by_replica: dict[int, list[str]] = {} for request_id in request_ids: @@ -1206,18 +1211,45 @@ class StagePool: continue request_ids_by_replica.setdefault(replica_id, []).append(request_id) + abort_outputs: list[tuple[str, Any]] = [] + is_diffusion = self.stage_type == "diffusion" for replica_id, replica_request_ids in request_ids_by_replica.items(): + # Orchestrator ids are OP external ids; EngineCore may use a + # different internal request_id assigned by InputProcessor. + # Always key collected outputs by the orchestrator id we passed + # (internal=False), never by RequestOutput.request_id — some stages + # still surface EngineCore/internal ids and would drop the prefix. + engine_abort_ids = list(replica_request_ids) + if not is_diffusion and self._output_processor is not None: + collect = getattr(self._output_processor, "abort_requests_collecting_outputs", None) + if collect is not None: + engine_abort_ids = [] + for orch_req_id in replica_request_ids: + collected_ids, stage_outputs = collect( + [orch_req_id], + internal=False, + commit_state=False, + ) + for req_out in stage_outputs: + abort_outputs.append((orch_req_id, req_out)) + if collected_ids: + engine_abort_ids.extend(collected_ids) + else: + engine_abort_ids.append(orch_req_id) + else: + aborted = self._output_processor.abort_requests(replica_request_ids, internal=False) + if aborted: + engine_abort_ids = list(aborted) client = self.clients[replica_id] if client is None: continue - await client.abort_requests_async(replica_request_ids) + await client.abort_requests_async(engine_abort_ids) + if not is_diffusion and self._output_processor is not None: + commit = getattr(self._output_processor, "commit_aborted_request_state", None) + if callable(commit): + commit(replica_request_ids, internal=False) - # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that - # would otherwise leak — aborted requests never produce a final - # EngineCoreOutput, so process_outputs() never fires its cleanup path. - all_aborted = [rid for ids in request_ids_by_replica.values() for rid in ids] - if all_aborted and self._output_processor is not None: - self._output_processor.abort_requests(all_aborted, internal=True) + return abort_outputs async def collective_rpc( self, @@ -1227,7 +1259,12 @@ class StagePool: args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> dict[str, Any] | Any: - """Dispatch a stage-scoped control-plane RPC to one physical route.""" + """Dispatch a stage-scoped control-plane RPC to one physical route. + + EngineCore control methods (``pause_scheduler``, ``resume_scheduler``, + ``sleep``, ``wake_up``) propagate worker exceptions. Capability probes + and other RPCs still return ``{"supported": False, "error": ...}``. + """ args = tuple(args or ()) kwargs = dict(kwargs or {}) client = self.clients[replica_id] @@ -1258,9 +1295,15 @@ class StagePool: replica_id, method, ) + if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS: + raise + if isinstance(exc, TimeoutError): + error = f"{type(exc).__name__}: {method} timed out after {timeout}s" + else: + error = str(exc) or repr(exc) return { "supported": False, - "error": str(exc), + "error": error, } def shutdown_replica(self, replica_id: int) -> None: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index a00554209..8ac0fde56 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -17,7 +17,7 @@ from vllm import TokensPrompt from vllm.engine.protocol import EngineClient, StreamingInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.outputs import PoolingRequestOutput +from vllm.outputs import CompletionOutput, PoolingRequestOutput from vllm.plugins.io_processors import get_io_processor from vllm.pooling_params import PoolingParams from vllm.renderers.inputs.preprocess import extract_prompt_components @@ -150,12 +150,19 @@ class AsyncOmni(EngineClient, OmniBase): OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False + # In-flight EngineCore submits (non-streaming add_request, or each + # streaming ADD/update/final marker). sleep() waits for this to hit + # zero so a pipelined request cannot race into EngineCore during + # drain/offload. Streaming generate does not hold a slot while + # waiting for the next client chunk. + self._admitting: int = 0 # True after pause_generation() or AR EngineCore sleep; wake_up must # not reopen generate() until resume_generation(). Diffusion-only # sleep uses _paused as a temporary admission gate and clears it # on wake so sleep → wake → generate keeps working. self._hold_admission_until_resume: bool = False self._sleeping_tags: set[str] = set() + self._stage_sleeping_tags: dict[int, set[str]] = {} self._level2_sleeping: bool = False self._duplex_request_client: DuplexRequestClient | None = None self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() @@ -513,31 +520,40 @@ class AsyncOmni(EngineClient, OmniBase): external_request_id = request_id request_id = self._get_unique_request_id(external_request_id) - # Wait until generation is resumed if the engine is paused + # Wait until generation is resumed if the engine is paused. Non-streaming + # generate holds an admission slot until add_request completes so sleep() + # cannot race a just-unblocked generate into EngineCore during offload. + # Streaming generate does **not** hold the slot while waiting for the + # first client chunk; the input pump acquires it immediately before each + # EngineCore ADD/update. + streaming_input = isinstance(prompt, AsyncGenerator) async with self._pause_cond: await self._pause_cond.wait_for(lambda: not self._paused) + if not streaming_input: + self._admitting = getattr(self, "_admitting", 0) + 1 + admitting = not streaming_input logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") - _sleeping_tags = getattr(self, "_sleeping_tags", None) - if _sleeping_tags: - raise RuntimeError( - f"Generation rejected: Engine is partially or fully asleep. " - f"Currently sleeping tags: {list(_sleeping_tags)}. " - f"Please perform a full wake_up before generating." - ) - - # Reject diffusion list-prompt early with a clear API error. - if isinstance(prompt, list) and any( - getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) - ): - raise ValueError( - "Diffusion stages accept only a single prompt per request. " - "Submit multiple independent requests to use scheduler batching." - ) - input_stream_task: asyncio.Task | None = None try: + _sleeping_tags = getattr(self, "_sleeping_tags", None) + if _sleeping_tags: + raise RuntimeError( + f"Generation rejected: Engine is partially or fully asleep. " + f"Currently sleeping tags: {list(_sleeping_tags)}. " + f"Please perform a full wake_up before generating." + ) + + # Reject diffusion list-prompt early with a clear API error. + if isinstance(prompt, list) and any( + getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) + ): + raise ValueError( + "Diffusion stages accept only a single prompt per request. " + "Submit multiple independent requests to use scheduler batching." + ) + # Start final output dispatcher on the first call to generate() self._final_output_handler() @@ -598,8 +614,11 @@ class AsyncOmni(EngineClient, OmniBase): req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id]) # Add request(s) to stage 0. For streaming inputs, submit - # chunks incrementally through streaming_update. - if isinstance(prompt, AsyncGenerator): + # chunks incrementally through streaming_update. The helper + # returns as soon as the pump task is created; each ADD/update + # takes its own admission slot inside the pump. + if streaming_input: + first_chunk_submitted = asyncio.get_running_loop().create_future() input_stream_task = await self._add_streaming_input_request( request_id=request_id, input_stream=prompt, @@ -608,7 +627,9 @@ class AsyncOmni(EngineClient, OmniBase): final_output_stage_ids=final_output_stage_ids, arrival_time=wall_start_ts, lora_request=lora_request, + first_chunk_submitted=first_chunk_submitted, ) + await first_chunk_submitted else: await self.engine.add_request_async( request_id=request_id, @@ -622,6 +643,9 @@ class AsyncOmni(EngineClient, OmniBase): submit_ts = time.time() req_state.metrics.stage_first_ts[0] = submit_ts req_start_ts[request_id] = submit_ts + if admitting: + await self._release_generate_admission() + admitting = False # Refresh gauges on arrival. self._publish_request_gauges(len(self.request_states)) @@ -640,18 +664,7 @@ class AsyncOmni(EngineClient, OmniBase): logger.debug(f"[AsyncOmni] Request {request_id} completed") - # The input pump can outlive a normally-completed generation - # (e.g. a realtime client that keeps streaming after the final - # output); a live pump keeps issuing streaming updates for a - # request the orchestrator no longer tracks. - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() - - self._log_summary_and_cleanup(request_id) - except (asyncio.CancelledError, GeneratorExit): - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() self._record_request_failure_once(request_id, reason="client_disconnect") await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} aborted.") @@ -661,6 +674,28 @@ class AsyncOmni(EngineClient, OmniBase): await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise + finally: + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() + if admitting: + await self._release_generate_admission() + self._log_summary_and_cleanup(request_id) + + async def _release_generate_admission(self) -> None: + """Drop one in-flight generate admission slot held across add_request.""" + async with self._pause_cond: + self._admitting = max(getattr(self, "_admitting", 1) - 1, 0) + self._pause_cond.notify_all() + + async def _submit_with_admission(self, awaitable): + """Wait for resume, hold one admission slot for a single EngineCore submit.""" + async with self._pause_cond: + await self._pause_cond.wait_for(lambda: not self._paused) + self._admitting = getattr(self, "_admitting", 0) + 1 + try: + return await awaitable + finally: + await self._release_generate_admission() async def _add_streaming_input_request( self, @@ -672,6 +707,7 @@ class AsyncOmni(EngineClient, OmniBase): final_output_stage_ids: Sequence[int], arrival_time: float, lora_request: Any = None, + first_chunk_submitted: asyncio.Future[None] | None = None, ) -> asyncio.Task: """Submit a streaming input generator as incremental stage-0 updates.""" if not sampling_params_list: @@ -688,6 +724,10 @@ class AsyncOmni(EngineClient, OmniBase): stage0_params = stage0_params.clone() stage0_params.skip_clone = True + def _mark_first_chunk_submitted() -> None: + if first_chunk_submitted is not None and not first_chunk_submitted.done(): + first_chunk_submitted.set_result(None) + async def handle_inputs() -> None: nonlocal has_submitted_first_chunk cancelled = False @@ -701,29 +741,34 @@ class AsyncOmni(EngineClient, OmniBase): prompt_text, _, _ = extract_prompt_components(self.model_config, chunk_prompt) if not has_submitted_first_chunk: - await self.engine.add_request_async( - request_id=request_id, - prompt=chunk_prompt, - prompt_text=prompt_text, - sampling_params_list=chunk_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=True, + await self._submit_with_admission( + self.engine.add_request_async( + request_id=request_id, + prompt=chunk_prompt, + prompt_text=prompt_text, + sampling_params_list=chunk_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=True, + ) ) has_submitted_first_chunk = True + _mark_first_chunk_submitted() else: - await self.engine.add_streaming_update_async( - request_id=request_id, - prompt=chunk_prompt, - prompt_text=prompt_text, - sampling_params_list=chunk_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=True, + await self._submit_with_admission( + self.engine.add_streaming_update_async( + request_id=request_id, + prompt=chunk_prompt, + prompt_text=prompt_text, + sampling_params_list=chunk_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=True, + ) ) except (asyncio.CancelledError, GeneratorExit): cancelled = True @@ -738,37 +783,47 @@ class AsyncOmni(EngineClient, OmniBase): ) ) finally: - if not cancelled: - # Send empty final request to indicate that inputs have - # finished. Don't send if canceled (session was aborted). - final_sampling_params_list = list(sampling_params_list) - final_sampling_params_list[0] = stage0_params - final_prompt = TokensPrompt(prompt_token_ids=[0]) + try: + if not cancelled: + # Send empty final request to indicate that inputs have + # finished. Don't send if canceled (session was aborted). + final_sampling_params_list = list(sampling_params_list) + final_sampling_params_list[0] = stage0_params + final_prompt = TokensPrompt(prompt_token_ids=[0]) - if has_submitted_first_chunk: - await self.engine.add_streaming_update_async( - request_id=request_id, - prompt=final_prompt, - prompt_text=None, - sampling_params_list=final_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=False, - ) - else: - await self.engine.add_request_async( - request_id=request_id, - prompt=final_prompt, - prompt_text=None, - sampling_params_list=final_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=False, - ) + if has_submitted_first_chunk: + await self._submit_with_admission( + self.engine.add_streaming_update_async( + request_id=request_id, + prompt=final_prompt, + prompt_text=None, + sampling_params_list=final_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=False, + ) + ) + else: + await self._submit_with_admission( + self.engine.add_request_async( + request_id=request_id, + prompt=final_prompt, + prompt_text=None, + sampling_params_list=final_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=False, + ) + ) + has_submitted_first_chunk = True + finally: + # Unblock generate() even on cancel / empty stream / submit + # failure so it can observe a terminal abort or empty result. + _mark_first_chunk_submitted() input_stream_task = asyncio.create_task(handle_inputs()) req_state.input_stream_task = input_stream_task @@ -1129,22 +1184,76 @@ class AsyncOmni(EngineClient, OmniBase): await self._abort(internal_req_ids) async def _abort(self, request_ids: list[str]) -> None: - """Abort request IDs via the engine and clean frontend state after ack. + """Abort request IDs via the engine and enqueue terminal abort outputs. - Waits for orchestrator abort acknowledgment before popping - ``request_states`` so generate() cleanup stays consistent with - backend binding/request teardown. Orchestrator abort errors propagate. + Waits for orchestrator abort acknowledgment, enqueues any AR terminal + abort outputs (partial tokens) into each request's asyncio queue, then + cancels the input pump. Frontend ``request_states`` stay registered so + ``generate()`` can consume the terminal message in + ``_process_orchestrator_results`` and run normal cleanup. + + When ``abort_async`` returns no output for an active request (OP not + registered yet, unbound replica, or orchestrator id drop), enqueue a + synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. """ - await self.engine.abort_async(request_ids) + abort_outputs = await self.engine.abort_async(request_ids) or [] + delivered: set[str] = set() + for output_msg in abort_outputs: + req_id = getattr(output_msg, "request_id", None) + if req_id is None: + continue + state = self.request_states.get(req_id) + if state is None: + logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) + continue + await state.queue.put(output_msg) + delivered.add(req_id) + for rid in request_ids: + state = self.request_states.get(rid) + if state is not None and rid not in delivered: + queue = getattr(state, "queue", None) + if queue is not None: + await state.queue.put(self._synthetic_abort_output_message(rid)) + delivered.add(rid) for rid in request_ids: self._record_request_failure_once(rid, reason="client_abort") - state = self.request_states.pop(rid, None) + state = self.request_states.get(rid) input_stream_task = getattr(state, "input_stream_task", None) if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) + @staticmethod + def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + """Terminal abort OutputMessage used when the engine returned none.""" + engine_output = OmniRequestOutput( + request_id=request_id, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return OutputMessage( + request_id=request_id, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" n_stages = len(self.engine.stage_clients) @@ -1168,6 +1277,47 @@ class AsyncOmni(EngineClient, OmniBase): ar_stage_ids.append(sid) return ar_stage_ids, diffusion_stage_ids + def _sleeping_tags_for_stages(self, stage_ids: list[int]) -> set[str]: + """Union of sleeping tags recorded for the given stages.""" + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + tags: set[str] = set() + for sid in stage_ids: + tags.update(per_stage.get(sid, set())) + return tags + + def _refresh_union_sleeping_tags(self) -> None: + """Keep ``_sleeping_tags`` as the engine-wide union of per-stage tags.""" + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + union: set[str] = set() + for stage_tags in per_stage.values(): + union.update(stage_tags) + self._sleeping_tags = union + + def _record_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None: + per_stage = getattr(self, "_stage_sleeping_tags", None) + if per_stage is None: + per_stage = {} + self._stage_sleeping_tags = per_stage + tag_set = set(tags) + for sid in stage_ids: + per_stage.setdefault(sid, set()).update(tag_set) + self._refresh_union_sleeping_tags() + + def _clear_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None: + per_stage = getattr(self, "_stage_sleeping_tags", None) + if per_stage is None: + self._sleeping_tags = set() + return + tag_set = set(tags) + for sid in stage_ids: + remaining = per_stage.get(sid) + if remaining is None: + continue + remaining.difference_update(tag_set) + if not remaining: + per_stage.pop(sid, None) + self._refresh_union_sleeping_tags() + async def pause_generation( self, *, @@ -1266,11 +1416,24 @@ class AsyncOmni(EngineClient, OmniBase): return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages) async def reset_mm_cache(self) -> None: - """Reset the multi-modal cache for all stages. + """Reset the frontend (P0) multimodal processor cache. - TODO: Forward to Orchestrator process via message. + ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache. + Clearing P0 avoids hash-only follow-up requests after that reset. """ - logger.warning("[AsyncOmni] reset_mm_cache not yet supported with Orchestrator process") + processor = getattr(self, "input_processor", None) + if processor is None: + processor = getattr(self.engine, "input_processor", None) + cache = getattr(processor, "mm_processor_cache", None) + if cache is None: + logger.debug("[AsyncOmni] reset_mm_cache: no frontend mm_processor_cache") + return + for name in ("clear", "reset", "clear_cache"): + fn = getattr(cache, name, None) + if callable(fn): + fn() + return + logger.debug("[AsyncOmni] reset_mm_cache: cache has no clear/reset method") async def reset_encoder_cache(self) -> None: """Reset the encoder cache for all stages. @@ -1314,10 +1477,15 @@ class AsyncOmni(EngineClient, OmniBase): admission and ``sleep → wake → generate`` keeps working. """ # Block admission before any sleep RPC so generate() waits on - # _pause_cond during the drain/offload window. EngineCore.sleep will - # pause the AR scheduler itself; avoid a second pause_scheduler here. + # _pause_cond during the drain/offload window. Wait until generate() + # coroutines that already passed the pause check have submitted (or + # failed) so EngineCore does not see ADD frames while sleeping. async with self._pause_cond: self._paused = True + await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0) + + # P0 sender cache must drop hashes before EngineCore.sleep clears P1. + await self.reset_mm_cache() self._final_output_handler() ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) @@ -1353,9 +1521,10 @@ class AsyncOmni(EngineClient, OmniBase): if diffusion_stage_ids: final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level)) - if not hasattr(self, "_sleeping_tags"): - self._sleeping_tags = set() - self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value]) + self._record_stage_sleep( + ar_stage_ids + diffusion_stage_ids, + [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value], + ) if level == 2: self._level2_sleeping = True return final_acks @@ -1399,7 +1568,12 @@ class AsyncOmni(EngineClient, OmniBase): "Use sleep(level=1) instead, which offloads weights to CPU RAM " "and supports fast DMA restore." ) - _current_tags = getattr(self, "_sleeping_tags", set()) + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + target_stage_ids = ar_stage_ids + diffusion_stage_ids + _current_tags = self._sleeping_tags_for_stages(target_stage_ids) + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + if not _current_tags and not per_stage: + _current_tags = set(getattr(self, "_sleeping_tags", set())) if tags is None: requested_tags = list(_current_tags) else: @@ -1408,8 +1582,6 @@ class AsyncOmni(EngineClient, OmniBase): logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.") return [] - ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) - final_acks: list[OmniACK] = [] if ar_stage_ids: logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids) @@ -1433,9 +1605,7 @@ class AsyncOmni(EngineClient, OmniBase): if diffusion_stage_ids: final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags)) - for t in requested_tags: - if hasattr(self, "_sleeping_tags"): - self._sleeping_tags.discard(t) + self._clear_stage_sleep(target_stage_ids, requested_tags) # Only clear the level-2 flag once all tags are warm, in case partial # wake support (e.g. tags=["kv_cache"] only) is added in the future. if not getattr(self, "_sleeping_tags", None): diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index c899b2697..281cbc55c 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -10,19 +10,26 @@ from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +# Default for ``visible_baseline``: read the live device-control env var. +# Explicit ``None`` means the caller captured an unset env and must not fall +# back to a later thread's temporary CUDA_VISIBLE_DEVICES. +_LIVE_DEVICE_ENV = object() + def resolve_stage_physical_devices( stage_id: int, devices: str | int | None, *, - visible_baseline: str | None = None, + visible_baseline: str | None | object = _LIVE_DEVICE_ENV, ) -> str | None: """Map logical stage devices to physical IDs without mutating process env. - When ``visible_baseline`` is provided it is used for logical-to-physical - mapping instead of the current device-control env var. This keeps parallel - stage initialization correct even if another thread has temporarily - narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. + ``visible_baseline`` is the device-control env captured at the start of + replica init. Pass the captured string, or ``None`` when that env was + unset. Omitting the argument reads the live env (single-threaded callers). + + Parallel stage init must pass the captured value: another thread may have + already narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. """ if devices in (None, "cpu"): return None @@ -32,7 +39,7 @@ def resolve_stage_physical_devices( device_list = _parse_device_list(devices) env_var = current_omni_platform.device_control_env_var - visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var) + visible_devices = os.environ.get(env_var) if visible_baseline is _LIVE_DEVICE_ENV else visible_baseline if visible_devices is not None: visible_device_list = _parse_device_list(visible_devices) device_list = _map_device_list(stage_id, device_list, visible_device_list) @@ -43,7 +50,7 @@ def set_stage_devices( stage_id: int, devices: str | int | None, *, - visible_baseline: str | None = None, + visible_baseline: str | None | object = _LIVE_DEVICE_ENV, ) -> str | None: """Configure per-stage device visibility and current device (CUDA or NPU). diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index c6ad6ca28..13844518e 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from collections.abc import Iterable from dataclasses import fields as dataclass_fields from typing import Any @@ -432,6 +433,116 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): return self._native_text_metrics_by_request.pop(request_id, {}) def abort_requests(self, request_ids, internal: bool) -> list[str]: + aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal) + return aborted_ids + + def abort_requests_collecting_outputs( + self, + request_ids: Iterable[str], + *, + internal: bool, + commit_state: bool = True, + ) -> tuple[list[str], list[RequestOutput | PoolingRequestOutput]]: + """Abort requests and return terminal abort outputs with partial tokens. + + Mirrors upstream ``OutputProcessor.abort_requests``, but also returns + abort ``RequestOutput`` objects when ``queue`` is ``None`` (Omni stage + pools register OP state without a collector queue). + + Terminal abort outputs use ``RequestOutputKind.CUMULATIVE`` so DELTA + streaming requests still surface the full prefix generated so far. + ``new_token_ids`` is the detokenizer prefix: Omni's + ``make_request_output`` writes that list onto ``CompletionOutput.token_ids``. + + When ``commit_state`` is False, processor mappings stay in place so a + failed physical EngineCore abort can retry with the same prefix. + """ + request_ids = list(request_ids) + if commit_state: + for request_id in request_ids: + if internal: + req_state = self.request_states.get(request_id) + if req_state is not None: + self._native_text_metrics_by_request.pop(req_state.external_req_id, None) + else: + self._native_text_metrics_by_request.pop(request_id, None) + + internal_req_ids: list[str] = [] + for request_id in request_ids: + if internal: + internal_req_ids.append(request_id) + if commit_state: + if req_state := self.request_states.get(request_id): + external_req_id = req_state.external_req_id + internal_ids = self.external_req_ids.get(external_req_id) + if internal_ids is not None: + try: + internal_ids.remove(request_id) + except ValueError: + pass + if not internal_ids: + self.external_req_ids.pop(external_req_id, None) + elif commit_state: + if internal_ids := self.external_req_ids.pop(request_id, []): + internal_req_ids.extend(internal_ids) + else: + internal_req_ids.extend(self.external_req_ids.get(request_id, [])) + + request_ids_to_abort: list[str] = [] + abort_outputs: list[RequestOutput | PoolingRequestOutput] = [] + for request_id in internal_req_ids: + req_state = ( + self.request_states.pop(request_id, None) if commit_state else self.request_states.get(request_id) + ) + if req_state is not None: + if commit_state: + self.lora_states.request_finished(request_id, req_state.lora_name) + request_ids_to_abort.append(request_id) + original_kind = req_state.output_kind + req_state.output_kind = RequestOutputKind.CUMULATIVE + detok = req_state.detokenizer + new_token_ids = list(detok.output_token_ids) if detok is not None else [] + saved_children = None + if not commit_state and req_state.parent_req is not None: + saved_children = set(req_state.parent_req.child_requests) + try: + request_output = req_state.make_request_output( + new_token_ids=new_token_ids, + pooling_output=None, + finish_reason=FinishReason.ABORT, + stop_reason=None, + kv_transfer_params=None, + ) + finally: + req_state.output_kind = original_kind + if saved_children is not None and req_state.parent_req is not None: + req_state.parent_req.child_requests = saved_children + if request_output is not None: + abort_outputs.append(request_output) + if commit_state and req_state.queue is not None: + req_state.queue.put(request_output) + if commit_state: + parent_req = req_state.parent_req + if parent_req is not None: + parent_req.child_requests.discard(request_id) + if not parent_req.child_requests: + self.parent_requests.pop(parent_req.request_id, None) + elif parent := self.parent_requests.get(request_id): + if parent.child_requests: + child_reqs = list(parent.child_requests) + child_aborted, child_outputs = self.abort_requests_collecting_outputs( + child_reqs, + internal=True, + commit_state=commit_state, + ) + request_ids_to_abort.extend(child_aborted) + abort_outputs.extend(child_outputs) + if commit_state: + self.parent_requests.pop(request_id, None) + return request_ids_to_abort, abort_outputs + + def commit_aborted_request_state(self, request_ids: Iterable[str], *, internal: bool) -> None: + """Drop processor mappings after a successful physical EngineCore abort.""" request_ids = list(request_ids) for request_id in request_ids: if internal: @@ -440,7 +551,37 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): self._native_text_metrics_by_request.pop(req_state.external_req_id, None) else: self._native_text_metrics_by_request.pop(request_id, None) - return super().abort_requests(request_ids, internal) + + internal_req_ids: list[str] = [] + for request_id in request_ids: + if internal: + internal_req_ids.append(request_id) + if req_state := self.request_states.get(request_id): + external_req_id = req_state.external_req_id + internal_ids = self.external_req_ids.get(external_req_id) + if internal_ids is not None: + try: + internal_ids.remove(request_id) + except ValueError: + pass + if not internal_ids: + self.external_req_ids.pop(external_req_id, None) + elif internal_ids := self.external_req_ids.pop(request_id, []): + internal_req_ids.extend(internal_ids) + + for request_id in internal_req_ids: + req_state = self.request_states.pop(request_id, None) + if req_state is not None: + self.lora_states.request_finished(request_id, req_state.lora_name) + parent_req = req_state.parent_req + if parent_req is not None: + parent_req.child_requests.discard(request_id) + if not parent_req.child_requests: + self.parent_requests.pop(parent_req.request_id, None) + elif parent := self.parent_requests.get(request_id): + if parent.child_requests: + self.commit_aborted_request_state(list(parent.child_requests), internal=True) + self.parent_requests.pop(request_id, None) def add_request( self, diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 320ac2cae..d616347d8 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -201,11 +201,19 @@ class OmniGPUWorkerBase(GPUWorker): """ from vllm.device_allocator.cumem import CuMemAllocator + # Drain in-flight kernels before CuMem unmap. Abort/decode can still + # have GPU work after the scheduler reports idle; unmapping then + # returns CUDA_ERROR_INVALID_VALUE (cudaErrorInvalidValue). + current_omni_platform.synchronize() + gc.collect() + mem_before = current_omni_platform.get_current_memory_usage(self.device) offload_tags = ("weights",) if level == 1 else tuple() allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=offload_tags) - current_omni_platform.empty_cache() + # allocator.sleep already gc.collects and empty_cache()s. A second + # empty_cache on the pluggable allocator is a known CUDA invalid + # argument (pytorch/pytorch#145168). Only sync for the readout. current_omni_platform.synchronize() mem_after = current_omni_platform.get_current_memory_usage(self.device) freed = max(0, mem_before - mem_after)