[Debug] Add trigger to concurrent stage init (#1274)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/entrypoints/omni_stage.py

@@ -49,6 +49,7 @@ from vllm_omni.entrypoints.stage_utils import (
     maybe_dump_to_shm,
     set_stage_devices,
 )
+from vllm_omni.entrypoints.utils import detect_pid_host
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
 from vllm_omni.metrics import count_tokens_from_outputs
 from vllm_omni.outputs import OmniRequestOutput
@@ -66,12 +67,17 @@ 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():
-        logger.debug("NVML process-scoped memory available — concurrent init is safe, skipping locks")
+    if is_process_scoped_memory_available() and detect_pid_host():
+        logger.debug(
+            "NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks"
+        )
         yield
         return
-
-    logger.debug("NVML unavailable — using sequential init locks")
+    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"
+        )
 
     from vllm_omni.platforms import current_omni_platform
 

vllm_omni/entrypoints/utils.py

@@ -280,3 +280,63 @@ def get_final_stage_id_for_e2e(
         final_stage_id_for_e2e = last_stage_id
 
     return final_stage_id_for_e2e
+
+
+# The following code detects if the process is running in a container and if
+# PID host is available. If so, we can use process-scoped memory tracking;
+# otherwise we need sequential init locks.
+
+
+def _read_text(path: str) -> str | None:
+    try:
+        with open(path, encoding="utf-8", errors="replace") as f:
+            return f.read()
+    except (FileNotFoundError, PermissionError, OSError):
+        return None
+
+
+def in_container() -> bool:
+    # Common Docker signal
+    if os.path.exists("/.dockerenv"):
+        return True
+
+    # cgroup markers (works for Docker/containerd/K8s/Podman in many setups)
+    cg = _read_text("/proc/1/cgroup") or ""
+    markers = ("docker", "containerd", "kubepods", "libpod", "podman")
+    return any(m in cg for m in markers)
+
+
+def has_pid_host() -> bool | None:
+    """
+    Returns:
+      True  -> very likely running with --pid=host (host PID namespace)
+      False -> very likely isolated PID namespace (default)
+      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
+
+
+def detect_pid_host() -> bool:
+    ic = in_container()
+    if not ic:
+        return True
+
+    return has_pid_host()

vllm_omni/worker/base.py

@@ -9,8 +9,10 @@ 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.entrypoints.utils import detect_pid_host
 from vllm_omni.worker.gpu_memory_utils import (
     get_process_gpu_memory,
+    is_process_scoped_memory_available,
 )
 
 logger = init_logger(__name__)
@@ -59,7 +61,11 @@ class OmniGPUWorkerBase(GPUWorker):
         self.non_torch_memory = profile_result.non_torch_increase
         self.peak_activation_memory = profile_result.torch_peak_increase
 
-        process_memory = get_process_gpu_memory(self.local_rank)
+        process_memory = (
+            get_process_gpu_memory(self.local_rank)
+            if is_process_scoped_memory_available() and detect_pid_host()
+            else None
+        )
 
         if process_memory is not None:
             # NVML available: use per-process memory