From 5a386d789f7646f42e7c97fe873848de9dd400aa Mon Sep 17 00:00:00 2001 From: summer <128961079+zhang-keliang@users.noreply.github.com> Date: Fri, 11 Sep 2026 08:04:08 +0800 Subject: [PATCH] [Bugfix] Fix host-memory leak on aborted /v1/images/generations (#6462) (#6561) Signed-off-by: summer <128961079+zhang-keliang@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_async_omni.py | 89 +++++++++++++++++++ tests/metrics/test_emit_calls.py | 3 +- tests/metrics/test_prometheus.py | 2 - vllm_omni/entrypoints/client_request_state.py | 6 ++ vllm_omni/entrypoints/omni_base.py | 18 +--- 5 files changed, 100 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 2549889b3..a69bcce09 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -329,6 +329,95 @@ def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): asyncio.run(run()) +@pytest.mark.cpu +def test_abort_drops_consumed_metric_message_ids_with_state(): + """Exercise the production leak path from #6462: a metric-bearing + ``OutputMessage`` goes through ``_handle_output_message`` (which populates + the per-request de-dup set), the request is aborted (which, since #6367, + keeps ``request_states`` registered and enqueues a terminal abort output), + and the set must be released with the state by ``generate()``'s normal + cleanup — with no independent request-keyed metric state surviving on the + instance at any point. A reintroduced + class-level ``_consumed_metric_messages``-style map populated by the + production handler fails the sweep below. + """ + import time as _time + from collections.abc import Mapping + + from vllm_omni.engine.messages import OutputMessage + from vllm_omni.entrypoints.client_request_state import ClientRequestState + from vllm_omni.metrics.stats import OrchestratorAggregator + + async def run(): + omni = get_async_omni_instance() + omni.engine.get_stage_metadata = lambda stage_id: SimpleNamespace( + stage_type="llm", + final_output=True, + final_output_type="text", + ) + + rid = "req-1-cccc" + state = ClientRequestState(request_id=rid, external_request_id="req-1") + state.metrics = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=_time.time(), + final_stage_id_for_e2e=0, + ) + state.input_stream_task = None + omni.request_states[rid] = state + + msg = OutputMessage( + request_id=rid, + stage_id=0, + engine_outputs=SimpleNamespace(final_output_type="text"), + metrics=SimpleNamespace( + num_tokens_in=3, + num_tokens_out=5, + stage_gen_time_ms=1.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + ), + finished=False, + ) + handled, out_rid, out_stage, out_state = omni._handle_output_message(msg) + assert handled is False and out_state is state + # The production handler recorded the de-dup entry on the state... + assert id(msg) in state.consumed_metric_message_ids + # ...and de-duplicates a replay of the same message object. + omni._handle_output_message(msg) + assert len(state.consumed_metric_message_ids) == 1 + + assert not hasattr(omni, "_consumed_metric_messages") + + # #6367 contract: abort() keeps the state registered and enqueues a + # terminal abort output; generate()'s ``finally`` owns the cleanup. + await omni.abort("req-1") + assert rid in omni.request_states + terminal = state.queue.get_nowait() + assert terminal.finished is True + assert terminal.engine_outputs.outputs[0].finish_reason == "abort" + # The de-dup set still lives only on the (still-registered) state. + for name, value in vars(omni).items(): + if name == "request_states": + continue + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed metric state kept outside request_states in {name!r}" + ) + + # generate()'s normal cleanup releases the state and the set with it. + omni._log_summary_and_cleanup(rid) + assert rid not in omni.request_states + for name, value in vars(omni).items(): + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed state survived cleanup in {name!r}" + ) + assert not hasattr(omni, "_consumed_metric_messages") + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index b46f4080f..02a98fc00 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -41,7 +41,6 @@ def _make_omni_base_with_mock_prom(mocker): obj = object.__new__(OmniBase) obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = True return obj, obj.prom_metrics @@ -270,7 +269,7 @@ class TestStageWorkloadMetricScope: obj = object.__new__(OmniBase) obj._enable_ar_profiler = False - obj._consumed_metric_messages = {} + obj.request_states = {"req-replay": SimpleNamespace(consumed_metric_message_ids=set())} obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.mod_metrics = mocker.Mock() obj.engine = SimpleNamespace( diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index e45365cf1..ca8dd8d5e 100644 --- a/tests/metrics/test_prometheus.py +++ b/tests/metrics/test_prometheus.py @@ -154,7 +154,6 @@ class TestRequestLifecycleGauges: obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Simulate request lifecycle: start (counter 0→1, dict {} → {req}), @@ -183,7 +182,6 @@ class TestRequestLifecycleGauges: obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Two in flight, one finalizes — running should report 1, waiting 0. diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 0aa0c7c5a..22f1c6d5e 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -39,3 +39,9 @@ class ClientRequestState: # without re-querying stage_pools. self.audio_emit_stage_id: int | None = None self.audio_emit_replica_id: int | None = None + # De-dup set for metric messages: OmniBase populates this in + # ``_handle_output_message`` / ``_process_single_result`` so the same + # ``id(msg)`` isn't counted twice into per-request metrics. Kept on + # the request state (not a class-level dict) so it is released with + # the state — see #6462 / #6561. + self.consumed_metric_message_ids: set[int] = set() diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 264b37fb5..899ff116d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -231,7 +231,6 @@ class OmniBase(PDDisaggregationMixin): self.async_chunk = bool(getattr(self.engine, "async_chunk", False)) self.request_states: dict[str, ClientRequestState] = {} - self._consumed_metric_messages: dict[str, set[int]] = {} self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list @@ -298,13 +297,6 @@ class OmniBase(PDDisaggregationMixin): """True when a non-empty stage pool has lost all of its replicas.""" return len(pool.clients) > 0 and self._live_replica_count(pool) == 0 - def _consumed_metric_message_ids(self, request_id: str) -> set[int]: - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is None: - consumed_by_request = {} - self._consumed_metric_messages = consumed_by_request - return consumed_by_request.setdefault(request_id, set()) - @property def is_running(self) -> bool: return self.engine.is_alive() @@ -449,9 +441,6 @@ class OmniBase(PDDisaggregationMixin): ) finally: self.request_states.pop(request_id, None) - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is not None: - consumed_by_request.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. @@ -532,7 +521,7 @@ class OmniBase(PDDisaggregationMixin): stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type) msg_id = id(msg) - consumed = self._consumed_metric_message_ids(req_id) + consumed = req_state.consumed_metric_message_ids if msg_id not in consumed: req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type) submit_ts = msg.stage_submit_ts @@ -642,8 +631,9 @@ class OmniBase(PDDisaggregationMixin): output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) if finished and _m is not None: msg_id = id(result) - consumed = self._consumed_metric_message_ids(req_id) - if msg_id not in consumed: + req_state = self.request_states.get(req_id) + consumed = req_state.consumed_metric_message_ids if req_state is not None else None + if consumed is not None and msg_id not in consumed: metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id)