diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 104d6df26..4162e60e8 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -124,6 +124,30 @@ steps: - 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 · Qwen3-Omni Async-Chunk SHM Leak" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py + - tests/helpers/runtime.py + - tests/helpers/stage_config.py + - vllm_omni/core/sched/omni_ar_scheduler.py + - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py + - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py + - vllm_omni/engine/orchestrator.py + - vllm_omni/engine/stage_engine_core_proc.py + - vllm_omni/engine/stage_pool.py + - vllm_omni/entrypoints/openai/serving_chat.py + - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py + - 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 + timeout_in_minutes: 40 + commands: + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model' + mirror_hardwares: h100_2 + - label: "Omni · MiniCPM-o 4.5 Test" source_file_dependencies: omni_minicpmo_4_5_function timeout_in_minutes: 50 diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py index cffe3f289..86102c451 100644 --- a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py +++ b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py @@ -1,10 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest from vllm import SamplingParams from vllm.v1.engine import EngineCoreRequest +from vllm.v1.request import RequestStatus +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata from vllm_omni.engine.serialization import deserialize_additional_information @@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest: ) -def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: +def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]: tagged = apply_omni_final_stage_metadata( _make_engine_request(), final_stage_id=0, @@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: request_id="req", additional_information=tagged.additional_information, ) - result = scheduler._request_omits_kv_transfer_to_next_stage(request) + omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request) + omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request) metadata = deserialize_additional_information(tagged.additional_information) - return result, metadata + return omits_kv, omits_chunk, metadata def test_stage_zero_request_omits_kv_transfer(): - omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False) + omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False) - assert omits_transfer + assert omits_kv + assert omits_chunk assert "omni_force_kv_transfer" not in metadata def test_cfg_companion_forces_kv_transfer_without_downstream_payload(): - omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True) + omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True) - assert not omits_transfer + assert not omits_kv + assert omits_chunk assert metadata["omni_final_stage_id"] == 0 assert metadata["omni_force_kv_transfer"] is True + + +class _ChunkRequest(SimpleNamespace): + def __hash__(self): + return hash(self.request_id) + + def __eq__(self, other): + return isinstance(other, _ChunkRequest) and other.request_id == self.request_id + + +def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False): + tagged = apply_omni_final_stage_metadata( + _make_engine_request(), + final_stage_id=final_stage_id, + force_kv_transfer=force_kv_transfer, + ) + request = _ChunkRequest( + request_id="req-shm-leak", + additional_information=tagged.additional_information, + num_in_flight_tokens=1, + num_stale_output_tokens=0, + sampling_params=SimpleNamespace(num_logprobs=None), + has_encoder_inputs=False, + pooling_params=None, + status=RequestStatus.RUNNING, + resumable=False, + num_output_placeholders=0, + spec_token_ids=[], + _output_token_ids=[], + num_computed_tokens=1, + stop_reason=None, + num_nans_in_logits=None, + client_index=0, + trace_headers=None, + ) + request.is_finished = lambda: False + request.get_finished_reason = lambda: None + request.take_prefill_stats = lambda: None + request.take_events = lambda: None + return request + + +def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None): + adapter = mocker.MagicMock() + adapter._confirmed_num_computed_tokens.return_value = 1 + + sched = mocker.MagicMock() + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.defer_block_free = False + sched.structured_output_manager.should_advance.return_value = False + sched._update_request_with_output.return_value = ([42], True) + sched._process_kv_transfer_trigger.return_value = False + sched._handle_stopped_request.return_value = True + + def _free_request(req, delay_free_blocks=False): + sched._omits_kv_transfer_cache.pop(req.request_id, None) + return None, None + + sched._free_request.side_effect = _free_request + sched.chunk_transfer_adapter = adapter + sched.running = [request] + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) + + scheduler_output = mocker.MagicMock() + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.total_num_scheduled_tokens = 1 + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = mocker.MagicMock() + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + model_runner_output.inter_stage_outputs = inter_stage_outputs + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + return adapter + + +@pytest.mark.parametrize( + ("final_stage_id", "force_kv_transfer", "expect_save"), + [ + pytest.param(0, False, False, id="stage-zero-final"), + pytest.param(1, False, True, id="downstream-final"), + ], +) +def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save): + request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer) + adapter = _run_finished_save_step(mocker, request) + + assert adapter.save_async.called is expect_save + + +def test_stage_zero_final_skips_inter_stage_output(mocker): + warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning") + request = _make_chunk_request(0) + adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}]) + + adapter.save_async.assert_not_called() + assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list) + + +def test_streaming_session_update_invalidates_omits_kv_transfer_cache(): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + scheduler.skipped_waiting = [] + scheduler.log_stats = False + tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0) + session = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + status=RequestStatus.WAITING, + ) + assert scheduler._request_omits_chunk_transfer_to_next_stage(session) + assert "req" in scheduler._omits_kv_transfer_cache + update = SimpleNamespace( + additional_information={"omni_final_stage_id": 1}, + model_intermediate_buffer=None, + arrival_time=0.0, + sampling_params=None, + ) + + OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update) + + assert scheduler._omits_kv_transfer_cache == {} + assert session.additional_information == {"omni_final_stage_id": 1} + assert not scheduler._request_omits_chunk_transfer_to_next_stage(session) + assert not scheduler._request_omits_kv_transfer_to_next_stage(session) + + +def test_omits_cache_misses_when_payload_object_is_replaced(): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0) + request = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + ) + assert scheduler._request_omits_chunk_transfer_to_next_stage(request) + request.additional_information = {"omni_final_stage_id": 1} + + assert not scheduler._request_omits_chunk_transfer_to_next_stage(request) diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index b4eb839a2..e5bd8dd4e 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers # isort: on @@ -103,6 +104,7 @@ def _make_drain_sched(session: Request) -> MagicMock: sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) return sched diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 94d796fd6..bdfbadba1 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers # isort: on @@ -175,6 +176,7 @@ def _run_resumable_segment_stop( sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} @@ -312,6 +314,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} @@ -400,6 +403,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None: sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index 008e52636..2ecd2772a 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from types import SimpleNamespace @@ -76,7 +79,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None: scheduler = OmniARScheduler.__new__(OmniARScheduler) scheduler.input_coordinator = coordinator scheduler.chunk_transfer_adapter = None - scheduler._omits_kv_transfer_cache = {"req-free": True} + scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)} scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None) scheduler.finished_req_ids = set() scheduler.finished_req_ids_dict = None diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 8f85366e3..438414854 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -15,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats from vllm.v1.request import Request, RequestStatus +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct @@ -95,6 +96,30 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc assert not adapter.code_prompt_token_ids +def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): + adapter, connector = shm_sender + ext_id = f"release_{uuid.uuid4().hex}" + request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id) + sibling_key = f"{ext_id}_1_0_0" + assert connector.put("0", "1", sibling_key, "sibling")[0] + for _ in range(3): + adapter.save_async(None, request) + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + # The consumer read only the first chunk before its own request ended. + assert connector.get("0", "1", f"{ext_id}_0_0") is not None + + adapter.release_shm_resources(ext_id) + # The release is queued behind sends; nothing is unlinked on the caller. + assert f"{ext_id}_0_1" in connector._pending_keys + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert adapter._reclaimed_shm_total == 2 + assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys) + assert connector.get("0", "1", f"{ext_id}_0_1") is None + assert connector.get("0", "1", sibling_key)[0] == "sibling" + + @pytest.mark.parametrize("fail_terminal", [False, True]) def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal): adapter, connector = shm_sender @@ -2620,10 +2645,9 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw)) adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw)) - from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler - scheduler = mocker.MagicMock() scheduler.chunk_transfer_adapter = adapter_mock + bind_omits_transfer_helpers(scheduler) scheduler.connector = None scheduler.perf_metrics = None scheduler.log_stats = False diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 39a0e8327..d0efd06f5 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -213,7 +213,7 @@ class TestCleanup: connector.put("s0", "s1", "cleanup_req_42", data) assert "cleanup_req_42" in connector._pending_keys - connector.cleanup("cleanup_req_42") + assert connector.cleanup("cleanup_req_42") is True assert "cleanup_req_42" not in connector._pending_keys result = connector.get("s0", "s1", "cleanup_req_42", metadata=None) @@ -224,7 +224,7 @@ class TestCleanup: connector.put("s0", "s1", "consumed_req_99", data) connector.get("s0", "s1", "consumed_req_99", metadata=None) - connector.cleanup("consumed_req_99") + assert connector.cleanup("consumed_req_99") is False assert "consumed_req_99" not in connector._pending_keys def test_close_cleans_all_pending(self, connector): @@ -296,6 +296,24 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch): assert key not in connector._pending_keys +@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"]) +def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix): + request_id = f"prefix_{uuid.uuid4().hex}" + own_keys = [f"{request_id}_0_{i}" for i in range(3)] + sibling_key = f"{request_id}{suffix}_0_0" + for key in own_keys: + assert connector.put("0", "1", key, key)[0] + assert connector.put("0", "1", sibling_key, "sibling")[0] + # A consumed chunk is no longer counted as reclaimed. + assert connector.get("0", "1", own_keys[0])[0] == own_keys[0] + + assert connector.cleanup_prefix(f"{request_id}_0_") == 2 + for key in own_keys: + assert key not in connector._pending_keys + assert connector.get("0", "1", key) is None + assert connector.get("0", "1", sibling_key)[0] == "sibling" + + # ── Arrival wakeups ─────────────────────────────────────────────────── diff --git a/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py b/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py new file mode 100644 index 000000000..2afd8bfd8 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final). + +The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with +``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the +thinker→talker SharedMemoryConnector path. + +A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream +ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still +``put()`` a finished marker on finish and nothing unlinked it, so every such +request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``. +Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any +undrained segment on request cleanup. + +This module starts the production async-chunk deploy and asserts the three +text-only request ids leave no connector lockfile or POSIX segment in +``/dev/shm``. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1" +_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct") +_SHM_DIR = Path("/dev/shm") +_NUM_TEXT_REQUESTS = 3 +_SETTLE_S = 15.0 + +# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0 +# still uses SharedMemoryConnector. Production yaml is already async_chunk:true. +_ASYNC_CHUNK_DEPLOY = modify_stage_config( + get_deploy_config_path("ci/qwen3_omni_moe.yaml"), + updates={"async_chunk": True}, +) + +test_params = [ + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=_ASYNC_CHUNK_DEPLOY, + use_stage_cli=True, + ), + id="async_chunk", + ) +] + + +def _system_prompt() -> dict: + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } + + +def _shm_used_mib() -> str: + try: + st = os.statvfs(_SHM_DIR) + used = (st.f_blocks - st.f_bavail) * st.f_frsize + return f"{used / (1024 * 1024):.1f}MiB" + except OSError: + return "unknown" + + +def _is_entry_for_request(name: str, request_id: str) -> bool: + """Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id.""" + return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_")) + + +def _entries_for_request_ids(request_ids: list[str]) -> list[str]: + if not _SHM_DIR.is_dir(): + pytest.skip("/dev/shm is not available") + found: list[str] = [] + with os.scandir(_SHM_DIR) as entries: + for entry in entries: + if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids): + found.append(entry.name) + return found + + +def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]: + """Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*.""" + deadline = time.monotonic() + timeout_s + time.sleep(1.0) + leftover: list[str] = _entries_for_request_ids(request_ids) + while leftover: + if time.monotonic() >= deadline: + return leftover + time.sleep(2.0) + leftover = _entries_for_request_ids(request_ids) + return leftover + + +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None: + """Stage-0-final (text-only) requests must not leave connector SHM behind.""" + print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True) + + messages = dummy_messages_from_mix_data( + system_prompt=_system_prompt(), + content_text="What is the capital of China? Answer in 20 words.", + ) + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + "key_words": {"text": ["beijing"]}, + } + responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS) + produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)] + assert len(produced) == _NUM_TEXT_REQUESTS, ( + f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; " + f"got {len(produced)} non-empty texts. A preprocess/engine miss would " + f"leave /dev/shm unchanged and hide a leak." + ) + request_ids = [resp.request_id for resp in responses] + assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, ( + f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}" + ) + + leftover = _wait_no_entries_for_request_ids(request_ids) + print( + f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: " + f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}", + flush=True, + ) + assert not leftover, ( + f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}" + ) diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index 6b907c3d2..7fa04d191 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Deterministic tests for the CFG-companion output lifecycle. Regression suite for the companion-output race behind the nightly Bagel @@ -94,6 +94,9 @@ class _FakePool: def release_bindings(self, request_ids): pass + async def release_request_resources(self, request_ids): + pass + def _make_orchestrator(num_stages: int = 2) -> Orchestrator: orch = Orchestrator.__new__(Orchestrator) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index 281b4b504..d15df2bdd 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -154,6 +154,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails(): asyncio.run(run()) +@pytest.mark.cpu +def test_release_request_resources_skips_without_async_chunk(): + async def run() -> None: + call = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=call)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)), + ) + + await pool.release_request_resources(["req-1"]) + call.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch): + async def run() -> None: + async def hang(*_args): + await asyncio.Event().wait() + + live = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)), + ) + monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2) + + release = asyncio.create_task(pool.release_request_resources(["req-1"])) + await asyncio.sleep(0.05) + live.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) + assert not release.done() + + await asyncio.wait_for(release, timeout=1.0) + + asyncio.run(run()) + + @pytest.mark.cpu @pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"]) @pytest.mark.parametrize("failure", ["error", "timeout", "missing"]) diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 8637e510d..883659565 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -97,6 +97,7 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[ class OmniResponse: """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``).""" + request_id: str | None = None text_content: str | None = None audio_data: list[str] | None = None audio_content: str | None = None @@ -421,6 +422,8 @@ class OnlineOmniClient: text_content = "" audio_data = [] for chunk in chat_completion: + if result.request_id is None: + result.request_id = getattr(chunk, "id", None) for choice in chunk.choices: content = getattr(getattr(choice, "delta", None), "content", None) modality = getattr(chunk, "modality", None) @@ -453,6 +456,7 @@ class OnlineOmniClient: """Wall clock from *before* ``chat.completions.create`` through response parse + local decode.""" result = OmniResponse() try: + result.request_id = getattr(chat_completion, "id", None) audio_data = None text_content = None for choice in chat_completion.choices: diff --git a/tests/helpers/omni_scheduler.py b/tests/helpers/omni_scheduler.py new file mode 100644 index 000000000..e5a4c58a5 --- /dev/null +++ b/tests/helpers/omni_scheduler.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Scheduler test helpers shared by AR / chunk-transfer unit tests.""" + +from __future__ import annotations + +from types import MethodType + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + +def bind_omits_transfer_helpers(sched) -> None: + """Bind the real omit helpers so a MagicMock scheduler does not skip save_async. + + MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage`` + silently skips the chunk put. + """ + sched._omits_kv_transfer_cache = {} + sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched) + sched._request_omits_kv_transfer_to_next_stage = MethodType( + OmniARScheduler._request_omits_kv_transfer_to_next_stage, + sched, + ) + sched._request_omits_chunk_transfer_to_next_stage = MethodType( + OmniARScheduler._request_omits_chunk_transfer_to_next_stage, + sched, + ) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7dc0a5d7c..9c6f2e6d2 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -144,8 +144,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Track requests that have already triggered prefill transfer to avoid duplicates self.transfer_triggered_requests: set[str] = set() - # Cache per-request flag to avoid repeated deserialization of additional_information - self._omits_kv_transfer_cache: dict[str, bool] = {} + # Cache per-request flags to avoid repeated deserialization of + # additional_information. Value is (payload_id, is_stage_zero_final, + # force_kv_transfer); payload_id invalidates the entry when the + # object is replaced. + self._omits_kv_transfer_cache: dict[str, tuple[int, bool, bool]] = {} # KV-wait start ts for the vllm_omni:kv_wait_s metric; see # _emit_kv_wait_output for the engine-core → orchestrator carry. @@ -264,26 +267,40 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return config.get(key, default) return getattr(config, key, default) if config is not None else default - def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: - """True when this stage-zero-final request does not need downstream KV. - - The result is cached per request to avoid repeated deserialization of - additional_information on every scheduler tick. - """ - rid = request.request_id - cached = self._omits_kv_transfer_cache.get(rid) - if cached is not None: - return cached - + def _omni_final_stage_flags(self, request: Request) -> tuple[bool, bool]: + """Return ``(is_stage_zero_final, force_kv_transfer)`` from request metadata.""" payload = getattr(request, "additional_information", None) if payload is None: - result = False - else: - info = deserialize_additional_information(payload) - result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False)) + return False, False - self._omits_kv_transfer_cache[rid] = result - return result + cache = getattr(self, "_omits_kv_transfer_cache", None) + rid = request.request_id + payload_id = id(payload) + if cache is not None: + cached = cache.get(rid) + if cached is not None and cached[0] == payload_id: + return cached[1], cached[2] + + info = deserialize_additional_information(payload) + is_final = info.get("omni_final_stage_id") == 0 + force_kv = bool(info.get("omni_force_kv_transfer", False)) + if cache is not None: + cache[rid] = (payload_id, is_final, force_kv) + return is_final, force_kv + + def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: + """True when this stage-zero-final request does not need downstream KV.""" + is_final, force_kv = self._omni_final_stage_flags(request) + return is_final and not force_kv + + def _request_omits_chunk_transfer_to_next_stage(self, request: Request) -> bool: + """True when this request has no downstream chunk consumer. + + CFG companions still force KV transfer but are stage-0-final for + ordinary inter-stage chunks, so they omit ``save_async`` here. + """ + is_final, _ = self._omni_final_stage_flags(request) + return is_final def _should_defer_waiting_admission(self) -> bool: return False @@ -669,6 +686,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): confirmed_num_computed_tokens = None boundary_generation = None + # Only when this step might save. Read additional_information + # before _free_request rewrites it. + omits_chunk_transfer = False + if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped): + omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request) # Capture before resumable stop handling can clear token history. output_token_ids: Any = getattr(request, "output_token_ids", None) if output_token_ids is None: @@ -787,8 +809,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors - if self.chunk_transfer_adapter is not None and ( - inter_stage_output is not None or is_segment_finished or finished + if omits_chunk_transfer and inter_stage_output is not None: + logger.warning( + "Skipping inter-stage chunk for request %s: " + "omni_final_stage_id=0 but inter_stage_output is present", + req_id, + ) + if ( + self.chunk_transfer_adapter is not None + and not omits_chunk_transfer + and (inter_stage_output is not None or is_segment_finished or finished) ): save_kwargs = { "new_token_ids": new_token_ids, diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 77a568d14..6f9f03d6c 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -229,6 +229,9 @@ class OmniSchedulerMixin(_SchedulerMixinBase): def _finish_streaming_session_update(self, session: Request, update: StreamingUpdate) -> None: """Install payload metadata and send an updated session to admission.""" + cache = getattr(self, "_omits_kv_transfer_cache", None) + if cache is not None: + cache.pop(session.request_id, None) session.additional_information = update.additional_information or None session.model_intermediate_buffer = getattr( update, diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 09854d22d..d0f2cdb1a 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -307,9 +307,13 @@ class SharedMemoryConnector(OmniConnectorBase): self._metrics["gets"] += 1 return result - def cleanup(self, request_id: str) -> None: - """Unlink the exact key passed to ``put()``, never a request-id prefix.""" + def cleanup(self, request_id: str) -> bool: + """Unlink the exact key passed to ``put()``, never a request-id prefix. + + Returns True when an unconsumed segment was actually unlinked. + """ key = request_id + unlinked = False with self._pending_keys_lock: self._pending_keys.pop(key, None) try: @@ -317,6 +321,7 @@ class SharedMemoryConnector(OmniConnectorBase): seg.close() seg.unlink() logger.debug("cleanup: unlinked unconsumed SHM segment %s", key) + unlinked = True except FileNotFoundError: pass except Exception as e: @@ -327,6 +332,18 @@ class SharedMemoryConnector(OmniConnectorBase): os.remove(lock_file) except OSError: pass + return unlinked + + def cleanup_prefix(self, key_prefix: str) -> int: + """Unlink every tracked key of the form ``{key_prefix}{chunk_id}``. + + Only keys still in ``_pending_keys`` are considered, and the suffix + must be a bare integer chunk id, so another request whose id merely + shares the prefix is never matched. Returns the unlinked count. + """ + with self._pending_keys_lock: + keys = [k for k in self._pending_keys if k.startswith(key_prefix) and k[len(key_prefix) :].isdigit()] + return sum(self.cleanup(key) for key in keys) def close(self) -> None: """Unlink all remaining tracked SHM segments.""" diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index bd6d7abf1..5e4ef6820 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -28,6 +28,8 @@ from .base import OmniTransferAdapterBase logger = get_connector_logger(__name__) +_RECLAIM_WARN_INTERVAL = 1000 + class _SenderGeneration: """Fence one external request generation without blocking cleanup.""" @@ -196,6 +198,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.request_payload = {} self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list) self.request_ids_mapping: dict[str, str] = {} + # Save-thread only: running count of segments reclaimed after finish. + self._reclaimed_shm_total = 0 self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() @@ -688,6 +692,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): for chunk_id in range(num_puts): self.connector.cleanup(f"{key_prefix}_{chunk_id}") return + if "release_shm" in task: + self._release_shm_prefix(task["release_shm"]) + return request = task["request"] external_req_id = request.external_req_id sender_token = task.get("sender_token") @@ -979,8 +986,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Called after a terminal chunk is sent or when the scheduler aborts the request before a terminal chunk can be produced. In-flight sends are - cancelled here and reclaim their own state in ``finally``; cleanup - never waits for connector I/O on the scheduler thread. + cancelled here and reclaim adapter state plus leftover SHM in + ``finally`` after ``put()`` returns; cleanup never waits for connector + I/O on the scheduler thread. Idempotent: calling with an already-cleaned or unknown id is safe. """ @@ -1042,6 +1050,42 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if cached_ic is not None: cached_ic.pop(external_req_id, None) + def release_shm_resources(self, request_id: str) -> None: + """Queue reclaim of inter-stage segments this request left unconsumed. + + Safe only once every stage has finished with the request: the producer + keeps writing until its own request ends, and a consumer that stopped + early never drains the remainder. Successful terminal chunks are kept + by ``_queue_shm_cleanup_locked`` for exactly that consumer, so the + orchestrator -- the only party that knows all stages are done -- drives + this (see ``Orchestrator._cleanup_request_ids``). + + Like abort cleanup, the unlink runs on the save thread, so the + scheduler never performs SHM I/O. + """ + if not isinstance(self.connector, SharedMemoryConnector): + return + external_req_id = self.request_ids_mapping.get(request_id, request_id) + self._pending_save_reqs.append({"release_shm": f"{external_req_id}_{self.connector.stage_id}_"}) + with self._save_cond: + self._save_cond.notify() + + def _release_shm_prefix(self, key_prefix: str) -> None: + reclaimed = self.connector.cleanup_prefix(key_prefix) + if not reclaimed: + return + # Non-zero means a consumer stopped before draining the producer. + # Common on audio requests, so warn only when the running total + # crosses another _RECLAIM_WARN_INTERVAL boundary. + prev = self._reclaimed_shm_total + self._reclaimed_shm_total = prev + reclaimed + if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_shm_total // _RECLAIM_WARN_INTERVAL: + logger.warning( + "Reclaimed %d unconsumed inter-stage segments so far; " + "a downstream stage finished before draining its producer", + self._reclaimed_shm_total, + ) + def cleanup( self, request_id: str, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8f55280cd..fad9a8f52 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -301,6 +301,7 @@ class OrchestratorBase: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None + _transfer_release_tasks: set[asyncio.Task] = set() def __init__( self, @@ -348,6 +349,8 @@ class OrchestratorBase: self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} + # Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC. + self._transfer_release_tasks: set[asyncio.Task] = set() self._init_metrics_state( stage_pools, running_counter, @@ -690,6 +693,36 @@ class OrchestratorBase: output_msg.finished = index == last_index_by_req[output_msg.request_id] return abort_outputs + def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: + """Drop each stage's inter-stage transfer resources for finished requests. + + This is the only point that knows every stage is done with the request, + so it is the only safe place to reclaim segments a consumer never + drained. Scheduled rather than awaited: reclaim is best-effort and must + not add RPC latency to request teardown. + """ + if not request_ids: + return + + async def _run() -> None: + results = await asyncio.gather( + *(pool.release_request_resources(request_ids) for pool in self.stage_pools), + return_exceptions=True, + ) + for result in results: + if isinstance(result, Exception): + logger.warning("[Orchestrator] release transfer resources failed: %s", result) + + try: + task = asyncio.get_running_loop().create_task(_run()) + self._transfer_release_tasks.add(task) + task.add_done_callback(self._transfer_release_tasks.discard) + except RuntimeError: + logger.warning( + "[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s", + request_ids, + ) + def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" for pool in self.stage_pools: @@ -1645,6 +1678,7 @@ class OrchestratorBase: if abort: abort_outputs = await self._abort_request_ids(cleanup_ids) self._release_request_bindings(cleanup_ids) + self._release_stage_transfer_resources(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index c19a5cf75..f9fcd34b9 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -131,6 +131,21 @@ class StageEngineCoreProc(EngineCoreProc): if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler): logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.") + def omni_release_request_resources(self, request_ids: list[str]) -> None: + """Release this stage's inter-stage transfer resources for *request_ids*. + + Invoked over the UTILITY channel by the orchestrator once every stage + has finished with the request. Idempotent and safe for unknown ids. + """ + adapter = getattr(getattr(self, "scheduler", None), "chunk_transfer_adapter", None) + if adapter is None: + return + for request_id in request_ids or (): + try: + adapter.release_shm_resources(request_id) + except Exception as e: + logger.debug("omni_release_request_resources(%s) failed: %s", request_id, e) + def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" scheduler_request, current_wave = super().preprocess_add_request(request) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 222cd3b5d..176eb9dc0 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -99,6 +99,8 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 + # A replica that dies mid-release never answers; without a bound its background release never finishes. + RELEASE_RPC_TIMEOUT_S: float = 5.0 # Only these EngineCore helpers may skip collective_rpc_async. A generic # ``{method}_async`` on AsyncMPClient must not silently drop timeout. _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"}) @@ -135,6 +137,9 @@ class StagePool: self.clients: list[StagePoolClient | None] = list(normalized_clients) self._output_processor = output_processor self._stage_vllm_config = stage_vllm_config + self._has_chunk_transfer_adapter = bool( + getattr(getattr(stage_vllm_config, "model_config", None), "async_chunk", False) + ) self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} self._unavailable_replicas: set[int] = set() @@ -1304,6 +1309,28 @@ class StagePool: return abort_outputs + async def release_request_resources(self, request_ids: list[str]) -> None: + """Ask every live replica to drop transfer resources for *request_ids*. + + Broadcast rather than binding-routed: the orchestrator releases route + bindings as part of the same teardown, so a binding lookup here would + race it. The engine-core handler is idempotent for unknown ids. + """ + if not request_ids or not self._has_chunk_transfer_adapter: + return + ids = list(request_ids) + + async def release(replica_id: int, call: Any) -> None: + try: + await asyncio.wait_for(call("omni_release_request_resources", ids), timeout=self.RELEASE_RPC_TIMEOUT_S) + except Exception as e: + logger.warning( + "[StagePool-%s] release_request_resources on replica %s failed: %r", self.stage_id, replica_id, e + ) + + calls = [(i, getattr(self.clients[i], "call_utility_async", None)) for i in self.live_replica_ids()] + await asyncio.gather(*(release(i, call) for i, call in calls if call is not None)) + async def collective_rpc( self, replica_id: int,