[Feature]: Diffusion sleep to use process level memory calculation (#1276)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

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."""

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,
         )

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

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