diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index ca5b2f43e..e32adf737 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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 diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 4ebff3290..0abe029d6 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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() diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 2e7d07858..77a42602b 100644 --- a/vllm_omni/worker/base.py +++ b/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