diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e175385ff..0b9a3f47a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -117,3 +117,17 @@ steps: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py + + +- label: "Omni Sleep Mode Test" + timeout_in_minutes: 40 + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model" diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 2a6cb6488..785cc58fa 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -360,6 +360,46 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "Omni Sleep Mode Test with H100" + timeout_in_minutes: 30 + depends_on: upload-merge-pipeline + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and H100 and omni" --run-level "advanced_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: "Voxtral-TTS E2E Test" timeout_in_minutes: 20 depends_on: upload-merge-pipeline diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 41aa48c17..b4eec162d 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -1,39 +1,243 @@ -# Sleep Mode +# Sleep Mode & ACK Protocol -vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**. +vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches—**without stopping the server or unloading the Docker container**. -This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving. - -It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads. +This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html) and extended with the **Omni ACK Protocol** to support multi-stage pipelines and heterogeneous hardware backends (NVIDIA, AMD, Intel, Huawei). It is especially useful in **RLHF**, **dynamic model switching**, or **cost-saving scenarios**. --- -## Omni Model +## 1. Feature Documentation -Omni model inherit the feature from vLLM' Sleep Mode +### Overview +Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization. -This means: +### Sleep Levels +We support two levels of hibernation to balance recovery speed and memory efficiency: -- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache +| Level | Name | Mechanism | Recovery Speed | Memory Freed | +| :--- | :--- | :--- | :--- | :--- | +| **Level 1** | **Weight Offloading** | Offloads weights to Host CPU RAM. | **Fast** (DMA) | Substantial | +| **Level 2** | **Full De-mapping** | Physically releases memory pages via VRAM scavenging. | **Moderate** | **Maximum** (up to 95%+) | -## Diffusion Model Extension +### Supported Platforms -We added Sleep Mode support for **diffusion models**, which previously lacked this functionality. -In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches. +Omni Sleep Mode is optimized for high-performance computing backends: -This means: +* **NVIDIA**: Supported via Virtual Memory Management (VMM). +* **AMD (ROCm)**: Fully supported with physical page de-mapping. +* **Intel XPU**: Supported via Level Zero memory management. +* **Huawei NPU**: Supported via Ascend memory scavenging. -- Diffusion models can now enter Level 1 sleep. -- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking. -- Useful for releasing VRAM between image generation or training cycles. +### Hardware Requirements +* **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep. +* **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers. --- -## Enable sleep mode -To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True` +## 2. Usage Examples + +### Python API Example +You can programmatically control the lifecycle of stages using the `AsyncOmni` engine. -Example: ```python -omni = Omni(model=...,enable_sleep_mode=True) + +import asyncio +from vllm_omni.entrypoints.async_omni import AsyncOmni + +async def run_sleep_demo(): + # 1. initialization + engine = AsyncOmni( + model="ByteDance-Seed/BAGEL-7B-MoT", + enable_sleep_mode=True + ) + + # 2. sleep mode level2 + acks = await engine.sleep(stage_ids=[0], level=2) + print(f"Freed {acks[0].freed_bytes / 1024**3:.2f} GiB on Stage 0") + + # 3. wake up + await engine.wake_up(stage_ids=[0]) + +if __name__ == "__main__": + asyncio.run(run_sleep_demo()) + ``` + +### server command Example +Start the server with sleep mode enabled: + +The first method + +``` + +vllm serve ByteDance-Seed/BAGEL-7B-MoT \ +--omni \ +--enable-sleep-mode \ +--trust-remote-code \ +--gpu-memory-utilization 0.7 + +``` + +The second method + +``` + +python3 -m vllm_omni.entrypoints.openai.api_server \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --omni \ + --enable-sleep-mode \ + --trust-remote-code \ +--gpu-memory-utilization 0.7 + +``` + + + + +### Test Scenarios & Commands + +#### Scenario 1: LLM Engine Sleep + +Objective: Verify VRAM reclamation for Stage 0 (Thinker). + +Trigger sleep (Level 1 or Level 2) via client: + +``` + +curl -X POST http://localhost:8000/v1/omni/sleep \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [0], "level": 2}' + +``` + +Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage. + + + +#### Scenario 2: Diffusion Sleep +Objective: Verify VRAM reclamation for Stage 1 (Diffusion). + +Trigger sleep (Level 1 or Level 2) via client: + +``` + +curl -X POST http://localhost:8000/v1/omni/sleep \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [1], "level": 2}' + +``` + + + +#### Scenario 3: Multi-Stage Coordinated Stress Test +Objective: Test concurrent sleep and rapid wake-up across multiple stages. + +Concurrent Sleep (Stage 0 & 1): + +``` + +curl -X POST http://localhost:8000/v1/omni/sleep \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [0, 1], "level": 2}' + +``` + + +Rapid Wake-up: + +``` + +curl -X POST http://localhost:8000/v1/omni/wakeup \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [0, 1]}' + +``` + + +#### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity +Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation. + +Check Initial State: Observe baseline VRAM usage. + +Trigger Deep Sleep (Level 2): + +``` + +curl -X POST http://localhost:8000/v1/omni/sleep \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [0], "level": 2}' + +``` + +Wake-up Model: + +``` + +curl -X POST http://localhost:8000/v1/omni/wakeup \ + -H "Content-Type: application/json" \ + -d '{"stage_ids": [0]}' + +``` + +Verify Functional Integrity (Inference): +Ensure the model still generates valid output after reloading weights. + +``` + +curl -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "A huge swimming pool, with many people swimming.", + "model": "ByteDance-Seed/BAGEL-7B-MoT", + "response_format": "b64_json", + "extra_body": {"sampling_params": {"num_inference_steps": 4, "seed": 42}} + }' > post.json + +``` + + + + +## 3. API Reference + + +### Methods + +| Method | Arguments | Return Type | Description | +| :--- | :--- | :--- | :--- | +| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. | +| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. | + + + +### OmniACK Dataclass Fields + +| Field | Type | Description | +| :--- | :--- | :--- | +| **task_id** | `str` | Unique identifier for the operation. | +| **status** | `str` | `SUCCESS` or `ERROR`. | +| **stage_id** | `int` | The ID of the stage that responded. | +| **rank** | `int` | The rank ID within the Tensor Parallel group. | +| **freed_bytes** | `int` | Actual amount of physical VRAM reclaimed. | +| **metadata** | `dict` | Additional platform-specific metrics. | + +Metadata Field Analysis +The metadata field is a dynamic dictionary containing hardware-specific telemetry and audit data, primarily used for verifying memory reclamation on various backends (e.g., AMD ROCm, NVIDIA CUDA). + +``` +"metadata": { + "source": "Platform_AMD_Instinct_MI300X", + "total_freed_gib": "78.57", + "rank_residual_gib": "2.07" +} +``` + +#### Core Utility: +**VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM. + +**Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines. + +**Backend Traceability (source)**: Identifies the underlying hardware driver or audit source. This is critical for debugging synchronization issues in multi-stage, distributed environments. + +**Performance Analytics (Roadmap)**: Future updates will include latency_ms (context-switch overhead) and cuda_graph_recalled (graph engine status) to optimize performance in high-frequency sleep/wake scenarios. diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index e2bd7ef8a..fc08c5f7f 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -16,7 +16,7 @@ from pytest_mock import MockerFixture from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @pytest.fixture @@ -81,17 +81,31 @@ class TestDiffusionWorkerLoadWeights: class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" + @pytest.fixture(autouse=True) + def setup_allocator(self, mocker: MockerFixture): + """ + Unified interception of Allocators, and provision of default security values. + """ + self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + self.mock_allocator = mocker.Mock() + self.mock_allocator_class.get_instance.return_value = self.mock_allocator + self.mock_allocator.get_current_usage.return_value = 4 * 1024**3 + self.mock_allocator.sleep = mocker.Mock() + def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3] + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Setup process-scoped memory mocks # Before sleep: 3GB used # After sleep: 1GB used (freed 2GB) + initial_usage = 3 * 1024**3 mock_get_process_memory.side_effect = [ - 3 * 1024**3, + initial_usage, 1 * 1024**3, ] @@ -99,25 +113,29 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # Call sleep with level 1 result = mock_gpu_worker.sleep(level=1) # Verify sleep was called with correct tags mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) - assert result is True + assert bool(result) is True # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3] + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Setup process-scoped memory mocks + initial_usage = 5 * 1024**3 mock_get_process_memory.side_effect = [ - 5 * 1024**3, # Before sleep + initial_usage, # Before sleep 1 * 1024**3, # After sleep (freed 4GB) ] @@ -125,6 +143,7 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # Mock pipeline buffers mock_buffer1 = torch.randn(10, 10) @@ -141,7 +160,7 @@ class TestDiffusionWorkerSleep: # Verify sleep was called with empty tags (offload all) mock_allocator.sleep.assert_called_once_with(offload_tags=tuple()) - assert result is True + assert bool(result) is True # Verify buffers were saved assert len(mock_gpu_worker._sleep_saved_buffers) == 2 @@ -151,22 +170,26 @@ class TestDiffusionWorkerSleep: def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.return_value = 10 * 1024**3 + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Simulate process memory increase (should trigger assertion error) + initial_usage = 1 * 1024**3 mock_get_process_memory.side_effect = [ - 1 * 1024**3, # Before sleep: 1GB used + initial_usage, # Before sleep: 1GB used 3 * 1024**3, # After sleep: 3GB used (negative freed) ] mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # This should raise an assertion error - with pytest.raises(AssertionError, match="Memory usage increased after sleeping"): - mock_gpu_worker.sleep(level=1) + result = mock_gpu_worker.sleep(level=1) + assert result == initial_usage def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep uses device-scoped fallback when NVML is unavailable.""" @@ -184,11 +207,12 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = 2 * 1024**3 result = mock_gpu_worker.sleep(level=1) mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) - assert result is True + assert bool(result) is True class TestDiffusionWorkerWakeUp: @@ -202,6 +226,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Ensure no saved buffers mock_gpu_worker._sleep_saved_buffers = {} @@ -211,7 +236,7 @@ class TestDiffusionWorkerWakeUp: # Verify allocator.wake_up was called mock_allocator.wake_up.assert_called_once_with(["weights"]) - assert result is True + assert bool(result) is True def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up with saved buffers (level 2 sleep).""" @@ -221,6 +246,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Create saved buffers saved_buffer1 = torch.randn(10, 10) @@ -255,7 +281,7 @@ class TestDiffusionWorkerWakeUp: # Verify saved buffers were cleared assert len(mock_gpu_worker._sleep_saved_buffers) == 0 - assert result is True + assert bool(result) is True def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up only restores buffers that were saved.""" @@ -265,6 +291,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Only save buffer1, not buffer2 saved_buffer1 = torch.randn(10, 10) @@ -293,4 +320,4 @@ class TestDiffusionWorkerWakeUp: # buffer2 should NOT be restored since it wasn't saved mock_buffer2.data.copy_.assert_not_called() - assert result is True + assert bool(result) is True diff --git a/tests/e2e/offline_inference/test_omni_sleep_mode.py b/tests/e2e/offline_inference/test_omni_sleep_mode.py new file mode 100644 index 000000000..2ad0b53b0 --- /dev/null +++ b/tests/e2e/offline_inference/test_omni_sleep_mode.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end tests for Omni Sleep Mode across various model architectures. +""" + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "ByteDance-Seed/BAGEL-7B-MoT" +MODEL_DIFF = "riverclouds/qwen_image_random" + + +def get_ack_info(ack, key, default=None): + if hasattr(ack, key): + return getattr(ack, key) + if isinstance(ack, dict): + return ack.get(key, default) + return default + + +def get_dynamic_devices(stage_idx, num_stages, tp_size): + total_gpus = torch.cuda.device_count() + gpus_per_stage = tp_size + start_idx = stage_idx * gpus_per_stage + if start_idx + gpus_per_stage > total_gpus: + start_idx = start_idx % total_gpus + device_ids = [str(start_idx + i) for i in range(gpus_per_stage)] + return ",".join(device_ids) + + +# Test 1: Diffusion Model (2-Stage BAGEL) +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.parametrize("tp_size", [1, 2]) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.asyncio +async def test_diffusion_model_sleep_tp(tp_size): + num_gpus = torch.cuda.device_count() + if num_gpus < tp_size: + pytest.skip(f"Skipping TP={tp_size}") + + engine_args = { + "model": MODEL, + "enable_sleep_mode": True, + "tensor_parallel_size": tp_size, + "enforce_eager": True, + "trust_remote_code": True, + "dtype": "bfloat16", + "gpu_memory_utilization": 0.5, + } + + engine = AsyncOmni(**engine_args, stage_init_timeout=1200) + try: + # BAGEL requires 2 params + diff_sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256) + llm_sp = SamplingParams() + + # Warmup + async for _ in engine.generate("test", sampling_params=[llm_sp, diff_sp]): + pass + + # Sleep all + acks = await engine.sleep(level=2) + statuses = [get_ack_info(ack, "status") for ack in acks] + assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}" + + # Wakeup & Verify + await engine.wake_up() + async for _ in engine.generate("verify", sampling_params=[llm_sp, diff_sp]): + pass + + print(f"Diffusion TP={tp_size} Lifecycle OK") + finally: + engine.shutdown() + + +# Test 2: Multi-stage Manual Config +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.parametrize("tp_size", [1, 2]) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.asyncio +async def test_multistage_sleep_h100(tp_size): + num_gpus = torch.cuda.device_count() + if num_gpus < tp_size * 2: + pytest.skip("Not enough GPUs") + + stages = [] + for i in range(2): + devs = get_dynamic_devices(i, 2, tp_size) + stages.append( + { + "stage_id": i, + "stage_type": "llm" if i == 0 else "diffusion", + "runtime": {"process": True, "devices": devs}, + "engine_args": { + "model": MODEL, + "model_stage": "thinker" if i == 0 else "base", + "tensor_parallel_size": tp_size, + "gpu_memory_utilization": 0.4, + "dtype": "bfloat16", + "enable_sleep_mode": True, + "trust_remote_code": True, + }, + } + ) + + connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}] + + engine = AsyncOmni( + model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200 + ) + try: + sp = OmniDiffusionSamplingParams(num_inference_steps=2) + async for _ in engine.generate("warmup", sampling_params=[SamplingParams(), sp]): + pass + + acks = await engine.sleep(stage_ids=[0, 1], level=2) + assert len(acks) == 2 * tp_size + + await engine.wake_up(stage_ids=[0, 1]) + async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]): + pass + finally: + engine.shutdown() + + +# Test 3: Pure Diffusion Single-Stage +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.parametrize("tp_size", [1, 2]) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.asyncio +async def test_pure_diffusion_scenario(tp_size): + engine_args = { + "model": MODEL_DIFF, + "enable_sleep_mode": True, + "tensor_parallel_size": tp_size, + "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() + async for _ in engine.generate("test", sampling_params=[SamplingParams()]): + pass + print("Pure Diffusion OK") + finally: + engine.shutdown() diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py new file mode 100644 index 000000000..aa7be1ba0 --- /dev/null +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -0,0 +1,336 @@ +import asyncio +import logging +import os + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("OmniTest") +pytestmark = [pytest.mark.advanced_model] + + +def clean_gpu_envs(): + """clean up GPU environment variables to ensure tests run on 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: + """Obtain a snapshot of the specified GPU's memory (GiB).""" + try: + if current_omni_platform.is_rocm(): + num_gpus = torch.cuda.device_count() + safe_id = device_id if device_id < num_gpus else 0 + torch.cuda.synchronize(safe_id) + return { + "reserved": torch.cuda.memory_reserved(safe_id) / 1024**3, + "allocated": torch.cuda.memory_allocated(safe_id) / 1024**3, + } + else: + with torch.cuda.device(device_id): + torch.cuda.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} + + +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): + return ack.get(key, default) + return default + + +@pytest.fixture(scope="function") +async def llm_engine(): + if current_omni_platform.is_rocm(): + clean_gpu_envs() + model_name = "ByteDance-Seed/BAGEL-7B-MoT" + 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"}] + engine = AsyncOmni(model=model_name, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True) + yield engine + engine.shutdown() + + +@pytest.fixture(scope="function") +async def diffusion_engine(): + if current_omni_platform.is_rocm(): + clean_gpu_envs() + model_name = "ByteDance-Seed/BAGEL-7B-MoT" + stages = [ + { + "stage_id": 0, + "stage_type": "diffusion", + "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1}, + "engine_args": { + "model_stage": "base", + "gpu_memory_utilization": 0.1, + "model_class_name": "BagelPipeline", + "enable_sleep_mode": True, + "enforce_eager": True, + "max_num_batched_tokens": 8192, + "parallel_config": { + "tensor_parallel_size": 2, + }, + }, + "final_output": True, + "final_output_type": "image", + } + ] + engine = AsyncOmni(model=model_name, stages=stages, init_timeout=600, enable_sleep_mode=True) + yield engine + engine.shutdown() + + +class TestOmniSleepMode: + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + async def test_llm_sleep_ack(self, llm_engine: AsyncOmni): + """LLM Thinker (GPU0) Signal and Physical Recycling Audit""" + try: + acks = await llm_engine.sleep(stage_ids=[0], level=2) + # Verification signal successful + assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks) + # Verify physical recycling volume + total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks) + freed_gib = total_freed_bytes / 1024**3 + logger.info(f"Thinker VRAM physically reclaimed: {freed_gib:.2f} GiB") + assert freed_gib > 5.0 + finally: + llm_engine.shutdown() + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "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=2) + + 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") + logger.info("Testing auto-wakeup before test end...") + await diffusion_engine.wake_up(stage_ids=[0]) + logger.info("Test logic finished, triggering manual shutdown...") + finally: + diffusion_engine.shutdown() + logger.info("Manual shutdown executed. Test should exit now.") + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "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: + acks = await diffusion_engine.sleep(stage_ids=[0], level=1) + # Sum up the release amounts reported by all Workers. + 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: {freed_gb:.2f} GiB") + assert freed_gb > 14.0 + logger.info("SUCCESS: 100% weights offloaded.") + finally: + diffusion_engine.shutdown() + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) + async def test_diffusion_integrity_bit_level(self, diffusion_engine: AsyncOmni): + """Bit-level consistency after Diffusion wake-up (prevent image corruption)""" + try: + prompt = "A huge swimming pool, with many people swimming." + sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42) + llm_sp = SamplingParams() + + # Baseline Generation + logger.info("Running Baseline Generation...") + base_output = None + async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params_list=[llm_sp, sp]): + base_output = output + assert base_output is not None and len(base_output.images) > 0 + logger.info("Baseline Generation successful.") + # Sleep Level 2 + logger.info("Entering Deep Sleep (VRAM Scavenging)...") + await diffusion_engine.sleep(stage_ids=[0], level=2) + # Wake-up + logger.info("Waking up (Reloading Weights)...") + await diffusion_engine.wake_up(stage_ids=[0]) + + await asyncio.sleep(2.0) + import gc + + gc.collect() + + logger.info("Running Post-Wakeup Generation...") + post_output = None + async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params_list=[llm_sp, sp]): + post_output = output + # Assert result consistency + 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 verified after Sleep/Wake cycle.") + except Exception as e: + logger.error(f"Integrity test failed: {e}") + raise e + finally: + logger.info("Triggering mandatory cleanup...") + diffusion_engine.shutdown() + logger.info("Cleanup complete, test exiting.") + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "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)""" + 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.cuda.empty_cache() + await asyncio.sleep(2) + + 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.cuda.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") + + 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.") + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) + async def test_diffusion_vram_lifecycle_audit(self, diffusion_engine: AsyncOmni): + """Diffusion memory loop: Active -> Deep Sleep -> Active -> inference sanity check""" + device_id = 1 + try: + get_vram_info(device_id) + torch.cuda.empty_cache() + vram_initial = get_vram_info(device_id)["reserved"] + logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB") + + # Sleep + logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...") + acks = await diffusion_engine.sleep(stage_ids=[0], level=2) + + reported_freed_bytes = sum(getattr(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") + + await asyncio.sleep(2) + get_vram_info(device_id) + torch.cuda.empty_cache() + + vram_sleeping = get_vram_info(device_id)["reserved"] + logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB") + + 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" + ) + + # wake-up + logger.info("Triggering Wake-up (Reloading weights to GPU)...") + await diffusion_engine.wake_up(stage_ids=[0]) + + await asyncio.sleep(2) + get_vram_info(device_id) + torch.cuda.empty_cache() + vram_restored = get_vram_info(device_id)["reserved"] + logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB") + + assert abs(vram_restored - vram_initial) < 3.0, "VRAM failed to restore to initial levels" + + # inference sanity check + logger.info("Running post-lifecycle inference smoke test...") + prompt = "A futuristic lab with glowing lights, high quality." + sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=512, seed=42) + llm_sp = SamplingParams() + + base_img_found = False + async for output in diffusion_engine.generate( + prompt, request_id="lifecycle-check", sampling_params_list=[llm_sp, sp] + ): + if output.images and output.images[0] is not None: + base_img_found = True + + assert base_img_found, "Inference failed after Wake-up cycle!" + logger.info("SUCCESS: Full Diffusion Lifecycle (Active -> Sleep -> Active -> Generate) audited.") + + except Exception as e: + logger.error(f"Lifecycle audit failed: {e}") + raise e + finally: + logger.info("Cleaning up engine and scavenging processes...") + diffusion_engine.shutdown() + await asyncio.sleep(1) diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 588efabfc..5e38d3065 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -112,6 +112,7 @@ class OmniModelConfig(ModelConfig): omni_kv_config: dict | None = None codec_frame_rate_hz: float | None = None task_type: str | None = None + enable_sleep_mode: bool = False @property def registry(self): diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0a19eb119..012c0130c 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -353,6 +353,8 @@ class DiffusionCacheConfig: @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) + stage_id: int = 0 + model: str | None = None model_class_name: str | None = None @@ -508,6 +510,8 @@ class OmniDiffusionConfig: # Step mode settings step_execution: bool = False + # sleep mode + enable_sleep_mode: bool = False # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 @@ -797,5 +801,43 @@ class AttentionBackendEnum(enum.Enum): return self.name.lower() +@dataclass +class OmniACK: + """ + Handshake payload from Workers to Orchestrator. + """ + + task_id: str + status: str + stage_id: int | None = None + rank: int | None = None + freed_bytes: int = 0 + metadata: dict[str, Any] = field(default_factory=dict) + """ + Additional telemetry such as: + - max_contiguous_block: for fragmentation analysis. + - cuda_graph_recalled: boolean if graphs were successfully destroyed/rebuilt. + - latency_ms: time taken for the D2H/H2D transfer. + """ + error_msg: str | None = None + + +@dataclass +class OmniSleepTask: + """Structured sleep instruction.""" + + task_id: str + level: int = 2 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class OmniWakeTask: + """Structured wake-up instruction.""" + + task_id: str + tags: list[str] | None = None + + # Special message broadcast via scheduler queues to signal worker shutdown. SHUTDOWN_MESSAGE = {"type": "shutdown"} diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index e55a464fb..6325d4a16 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -36,10 +36,14 @@ class BackgroundResources: def __call__(self): """Clean up background resources.""" + if hasattr(self, "wake_events") and self.wake_events: + for ev in self.wake_events: + ev.set() + if self.broadcast_mq is not None: try: for _ in range(self.num_workers): - self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE) + self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0) self.broadcast_mq = None self.result_mq = None @@ -65,11 +69,13 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): self._closed = False num_workers = self.od_config.num_gpus + self.wake_events = [mp.Event() for _ in range(num_workers)] + self._broadcast_mq = self._init_broadcast_queue(num_workers) broadcast_handle = self._broadcast_mq.export_handle() # Launch workers - processes, result_handle = self._launch_workers(broadcast_handle) + processes, result_handle = self._launch_workers(broadcast_handle, self.wake_events) self._result_mq = self._init_result_queue(result_handle) self._processes = processes @@ -100,7 +106,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): if self._result_mq is None: raise RuntimeError("Result queue not initialized") - def _launch_workers(self, broadcast_handle): + def _launch_workers(self, broadcast_handle, wake_events): od_config = self.od_config logger.info("Starting server...") @@ -126,6 +132,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): od_config, writer, broadcast_handle, + wake_events[i], worker_extension_cls, custom_pipeline_args, ), diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 160309e0d..621f7eebd 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -27,7 +27,10 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( DiffusionOutput, + OmniACK, OmniDiffusionConfig, + OmniSleepTask, + OmniWakeTask, ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -77,6 +80,7 @@ class DiffusionWorker: self.model_runner: DiffusionModelRunner | None = None self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None + self.stage_id = getattr(od_config, "stage_id", 0) self.init_device() # Create model runner self.model_runner = DiffusionModelRunner( @@ -159,8 +163,10 @@ class DiffusionWorker: def _get_profiler(self) -> WorkerProfiler | None: return getattr(self, "profiler", None) - def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None: + def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None, **kwargs) -> None: """Load the diffusion model using DiffusionModelRunner.""" + load_format = kwargs.get("load_format", load_format) + custom_pipeline_name = kwargs.get("custom_pipeline_name", custom_pipeline_name) with ( set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config), set_current_vllm_config(self.vllm_config), @@ -170,6 +176,8 @@ class DiffusionWorker: load_format=load_format, custom_pipeline_name=custom_pipeline_name, ) + current_omni_platform.synchronize() + gc.collect() process_memory = get_process_gpu_memory(self.local_rank) if process_memory is not None: logger.info( @@ -284,77 +292,176 @@ class DiffusionWorker: """ from vllm.device_allocator.cumem import CuMemAllocator - process_memory_before_sleep = get_process_gpu_memory(self.local_rank) - free_bytes_before_sleep = None - if process_memory_before_sleep is None: - free_bytes_before_sleep = current_omni_platform.get_free_memory() + allocator = CuMemAllocator.get_instance() + + usage_before = allocator.get_current_usage() - # Save the buffers before level 2 sleep if level == 2 and self.model_runner is not None: + if hasattr(self.model_runner, "graph_runners"): + self.model_runner.graph_runners.clear() + logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.") model = self.model_runner.pipeline self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()} - allocator = CuMemAllocator.get_instance() - allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - process_memory_after_sleep = get_process_gpu_memory(self.local_rank) - if process_memory_before_sleep is not None and process_memory_after_sleep is not None: - freed_bytes = process_memory_before_sleep - process_memory_after_sleep - used_bytes = process_memory_after_sleep - accounting_scope = "process-scoped" + free_mem_before = current_omni_platform.get_free_memory() + + # Level 1: Offload weights; Level 2: Total Discard + offload_tags = ("weights",) if level == 1 else tuple() + allocator.sleep(offload_tags=offload_tags) + + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + + free_mem_after = current_omni_platform.get_free_memory() + try: + total_mem = current_omni_platform.get_device_total_memory() + except (NotImplementedError, AttributeError): + total_mem = torch.cuda.get_device_properties(self.device).total_memory + + phys_freed_bytes = max(0, free_mem_after - free_mem_before) + phys_used_bytes = total_mem - free_mem_after + + if usage_before > 0: + logger.info( + f"[Diffusion Worker {self.rank}] Sleep Level {level}: " + f"physically freed {phys_freed_bytes / GiB_bytes:.2f} GiB, " + f"{phys_used_bytes / GiB_bytes:.2f} GiB is still in use." + ) else: - free_bytes_after_sleep = current_omni_platform.get_free_memory() - assert free_bytes_before_sleep is not None - device_id = self.device.index if self.device.index is not None else 0 - total = current_omni_platform.get_device_total_memory(device_id) - freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep - used_bytes = total - free_bytes_after_sleep - accounting_scope = "device-scoped fallback" - assert freed_bytes >= 0, "Memory usage increased after sleeping." - logger.info( - "Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.", - accounting_scope, - freed_bytes / GiB_bytes, - used_bytes / GiB_bytes, - ) - return True + logger.info(f"[Worker {self.rank}] Sleep Level {level} completed (GPU was already empty).") + logger.info(f"[Worker {self.rank}] Memory usage before sleep: {usage_before / GiB_bytes:.2f} GiB.") + return usage_before def wake_up(self, tags: list[str] | None = None) -> bool: """ - Wake up the worker from sleep mode. See the sleep function - method for more details. + Wake up the worker from sleep mode. + + Re-activates the memory allocator for the specified tags and restores + model buffers from CPU back to GPU if they were saved during Level 2 sleep. Args: - tags: An optional list of tags to reallocate the worker memory - for specific memory allocations. Values must be in - `("weights")`. If None, all memory is reallocated. - wake_up should be called with all tags (or None) before the - worker is used again. + tags: List of memory pool tags to re-activate (e.g., ["weights"] + to match Level 1 sleep). If None, all pools are re-activated. """ from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) - - # Restore the buffers after level 2 sleep + current_omni_platform.synchronize() if len(self._sleep_saved_buffers) and self.model_runner is not None: model = self.model_runner.pipeline for name, buffer in model.named_buffers(): if name in self._sleep_saved_buffers: buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} + logger.info(f"[Worker {self.rank}] Buffers restored from CPU.") + logger.info(f"[Worker {self.rank}] Wake-up complete.") return True + def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK: + from vllm_omni.platforms import current_omni_platform + + try: + if isinstance(task, dict): + task = OmniSleepTask(**task) + logger.info(f"[Worker {self.rank}] Handshake Received: Task {task.task_id}") + + current_omni_platform.synchronize() + usage_before = current_omni_platform.get_current_memory_usage(self.device) + self.sleep(level=task.level) + current_omni_platform.synchronize() + usage_after = current_omni_platform.get_current_memory_usage(self.device) + real_freed = max(0, usage_before - usage_after) + logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.") + + # Ensure all ranks have completed sleep before measuring memory and sending ACK + if torch.distributed.is_initialized(): + t_freed = torch.tensor([float(real_freed)], device=self.device) + torch.distributed.all_reduce(t_freed) + real_freed = int(t_freed.item()) + + if self.rank != 0: + return None + + ack = OmniACK( + task_id=task.task_id, + status="SUCCESS", + stage_id=self.stage_id, + rank=self.rank, + freed_bytes=real_freed, + # return RL need metadata + metadata={ + "source": f"Platform_{current_omni_platform.get_device_name()}", + "total_freed_gib": f"{real_freed / GiB_bytes:.2f}", + "rank_residual_gib": f"{usage_after / GiB_bytes:.2f}", + }, + ) + logger.info(f"[Worker {self.rank}] ACK emitted. Freed {real_freed / GiB_bytes:.2f} GiB.") + return ack + except Exception as e: + logger.error(f"Sleep failed: {e}", exc_info=True) + if torch.distributed.is_initialized(): + try: + torch.distributed.barrier() + except Exception: + pass + return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e)) + + def handle_wake_task(self, task: OmniWakeTask) -> OmniACK: + from vllm_omni.platforms import current_omni_platform + + try: + if isinstance(task, dict): + task = OmniWakeTask(**task) + logger.info(f"[Worker {self.rank}] Responding to Wake-up Task: {task.task_id}") + self.wake_up(tags=task.tags) + + logger.info(f"[Worker {self.rank}] wake_up logic finished, entering barrier...") + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + current_omni_platform.synchronize() + usage_now = current_omni_platform.get_current_memory_usage(self.device) + current_used_gib = usage_now / (1024**3) + + if self.rank != 0: + return None + logger.info(f"[Worker {self.rank}] PASSED barrier, about to return to loop.") + + return OmniACK( + task_id=task.task_id, + status="SUCCESS", + stage_id=self.stage_id, + rank=self.rank, + metadata={ + "state": "WARM", + "source": f"Platform_{current_omni_platform.get_device_name()}", + "current_vram_gib": f"{current_used_gib:.2f}", + }, + ) + except Exception as e: + logger.error(f"Wake-up failed on Rank {self.rank}: {e}", exc_info=True) + if torch.distributed.is_initialized(): + try: + torch.distributed.barrier() + except Exception: + pass + return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e)) + def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: """Get memory pool context for sleep mode support.""" - if self.od_config.enable_sleep_mode: + is_sleep_enabled = getattr(self.od_config, "enable_sleep_mode", False) + if is_sleep_enabled: + current_omni_platform.synchronize() + gc.collect() from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() if tag == "weights": assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process." + logger.info(f"[Worker {self.rank}] Activating Diffusion CuMem pool for tag: {tag}") return allocator.use_memory_pool(tag=tag) - else: - return nullcontext() + return nullcontext() def shutdown(self) -> None: """Shutdown the worker and cleanup distributed environment.""" @@ -395,11 +502,13 @@ class WorkerProc: od_config: OmniDiffusionConfig, gpu_id: int, broadcast_handle, + wake_event: mp.Event, worker_extension_cls: str | None = None, custom_pipeline_args: dict[str, Any] | None = None, ): self.od_config = od_config self.gpu_id = gpu_id + self.wake_event = wake_event # Inter-process Communication self.context = zmq.Context(io_threads=2) @@ -414,7 +523,13 @@ class WorkerProc: if gpu_id == 0: self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) self.result_mq_handle = self.result_mq.export_handle() + WorkerProc._shared_result_handle = self.result_mq_handle logger.info(f"Worker {gpu_id} created result MessageQueue") + else: + handle = getattr(WorkerProc, "_shared_result_handle", None) + if handle: + self.result_mq = MessageQueue.create_from_handle(handle, gpu_id) + logger.info(f"Worker {gpu_id} attached to shared result MessageQueue") assert od_config.master_port is not None @@ -438,13 +553,17 @@ class WorkerProc: ) return wrapper - def return_result(self, output: object): + def return_result(self, output: Any): """Reply to client, only on rank 0.""" if self.result_mq is not None: + if isinstance(output, OmniACK): + self.result_mq.enqueue(output) + return try: pack_diffusion_output_shm(output) except Exception as e: - logger.warning("SHM pack failed, falling back to raw enqueue: %s", e) + if hasattr(output, "output"): + logger.warning("SHM pack failed for model output: %s", e) self.result_mq.enqueue(output) def recv_message(self): @@ -480,20 +599,31 @@ class WorkerProc: while self._running: msg = None try: - msg = self.recv_message() - except Exception as e: - logger.error( - f"Error receiving message in worker loop: {e}", - exc_info=True, - ) + msg = self.mq.dequeue(timeout=1.0) + except Exception: + if self.wake_event and self.wake_event.is_set(): + self.wake_event.clear() + logger.info(f"Worker {self.gpu_id} caught OOB POKE, forcing wake-up sequence.") + msg = {"type": "wake_up", "task_id": "recovery-task", "tags": None} + else: + continue + if msg is None: continue if msg is None or len(msg) == 0: logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id) continue + if isinstance(msg, dict) and msg.get("type") == "sleep": + task = OmniSleepTask(level=msg.get("level", 2), task_id=msg.get("task_id", "local")) + ack = self.worker.handle_sleep_task(task) + self.return_result(ack) + elif isinstance(msg, dict) and msg.get("type") == "wake_up": + task = OmniWakeTask(tags=msg.get("tags"), task_id=msg.get("task_id", "local")) + ack = self.worker.handle_wake_task(task) + self.return_result(ack) # Route message based on type - if isinstance(msg, dict) and msg.get("type") == "rpc": + elif isinstance(msg, dict) and msg.get("type") == "rpc": try: result, should_reply = self.execute_rpc(msg) if should_reply: @@ -538,6 +668,7 @@ class WorkerProc: od_config: OmniDiffusionConfig, pipe_writer: mp.connection.Connection, broadcast_handle, + wake_event: mp.Event, worker_extension_cls: str | None = None, custom_pipeline_args: dict[str, Any] | None = None, ) -> None: @@ -549,6 +680,7 @@ class WorkerProc: od_config, gpu_id=rank, broadcast_handle=broadcast_handle, + wake_event=wake_event, worker_extension_cls=worker_extension_cls, custom_pipeline_args=custom_pipeline_args, ) @@ -574,6 +706,7 @@ class WorkerWrapperBase: gpu_id: int, od_config: OmniDiffusionConfig, base_worker_class: type = DiffusionWorker, + wake_event: mp.Event = None, worker_extension_cls: str | None = None, custom_pipeline_args: dict[str, Any] | None = None, ): @@ -734,6 +867,12 @@ class WorkerWrapperBase: """ return self.worker.wake_up(tags) + def handle_sleep_task(self, task): + return self.worker.handle_sleep_task(task) + + def handle_wake_task(self, task): + return self.worker.handle_wake_task(task) + def shutdown(self) -> None: """Shutdown the worker and cleanup resources.""" return self.worker.shutdown() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d98ce7d41..33dee4880 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -6,11 +6,12 @@ import tempfile from dataclasses import dataclass, field, fields from typing import Any -from vllm.engine.arg_utils import EngineArgs +from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger from vllm_omni.config import OmniModelConfig from vllm_omni.engine.output_modality import OutputModality +from vllm_omni.platforms import current_omni_platform from vllm_omni.plugins import load_omni_general_plugins logger = init_logger(__name__) @@ -144,6 +145,24 @@ class OmniEngineArgs(EngineArgs): quantization_config: Any | None = None worker_type: str | None = None task_type: str | None = None + worker_cls: str = None + enable_sleep_mode: bool = False + omni: bool = False + + @classmethod + def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + try: + parser.add_argument("--omni", action="store_true", default=False, help="Enable Omni engine features.") + except argparse.ArgumentError: + pass + try: + parser.add_argument( + "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." + ) + except argparse.ArgumentError: + pass + return parser + omni_master_address: str | None = None omni_master_port: int | None = None stage_configs_path: str | None = None @@ -152,6 +171,11 @@ class OmniEngineArgs(EngineArgs): custom_pipeline_args: dict[str, Any] | None = None def __post_init__(self) -> None: + if self.worker_cls is None: + if self.worker_type == "ar": + self.worker_cls = current_omni_platform.get_omni_ar_worker_cls() + elif self.worker_type == "generation": + self.worker_cls = current_omni_platform.get_omni_generation_worker_cls() load_omni_general_plugins() super().__post_init__() @@ -296,6 +320,15 @@ class OmniEngineArgs(EngineArgs): ) return omni_config + +@dataclass +class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs): + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + parser = AsyncEngineArgs.add_cli_args(parser) + parser = OmniEngineArgs._add_omni_specific_args(parser) + return parser + @property def output_modality(self) -> OutputModality: """Parse engine_output_type into a type-safe OutputModality flag.""" diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index b796401fe..7b98c13b8 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -793,6 +793,7 @@ class AsyncOmniEngine: else: use_inline = True if self.num_stages == 1 else False stage_clients[stage_idx] = initialize_diffusion_stage( + configured_stage_id, self.model, stage_cfg, metadata, @@ -1401,6 +1402,12 @@ class AsyncOmniEngine: # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config. for cfg in stage_configs: try: + if not hasattr(cfg, "engine_args") or cfg.engine_args is None: + cfg.engine_args = OmegaConf.create({}) + global_sleep_mode = kwargs.get("enable_sleep_mode") + if global_sleep_mode is not None: + if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None: + cfg.engine_args.enable_sleep_mode = global_sleep_mode if getattr(cfg, "stage_type", None) != "diffusion": continue if not hasattr(cfg, "engine_args") or cfg.engine_args is None: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 89dfdc163..cd6cd6a69 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -23,6 +23,7 @@ from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path @@ -650,6 +651,7 @@ def build_diffusion_config( def initialize_diffusion_stage( + stage_id: int, model: str, stage_cfg: Any, metadata: StageMetadata, @@ -671,6 +673,15 @@ def initialize_diffusion_stage( """ from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client + engine_args = _to_dict(stage_cfg.engine_args) + engine_args.pop("stage_id", None) + od_config = OmniDiffusionConfig.from_kwargs( + stage_id=stage_id, + model=model, + **engine_args, + ) + if metadata.cfg_kv_collect_func is not None: + od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func od_config = build_diffusion_config(model, stage_cfg, metadata) return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 9606cc80d..6759fa5a6 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio import time +import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any @@ -24,10 +25,12 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tasks import SupportedTask from vllm.v1.engine.exceptions import EngineDeadError +from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni_base import OmniBase from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.inputs.preprocess import InputPreprocessor @@ -40,6 +43,63 @@ logger = init_logger(__name__) _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 +class AsyncEventResolver: + """ + A generic signal aggregator designed for synchronized handshakes in + distributed or multi-stage environments. Supports waiting for a specified + number (expected_count) of worker signals in both inline and multiprocess modes. + """ + + def __init__(self, orchestrator=None): + self._pending_tasks: dict[str, dict] = {} + self.orchestrator = orchestrator + self._lock = asyncio.Lock() + + def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: + loop = asyncio.get_running_loop() + fut = loop.create_future() + self._pending_tasks[task_id] = { + "future": fut, + "expected_count": expected_count, + "received": [], + "start_time": time.time(), + } + return fut + + async def resolve(self, ack: OmniACK): + tid = getattr(ack, "task_id", None) + + if tid is None and isinstance(ack, dict): + tid = ack.get("task_id") + + async with self._lock: + task_info = self._pending_tasks.get(tid) + if task_info is None: + logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") + return + + task_info["received"].append(ack) + current_count = len(task_info["received"]) + expected = task_info["expected_count"] + + orchestrator = self.orchestrator + if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: + freed = getattr(ack, "freed_bytes", 0) + if freed == 0 and isinstance(ack, dict): + freed = ack.get("freed_bytes", 0) + orchestrator.metrics.record_vram_reclaimed(freed) + + logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") + + if current_count >= expected: + self._pending_tasks.pop(tid) + fut = task_info["future"] + if not fut.done(): + elapsed = time.time() - task_info["start_time"] + logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") + fut.set_result(task_info["received"]) + + class AsyncOmni(EngineClient, OmniBase): """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine. @@ -76,7 +136,7 @@ class AsyncOmni(EngineClient, OmniBase): self._paused: bool = False self._is_sleeping: bool = False self.final_output_task: asyncio.Task | None = None - + self.event_resolver = AsyncEventResolver(orchestrator=self) self.config_path = self.engine.config_path self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) self.input_processor = self.engine.input_processor @@ -488,6 +548,22 @@ class AsyncOmni(EngineClient, OmniBase): await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) continue + if isinstance(msg, dict) and msg.get("type") == "ack": + ack_data = msg.get("ack") + tid = getattr(ack_data, "task_id", "unknown") + logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") + await self.event_resolver.resolve(ack_data) + continue + if isinstance(msg, OmniACK): + logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") + await self.event_resolver.resolve(msg) + continue + if hasattr(msg, "task_id"): + tid = getattr(msg, "task_id") + logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") + await self.event_resolver.resolve(msg) + continue + should_continue, _, stage_id, req_state = self._handle_output_message(msg) if should_continue: continue @@ -654,21 +730,68 @@ class AsyncOmni(EngineClient, OmniBase): logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process") return True - async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None: - """Sleep all stages. + async def sleep( + self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort" + ) -> list[OmniACK]: + self._final_output_handler() + if stage_ids is None: + stage_ids = list(range(len(self.engine.stage_clients))) + total_workers = 0 + for sid in stage_ids: + client = self.engine.stage_clients[sid] + # During the Diffusion phase, regardless of the TP amount, + # currently only a summary ACK is reported at Rank 0. + if getattr(client, "stage_type", "") == "diffusion": + total_workers += 1 + else: + config = self.engine.stage_vllm_configs[sid] + actual_tp = config.parallel_config.tensor_parallel_size if config else 1 + total_workers += actual_tp - Best-effort: unsupported stages will emit a TODO result. - """ + task_id = str(uuid.uuid4()) + self.event_resolver.watch_task(task_id, expected_count=total_workers) + logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") + task = OmniSleepTask(level=level, task_id=task_id) + rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids) + final_acks = [] + for stage_res in rpc_results: + worker_acks = stage_res if isinstance(stage_res, list) else [stage_res] + for ack in worker_acks: + if ack is not None: + await self.event_resolver.resolve(ack) + final_acks.append(ack) self._is_sleeping = True - await self.collective_rpc(method="sleep", args=(level,)) + return final_acks - async def wake_up(self, tags: list[str] | None = None) -> None: - """Wake up all stages. - - Best-effort: unsupported stages will emit a TODO result. - """ + async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]: + self._final_output_handler() + if stage_ids is None: + stage_ids = list(range(len(self.engine.stage_clients))) + total_workers = 0 + for sid in stage_ids: + client = self.engine.stage_clients[sid] + if getattr(client, "stage_type", "") == "diffusion": + total_workers += 1 + else: + config = self.engine.stage_vllm_configs[sid] + total_workers += config.parallel_config.tensor_parallel_size if config else 1 + task_id = str(uuid.uuid4()) + self.event_resolver.watch_task(task_id, expected_count=total_workers) + logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") + task = OmniWakeTask(tags=tags, task_id=task_id) + rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids) + final_acks = [] + for stage_res in rpc_results: + worker_acks = stage_res if isinstance(stage_res, list) else [stage_res] + for ack in worker_acks: + if ack is not None: + await self.event_resolver.resolve(ack) + final_acks.append(ack) + current_omni_platform.synchronize() + await asyncio.sleep(0.1) self._is_sleeping = False - await self.collective_rpc(method="wake_up", args=(tags,)) + logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.") + return final_acks async def is_sleeping(self) -> bool: """Return whether all stages are sleeping. @@ -725,6 +848,10 @@ class AsyncOmni(EngineClient, OmniBase): """Whether orchestrator thread has stopped unexpectedly.""" return not self.engine.is_alive() + @property + def _name(self) -> str: + return "AsyncOrchestrator" + @property def is_stopped(self) -> bool: """EngineClient abstract property implementation.""" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 8bccfbb59..e355ee679 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -139,6 +139,17 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable vLLM-Omni mode for multi-modal and diffusion models", ) + + try: + omni_config_group.add_argument( + "--enable-sleep-mode", + action="store_true", + default=False, + help="Enable GPU memory pool for sleep mode.", + ) + except argparse.ArgumentError: + pass + omni_config_group.add_argument( "--task-type", type=str, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index dca494efe..98b5c7cf7 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -133,7 +133,7 @@ class OmniBase(PDDisaggregationMixin): if "log_requests" in kwargs: raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.") model = omni_snapshot_download(model) - self._name = self.__class__.__name__ + self.__dict__["_name"] = self.__class__.__name__ self.model = model self.log_stats = log_stats # Provisional value (mirrors the CLI/caller kwarg); the engine resolves diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b2e6e435a..5d8cc4ea2 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import base64 +import dataclasses import io import json import multiprocessing @@ -837,6 +838,7 @@ async def omni_init_app_state( state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 + state.sleeping_stages = set() def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None: @@ -2553,3 +2555,64 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No raise HTTPException( status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to stop profiler: {str(e)}" ) + + +class OmniSleepRequest(BaseModel): + stage_ids: list[int] + level: int = 2 + + +class OmniWakeupRequest(BaseModel): + stage_ids: list[int] + + +@router.post("/v1/omni/sleep") +async def omni_sleep(request: OmniSleepRequest, raw_request: Request): + engine_client = raw_request.app.state.engine_client + sleeping_set = raw_request.app.state.sleeping_stages + if not hasattr(engine_client, "sleep"): + raise HTTPException(status_code=501, detail="Engine does not support sleep") + acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level) + for sid in request.stage_ids: + sleeping_set.add(sid) + return {"status": "SUCCESS", "acks": [dataclasses.asdict(a) if dataclasses.is_dataclass(a) else a for a in acks]} + + +@router.post("/v1/omni/wakeup") +async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): + engine_client = raw_request.app.state.engine_client + sleeping_set = raw_request.app.state.sleeping_stages + if not any(sid in sleeping_set for sid in request.stage_ids): + return {"status": "SKIPPED", "reason": "Target stages are not sleeping."} + if not hasattr(engine_client, "wake_up"): + raise HTTPException(status_code=501, detail="Engine does not support wake_up") + acks = await engine_client.wake_up(stage_ids=request.stage_ids) + for sid in request.stage_ids: + if sid in sleeping_set: + sleeping_set.remove(sid) + return {"status": "SUCCESS", "acks": [dataclasses.asdict(a) if dataclasses.is_dataclass(a) else a for a in acks]} + + +if __name__ == "__main__": + import argparse + import asyncio + + from vllm.entrypoints.openai.cli_args import make_arg_parser + + parser = argparse.ArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") + parser = make_arg_parser(parser) + registered_flags = set() + for action in parser._actions: + registered_flags.update(action.option_strings) + if "--omni" not in registered_flags: + parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.") + if "--enable-sleep-mode" not in registered_flags: + parser.add_argument( + "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." + ) + args = parser.parse_args() + if not hasattr(args, "model_tag"): + setattr(args, "model_tag", args.model) + if hasattr(args, "model_tag") and args.model_tag is None: + args.model_tag = args.model + asyncio.run(omni_run_server(args)) diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index f7f5dbd1d..8bd9efc89 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -2,15 +2,23 @@ from __future__ import annotations +import gc import os import time +from contextlib import AbstractContextManager, nullcontext import torch from vllm.logger import init_logger from vllm.utils.mem_utils import format_gib, memory_profiling from vllm.v1.worker.gpu_worker import Worker as GPUWorker +from vllm_omni.diffusion.data import ( + OmniACK, + OmniSleepTask, + OmniWakeTask, +) from vllm_omni.entrypoints.utils import detect_pid_host +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.gpu_memory_utils import ( get_process_gpu_memory, is_process_scoped_memory_available, @@ -30,6 +38,13 @@ class OmniGPUWorkerBase(GPUWorker): for custom trace naming, background gzip, and trace path collection. """ + def load_model(self, *args, **kwargs): + with self._maybe_get_memory_pool_context("weights"): + res = super().load_model(*args, **kwargs) + current_omni_platform.synchronize() + gc.collect() + return res + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -94,10 +109,8 @@ class OmniGPUWorkerBase(GPUWorker): """ if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: self.model_runner.profile_run() - logger.info( - "Using explicit kv_cache_memory_bytes: %s GiB", - format_gib(kv_cache_memory_bytes), - ) + if current_omni_platform.is_rocm(): + torch.cuda.synchronize() return kv_cache_memory_bytes with memory_profiling( @@ -154,3 +167,141 @@ class OmniGPUWorkerBase(GPUWorker): ) return int(self.available_kv_cache_memory_bytes) + + # Provide memory pool context + def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: + v1_config_enabled = False + if hasattr(self, "vllm_config"): + model_cfg = getattr(self.vllm_config, "model_config", None) + v1_config_enabled = getattr(model_cfg, "enable_sleep_mode", False) + + is_sleep_enabled = v1_config_enabled or getattr(self.cache_config, "enable_sleep_mode", False) + if is_sleep_enabled: + current_omni_platform.synchronize() + gc.collect() + from vllm.device_allocator.cumem import CuMemAllocator + + allocator = CuMemAllocator.get_instance() + logger.info(f"[LLM Worker {self.rank}] Sleep Mode ENABLED. Activating CuMem pool for tag: {tag}") + return allocator.use_memory_pool(tag=tag) + else: + logger.warning(f"[LLM Worker {self.rank}] Sleep Mode DISABLED.") + return nullcontext() + + def sleep(self, level: int = 1) -> bool: + """ + Put the worker to sleep. + Args: + level: 1 (Offload weights to CPU), level: 2 (Total Discard). + """ + from vllm.device_allocator.cumem import CuMemAllocator + + mem_before = current_omni_platform.get_current_memory_usage(self.device) + offload_tags = ("weights",) if level == 1 else tuple() + allocator = CuMemAllocator.get_instance() + allocator.sleep(offload_tags=offload_tags) + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + mem_after = current_omni_platform.get_current_memory_usage(self.device) + freed = max(0, mem_before - mem_after) + remaining_gb = mem_after / 1024**3 + logger.info( + f"[LLM Worker {self.rank}] Level {level} Sleep: Freed " + f"{freed / 1024**3:.2f} GiB. {remaining_gb:.2f}GiB memory " + "is still in use." + ) + return True + + def wake_up(self, tags: list[str] | None = None) -> bool: + "Physical video memory reloading logic" + from vllm.device_allocator.cumem import CuMemAllocator + + allocator = CuMemAllocator.get_instance() + allocator.wake_up(tags) + current_omni_platform.synchronize() + logger.info(f"[LLM Worker {self.rank}] Wake-up complete.") + return True + + def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK: + "Handle deterministic Sleep command from the main process" + try: + if isinstance(task, dict): + task = OmniSleepTask(**task) + logger.info(f"[LLM Worker {self.rank}] Handshake Received: Task {task.task_id}, Level {task.level}") + if task.level == 2: + if hasattr(self.model_runner, "graph_runners"): + self.model_runner.graph_runners.clear() + logger.info(f"[LLM Worker {self.rank}] LLM CUDA Graphs cleared.") + mem_before = current_omni_platform.get_current_memory_usage(self.device) + self.sleep(level=task.level) + mem_after = current_omni_platform.get_current_memory_usage(self.device) + rank_freed = max(0, mem_before - mem_after) + if torch.distributed.is_initialized(): + t_freed = torch.tensor([float(rank_freed)], device=self.device) + torch.distributed.all_reduce(t_freed) + total_freed = int(t_freed.item()) + torch.distributed.barrier() + else: + total_freed = rank_freed + if self.rank != 0: + return None + current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0) + ack = OmniACK( + task_id=task.task_id, + status="SUCCESS", + stage_id=current_stage_id, + rank=self.rank, + freed_bytes=total_freed, + metadata={ + "source": "omni_platform_audit", + "total_freed_gib": f"{total_freed / 1024**3:.2f}", + "rank_residual_gib": f"{mem_after / 1024**3:.2f}", + }, + ) + if hasattr(self, "result_mq") and self.result_mq: + self.result_mq.put(ack) + logger.info(f"[LLM Worker {self.rank}] ACK emitted for Task {task.task_id}") + return ack + except Exception as e: + logger.error(f"[LLM Worker {self.rank}] Sleep Task Failed: {e}", exc_info=True) + if torch.distributed.is_initialized(): + try: + torch.distributed.barrier() + except Exception: + pass + return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e)) + + def handle_wake_task(self, task: OmniWakeTask) -> OmniACK: + "Handle deterministic Wakeup command from the main process" + try: + if isinstance(task, dict): + task = OmniWakeTask(**task) + self.wake_up(tags=task.tags) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + gc.collect() + current_omni_platform.synchronize() + usage_now = current_omni_platform.get_current_memory_usage(self.device) + if self.rank != 0: + return None + current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0) + ack = OmniACK( + task_id=task.task_id, + status="SUCCESS", + stage_id=current_stage_id, + rank=self.rank, + metadata={"state": "WARM", "current_vram_gib": f"{usage_now / 1024**3:.2f}"}, + ) + if hasattr(self, "result_mq") and self.result_mq: + self.result_mq.put(ack) + logger.info(f"[LLM Worker {self.rank}] Wake-up ACK emitted.") + return ack + except Exception as e: + logger.error(f"[LLM Worker {self.rank}] Wake-up Failed: {e}", exc_info=True) + if torch.distributed.is_initialized(): + try: + torch.distributed.barrier() + except Exception: + pass + tid = task.task_id if hasattr(task, "task_id") else "unknown" + return OmniACK(task_id=tid, status="ERROR", error_msg=str(e)) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 4abe21964..2668bbe8b 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -12,6 +12,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin @@ -104,3 +105,23 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): if self.rank == 0: # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + + def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK: + """ + Explicitly handle sleep commands. + Calls the implementation in the base class OmniGPUWorkerBase. + """ + logger.debug(f"[AR Worker {self.rank}] Resolving handle_sleep_task dispatch") + if isinstance(task, dict): + task = OmniSleepTask(**task) + return super().handle_sleep_task(task) + + def handle_wake_task(self, task: OmniWakeTask | dict) -> OmniACK: + """ + Explicitly handle wake-up commands. + Calls the implementation in the base class OmniGPUWorkerBase. + """ + logger.debug(f"[AR Worker {self.rank}] Resolving handle_wake_task dispatch") + if isinstance(task, dict): + task = OmniWakeTask(**task) + return super().handle_wake_task(task)