[CI/Build] Cut L3 entrypoint BAGEL cold starts via L4 small-model sleep suites (#6851) (#7359)

Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

.buildkite/common/ci_source_file_dependencies.yml

@@ -188,6 +188,10 @@ source_file_dependencies:
 
   # --- models ---
   # --- omni / tts jobs ---
+  entrypoints_function:
+    - tests/entrypoints/
+    - vllm_omni/entrypoints/
+
   omni_qwen3_omni_function:
     - *qwen3_omni
     - tests/e2e/offline_inference/test_qwen3_omni.py

.buildkite/cuda/test-merge.yml

@@ -60,11 +60,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        # TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per
-        # run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds
-        # 2915 vs 2916), which pushed this suite past 60 min. Revert once the
-        # 0.26 startup regression is understood and fixed.
-        timeout_in_minutes: 90
+        timeout_in_minutes: 60
         commands:
           # Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5
           # (see the P0 thread on #7413); the test stays in the tree and runs again there.
@@ -73,10 +69,25 @@ steps:
 
       - label: "Entrypoint Test with L4"
         timeout_in_minutes: 60
+        # Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
+        # leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
+        # This job does not collect the ready leftover
+        # (test_diffusion_generate_request_id, core_model); isolate the
+        # advanced_model DiT handshake instead. Revert to a single sweep after
+        # the product shutdown fix.
         commands:
-          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
+          - |
+            set -eu
+            pytest -s -v tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model"
+            pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model" --deselect tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake
         mirror_hardwares: l4_1
 
+      - label: "Entrypoint Test with L4 · 2-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_2" --run-level "advanced_model"
+        mirror_hardwares: l4_2
+
   - group: ":card_index_dividers: E2E Test"
     depends_on: upload-merge-pipeline
     steps:

.buildkite/cuda/test-nightly.yml

@@ -132,6 +132,20 @@ steps:
         commands:
           - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model and H100 and B200 and cards_4" --run-level "core_model"
 
+  - group: ":card_index_dividers: Entrypoints Test"
+    key: nightly-entrypoints-test-group
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    steps:
+
+      - label: ":full_moon: Entrypoint Test with H100"
+        source_file_dependencies: entrypoints_function
+        # Merge collects advanced_model (5 duplex interruptions). Nightly
+        # collects full_model (20 interruptions + BAGEL sleep).
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "full_model and cuda and H100 and B200 and cards_2" --run-level "full_model"
+
   - group: ":card_index_dividers: TTS Model Test"
     key: nightly-tts-test-group
     depends_on: upload-nightly-pipeline

.buildkite/cuda/test-ready.yml

@@ -117,11 +117,18 @@ steps:
 
   - label: "Entrypoints Test"
     depends_on: upload-ready-pipeline
+    # Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
+    # leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
+    # Isolate that leftover case in its own process so the driver drops the
+    # mapping; the rest stays a directory sweep with --deselect. Revert to a
+    # single sweep after the product shutdown fix.
     commands:
       - |
         timeout 30m bash -c "
+          set -eu
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model'
+          pytest -sv tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id -m 'core_model and cuda' --run-level 'core_model'
+          pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model' --deselect tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id
         "
     mirror_hardwares: l4_1
 

.buildkite/cuda/test-weekly.yml

@@ -416,6 +416,7 @@ steps:
           - >-
             pytest -sv
             tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/offline_inference/test_bagel_expansion.py
             -m "slow and diffusion and H100 and cards_2"
             --run-level "full_model"
         mirror_hardwares: h100_2

tests/config/test_sleep_mode_deploy.py

@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-"""Check the sleep/wake E2E topology before loading large BAGEL weights."""
-
-import pytest
-
-from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config
-from vllm_omni.config.omni_config import (
-    VllmOmniConfig,
-    VllmOmniDiffusionStageConfig,
-    extract_diffusion_stage_config_kwargs,
-)
-from vllm_omni.config.pipeline_registry import resolve_pipeline_config
-from vllm_omni.config.stage_config import load_deploy_config
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)])
-def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size):
-    deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size))
-    pipeline = resolve_pipeline_config(deploy.pipeline or "bagel")
-    config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
-
-    assert len(config.stage_configs) == (1 if diffusion_only else 2)
-    for stage in config.stage_configs:
-        is_diffusion = diffusion_only or stage.stage_id == 1
-        assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion
-        # Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens.
-        assert stage.scheduler_config.max_num_batched_tokens == 32768
-        assert stage.model_config.enable_sleep_mode is True
-        assert stage.parallel_config.tensor_parallel_size == tp_size
-        assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size))
-        if not is_diffusion:
-            assert stage.cache_config.gpu_memory_utilization == 0.8
-        if is_diffusion:
-            engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel")
-            terminal = OmniDiffusionConfig.from_kwargs(
-                **extract_diffusion_stage_config_kwargs(
-                    engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
-                )
-            )
-            assert terminal.enable_sleep_mode is True
-            assert terminal.parallel_config.tensor_parallel_size == tp_size
-
-    if not diffusion_only:
-        assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,)
-        assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector"

tests/e2e/offline_inference/test_bagel_expansion.py

@@ -1,20 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
-Expanded end-to-end test for BAGEL in offline mode.
+Expanded end-to-end tests for BAGEL in offline mode.
 
-This test file primarily covers end-to-end tests for LoRA support (Stage 1 / DiT).
+Coverage:
+- LoRA scale / deactivation on Stage 1 DiT
 
-Validates that LoRA adapters are correctly loaded, applied with controllable
-scale, and cleanly deactivated.  Uses a synthetic rank-1 adapter targeting the
-first decoder layer's QKV projection.
-
-Assertions:
-  (a) LoRA at scale=1.0 visibly changes the output  (diff > 0.5)
-  (b) scale=2.0 produces a larger delta than scale=1.0  (linearity)
-  (c) The delta is bounded  (diff < 80, not corrupted)
-  (d) Deactivating LoRA exactly restores the baseline  (diff == 0)
+Sleep / wake (BagelPipeline TP=2 and coordinated dual-engine) lives in
+``tests/entrypoints/test_omni_sleep_mode.py`` as ``full_model`` H100 cases.
 """
 
 import json
@@ -133,17 +127,13 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
 
 
 # ---------------------------------------------------------------------------
-# Test
+# LoRA scale / deactivation
 # ---------------------------------------------------------------------------
 
 
-pytestmark = [
-    pytest.mark.slow,
-    pytest.mark.diffusion,
-    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
-]
-
-
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_lora_scale_and_deactivation(omni_runner: OmniRunner, tmp_path) -> None:
     """Validate LoRA effect, bounded perturbation, and clean deactivation."""

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -537,7 +537,10 @@ class TestQwen3OmniRealtimeWebSocket:
 
 # These regressions use the engine-owned Qwen duplex plugin. The older class
 # above intentionally exercises the separate turn/STT deployment.
-_DUPLEX_REPEATS = int(os.environ.get("VLLM_OMNI_TEST_QWEN_DUPLEX_REPEATS", "20"))
+# ``advanced_model`` (merge) uses 5 interruptions; ``full_model`` (nightly)
+# keeps the 20-repeat race hunt. Marks select the case, not ``--run-level``.
+_DUPLEX_REPEATS_ADVANCED = 5
+_DUPLEX_REPEATS_FULL = 20
 _DUPLEX_LONG_INSTRUCTIONS = "Answer the user's question in English using at least eight complete sentences."
 _DUPLEX_SHORT_INSTRUCTIONS = "Answer the user's question in one short English sentence."
 
@@ -906,12 +909,12 @@ async def _run_duplex_history_pruning(runtime, log_path: Path) -> None:
         _duplex_assert_completed_audio(events, response_id)
 
 
-@pytest.mark.advanced_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
 class TestQwen3OmniDuplexPlayback:
+    @pytest.mark.advanced_model
     @pytest.mark.parametrize("generation_finished", [False, True], ids=["generating", "queued-playback"])
-    @pytest.mark.parametrize("repetition", range(_DUPLEX_REPEATS))
+    @pytest.mark.parametrize("repetition", range(_DUPLEX_REPEATS_ADVANCED))
     def test_speech_interrupts_playback(
         self, qwen_duplex_server, tmp_path: Path, generation_finished: bool, repetition: int
     ):
@@ -922,5 +925,19 @@ class TestQwen3OmniDuplexPlayback:
             )
         )
 
+    @pytest.mark.full_model
+    @pytest.mark.parametrize("generation_finished", [False, True], ids=["generating", "queued-playback"])
+    @pytest.mark.parametrize("repetition", range(_DUPLEX_REPEATS_FULL))
+    def test_speech_interrupts_playback_full(
+        self, qwen_duplex_server, tmp_path: Path, generation_finished: bool, repetition: int
+    ):
+        """Same contract as ``test_speech_interrupts_playback``, 20 nightly repeats."""
+        asyncio.run(
+            _run_duplex_playback_interruption(
+                qwen_duplex_server, tmp_path / f"events-{repetition}.json", generation_finished=generation_finished
+            )
+        )
+
+    @pytest.mark.advanced_model
     def test_mixed_image_audio_history_pruning(self, qwen_duplex_server, tmp_path: Path):
         asyncio.run(_run_duplex_history_pruning(qwen_duplex_server, tmp_path / "events.json"))

tests/helpers/clean.py

@@ -60,6 +60,12 @@ def pick_least_used_device_indices(num_devices: int) -> list[int]:
     return get_physical_device_indices(logical_indices)
 
 
+def _logical_to_physical(logical: int) -> int:
+    """Physical id for logs / nvidia-smi; ``logical`` if the env does not map it."""
+    mapped = get_physical_device_indices([logical])
+    return mapped[0] if mapped else logical
+
+
 def wait_for_gpu_memory_to_clear(
     *,
     devices: list[int],
@@ -67,54 +73,80 @@ def wait_for_gpu_memory_to_clear(
     threshold_ratio: float | None = None,
     timeout_s: float = 120,
 ) -> None:
+    """Wait until *logical* device ordinals are under the memory threshold.
+
+    ``devices`` are CUDA/NPU ordinals in the current process (``0..N-1``),
+    the same space as ``current_omni_platform.device()`` / ``mem_get_info()``.
+    Physical ids (``CUDA_VISIBLE_DEVICES``) are only used for log labels and
+    Whisper-worker matching — remapping them into ``device()`` raises
+    ``invalid device ordinal`` when the visible set is a subset (e.g. ``4,5``).
+    """
     assert threshold_bytes is not None or threshold_ratio is not None
-    devices = get_physical_device_indices(devices)
+    logical_devices = list(devices)
+    physical_by_logical = {logical: _logical_to_physical(logical) for logical in logical_devices}
     start_time = time.time()
 
-    device_list = ", ".join(str(d) for d in devices)
+    device_list = ", ".join(str(physical_by_logical[d]) for d in logical_devices)
     if threshold_bytes is not None:
-        condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB"
+        condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB (excl. Whisper worker)"
 
         def is_free(used, total):
             return used <= threshold_bytes / 2**30
     else:
         ratio = threshold_ratio
         assert ratio is not None
-        condition_str = f"Memory usage ratio ≤ {ratio * 100:.1f}%"
+        condition_str = f"Memory usage ratio ≤ {ratio * 100:.1f}% (excl. Whisper worker)"
 
         def is_free(used, total):
             return used / total <= ratio
 
     print(f"[Device Memory Monitor] Waiting for device(s) {device_list} to free memory, Condition: {condition_str}")
-
-    def get_mem_gib(device: int) -> tuple[float, float]:
-        with current_omni_platform.device(device):
-            free_bytes, total_bytes = current_omni_platform.mem_get_info()
-        return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30
+    whisper_allowance_gib, whisper_physical = _whisper_vram_allowance()
+    if whisper_allowance_gib > 0 and whisper_physical is not None:
+        print(
+            f"[Device Memory Monitor] Whisper worker allowance {whisper_allowance_gib:.1f}GiB "
+            f"on device {whisper_physical}"
+        )
 
     while True:
-        output_raw = {d: get_mem_gib(d) for d in devices}
-        output = {
-            d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)"
-            for d, (used, total) in output_raw.items()
-        }
-
-        print("[Device Memory Status] Current usage:")
-        for device_id, mem_info in output.items():
-            print(f"  Device {device_id}: {mem_info}")
+        output_raw: dict[int, tuple[float, float, float]] = {}
+        for logical in logical_devices:
+            with current_omni_platform.device(logical):
+                free_bytes, total_bytes = current_omni_platform.mem_get_info()
+            used_gib = (total_bytes - free_bytes) / 2**30
+            total_gib = total_bytes / 2**30
+            physical = physical_by_logical[logical]
+            whisper_gib = (
+                min(used_gib, whisper_allowance_gib)
+                if whisper_physical is not None and physical == whisper_physical
+                else 0.0
+            )
+            output_raw[physical] = (max(0.0, used_gib - whisper_gib), total_gib, whisper_gib)
+        print("[Device Memory Status] Current usage (engine view excludes Whisper worker):")
+        for device_id, (used, total, whisper) in output_raw.items():
+            ratio_pct = (used / total) * 100 if total > 0 else 0.0
+            line = f"  Device {device_id}: {used:.1f}GiB/{total:.1f}GiB ({ratio_pct:.1f}%)"
+            if whisper > 0:
+                line += f" [Whisper worker {whisper:.1f}GiB excluded]"
+            print(line)
 
         dur_s = time.time() - start_time
-        if all(is_free(used, total) for used, total in output_raw.values()):
+        if all(is_free(used, total) for used, total, _whisper in output_raw.values()):
             print(f"[Device Memory Freed] Device(s) {device_list} meet memory condition")
             print(f"   Condition: {condition_str}")
             print(f"   Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)")
             break
 
         if dur_s >= timeout_s:
+            status = "\n".join(
+                f"  Device {d}: {used:.1f}GiB/{total:.1f}GiB"
+                + (f" (+Whisper {whisper:.1f}GiB excluded)" if whisper > 0 else "")
+                for d, (used, total, whisper) in output_raw.items()
+            )
             raise ValueError(
                 f"[Device Memory Timeout] Device(s) {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
                 f"Condition: {condition_str}\n"
-                f"Current status:\n" + "\n".join(f"  Device {d}: {output[d]}" for d in devices)
+                f"Current status:\n{status}"
             )
 
         gc.collect()
@@ -122,6 +154,27 @@ def wait_for_gpu_memory_to_clear(
         time.sleep(5)
 
 
+def _whisper_vram_allowance() -> tuple[float, int | None]:
+    """Living Whisper worker's GPU footprint and physical device.
+
+    Orthogonal to RFC #6851 leftover-engine PID reap / a raised wait ratio:
+    this only excludes the transcriber's own allocator (capped by current used
+    on that device). CPU / unknown device → ``(0.0, None)``.
+    """
+    try:
+        from tests.helpers.media import whisper_resident_device_index, whisper_resident_vram_gib
+    except Exception:
+        return 0.0, None
+    gib = whisper_resident_vram_gib()
+    logical = whisper_resident_device_index()
+    if gib <= 0.0 or logical is None:
+        return 0.0, None
+    mapped = get_physical_device_indices([logical])
+    if not mapped:
+        return 0.0, None
+    return gib, mapped[0]
+
+
 def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None:
     print("\n" + "=" * 80)
     print(label)
@@ -194,6 +247,16 @@ def cleanup_test_environment(*, shutdown_ray: bool = False) -> None:
         ray.shutdown()
 
     print("Pre-test device status:")
+    gc.collect()
+    if not current_omni_platform.is_cpu():
+        current_omni_platform.empty_cache()
+        try:
+            import torch
+
+            torch._C._host_emptyCache()
+        except AttributeError:
+            logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
+
     num_devices = current_omni_platform.device_count()
     if num_devices > 0:
         try:
@@ -205,16 +268,6 @@ def cleanup_test_environment(*, shutdown_ray: bool = False) -> None:
         except Exception as e:
             print(f"Device cleanup note: {e}")
 
-    gc.collect()
-    if not current_omni_platform.is_cpu():
-        current_omni_platform.empty_cache()
-        try:
-            import torch
-
-            torch._C._host_emptyCache()
-        except AttributeError:
-            logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
-
     if current_omni_platform.is_available():
         print("Post-test device status:")
         _print_device_processes()

tests/helpers/media.py

@@ -734,9 +734,13 @@ def _select_whisper_device() -> str:
 
 # Populated in the transcription worker, not in the pytest process.
 _WHISPER_MODELS: dict[str, Any] = {}
+# Pinned on the first ``whisper.load_model`` in this worker. Later sizes reuse
+# it so a CPU fallback cannot be mixed with a later GPU load in the same pool.
+_WHISPER_LOADED_DEVICE: str | None = None
 
 
 def _get_whisper_model(model_size: str) -> Any:
+    global _WHISPER_LOADED_DEVICE
     model = _WHISPER_MODELS.get(model_size)
     if model is None:
         import whisper
@@ -744,16 +748,41 @@ def _get_whisper_model(model_size: str) -> Any:
         # The device is picked on first load and the model stays on it for the
         # worker's lifetime: the current server or runner fixture instance, or
         # the test module for callers that transcribe without those fixtures.
-        device = _select_whisper_device()
+        device = _WHISPER_LOADED_DEVICE
+        if device is None:
+            device = _select_whisper_device()
+            _WHISPER_LOADED_DEVICE = device
         with _serialize_whisper_model_download(model_size):
             model = whisper.load_model(model_size, device=device)
         _WHISPER_MODELS[model_size] = model
     return model
 
 
+def _whisper_process_reserved_gib() -> float:
+    """GiB held by this worker's caching allocator, or 0 on CPU / if unreadable.
+
+    Parent-side ``nvidia-smi --query-compute-apps`` is empty on some CI
+    drivers (pmon shows ``-`` while device-wide used is tens of GiB), so the
+    pytest process cannot map the worker PID to VRAM. The allocator in *this*
+    process still knows what it reserved.
+    """
+    index = _accelerator_index_from_device(_WHISPER_LOADED_DEVICE)
+    if index is None:
+        return 0.0
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        if not current_omni_platform.is_available():
+            return 0.0
+        device = current_omni_platform.get_torch_device(index)
+        return current_omni_platform.memory_reserved(device) / 1024**3
+    except Exception:
+        return 0.0
+
+
 def _whisper_transcribe_in_current_process(
     output_path: str, model_size: str = "small", language: str | None = None
-) -> str:
+) -> tuple[str, str, float]:
     model = _get_whisper_model(model_size)
     text = model.transcribe(
         output_path,
@@ -764,7 +793,9 @@ def _whisper_transcribe_in_current_process(
         # language: callers include non-English audio tests.
         language=language,
     )["text"]
-    return text or ""
+    device = _WHISPER_LOADED_DEVICE or "cpu"
+    reserved = _whisper_process_reserved_gib() if _accelerator_index_from_device(device) is not None else 0.0
+    return text or "", device, reserved
 
 
 # Serializes a whole submit->result->cleanup on the parent side, so at most one
@@ -775,6 +806,52 @@ _TRANSCRIBER_CALL_LOCK = threading.Lock()
 # Guards the _TRANSCRIBER pointer itself.
 _TRANSCRIBER_LOCK = threading.Lock()
 _TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None
+# Parent-side record of the device the child reported and its allocator
+# reserved GiB. Credit only after a GPU result (not while the model might
+# still be on CPU). Reserved is the child's caching-allocator reading, not
+# nvidia-smi PID memory (N/A on some CI drivers) and not a size table.
+_TRANSCRIBER_DEVICE: str | None = None
+_TRANSCRIBER_RESERVED_GIB: float = 0.0
+
+
+def _accelerator_index_from_device(device: str | None) -> int | None:
+    """Logical device index for ``cuda:1`` / ``npu:0``; ``None`` for CPU or unknown."""
+    if device is None:
+        return None
+    name = device.strip().lower()
+    if not name or name == "cpu":
+        return None
+    if ":" in name:
+        suffix = name.rsplit(":", 1)[1]
+        if suffix.isdigit():
+            return int(suffix)
+    return None
+
+
+def whisper_resident_device_index() -> int | None:
+    """Logical accelerator index holding Whisper, or ``None`` if it is not on GPU."""
+    with _TRANSCRIBER_LOCK:
+        if _TRANSCRIBER is None:
+            return None
+        device = _TRANSCRIBER_DEVICE
+    return _accelerator_index_from_device(device)
+
+
+def whisper_resident_vram_gib() -> float:
+    """VRAM (GiB) held by the living Whisper worker, or 0 if none / CPU.
+
+    Only the child's caching-allocator reading. Unmeasured or CPU → 0 (fail
+    closed: do not invent a size-table credit). This is not a substitute for
+    engine PID reap or a raised 5% wait threshold (RFC #6851).
+    """
+    with _TRANSCRIBER_LOCK:
+        if _TRANSCRIBER is None:
+            return 0.0
+        device = _TRANSCRIBER_DEVICE
+        measured = _TRANSCRIBER_RESERVED_GIB
+    if _accelerator_index_from_device(device) is None:
+        return 0.0
+    return measured if measured > 0.0 else 0.0
 
 
 def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor:
@@ -792,11 +869,13 @@ def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> No
     Identity-checked so a stale reference can never shut down a newer worker that
     was installed after ``executor`` was replaced.
     """
-    global _TRANSCRIBER
+    global _TRANSCRIBER, _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
     with _TRANSCRIBER_LOCK:
         if _TRANSCRIBER is not executor:
             return
         _TRANSCRIBER = None
+        _TRANSCRIBER_DEVICE = None
+        _TRANSCRIBER_RESERVED_GIB = 0.0
     # Joining the worker can block; do it outside the lock.
     executor.shutdown(wait=True)
 
@@ -812,14 +891,25 @@ def release_audio_transcriber() -> None:
     Takes the call lock, so it waits for any in-flight transcription to finish
     rather than shutting the worker down underneath it.
     """
-    global _TRANSCRIBER
+    global _TRANSCRIBER, _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
     with _TRANSCRIBER_CALL_LOCK:
         with _TRANSCRIBER_LOCK:
             executor, _TRANSCRIBER = _TRANSCRIBER, None
+            _TRANSCRIBER_DEVICE = None
+            _TRANSCRIBER_RESERVED_GIB = 0.0
         if executor is not None:
             executor.shutdown(wait=True)
 
 
+def _unpack_transcribe_worker_result(result: object) -> tuple[str, str, float]:
+    """Accept ``(text, device)`` from tests or ``(text, device, reserved_gib)`` from the worker."""
+    if not isinstance(result, tuple) or len(result) not in (2, 3):
+        raise TypeError(f"whisper worker returned {type(result)!r}, expected a 2- or 3-tuple")
+    text, device = result[0], result[1]
+    reserved = float(result[2]) if len(result) == 3 else 0.0
+    return str(text), str(device), reserved
+
+
 def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
     """Convert an audio file to text in a reused, isolated subprocess.
 
@@ -837,13 +927,21 @@ def convert_audio_file_to_text(output_path: str, model_size: str = "small", lang
     and its resident model -- is torn down, and a dead worker
     (``BrokenProcessPool``) is additionally retried once.
     """
+    global _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
     with _TRANSCRIBER_CALL_LOCK:
         for attempt in range(2):
             executor = _get_transcriber()
             try:
-                return executor.submit(
-                    _whisper_transcribe_in_current_process, output_path, model_size, language
-                ).result()
+                text, device, reserved_gib = _unpack_transcribe_worker_result(
+                    executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language).result()
+                )
+                with _TRANSCRIBER_LOCK:
+                    _TRANSCRIBER_DEVICE = device
+                    if _accelerator_index_from_device(device) is None:
+                        _TRANSCRIBER_RESERVED_GIB = 0.0
+                    elif reserved_gib > 0.0:
+                        _TRANSCRIBER_RESERVED_GIB = reserved_gib
+                return text
             except BrokenProcessPool:
                 _discard_transcriber(executor)
                 if attempt == 1:
@@ -882,4 +980,6 @@ __all__ = [
     "get_asset_path",
     "preprocess_text",
     "release_audio_transcriber",
+    "whisper_resident_device_index",
+    "whisper_resident_vram_gib",
 ]

tests/helpers/runtime.py

@@ -169,9 +169,8 @@ class OmniServer:
         cleanup_test_environment()
         self.startup_timeout = startup_timeout
         self.model = model
-        args = list(serve_args)
-        self.serve_args = args
-        self.log_stats = "--disable-log-stats" not in args and "--log-stats" in args
+        self.serve_args = list(serve_args)
+        self.log_stats = "--disable-log-stats" not in self.serve_args and "--log-stats" in self.serve_args
         self.env_dict = env_dict
         self.use_omni = use_omni
         self.proc: subprocess.Popen | None = None

tests/helpers/tests/test_clean.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+
+import pytest
+
+from tests.helpers import clean
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_GIB = 2**30
+
+
+def _patch_wait_platform(mocker, monkeypatch, *, queried: list[int], mem_get_info):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5")
+    monkeypatch.setattr(clean.current_omni_platform, "device_control_env_var", "CUDA_VISIBLE_DEVICES")
+
+    def fake_device(idx):
+        queried.append(idx)
+        return nullcontext()
+
+    mocker.patch.object(clean.current_omni_platform, "device", side_effect=fake_device)
+    mocker.patch.object(clean.current_omni_platform, "mem_get_info", side_effect=mem_get_info)
+    mocker.patch.object(clean.current_omni_platform, "empty_cache")
+    mocker.patch.object(clean.time, "sleep")
+
+
+def test_wait_queries_logical_ordinals_when_visible_set_is_physical(mocker, monkeypatch):
+    """#6523 kept a physical remap after switching off NVML; device(4) is invalid."""
+    queried: list[int] = []
+
+    def mem_get_info():
+        return (99 * _GIB, 100 * _GIB)
+
+    _patch_wait_platform(mocker, monkeypatch, queried=queried, mem_get_info=mem_get_info)
+    monkeypatch.setattr(clean, "_whisper_vram_allowance", lambda: (0.0, None))
+
+    clean.wait_for_gpu_memory_to_clear(devices=[0, 1], threshold_ratio=0.05, timeout_s=1)
+
+    assert queried == [0, 1]
+    assert 4 not in queried
+    assert 5 not in queried
+
+
+def test_wait_matches_whisper_on_physical_id_not_logical(mocker, monkeypatch):
+    queried: list[int] = []
+
+    def mem_get_info():
+        logical = queried[-1]
+        if logical == 0:
+            return (99 * _GIB, 100 * _GIB)
+        return (50 * _GIB, 100 * _GIB)
+
+    _patch_wait_platform(mocker, monkeypatch, queried=queried, mem_get_info=mem_get_info)
+    # logical 1 → physical 5; without physical matching this 50% used card fails 15%.
+    monkeypatch.setattr(clean, "_whisper_vram_allowance", lambda: (40.0, 5))
+
+    clean.wait_for_gpu_memory_to_clear(devices=[0, 1], threshold_ratio=0.15, timeout_s=1)
+
+    assert queried == [0, 1]

tests/helpers/tests/test_media.py

@@ -76,7 +76,7 @@ class _FakeExecutor:
 
     def submit(self, fn, *args):
         self.submitted.append((fn, args))
-        outcome = self._outcomes.pop(0) if self._outcomes else "London"
+        outcome = self._outcomes.pop(0) if self._outcomes else ("London", "cpu")
 
         def _result():
             if isinstance(outcome, BaseException):
@@ -110,9 +110,13 @@ def _reset_transcriber_singletons():
     """The worker and its model cache are module-level state; do not leak them across tests."""
     media.release_audio_transcriber()
     media._WHISPER_MODELS.clear()
+    media._WHISPER_LOADED_DEVICE = None
+    media._TRANSCRIBER_RESERVED_GIB = 0.0
     yield
     media.release_audio_transcriber()
     media._WHISPER_MODELS.clear()
+    media._WHISPER_LOADED_DEVICE = None
+    media._TRANSCRIBER_RESERVED_GIB = 0.0
 
 
 @pytest.mark.parametrize(
@@ -236,7 +240,10 @@ def test_worker_loads_each_model_size_once(monkeypatch):
 
 @pytest.mark.parametrize(
     ("second_outcome", "recovers"),
-    [(["London"], True), ([BrokenProcessPool("second")], False)],
+    [
+        ([("London", "cpu")], True),
+        ([BrokenProcessPool("second")], False),
+    ],
 )
 def test_broken_pool_is_discarded_and_retried_once(monkeypatch, second_outcome, recovers):
     """A dead worker is discarded and the call retried exactly once.
@@ -286,3 +293,103 @@ def test_transcription_failure_propagates_and_discards_the_worker(monkeypatch, e
 
     assert media.convert_audio_file_to_text("/tmp/b.wav") == "London"
     assert len(created) == 2  # a fresh worker for the next call
+
+
+def test_whisper_resident_vram_is_zero_until_a_gpu_result(monkeypatch):
+    """CPU fallback and in-flight loads must not hide engine leaks on the last GPU."""
+    assert media.whisper_resident_vram_gib() == 0.0
+    assert media.whisper_resident_device_index() is None
+
+    created = _patch_executors(monkeypatch, outcomes_per_executor=([("London", "cpu")],))
+    media.convert_audio_file_to_text("/tmp/a.wav", "small")
+
+    assert created[0].submitted
+    assert media.whisper_resident_vram_gib() == 0.0
+    assert media.whisper_resident_device_index() is None
+
+
+def test_whisper_resident_vram_tracks_gpu_device_after_success(monkeypatch):
+    _patch_executors(
+        monkeypatch,
+        outcomes_per_executor=([("London", "cuda:1", 2.5), ("Paris", "cuda:1", 9.6)],),
+    )
+    media.convert_audio_file_to_text("/tmp/a.wav", "small")
+    assert media.whisper_resident_device_index() == 1
+    assert media.whisper_resident_vram_gib() == pytest.approx(2.5)
+
+    media.convert_audio_file_to_text("/tmp/b.wav", "large-v3")
+    assert media.whisper_resident_device_index() == 1
+    assert media.whisper_resident_vram_gib() == pytest.approx(9.6)
+
+    media.release_audio_transcriber()
+    assert media.whisper_resident_vram_gib() == 0.0
+    assert media.whisper_resident_device_index() is None
+
+
+def test_whisper_unmeasured_gpu_result_credits_zero(monkeypatch):
+    """A GPU device without a reserved reading must not invent a size-table credit."""
+    _patch_executors(monkeypatch, outcomes_per_executor=([("London", "cuda:0")],))
+    media.convert_audio_file_to_text("/tmp/a.wav", "not-a-whisper-size")
+    assert media.whisper_resident_device_index() == 0
+    assert media.whisper_resident_vram_gib() == 0.0
+
+
+def test_whisper_resident_vram_uses_child_reserved(monkeypatch):
+    _patch_executors(monkeypatch, outcomes_per_executor=([("London", "cuda:0", 1.0)],))
+    media.convert_audio_file_to_text("/tmp/a.wav", "small")
+    assert media.whisper_resident_vram_gib() == pytest.approx(1.0)
+
+
+def test_whisper_resident_vram_does_not_cap_measured_reserved(monkeypatch):
+    _patch_executors(monkeypatch, outcomes_per_executor=([("London", "cuda:0", 20.0)],))
+    media.convert_audio_file_to_text("/tmp/a.wav", "small")
+    assert media.whisper_resident_vram_gib() == pytest.approx(20.0)
+
+
+def test_whisper_unknown_size_uses_measured_reserved(monkeypatch):
+    _patch_executors(monkeypatch, outcomes_per_executor=([("London", "cuda:0", 3.0)],))
+    media.convert_audio_file_to_text("/tmp/a.wav", "not-a-whisper-size")
+    assert media.whisper_resident_vram_gib() == pytest.approx(3.0)
+
+
+def test_later_whisper_sizes_reuse_the_first_selected_device(monkeypatch):
+    selected: list[str] = []
+
+    def fake_model(_size, device=None):
+        return SimpleNamespace(transcribe=lambda *_a, **_k: {"text": "x"})
+
+    def select_device() -> str:
+        selected.append("cuda:1")
+        return "cuda:1"
+
+    monkeypatch.setattr(media, "_select_whisper_device", select_device)
+    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=fake_model))
+    monkeypatch.setattr(media, "_serialize_whisper_model_download", lambda model_size: nullcontext())
+
+    text, device, _reserved = media._whisper_transcribe_in_current_process("/tmp/a.wav", "small")
+    assert (text, device) == ("x", "cuda:1")
+    media._whisper_transcribe_in_current_process("/tmp/b.wav", "large-v3")
+    assert selected == ["cuda:1"]
+    assert media._WHISPER_LOADED_DEVICE == "cuda:1"
+
+
+def test_cleanup_whisper_allowance_is_zero_on_cpu(monkeypatch):
+    from tests.helpers import clean
+
+    monkeypatch.setattr(media, "whisper_resident_vram_gib", lambda: 11.0)
+    monkeypatch.setattr(media, "whisper_resident_device_index", lambda: None)
+    assert clean._whisper_vram_allowance() == (0.0, None)
+
+
+def test_cleanup_whisper_allowance_maps_logical_gpu(monkeypatch):
+    from tests.helpers import clean
+
+    monkeypatch.setattr(media, "whisper_resident_vram_gib", lambda: 2.5)
+    monkeypatch.setattr(media, "whisper_resident_device_index", lambda: 1)
+    monkeypatch.setattr(clean, "get_physical_device_indices", lambda devices: [7] if devices == [1] else [])
+    assert clean._whisper_vram_allowance() == (2.5, 7)
+
+
+def test_whisper_has_no_size_table():
+    assert not hasattr(media, "_WHISPER_VRAM_GIB")
+    assert not hasattr(media, "_WHISPER_VRAM_GIB_DEFAULT")

tools/pre_commit/check_torch_cuda.py

@@ -42,7 +42,6 @@ ALLOWED_FILES = {
     "tests/e2e/accuracy/wan22_i2v/run_wan22_i2v_diffusers_cp.py",
     "tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py",
     "tests/e2e/features/fullduplex/test_personaplex_elastic.py",
-    "tests/entrypoints/test_omni_sleep_mode.py",
     "tools/configure_stage_memory.py",
     "vllm_omni/diffusion/models/bagel/pipeline_bagel.py",
     "vllm_omni/diffusion/models/lance/pipeline_lance.py",

vllm_omni/platforms/cuda/platform.py

@@ -423,6 +423,10 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         free, total = torch.cuda.mem_get_info(device)
         return free, total
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        return int(torch.cuda.memory_reserved(device))
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)

vllm_omni/platforms/interface.py

@@ -300,6 +300,15 @@ class OmniPlatform(Platform):
     def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
         raise NotImplementedError
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        """Bytes reserved by this process's caching allocator on ``device``.
+
+        Device-wide ``get_device_memory`` counts every process; this is the
+        portable equivalent of ``torch.cuda.memory_reserved``.
+        """
+        return 0
+
     @classmethod
     def create_autocast_context(
         cls,
@@ -329,7 +338,7 @@ class OmniPlatform(Platform):
     def set_device_control_env_var(cls, devices: str | int | None) -> None:
         import os
 
-        os.environ[cls.device_control_env_var] = devices
+        os.environ[cls.device_control_env_var] = "" if devices is None else str(devices)
 
     @classmethod
     def unset_device_control_env_var(cls) -> None:

vllm_omni/platforms/musa/platform.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
 
 
 import torch
@@ -182,6 +182,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         free, total = torch.musa.mem_get_info(device)
         return free, total
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        return int(torch.musa.memory_reserved(device))
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.musa.get_device_name(device_id)
@@ -190,7 +194,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     def set_device_control_env_var(cls, devices: str | int | None) -> None:
         import os
 
-        os.environ["MUSA_VISIBLE_DEVICES"] = devices
+        os.environ["MUSA_VISIBLE_DEVICES"] = "" if devices is None else str(devices)
 
     @classmethod
     def unset_device_control_env_var(cls) -> None:

vllm_omni/platforms/npu/platform.py

@@ -344,6 +344,10 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         free, total = torch.npu.mem_get_info(device)
         return free, total
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        return int(torch.npu.memory_reserved(device))
+
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         device_props = torch.npu.get_device_properties(device_id)

vllm_omni/platforms/rocm/platform.py

@@ -180,12 +180,17 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
         free, total = torch.cuda.mem_get_info(device)
         return free, total
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        return int(torch.cuda.memory_reserved(device))
+
     @classmethod
     def set_device_control_env_var(cls, devices: str | int | None) -> None:
         import os
 
-        os.environ["HIP_VISIBLE_DEVICES"] = devices
-        os.environ["CUDA_VISIBLE_DEVICES"] = devices
+        visible = "" if devices is None else str(devices)
+        os.environ["HIP_VISIBLE_DEVICES"] = visible
+        os.environ["CUDA_VISIBLE_DEVICES"] = visible
 
     @classmethod
     def unset_device_control_env_var(cls) -> None:

vllm_omni/platforms/xpu/platform.py

@@ -137,6 +137,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
         free, total = torch.xpu.mem_get_info(device)
         return free, total
 
+    @classmethod
+    def memory_reserved(cls, device: torch.device | int | None = None) -> int:
+        return int(torch.xpu.memory_reserved(device))
+
     @classmethod
     def get_profiler_cls(cls) -> str:
         """Return XPU-specific profiler that handles XPU events."""