From 9225039d170607954e23ff32153bd0121ba3ce57 Mon Sep 17 00:00:00 2001
From: LiBai <91311486+RGB-loop@users.noreply.github.com>
Date: Thu, 9 Apr 2026 15:35:22 +0800
Subject: [PATCH] [Bug] Lazy-import entrypoints to fix subprocess pynvml crash
 (#2187)

Signed-off-by: Meng Jianwen <mengjianwen@liblib.ai>
---
 vllm_omni/__init__.py             | 18 +++++++++++++++++-
 vllm_omni/entrypoints/__init__.py | 17 +++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py
index b093272d2..cec8b0af7 100644
--- a/vllm_omni/__init__.py
+++ b/vllm_omni/__init__.py
@@ -24,11 +24,27 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
 
 from .config import OmniModelConfig
-from .entrypoints import AsyncOmni, Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 
 
+def __getattr__(name: str):
+    # Lazy import for AsyncOmni and Omni to avoid pulling in heavy
+    # dependencies (vllm model_loader → fused_moe → pynvml) at package
+    # import time.  This prevents crashes in lightweight subprocesses
+    # (e.g. model-architecture inspection) that lack a CUDA context.
+    # See: https://github.com/vllm-project/vllm-omni/issues/1793
+    if name == "AsyncOmni":
+        from .entrypoints.async_omni import AsyncOmni
+
+        return AsyncOmni
+    if name == "Omni":
+        from .entrypoints.omni import Omni
+
+        return Omni
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "__version__",
     "__version_tuple__",
diff --git a/vllm_omni/entrypoints/__init__.py b/vllm_omni/entrypoints/__init__.py
index 7b09adf93..b273929a8 100644
--- a/vllm_omni/entrypoints/__init__.py
+++ b/vllm_omni/entrypoints/__init__.py
@@ -5,8 +5,21 @@
 vLLM-Omni entrypoints module.
 """
 
-from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.omni import Omni
+
+def __getattr__(name: str):
+    # Lazy imports to avoid eagerly loading heavy modules (engine,
+    # model_loader, pynvml) when the package is imported in lightweight
+    # contexts such as model-architecture inspection subprocesses.
+    if name == "AsyncOmni":
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        return AsyncOmni
+    if name == "Omni":
+        from vllm_omni.entrypoints.omni import Omni
+
+        return Omni
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
 
 __all__ = [
     "AsyncOmni",