fix(multimodal): stop the media loaders from importing torch

components/src/dynamo/common/multimodal/__init__.py

@@ -1,44 +1,138 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Multimodal utilities for Dynamo components."""
+"""Multimodal utilities for Dynamo components.
 
-from collections.abc import Callable
+`AsyncEncoderCache` and the embedding-transfer members import `torch`, while the
+image, audio and video loaders do not. Re-exporting all of them eagerly meant
+`from dynamo.common.multimodal import ImageLoader` pulled in the whole ML stack,
+so the loaders could not be imported or unit-tested without it. The heavy names
+resolve on first access instead (PEP 562); the loaders stay eager because they
+cost nothing.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
-from dynamo.common.constants import EmbeddingTransferMode
-from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
 from dynamo.common.multimodal.audio_loader import AudioLoader
-from dynamo.common.multimodal.embedding_transfer import (
-    AbstractEmbeddingReceiver,
-    AbstractEmbeddingSender,
-    LocalEmbeddingReceiver,
-    LocalEmbeddingSender,
-    NixlReadEmbeddingReceiver,
-    NixlReadEmbeddingSender,
-    NixlWriteEmbeddingReceiver,
-    NixlWriteEmbeddingSender,
-    TransferRequest,
-)
 from dynamo.common.multimodal.image_loader import ImageLoader
 from dynamo.common.multimodal.video_loader import VideoLoader
 
-EMBEDDING_SENDER_FACTORIES: dict[
-    EmbeddingTransferMode, Callable[[], AbstractEmbeddingSender]
-] = {
-    EmbeddingTransferMode.LOCAL: LocalEmbeddingSender,
-    EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingSender,
-    EmbeddingTransferMode.NIXL_READ: NixlReadEmbeddingSender,
-}
-
-EMBEDDING_RECEIVER_FACTORIES: dict[
-    EmbeddingTransferMode, Callable[[], AbstractEmbeddingReceiver]
-] = {
-    EmbeddingTransferMode.LOCAL: LocalEmbeddingReceiver,
-    EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingReceiver,
-    # [gluo FIXME] can't use pre-registered tensor as NIXL requires descriptors
-    # to be at matching size, need to overwrite nixl connect library
-    EmbeddingTransferMode.NIXL_READ: lambda: NixlReadEmbeddingReceiver(max_items=0),
-}
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
+    from dynamo.common.multimodal.embedding_transfer import (
+        AbstractEmbeddingReceiver,
+        AbstractEmbeddingSender,
+        LocalEmbeddingReceiver,
+        LocalEmbeddingSender,
+        NixlReadEmbeddingReceiver,
+        NixlReadEmbeddingSender,
+        NixlWriteEmbeddingReceiver,
+        NixlWriteEmbeddingSender,
+        TransferRequest,
+    )
+
+    # Declared, not assigned: the runtime values come from `__getattr__` below.
+    # Without these mypy widens both to `Any` and callers lose the key and value
+    # types.
+    EMBEDDING_SENDER_FACTORIES: dict[
+        EmbeddingTransferMode, Callable[[], AbstractEmbeddingSender]
+    ]
+    EMBEDDING_RECEIVER_FACTORIES: dict[
+        EmbeddingTransferMode, Callable[[], AbstractEmbeddingReceiver]
+    ]
+
+_EMBEDDING_TRANSFER_EXPORTS = frozenset(
+    {
+        "AbstractEmbeddingReceiver",
+        "AbstractEmbeddingSender",
+        "LocalEmbeddingReceiver",
+        "LocalEmbeddingSender",
+        "NixlReadEmbeddingReceiver",
+        "NixlReadEmbeddingSender",
+        "NixlWriteEmbeddingReceiver",
+        "NixlWriteEmbeddingSender",
+        "TransferRequest",
+    }
+)
+
+
+def _build_sender_factories() -> (
+    dict[EmbeddingTransferMode, Callable[[], AbstractEmbeddingSender]]
+):
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.embedding_transfer import (
+        LocalEmbeddingSender,
+        NixlReadEmbeddingSender,
+        NixlWriteEmbeddingSender,
+    )
+
+    return {
+        EmbeddingTransferMode.LOCAL: LocalEmbeddingSender,
+        EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingSender,
+        EmbeddingTransferMode.NIXL_READ: NixlReadEmbeddingSender,
+    }
+
+
+def _build_receiver_factories() -> (
+    dict[EmbeddingTransferMode, Callable[[], AbstractEmbeddingReceiver]]
+):
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.embedding_transfer import (
+        LocalEmbeddingReceiver,
+        NixlReadEmbeddingReceiver,
+        NixlWriteEmbeddingReceiver,
+    )
+
+    return {
+        EmbeddingTransferMode.LOCAL: LocalEmbeddingReceiver,
+        EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingReceiver,
+        # [gluo FIXME] can't use pre-registered tensor as NIXL requires descriptors
+        # to be at matching size, need to overwrite nixl connect library
+        EmbeddingTransferMode.NIXL_READ: lambda: NixlReadEmbeddingReceiver(max_items=0),
+    }
+
+
+# Hidden from the type checker on purpose. A module-level `__getattr__` tells
+# mypy the package may have any attribute, which would silently accept a typo
+# like `from dynamo.common.multimodal import DoesNotExistAtAll` in every module
+# that imports from here. The `TYPE_CHECKING` block above already declares the
+# real names, so the checker has what it needs without the widening.
+if not TYPE_CHECKING:
+
+    def __getattr__(name: str) -> Any:
+        if name == "AsyncEncoderCache":
+            from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
+
+            value: Any = AsyncEncoderCache
+        elif name in _EMBEDDING_TRANSFER_EXPORTS:
+            from dynamo.common.multimodal import embedding_transfer
+
+            value = getattr(embedding_transfer, name)
+        elif name == "EMBEDDING_SENDER_FACTORIES":
+            value = _build_sender_factories()
+        elif name == "EMBEDDING_RECEIVER_FACTORIES":
+            value = _build_receiver_factories()
+        else:
+            raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+        # Cache on the module so repeat access skips this function entirely.
+        globals()[name] = value
+        return value
+
+
+def __dir__() -> list[str]:
+    # `__all__` alone under-reports: it omits the submodules, and it omits
+    # `AbstractEmbeddingSender` / `AbstractEmbeddingReceiver`, which resolve
+    # through `__getattr__` but were never public names. Union in the lazy
+    # exports and whatever is already materialised so `dir()` and REPL
+    # completion do not hide things that import fine.
+    return sorted(set(__all__) | _EMBEDDING_TRANSFER_EXPORTS | set(globals()))
+
 
 __all__ = [
     "AsyncEncoderCache",

components/src/dynamo/common/tests/multimodal/test_lazy_imports.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""The package must not drag `torch` in behind the lightweight media loaders."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+
+pytestmark = [
+    pytest.mark.pre_merge,
+    pytest.mark.unit,
+    pytest.mark.gpu_0,
+    pytest.mark.timeout(60),
+]
+
+# A fresh interpreter is the only honest way to ask this: once any earlier test
+# has imported torch, an eager import here would go unnoticed.
+#
+# The answer is tagged and parsed off its own line rather than read from the
+# whole of stdout. Imports here are free to print (a CuPy fallback notice from
+# `dynamo.nixl_connect` is one that really happens), and a bare
+# `stdout.strip() == "True"` would read that extra output as False, turning the
+# regression this test exists to catch into a pass.
+_PROBE = """
+import sys
+from dynamo.common.multimodal import {names}
+print("RESULT:" + str("torch" in sys.modules))
+"""
+
+
+def _torch_imported_by(names: str) -> bool:
+    result = subprocess.run(
+        [sys.executable, "-c", _PROBE.format(names=names)],
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+    # Checked here rather than with check=True: that raises CalledProcessError,
+    # whose message names the command and the exit status only. The reason the
+    # child died is in stderr, which pytest never shows, so a missing torch or
+    # PIL or a raise during package import reads as a bare non-zero exit.
+    assert (
+        result.returncode == 0
+    ), f"probe for {names} exited {result.returncode}\nstderr:\n{result.stderr}"
+    answers = [
+        line.removeprefix("RESULT:")
+        for line in result.stdout.splitlines()
+        if line.startswith("RESULT:")
+    ]
+    assert len(answers) == 1, (
+        f"probe for {names} did not report exactly one result; "
+        f"stdout={result.stdout!r} stderr={result.stderr!r}"
+    )
+    assert answers[0] in ("True", "False"), f"unexpected probe answer {answers[0]!r}"
+    return answers[0] == "True"
+
+
+@pytest.mark.parametrize("name", ["ImageLoader", "AudioLoader", "VideoLoader"])
+def test_media_loader_import_does_not_pull_torch(name: str) -> None:
+    assert not _torch_imported_by(
+        name
+    ), f"importing {name} pulled in torch; the loaders must stay usable without the ML stack"
+
+
+def test_heavy_members_still_import() -> None:
+    assert _torch_imported_by(
+        "AsyncEncoderCache"
+    ), "AsyncEncoderCache should still resolve, and it legitimately needs torch"