[Bugfix] Fix host-memory leak on aborted /v1/images/generations (#6462)

[zhang-keliang]

Purpose

Fixes #6462 (host-memory leak in api_server on aborted /v1/images/generations).

Every request keeps self._consumed_metric_messages[request_id] — a per-request set[int] populated by _handle_output_message (omni_base.py) to de-dup metrics. The completion path (_log_summary_and_cleanup) pops both request_states[rid] and _consumed_metric_messages[rid]. The abort path — AsyncOmni._abort and the sync mirror Omni.abort — popped only request_states, so every aborted request retained one set[int] entry for the process lifetime.

Reporter's fingerprint matches: abort-only, diffusion-only (audio has its own dict with explicit .pop() sites in serving_speech.py), linear, retained after load stops. Related to #6413 / #6439 which fixed the executor-side abort leak (_completed_outputs); this covers the frontend side.

Fix: _abort (async) and Omni.abort (sync) now pop _consumed_metric_messages[req_id] in lockstep with request_states. Uses getattr(self, "_consumed_metric_messages", None) so existing subclasses / test doubles that don't initialise the dict still work.

Test Plan

Added tests/entrypoints/test_async_omni.py::test_abort_evicts_consumed_metric_messages_entry alongside the existing test_abort_handles_internal_request_mapping. Uses the pre-existing get_async_omni_instance scaffold, seeds both dicts, calls omni._abort([req_id]), asserts both dicts drop the key. Fails on main, passes on this branch.

CPU-only tracemalloc reproduction (200 abort cycles × 40 metrics messages per request):

• Before fix: 71.1 KB retained (364 B/req); _consumed_metric_messages size = 200.
• After fix: 3232 B retained (16 B/req, all tracemalloc bookkeeping); dict empty.
• Completion control (both before and after): ~4 B/req, dict empty.

vLLM Version: N/A for this fix (pure frontend Python; no vLLM version bump)

vLLM-Omni Commit: off 3d035bf (main at the time of the fix)

Test Result

• Regression test passes on this branch (verified in isolation via a standalone driver that extracts _abort source, since the full pytest suite has heavy runtime deps not present in my minimal env).
• ruff check clean; ruff format --check clean (ruff==0.14.10 per .pre-commit-config.yaml).
• No behavior change on the non-abort paths; no public API change.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/entrypoints.md.

Module owners: @alex-jw-brooks @linyueqian @NickCao

@zhang-keliang, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[NickCao]

Could _consumed_metric_messages, instead of being a dict keyed by request_id, be an attr attached to the request object itself? That way it can never leak.

[zhang-keliang]

Adopted, thanks — moved the dedup set to ClientRequestState.consumed_metric_message_ids. Dropping the class-level OmniBase._consumed_metric_messages dict + its setdefault helper + the explicit dict pop in _log_summary_and_cleanup; the set now drops with the state via GC on both completion and abort. Net -14 LoC. New head 2868c44.

[zhang-keliang]

Self-review:

• Scope: 4 files, +35/-16, all Python host-side. No CUDA / kernel / config touched.
• Root cause: _consumed_metric_messages[rid] populated on every stage metric emit but only popped in _log_summary_and_cleanup; abort skipped it → linear leak.
• After @NickCao's suggestion the set now lives on ClientRequestState, so completion and abort share one cleanup site (request_states.pop(rid)); no per-callsite dict maintenance.
• Regression test: test_abort_drops_consumed_metric_message_ids_with_state — populated set on the state → abort → state pops → set gone (also guards against re-introducing a class-level dict via hasattr).
• Ran locally: pytest tests/entrypoints/test_async_omni.py tests/metrics/test_prometheus.py -m cpu → 13 passed. ruff check and ruff format --check on all 6 modified files → clean (only pre-existing UP038 warnings on unrelated lines).
• No GPU / model-weight reproduction on this branch (host-side lifecycle only; leak measured by reporter with tracemalloc in api_server PID 1, and my CPU repro in the linked issue does the same). No perf benchmarks claimed.

[hsliuustc0106]

resolve conflicts

[zhang-keliang]

Rebased onto latest main and resolved the conflict — main had moved prom_metrics initialization earlier in __init__ (post-0.28 rebase), and the new test_same_finished_image_message_is_observed_exactly_once in test_emit_calls.py stubbed the old class-level dict, so its fixture now seeds request_states with a per-request state instead. Branch is conflict-free (head 5abddb2); pytest tests/metrics/test_emit_calls.py tests/entrypoints/test_async_omni.py tests/metrics/test_prometheus.py -m cpu → 54 passed locally.

[hsliuustc0106]

[P2] Exercise the production leak path in this regression

This helper uses object.__new__, manually seeds the new field, and never runs OmniBase.__init__ or a metric-bearing output through _handle_output_message. Reintroducing and populating the old request-keyed map in production could therefore still pass this test. Drive a metric message through the production handler, abort the request, and assert that no independent request-keyed metric state survives.

[zhang-keliang]

Done in d9d2b69 — the regression now drives a metric-bearing OutputMessage through _handle_output_message (asserting the handler records the de-dup entry on the state and de-duplicates a replay of the same message object), aborts via the public abort(), and then sweeps vars(omni) so no mapping on the instance still keys anything by the request id. A reintroduced request-keyed map populated by the production handler fails either the sweep or the hasattr guard.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit c3d8baaf6503 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[hsliuustc0106]

I found one blocking integration issue in the regression test; see the inline comment. No additional high-confidence production-code findings.

[hsliuustc0106]

P1 — Please rebase this regression test onto current main's abort lifecycle. Since #6367, AsyncOmni._abort() keeps request_states registered, queues a terminal abort output, and relies on generate() cleanup to remove the state. Current main therefore conflicts with this test, and if this assertion is carried over it fails immediately after await omni.abort("req-1"). Please preserve the newer abort tests and assert that the per-request set is released after the normal generate()/finally cleanup, or update this test to the terminal-message contract.

[zhang-keliang]

Done — rebased onto current main (973d20f) and moved the test to the #6367 contract. Both new abort tests from main (test_abort_enqueues_prefix_tokens_from_engine, test_abort_enqueues_synthetic_finished_when_engine_returns_empty) are preserved. The regression now: drives a metric-bearing OutputMessage through _handle_output_message → abort() → asserts the state is still registered and the terminal abort output is on its queue (with the de-dup set living only on that state) → runs generate()'s cleanup via _log_summary_and_cleanup → asserts the state and the set are released together, with a vars(omni) sweep for any surviving request-keyed mapping. pytest tests/entrypoints/test_async_omni.py tests/metrics/test_prometheus.py tests/metrics/test_emit_calls.py -m cpu → 58 passed.

[hsliuustc0106]

lgtm