From 3d43571b94f1683412023c45bd886f9ce30f76bd Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:06:58 +0800 Subject: [PATCH] [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] Co-authored-by: Cursor --- .../common/ci_source_file_dependencies.yml | 4 + .buildkite/cuda/test-merge.yml | 23 +- .buildkite/cuda/test-nightly.yml | 14 + .buildkite/cuda/test-ready.yml | 9 +- .buildkite/cuda/test-weekly.yml | 1 + tests/config/test_sleep_mode_deploy.py | 51 - .../offline_inference/test_bagel_expansion.py | 30 +- .../test_qwen3_omni_realtime_websocket.py | 25 +- tests/entrypoints/test_omni_sleep_mode.py | 895 +++++++++++------- tests/helpers/clean.py | 113 ++- tests/helpers/media.py | 116 ++- tests/helpers/runtime.py | 5 +- tests/helpers/tests/test_clean.py | 61 ++ tests/helpers/tests/test_media.py | 111 ++- tools/pre_commit/check_torch_cuda.py | 1 - vllm_omni/platforms/cuda/platform.py | 4 + vllm_omni/platforms/interface.py | 11 +- vllm_omni/platforms/musa/platform.py | 8 +- vllm_omni/platforms/npu/platform.py | 4 + vllm_omni/platforms/rocm/platform.py | 9 +- vllm_omni/platforms/xpu/platform.py | 4 + 21 files changed, 1012 insertions(+), 487 deletions(-) delete mode 100644 tests/config/test_sleep_mode_deploy.py create mode 100644 tests/helpers/tests/test_clean.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 67231eb8a..d1336e66c 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.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 diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 021f5d678..a14a57c6c 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.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: diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 36df33a41..4f6a88f52 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.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 diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 5fe28722f..c474e64a2 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.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 diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index 5c193945e..1d864c176 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.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 diff --git a/tests/config/test_sleep_mode_deploy.py b/tests/config/test_sleep_mode_deploy.py deleted file mode 100644 index 13040ebd9..000000000 --- a/tests/config/test_sleep_mode_deploy.py +++ /dev/null @@ -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" diff --git a/tests/e2e/offline_inference/test_bagel_expansion.py b/tests/e2e/offline_inference/test_bagel_expansion.py index 4c31be722..792988e57 100644 --- a/tests/e2e/offline_inference/test_bagel_expansion.py +++ b/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.""" diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 7aaa975b5..4aea42159 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/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")) diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index c4900e9ca..1fb6c7512 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -1,20 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Omni sleep mode: entrypoint-level VRAM/ACK tests plus multi-stage e2e. +"""Entrypoint sleep-mode coverage on small models plus H100 BAGEL. -Class-scoped engines amortize BAGEL load across sleep/wake assertions within -each suite (LLM vs diffusion). Suites are separate classes so only one heavy -engine is alive at a time. Device cleanup runs once per module. +Layering (AR before tiny DiT on the same L4). Thinker lives in a child +``StageEngineCoreProc`` that shutdown can reap; tiny DiT is an inline +UniProc ``DiffusionWorker`` in pytest and leftover VRAM cannot be killed +(#8016). DiT-then-7B OOMs (6.78 GiB leftover + 16.78 GiB weights > 22 GiB). +1. AR protocol (#4473) — ``Qwen/Qwen2.5-Omni-7B`` thinker-only on L4 +2. Diffusion sleep/wake/generate — ``riverclouds/qwen_image_random`` on L4 +3. Light multistage orchestration — thinker-only AR + tiny DiT on L4×2 +4. BAGEL BagelPipeline TP=2 / coordinated dual-engine — H100, ``full_model`` """ +from __future__ import annotations + import asyncio import logging -import os +from typing import Any, NamedTuple import pytest import pytest_asyncio -import torch from vllm import SamplingParams from tests.helpers.mark import hardware_test @@ -25,69 +31,122 @@ from vllm_omni.platforms import current_omni_platform logging.basicConfig(level=logging.INFO) logger = logging.getLogger("OmniTest") -pytestmark = [ - pytest.mark.advanced_model, -] - -def clean_device_envs(): - """Clear device-visibility env vars so tests see all available devices.""" - device_visibility_vars = [ - "CUDA_VISIBLE_DEVICES", # NVIDIA - "HIP_VISIBLE_DEVICES", # AMD ROCm - "ZE_AFFINITY_MASK", # Intel XPU - "ONEAPI_DEVICE_SELECTOR", # Intel OneAPI - "ASCEND_RT_VISIBLE_DEVICES", # Huawei NPU (CAN) - ] - for key in device_visibility_vars: - os.environ.pop(key, None) - - -def get_vram_info(device_id: int) -> dict: - """Per-**process** CUDA allocator stats (GiB). Does not see other processes' usage.""" - try: - if current_omni_platform.is_rocm(): - num_gpus = torch.accelerator.device_count() - safe_id = device_id if device_id < num_gpus else 0 - torch.accelerator.synchronize(safe_id) - return { - "reserved": torch.cuda.memory_reserved(safe_id) / 1024**3, - "allocated": torch.cuda.memory_allocated(safe_id) / 1024**3, +MODEL_DIFF = "riverclouds/qwen_image_random" +MODEL_AR = "Qwen/Qwen2.5-Omni-7B" +MODEL_BAGEL = "ByteDance-Seed/BAGEL-7B-MoT" +# Thinker weights are 16.78 GiB. On L4 (22 GiB) util 0.45 only budgets 9.9 GiB +# so device-level KV profiling reports 0.0 GiB and StageEngineCoreProc dies. +# Keep abort-test util 0.90; cap len/batch at 2048 so the leftover ~3 GiB KV fits. +_AR_SLEEP_GPU_MEMORY_UTILIZATION = 0.90 +_AR_SLEEP_MAX_MODEL_LEN = 2048 +_AR_SLEEP_MAX_NUM_BATCHED_TOKENS = 2048 +AR_STAGE_CONFIG = modify_stage_config( + get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml"), + updates={ + "stages": { + 0: { + "enable_sleep_mode": True, + "gpu_memory_utilization": _AR_SLEEP_GPU_MEMORY_UTILIZATION, + "max_model_len": _AR_SLEEP_MAX_MODEL_LEN, + "max_num_batched_tokens": _AR_SLEEP_MAX_NUM_BATCHED_TOKENS, } - else: - with torch.cuda.device(device_id): - torch.accelerator.synchronize() - return { - "reserved": torch.cuda.memory_reserved() / 1024**3, - "allocated": torch.cuda.memory_allocated() / 1024**3, - } - except Exception as e: - logger.warning(f"memory skip ({device_id}): {e}") - return {"reserved": 0.0, "allocated": 0.0} + } + }, +) +_MULTISTAGE_SLEEP_PIPELINE = "qwen2_5_omni_thinker_tiny_dit" -def get_device_global_memory_used_gib(device_id: int) -> float: - """GPU-wide memory in use (GiB), includes all processes (driver view). +def _tiny_dit_multistage_deploy_config() -> str: + """Thinker-only AR + tiny DiT deploy for joint sleep/wake. - Uses ``torch.cuda.mem_get_info`` (CUDA / ROCm when available). Use this when stages run in - **separate worker processes**; ``memory_reserved`` in the test process does not reflect - child workers' allocations. + ``AsyncOmni(stages=...)`` is ignored by ``resolve_omni_config``: a + ``model=Qwen2.5-Omni-7B`` call otherwise loads the native 3-stage + thinker/talker/code2wav pipeline (``Expected 3 sampling params, got 2``). """ - try: - with torch.cuda.device(device_id): - torch.accelerator.synchronize() - free_b, total_b = torch.cuda.mem_get_info() - return (total_b - free_b) / 1024**3 - except Exception as e: - logger.warning("get_device_global_memory_used_gib(%s): %s", device_id, e) - return 0.0 + import atexit + import os + import tempfile + from pathlib import Path + + import yaml + + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline + from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig + + if _MULTISTAGE_SLEEP_PIPELINE not in OMNI_PIPELINES: + register_pipeline( + PipelineConfig( + model_type=_MULTISTAGE_SLEEP_PIPELINE, + hf_architectures=(), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + model_arch="Qwen2_5OmniForConditionalGeneration", + engine_output_type="latent", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="image", + ), + ), + ) + ) + + config = { + "pipeline": _MULTISTAGE_SLEEP_PIPELINE, + "async_chunk": False, + "dtype": "bfloat16", + "trust_remote_code": True, + "enable_prefix_caching": False, + "connectors": { + "shared_memory_connector": {"name": "SharedMemoryConnector"}, + }, + "stages": [ + { + "stage_id": 0, + "devices": "0", + "enable_sleep_mode": True, + "gpu_memory_utilization": _AR_SLEEP_GPU_MEMORY_UTILIZATION, + "enforce_eager": True, + "max_model_len": _AR_SLEEP_MAX_MODEL_LEN, + "max_num_batched_tokens": _AR_SLEEP_MAX_NUM_BATCHED_TOKENS, + "skip_mm_profiling": True, + "mm_processor_cache_gb": 0, + "max_num_seqs": 1, + }, + { + "stage_id": 1, + "devices": "1", + "model": MODEL_DIFF, + "enable_sleep_mode": True, + "gpu_memory_utilization": 0.4, + "enforce_eager": True, + "tensor_parallel_size": 1, + "max_num_seqs": 1, + "input_connectors": {"from_stage_0": "shared_memory_connector"}, + }, + ], + } + fd, path = tempfile.mkstemp(prefix="sleep_thinker_tiny_dit_", suffix=".yaml") + atexit.register(Path(path).unlink, missing_ok=True) + with os.fdopen(fd, "w", encoding="utf-8") as f: + yaml.safe_dump(config, f, sort_keys=False) + return path def get_ack_info(ack, key, default=None): - """ - Since ACKs in a distributed environment can be either objects or dictionaries, - this tool ensures compatibility. - """ if hasattr(ack, key): return getattr(ack, key) if isinstance(ack, dict): @@ -95,12 +154,7 @@ def get_ack_info(ack, key, default=None): return default -MODEL = "ByteDance-Seed/BAGEL-7B-MoT" -MODEL_DIFF = "riverclouds/qwen_image_random" - - -async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: - """Best-effort full wake + resume so the next shared-engine test starts active.""" +async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int] | None = None) -> None: try: await engine.wake_up(stage_ids=stage_ids) except Exception as e: @@ -111,383 +165,512 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e) -def _sleep_deploy_config(*, diffusion_only: bool = False, tp_size: int = 1) -> str: - """Use registered BAGEL topology and override only deployment settings.""" - filename = "bagel_single_stage.yaml" if diffusion_only else "bagel.yaml" - stage_ids = (0,) if diffusion_only else (0, 1) - return modify_stage_config( - get_deploy_config_path(filename), - updates={ - "stages": { - stage_id: { - "devices": ",".join(str(stage_id * tp_size + rank) for rank in range(tp_size)), - "tensor_parallel_size": tp_size, - "enable_sleep_mode": True, - "enforce_eager": True, - "dtype": "bfloat16", - # Stages use disjoint GPUs. On an 80 GiB H100, 40% cannot - # fit TP=1 Thinker weights, activation peaks and KV cache. - "gpu_memory_utilization": 0.8 if not diffusion_only and stage_id == 0 else 0.4, - } - for stage_id in stage_ids - } - }, +async def _tiny_dit_multistage_generate(engine: AsyncOmni, prompt: str, request_id: str): + """2-step tiny-DiT generate through the joint AR + DiT pipeline. + + Restores the old H100 ``test_multistage_llm_diffusion_sleep_wake`` contract + (generate before and after ``sleep(stage_ids=[0, 1])``) without BAGEL. + AR is capped so this stays a sleep/wake orchestration check, not a 7B + completion. + """ + params_list = [ + SamplingParams(max_tokens=4), + OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), + ] + output = None + async for item in engine.generate(prompt, request_id=request_id, sampling_params_list=params_list): + output = item + assert output is not None, f"generate({request_id!r}) produced no output" + return output + + +_ENGINE_CHILD_MARKERS = ( + "stagediffusionproc", + "enginecore", + "diffusionworker", + "vllm::worker", + "vllm-omni:", +) + + +def _is_engine_worker_proc(proc) -> bool: + blob = " ".join([proc.name(), *proc.cmdline()]).lower() + return any(marker in blob for marker in _ENGINE_CHILD_MARKERS) + + +class _EngineWorkerSnap(NamedTuple): + """Pre-shutdown worker handle. ``proc`` already caches identity.""" + + proc: Any + pid: int + create_time: float + + +def _snapshot_engine_workers() -> list[_EngineWorkerSnap]: + """Engine workers under this pytest process (before shutdown). + + ``DiffusionWorker`` is started ``daemon=True``. After StagePool joins + ``StageDiffusionProc``, that worker is reparented to PID 1 and a later + ``children(recursive=True)`` scan misses it while it still holds VRAM. + + Keep the original ``psutil.Process`` (and its cached ``create_time``) + across ``engine.shutdown()``. Rebuilding ``Process(pid)`` after exit + can bind to a reused PID and ``kill()`` an unrelated process. + """ + try: + import psutil + except ImportError: + return [] + try: + children = psutil.Process().children(recursive=True) + except psutil.Error: + return [] + snaps: list[_EngineWorkerSnap] = [] + for proc in children: + try: + if _is_engine_worker_proc(proc): + snaps.append(_EngineWorkerSnap(proc, proc.pid, proc.create_time())) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + return snaps + + +def _same_engine_worker(snap: _EngineWorkerSnap) -> bool: + """True only if ``snap.proc`` is still the original worker, not a reused PID.""" + try: + import psutil + except ImportError: + return False + try: + if not snap.proc.is_running(): + return False + return snap.proc.pid == snap.pid and snap.proc.create_time() == snap.create_time + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + + +def _reap_engine_workers(snaps: list[_EngineWorkerSnap]) -> None: + """SIGKILL leftover snapshot workers even if they were reparented.""" + try: + import psutil + except ImportError: + return + + procs = [snap.proc for snap in snaps if _same_engine_worker(snap)] + print( + f"[sleep_mode] engine worker snapshot pids={[snap.pid for snap in snaps]} still_alive={[p.pid for p in procs]}" ) + if not procs: + return + # shutdown() already tried a graceful exit. SIGTERM + 5s here just + # delays SIGKILL on CUDA workers that ignore it. + for proc in procs: + try: + if not proc.is_running(): + continue + proc.kill() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + psutil.wait_procs(procs, timeout=3) + + +async def _shutdown_engine_and_clear_gpu(engine: AsyncOmni) -> None: + """Shut down the engine, reap leftover workers, then run shared cleanup.""" + from tests.helpers.clean import cleanup_test_environment + + leftover_workers = _snapshot_engine_workers() + engine.shutdown() + _reap_engine_workers(leftover_workers) + cleanup_test_environment() @pytest.fixture(scope="module", autouse=True) def _module_device_cleanup(): - """One device cleanup pass around the module (engines are class-scoped).""" from tests.helpers.clean import cleanup_test_environment print("\n=== PRE-MODULE DEVICE CLEANUP (sleep_mode) ===") cleanup_test_environment() yield - print("\n=== POST-MODULE DEVICE CLEANUP (sleep_mode) ===") - cleanup_test_environment() + # Fixture / test teardown already ran cleanup_test_environment(); a second + # POST wait just repeats the same 5%×60s on leftover worker VRAM. + + +# --------------------------------------------------------------------------- +# 1) AR protocol — Omni thinker-only (L4) +# --------------------------------------------------------------------------- @pytest_asyncio.fixture(scope="class", loop_scope="class") -async def llm_engine(): - """Shared BAGEL Thinker + DiT engine for LLM sleep/wake protocol tests.""" - if current_omni_platform.is_rocm(): - clean_device_envs() - engine = AsyncOmni(model=MODEL, deploy_config=_sleep_deploy_config(), init_timeout=600, enable_sleep_mode=True) +async def ar_engine(): + """Shared thinker-only AR engine for #4473 protocol checks on L4.""" + engine = AsyncOmni( + model=MODEL_AR, + deploy_config=AR_STAGE_CONFIG, + enable_sleep_mode=True, + stage_init_timeout=1200, + ) yield engine - engine.shutdown() - await asyncio.sleep(1.5) + await _shutdown_engine_and_clear_gpu(engine) + + +class TestOmniArSleepMode: + """AR sleep protocol on Qwen2.5-Omni thinker-only.""" + + @pytest.mark.advanced_model + @pytest.mark.omni + @pytest.mark.asyncio(loop_scope="class") + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_llm_sleep_ack(self, ar_engine: AsyncOmni): + """AR sleep reports EngineCore SUCCESS; generate works after a full wake. + + ``AsyncOmni.sleep()`` synthesizes OmniACK with ``freed_bytes=0`` (no + worker handshake), so VRAM / ``freed_bytes`` is not a real contract. + """ + try: + acks = await ar_engine.sleep(stage_ids=[0], level=1) + assert acks + assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) + for ack in acks: + meta = get_ack_info(ack, "metadata") or {} + assert meta.get("path") == "engine_core", f"expected EngineCore ACK, got metadata={meta}" + + await ar_engine.wake_up(stage_ids=[0]) + await ar_engine.resume_generation(stage_ids=[0]) + output = None + async for item in ar_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): + output = item + assert output is not None, "generate after full wake produced no output" + finally: + await _ensure_awake(ar_engine, [0]) + + @pytest.mark.advanced_model + @pytest.mark.omni + @pytest.mark.asyncio(loop_scope="class") + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_partial_wake_blocks_generate(self, ar_engine: AsyncOmni): + """#4473 Repro B: generate() rejected while kv_cache stays asleep.""" + try: + await ar_engine.sleep(stage_ids=[0], level=1) + await ar_engine.wake_up(stage_ids=[0], tags=["weights"]) + await ar_engine.resume_generation(stage_ids=[0]) + with pytest.raises(RuntimeError, match="partially or fully asleep"): + async for _ in ar_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): + pass + finally: + await _ensure_awake(ar_engine, [0]) + + @pytest.mark.advanced_model + @pytest.mark.omni + @pytest.mark.asyncio(loop_scope="class") + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_duplicate_wake_is_idempotent(self, ar_engine: AsyncOmni): + """#4473 Repro C: duplicate wake_up(tags=None) is a safe no-op.""" + try: + await ar_engine.sleep(stage_ids=[0], level=1) + first_acks = await ar_engine.wake_up(stage_ids=[0]) + assert len(first_acks) > 0, "First wake_up() should return ACKs" + second_acks = await ar_engine.wake_up(stage_ids=[0]) + assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}" + finally: + await _ensure_awake(ar_engine, [0]) + + +# --------------------------------------------------------------------------- +# 2) Diffusion sleep/wake/generate — qwen_image_random (L4) +# --------------------------------------------------------------------------- @pytest_asyncio.fixture(scope="class", loop_scope="class") async def diffusion_engine(): - """Shared BAGEL diffusion TP=2 engine for sleep/wake + generate checks.""" - if current_omni_platform.is_rocm(): - clean_device_envs() + """Shared tiny diffusion engine on L4.""" engine = AsyncOmni( - model=MODEL, - deploy_config=_sleep_deploy_config(diffusion_only=True, tp_size=2), + model=MODEL_DIFF, + enable_sleep_mode=True, + tensor_parallel_size=1, + enforce_eager=True, + dtype="bfloat16", + gpu_memory_utilization=0.5, + stage_init_timeout=1200, + ) + yield engine + await _shutdown_engine_and_clear_gpu(engine) + + +class TestOmniDiffusionSleepMode: + """Diffusion worker sleep/wake on ``qwen_image_random`` (TP=1).""" + + @pytest.mark.advanced_model + @pytest.mark.omni + @pytest.mark.asyncio(loop_scope="class") + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_diffusion_sleep_handshake(self, diffusion_engine: AsyncOmni): + try: + acks = await diffusion_engine.sleep(level=1) + assert acks is not None + assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) + await diffusion_engine.wake_up() + finally: + await _ensure_awake(diffusion_engine) + + @pytest.mark.omni + @pytest.mark.core_model + @pytest.mark.asyncio(loop_scope="class") + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_diffusion_sleep_wake_generate(self, diffusion_engine: AsyncOmni): + try: + acks = await diffusion_engine.sleep(level=1) + assert acks is not None + assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) + await diffusion_engine.wake_up() + await diffusion_engine.resume_generation() + output = None + async for item in diffusion_engine.generate( + "test", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), + ): + output = item + assert output is not None, "generate after sleep/wake produced no output" + assert output.images and output.images[0] is not None, "generate after sleep/wake produced no image" + finally: + await _ensure_awake(diffusion_engine) + + +# --------------------------------------------------------------------------- +# 3) Light multistage — small AR + small DiT (L4×2) +# --------------------------------------------------------------------------- + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) +@pytest.mark.asyncio +async def test_multistage_ar_diffusion_sleep_wake(): + """Orchestration: generate → joint sleep/wake/resume → generate. + + Covers the L4 path that ``TestBagelCoordinatedSleepMode`` (still skipped) + and BAGEL TP=2 (diffusion-only, ``full_model``) do not: a 2-step tiny-DiT + generate after ``sleep(stage_ids=[0, 1])`` + ``resume_generation``. + """ + engine = AsyncOmni( + model=MODEL_AR, + deploy_config=_tiny_dit_multistage_deploy_config(), + enable_sleep_mode=True, + stage_init_timeout=1200, + ) + try: + await _tiny_dit_multistage_generate(engine, "warmup", "warmup") + + acks = await engine.sleep(stage_ids=[0, 1], level=1) + assert len(acks) == 2 + assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) + + await engine.wake_up(stage_ids=[0, 1]) + await engine.resume_generation(stage_ids=[0, 1]) + await _tiny_dit_multistage_generate(engine, "verify", "verify") + logger.info("Light multistage joint sleep/wake/resume generate OK") + finally: + await _shutdown_engine_and_clear_gpu(engine) + + +# --------------------------------------------------------------------------- +# 4) BAGEL — H100/B200 ``full_model`` (nightly Entrypoints Test) +# --------------------------------------------------------------------------- + + +def _get_device_global_memory_used_gib(device_id: int) -> float: + """GPU-wide memory in use (GiB), including all processes (driver view). + + Fail closed: a swallowed query that returned 0.0 used to inflate + ``drop_gib`` and false-pass VRAM assertions. + """ + with current_omni_platform.device(device_id): + current_omni_platform.synchronize() + free_b, total_b = current_omni_platform.get_device_memory() + return (total_b - free_b) / 1024**3 + + +def _bagel_diffusion_deploy_config() -> str: + """Official single-stage BAGEL DiT, TP=2 on cards 0,1. + + ``AsyncOmni(stages=...)`` is ignored and leaked as an unknown diffusion + field (``Unknown diffusion config field(s) for stage 1: 'stages'``). + """ + return modify_stage_config( + get_deploy_config_path("bagel_single_stage.yaml"), + updates={ + "stages": { + 0: { + "devices": "0,1", + "tensor_parallel_size": 2, + "enable_sleep_mode": True, + "enforce_eager": True, + "dtype": "bfloat16", + "gpu_memory_utilization": 0.4, + } + } + }, + ) + + +@pytest_asyncio.fixture(scope="class", loop_scope="class") +async def bagel_diffusion_engine(): + """Shared BAGEL single-stage DiT TP=2 engine for sleep/wake + generate.""" + engine = AsyncOmni( + model=MODEL_BAGEL, + deploy_config=_bagel_diffusion_deploy_config(), init_timeout=600, enable_sleep_mode=True, ) yield engine - engine.shutdown() - await asyncio.sleep(1.5) + await _shutdown_engine_and_clear_gpu(engine) -class TestOmniLlmSleepMode: - """LLM sleep/wake protocol tests sharing one class-scoped engine.""" +@pytest.mark.full_model +@pytest.mark.omni +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) +class TestBagelDiffusionSleepMode: + """BAGEL diffusion sleep/wake on a class-scoped TP=2 BagelPipeline.""" @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_llm_sleep_ack(self, llm_engine: AsyncOmni): - """LLM Thinker (GPU0) Signal and Physical Recycling Audit""" - device_id = 0 + async def test_diffusion_sleep_handshake(self, bagel_diffusion_engine: AsyncOmni): try: - used_before = get_device_global_memory_used_gib(device_id) - acks = await llm_engine.sleep(stage_ids=[0], level=1) - await asyncio.sleep(1.5) - used_after = get_device_global_memory_used_gib(device_id) - drop_gib = used_before - used_after + acks = await bagel_diffusion_engine.sleep(stage_ids=[0], level=1) + assert len(acks) >= 1 assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) - total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) - freed_gib = total_freed_bytes / 1024**3 - logger.info( - "Thinker: ACK freed=%.2f GiB, global GPU used drop=%.2f GiB (before=%.2f, after=%.2f)", - freed_gib, - drop_gib, - used_before, - used_after, - ) - assert freed_gib > 5.0 or drop_gib > 3.0, ( - "Expected either ACK freed_bytes or global VRAM drop after sleep. " - f"ACK={freed_gib:.2f} GiB, global_drop={drop_gib:.2f} GiB" - ) + await bagel_diffusion_engine.wake_up(stage_ids=[0]) finally: - await _ensure_awake(llm_engine, [0]) + await _ensure_awake(bagel_diffusion_engine, [0]) @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni): - """Regression for #4473 Repro B: generate() must be rejected if kv_cache - is still asleep after wake_up(tags=["weights"]), instead of crashing with - CUDA illegal memory access.""" + async def test_cross_device_cleanup(self, bagel_diffusion_engine: AsyncOmni): try: - await llm_engine.sleep(stage_ids=[0], level=1) - await llm_engine.wake_up(stage_ids=[0], tags=["weights"]) - # wake_up restores weights but keeps the AR admission hold. - # Resume so generate() can reach the sleeping-tags rejection. - await llm_engine.resume_generation(stage_ids=[0]) - with pytest.raises(RuntimeError, match="partially or fully asleep"): - async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): - pass - finally: - await _ensure_awake(llm_engine, [0]) - - @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni): - """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a - safe no-op instead of raising a cumem CUDA invalid argument error.""" - try: - await llm_engine.sleep(stage_ids=[0], level=1) - first_acks = await llm_engine.wake_up(stage_ids=[0]) - assert len(first_acks) > 0, "First wake_up() should return ACKs" - second_acks = await llm_engine.wake_up(stage_ids=[0]) - assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}" - finally: - await _ensure_awake(llm_engine, [0]) - - -class TestOmniDiffusionSleepMode: - """Diffusion sleep/wake tests sharing one class-scoped TP=2 engine.""" - - @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_diffusion_sleep_handshake(self, diffusion_engine: AsyncOmni): - """Diffusion Worker stage signal loop""" - try: - logger.info("Starting Diffusion Worker Handshake Test") - acks = await diffusion_engine.sleep(stage_ids=[0], level=1) - - def _get_status(ack): - return ack.status if hasattr(ack, "status") else ack.get("status") - - assert len(acks) >= 1, "Expected at least 1 ACK from Diffusion Workers" - assert all(_get_status(ack) == "SUCCESS" for ack in acks) - logger.info(f"Success: Received {len(acks)} Diffusion Worker ACKs") - await diffusion_engine.wake_up(stage_ids=[0]) - finally: - await _ensure_awake(diffusion_engine, [0]) - - @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_cross_device_cleanup(self, diffusion_engine: AsyncOmni): - """Physical recycling audit: leveraging deterministic data returned by Workers""" - try: - used_before = get_device_global_memory_used_gib(0) + get_device_global_memory_used_gib(1) - acks = await diffusion_engine.sleep(stage_ids=[0], level=1) + used_before = _get_device_global_memory_used_gib(0) + _get_device_global_memory_used_gib(1) + acks = await bagel_diffusion_engine.sleep(stage_ids=[0], level=1) await asyncio.sleep(1.5) - used_after = get_device_global_memory_used_gib(0) + get_device_global_memory_used_gib(1) + used_after = _get_device_global_memory_used_gib(0) + _get_device_global_memory_used_gib(1) drop_gib = used_before - used_after - total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) - freed_gb = total_freed_bytes / 1024**3 - logger.info("Physical reclamation summary from workers:") - logger.info(f"- Total Workers: {len(acks)}") - logger.info(f"- Total Freed (ACK): {freed_gb:.2f} GiB, global used drop: {drop_gib:.2f} GiB") - assert freed_gb > 14.0 or drop_gib > 8.0, ( - "Expected either ACK freed_bytes or global VRAM drop on GPUs 0+1. " - f"ACK={freed_gb:.2f} GiB, global_drop={drop_gib:.2f} GiB" - ) - logger.info("SUCCESS: 100% weights offloaded.") + freed_gb = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) / 1024**3 + assert freed_gb > 14.0 or drop_gib > 8.0, f"ACK={freed_gb:.2f} GiB, global_drop={drop_gib:.2f} GiB" finally: - await _ensure_awake(diffusion_engine, [0]) + await _ensure_awake(bagel_diffusion_engine, [0]) @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_diffusion_sleep_wake_generate(self, diffusion_engine: AsyncOmni): - """Merged diffusion lifecycle with test points from the former trio: - - - ``test_diffusion_integrity_bit_level``: 4-step/512 baseline generate → - sleep → wake → post generate (images present, count unchanged). - - ``test_diffusion_vram_lifecycle_audit``: ACK/process VRAM reclaim on - sleep; after wake, process VRAM restores to within 3 GiB of initial; - post-wake generate smoke. - - ``test_diffusion_model_sleep_tp``: sleep ACK SUCCESS → wake → verify - generate (same Active→Sleep→Active→Generate path; runs on the shared - class TP=2 BagelPipeline engine rather than a second default-config load). - """ + async def test_diffusion_sleep_wake_generate(self, bagel_diffusion_engine: AsyncOmni): import gc device_id = 1 try: - # --- integrity: baseline generate at original resolution/steps --- prompt = "A huge swimming pool, with many people swimming." sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42) - logger.info("Running Baseline Generation...") base_output = None - async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params=sp): + async for output in bagel_diffusion_engine.generate(prompt, request_id="base", sampling_params=sp): base_output = output assert base_output is not None and len(base_output.images) > 0 - logger.info("Baseline Generation successful.") - # Snapshot Active VRAM after generate (lifecycle compared Sleep/Wake - # against the pre-sleep active watermark). - get_vram_info(device_id) - torch.accelerator.empty_cache() - vram_initial = get_vram_info(device_id)["reserved"] - logger.info(f"Diffusion Active VRAM before sleep: {vram_initial:.2f} GiB") + current_omni_platform.empty_cache() + vram_initial = _get_device_global_memory_used_gib(device_id) - # --- lifecycle + model_sleep_tp: level-1 sleep reclaim --- - logger.info("Entering Deep Sleep (VRAM Scavenging)...") - acks = await diffusion_engine.sleep(stage_ids=[0], level=1) + acks = await bagel_diffusion_engine.sleep(stage_ids=[0], level=1) statuses = [get_ack_info(ack, "status") for ack in acks] assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}" - reported_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) - reported_freed_gib = reported_freed_bytes / 1024**3 - logger.info(f"Worker internally reported freed: {reported_freed_gib:.2f} GiB") - + reported_freed_gib = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) / 1024**3 await asyncio.sleep(2) - get_vram_info(device_id) - torch.accelerator.empty_cache() - vram_sleeping = get_vram_info(device_id)["reserved"] - logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB") + current_omni_platform.empty_cache() + vram_sleeping = _get_device_global_memory_used_gib(device_id) assert reported_freed_gib > 14.0 or vram_sleeping < 5.0, ( - f"Reclamation failed. Reported: {reported_freed_gib:.2f}G, Measured: {vram_sleeping:.2f}G" + f"Reported: {reported_freed_gib:.2f}G, Measured: {vram_sleeping:.2f}G" ) - # --- wake + VRAM restore to initial (±3 GiB) --- - logger.info("Waking up (Reloading Weights)...") - await diffusion_engine.wake_up(stage_ids=[0]) - # AR stages keep the frontend admission gate held after wake. - await diffusion_engine.resume_generation(stage_ids=[0]) + await bagel_diffusion_engine.wake_up(stage_ids=[0]) + await bagel_diffusion_engine.resume_generation(stage_ids=[0]) await asyncio.sleep(2.0) gc.collect() - get_vram_info(device_id) - torch.accelerator.empty_cache() - vram_restored = get_vram_info(device_id)["reserved"] - logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB") + current_omni_platform.empty_cache() + vram_restored = _get_device_global_memory_used_gib(device_id) assert abs(vram_restored - vram_initial) < 3.0, "VRAM failed to restore to initial levels" - # --- integrity + lifecycle: post-wake generate --- - logger.info("Running Post-Wakeup Generation...") post_output = None - async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params=sp): + async for output in bagel_diffusion_engine.generate(prompt, request_id="post", sampling_params=sp): post_output = output assert post_output is not None assert len(base_output.images) == len(post_output.images) assert post_output.images[0] is not None - logger.info("SUCCESS: Diffusion integrity + VRAM lifecycle + sleep/wake generate verified.") finally: - await _ensure_awake(diffusion_engine, [0]) + await _ensure_awake(bagel_diffusion_engine, [0]) -class TestOmniCoordinatedSleepMode: +def _build_bagel_llm_stages() -> tuple[list[dict], list[dict]]: + common_args = { + "worker_type": "ar", + "enable_sleep_mode": True, + "dtype": "bfloat16", + "trust_remote_code": True, + "max_model_len": 2048, + "max_num_batched_tokens": 8192, + "enforce_eager": True, + } + stages = [ + { + "stage_id": 0, + "stage_type": "llm", + "runtime": {"process": True, "devices": "0", "max_batch_size": 1}, + "engine_args": {**common_args, "model_stage": "thinker", "gpu_memory_utilization": 0.1}, + }, + { + "stage_id": 1, + "stage_type": "llm", + "engine_input_source": [0], + "runtime": {"process": True, "devices": "1", "max_batch_size": 1, "connector_type": "queue"}, + "engine_args": {**common_args, "model_stage": "talker", "gpu_memory_utilization": 0.1}, + }, + ] + connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}] + return stages, connectors + + +@pytest.mark.full_model +@pytest.mark.omni +@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2) +class TestBagelCoordinatedSleepMode: """Dual-engine coordination (kept skipped; do not delete).""" @pytest.mark.skip( reason=( "Flaky/CI: dual AsyncOmni can fail with " "RuntimeError: Orchestrator init failed, StageDiffusionProc died during handshake. " - "Re-enable when stable (no OOM on coordinated talker+diffusion). " - "Note: this test is fixture-free so skip does not init llm/diffusion engines." + "Re-enable when stable (no OOM on coordinated talker+diffusion)." ) ) @pytest.mark.asyncio(loop_scope="class") - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - async def test_coordinated_cross_device(self, llm_engine: AsyncOmni, diffusion_engine: AsyncOmni): - """Heterogeneous Coordinated Cleanup Test (Talker and Diffusion on GPU 1)""" + async def test_coordinated_cross_device(self): + """Heterogeneous coordinated cleanup (talker + diffusion on GPU 1).""" + llm_stages, llm_connectors = _build_bagel_llm_stages() + llm_engine = AsyncOmni( + model=MODEL_BAGEL, stages=llm_stages, connectors=llm_connectors, init_timeout=600, enable_sleep_mode=True + ) + diffusion_engine = AsyncOmni( + model=MODEL_BAGEL, + deploy_config=_bagel_diffusion_deploy_config(), + init_timeout=600, + enable_sleep_mode=True, + ) device_id = 1 try: - logger.info(f"Waking up both engines on GPU {device_id}...") await llm_engine.wake_up(stage_ids=[1]) await diffusion_engine.wake_up(stage_ids=[0]) - - get_vram_info(device_id) - torch.accelerator.empty_cache() + current_omni_platform.empty_cache() await asyncio.sleep(2) + initial_vram = _get_device_global_memory_used_gib(device_id) - initial_vram = get_vram_info(device_id)["reserved"] - logger.info(f"GPU {device_id} Peak Pressure: {initial_vram:.2f} GiB") - - # coordinated sleep - logger.info("Issuing concurrent SLEEP commands...") await llm_engine.sleep(stage_ids=[1], level=2) await asyncio.sleep(1.0) await diffusion_engine.sleep(stage_ids=[0], level=2) - await asyncio.sleep(3.0) - torch.accelerator.empty_cache() - - final_vram = get_vram_info(device_id)["reserved"] - logger.info(f"GPU {device_id} Final VRAM after coordinated sleep: {final_vram:.2f} GiB") - + current_omni_platform.empty_cache() + final_vram = _get_device_global_memory_used_gib(device_id) assert initial_vram - final_vram > 15.0 or final_vram < 8.0 - logger.info(f"SUCCESS: Heterogeneous VRAM drop verified on GPU {device_id}.") - except Exception as e: - logger.error(f"Coordinated test failed: {e}") - raise e finally: - logger.info("Triggering mandatory cleanup for both engines...") - llm_engine.shutdown() - diffusion_engine.shutdown() - logger.info("All engines scavenged. Ready for next test.") - - -# --------------------------------------------------------------------------- -# Multi-stage (H100) / pure diffusion TP=1 (2×L4): sleep→wake→generate -# (BAGEL TP=2 diffusion path covered by TestOmniDiffusionSleepMode) -# --------------------------------------------------------------------------- - - -@pytest.mark.omni -@pytest.mark.advanced_model -@pytest.mark.parametrize("tp_size", [1, 2]) -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.asyncio -async def test_multistage_llm_diffusion_sleep_wake(tp_size: int): - """Registered BAGEL Thinker + DiT topology; sleep/wake both stages.""" - if current_omni_platform.is_rocm(): - clean_device_envs() - num_gpus = torch.accelerator.device_count() - if num_gpus < tp_size * 2: - pytest.skip("Not enough GPUs") - - engine = AsyncOmni( - model=MODEL, - deploy_config=_sleep_deploy_config(tp_size=tp_size), - enable_sleep_mode=True, - stage_init_timeout=1200, - ) - try: - sp = OmniDiffusionSamplingParams(num_inference_steps=2) - async for _ in engine.generate("warmup", sampling_params_list=[engine.default_sampling_params_list[0], sp]): - pass - - acks = await engine.sleep(stage_ids=[0, 1], level=1) - assert len(acks) == 2 - - await engine.wake_up(stage_ids=[0, 1]) - await engine.resume_generation(stage_ids=[0, 1]) - async for _ in engine.generate("verify", sampling_params_list=[engine.default_sampling_params_list[0], sp]): - pass - finally: - engine.shutdown() - - -@pytest.mark.omni -@pytest.mark.core_model -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) -@pytest.mark.asyncio -async def test_pure_diffusion_sleep_wake(): - """Single-stage random diffusion (TP=1): sleep → wake → generate. - - Diffusion TP=2 sleep/wake is covered by TestOmniDiffusionSleepMode - (shared BAGEL TP=2 engine); this case keeps the default AsyncOmni + - qwen_image_random entrypoint path on L4. - """ - if current_omni_platform.is_rocm(): - clean_device_envs() - engine_args = { - "model": MODEL_DIFF, - "enable_sleep_mode": True, - "tensor_parallel_size": 1, - "enforce_eager": True, - "dtype": "bfloat16", - "gpu_memory_utilization": 0.5, - } - - engine = AsyncOmni(**engine_args, stage_init_timeout=1200) - try: - await engine.sleep(level=1) - await engine.wake_up() - await engine.resume_generation() - async for _ in engine.generate( - "test", - sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), - ): - pass - logger.info("Pure diffusion sleep/wake OK") - finally: - engine.shutdown() + await _shutdown_engine_and_clear_gpu(llm_engine) + await _shutdown_engine_and_clear_gpu(diffusion_engine) diff --git a/tests/helpers/clean.py b/tests/helpers/clean.py index 5cd3f5b39..39e24f10d 100644 --- a/tests/helpers/clean.py +++ b/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() diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 2a12a6038..18da70fa4 100644 --- a/tests/helpers/media.py +++ b/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", ] diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index c5ba1ada2..2a6addd93 100644 --- a/tests/helpers/runtime.py +++ b/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 diff --git a/tests/helpers/tests/test_clean.py b/tests/helpers/tests/test_clean.py new file mode 100644 index 000000000..0eb9198ab --- /dev/null +++ b/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] diff --git a/tests/helpers/tests/test_media.py b/tests/helpers/tests/test_media.py index 4683dae12..228251edc 100644 --- a/tests/helpers/tests/test_media.py +++ b/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") diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index a298baae5..9a422ee9e 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/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", diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index b7c0a1838..7c6cf55c6 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 7ab4dd084..d0de42902 100644 --- a/vllm_omni/platforms/interface.py +++ b/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: diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index d667c9bf1..3ea177349 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/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: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 2c8d42f60..24e29f342 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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) diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 4eaf8c347..d29e90437 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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: diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index cedbf52ad..73d8324f3 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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."""