From c6f314df4f7847f31bdbd2e6439f650414575230 Mon Sep 17 00:00:00 2001 From: dsinghvi Date: Thu, 12 Feb 2026 12:17:54 +0000 Subject: [PATCH] [Feature]: Diffusion sleep to use process level memory calculation (#1276) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Divyansh Singhvi Signed-off-by: 汪志鹏 Signed-off-by: dsinghvi Co-authored-by: 汪志鹏 --- tests/diffusion/test_diffusion_worker.py | 73 ++++++++++++++----- .../diffusion/worker/diffusion_worker.py | 35 +++++++-- vllm_omni/entrypoints/omni_stage.py | 14 ++-- vllm_omni/entrypoints/utils.py | 57 ++++++++------- 4 files changed, 121 insertions(+), 58 deletions(-) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 1c767b98f..73cfa104f 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -81,18 +81,18 @@ class TestDiffusionWorkerLoadWeights: class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" - # Setup memory info mocks - # Before sleep: 1GB free - # After sleep: 3GB free (freed 2GB) - mock_platform.get_free_memory.side_effect = [ - 1 * 1024**3, # Before sleep - 3 * 1024**3, # After sleep + # Setup process-scoped memory mocks + # Before sleep: 3GB used + # After sleep: 1GB used (freed 2GB) + mock_get_process_memory.side_effect = [ + 3 * 1024**3, + 1 * 1024**3, ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -108,16 +108,16 @@ class TestDiffusionWorkerSleep: # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" - # Setup memory info mocks - mock_platform.get_free_memory.side_effect = [ - 1 * 1024**3, # Before sleep - 5 * 1024**3, # After sleep (freed 4GB) + # Setup process-scoped memory mocks + mock_get_process_memory.side_effect = [ + 5 * 1024**3, # Before sleep + 1 * 1024**3, # After sleep (freed 4GB) ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -146,16 +146,22 @@ class TestDiffusionWorkerSleep: assert "buffer1" in mock_gpu_worker._sleep_saved_buffers assert "buffer2" in mock_gpu_worker._sleep_saved_buffers + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_memory_freed_validation( + self, + mock_allocator_class, + mock_platform, + mock_get_process_memory, + mock_gpu_worker, + ): """Test that sleep validates memory was actually freed.""" - # Simulate memory increase (should trigger assertion error) - mock_platform.get_free_memory.side_effect = [ - 3 * 1024**3, # Before sleep: 3GB free - 1 * 1024**3, # After sleep: 1GB free (negative freed!) + # Simulate process memory increase (should trigger assertion error) + mock_get_process_memory.side_effect = [ + 1 * 1024**3, # Before sleep: 1GB used + 3 * 1024**3, # After sleep: 3GB used (negative freed) ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 mock_allocator = Mock() mock_allocator_class.get_instance = Mock(return_value=mock_allocator) @@ -165,6 +171,33 @@ class TestDiffusionWorkerSleep: with pytest.raises(AssertionError, match="Memory usage increased after sleeping"): mock_gpu_worker.sleep(level=1) + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") + @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + @patch("vllm.device_allocator.cumem.CuMemAllocator") + def test_sleep_falls_back_to_device_memory_when_nvml_unavailable( + self, + mock_allocator_class, + mock_platform, + mock_get_process_memory, + mock_gpu_worker, + ): + """Test sleep uses device-scoped fallback when NVML is unavailable.""" + mock_get_process_memory.side_effect = [None, None] + mock_platform.get_free_memory.side_effect = [ + 1 * 1024**3, # Before sleep + 3 * 1024**3, # After sleep + ] + mock_platform.get_device_total_memory.return_value = 8 * 1024**3 + + mock_allocator = Mock() + mock_allocator_class.get_instance = Mock(return_value=mock_allocator) + mock_allocator.sleep = Mock() + + result = mock_gpu_worker.sleep(level=1) + + mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) + assert result is True + class TestDiffusionWorkerWakeUp: """Test DiffusionWorker.wake_up method.""" diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 90fff548e..6a8494991 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -40,6 +40,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.lora.request import LoRARequest from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory logger = init_logger(__name__) @@ -135,6 +136,14 @@ class DiffusionWorker: load_format=load_format, custom_pipeline_name=custom_pipeline_name, ) + process_memory = get_process_gpu_memory(self.local_rank) + if process_memory is not None: + logger.info( + "Worker %d: Process-scoped GPU memory after model loading: %.2f GiB.", + self.rank, + process_memory / GiB_bytes, + ) + assert self.model_runner.pipeline is not None def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker.""" @@ -201,7 +210,10 @@ class DiffusionWorker: """ from vllm.device_allocator.cumem import CuMemAllocator - free_bytes_before_sleep = current_omni_platform.get_free_memory() + 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() # Save the buffers before level 2 sleep if level == 2 and self.model_runner is not None: @@ -210,14 +222,23 @@ class DiffusionWorker: allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - free_bytes_after_sleep = current_omni_platform.get_free_memory() - 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 + 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" + 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 freed %.2f GiB memory, %.2f GiB memory is still in use.", + "Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.", + accounting_scope, freed_bytes / GiB_bytes, used_bytes / GiB_bytes, ) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 92ca363ff..f53177ce3 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -67,16 +67,20 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = """ from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available - if is_process_scoped_memory_available() and detect_pid_host(): - logger.debug( + nvml_available = is_process_scoped_memory_available() + pid_host = detect_pid_host() + + if nvml_available and pid_host: + logger.info( "NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks" ) yield return else: - logger.debug( - "NVML unavailable or PID host is not available (usually inside a container, " - "--pid=host is not set in docker run command) — using sequential init locks" + logger.info( + "Using sequential init locks (nvml_available=%s, pid_host=%s)", + nvml_available, + pid_host, ) from vllm_omni.platforms import current_omni_platform diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 0abe029d6..58d44919d 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -306,37 +306,42 @@ def in_container() -> bool: return any(m in cg for m in markers) -def has_pid_host() -> bool | None: - """ +def _nvml_pid_matches() -> bool | None: + """Check if os.getpid() matches what NVML would report via NSpid. + + If /proc/self/status has a single NSpid entry, the container PID equals + the host PID, so NVML process-scoped memory tracking works correctly. + Returns: - True -> very likely running with --pid=host (host PID namespace) - False -> very likely isolated PID namespace (default) + True -> PID namespaces match (single NSpid entry) + False -> PID namespaces differ (multiple NSpid entries) None -> cannot determine """ - # Strong signal: in host pid namespace, PID 2 is usually kthreadd - comm2 = _read_text("/proc/2/comm") - if comm2 is not None: - comm2 = comm2.strip() - if comm2 == "kthreadd": - return True - # If PID 2 exists and is NOT kthreadd, we're almost certainly not in host pid ns - return False - - # Fallback: check for other low-numbered kernel threads (best-effort) - for pid, name in [(3, "rcu_gp"), (4, "rcu_par_gp"), (10, "ksoftirqd/0")]: - comm = _read_text(f"/proc/{pid}/comm") - if comm is not None: - if comm.strip() == name: - return True - else: - return False - - return False + status = _read_text("/proc/self/status") + if status is None: + return None + for line in status.splitlines(): + if line.startswith("NSpid:"): + pids = line.split()[1:] + # Single PID means container PID == host PID + return len(pids) == 1 + return None def detect_pid_host() -> bool: - ic = in_container() - if not ic: + """Check if os.getpid() matches the PID that NVML will report. + + NVML always reports host-namespace PIDs. If this process lives in a + different PID namespace the match will fail and process-scoped memory + tracking won't work. We detect this via the NSpid field in + /proc/self/status: a single entry means host PID == container PID. + """ + if not in_container(): return True - return has_pid_host() + result = _nvml_pid_matches() + if result is not None: + return result + + # Cannot determine — assume mismatch to be safe. + return False