feat(lammps): load installed backends dynamically

deepmd/lmp.py

@@ -10,25 +10,13 @@
     Path,
 )
 
-import torch  # noqa: TID253
 from packaging.version import (
     Version,
 )
 
 from deepmd.env import (
     SHARED_LIB_DIR,
 )
-from deepmd.tf.env import (  # noqa: TID253
-    TF_VERSION,
-    tf,
-)
-
-if Version(TF_VERSION) < Version("2.12"):
-    from find_libpython import (
-        find_libpython,
-    )
-else:
-    find_libpython = None
 
 
 def get_env(paths: list[str | None]) -> str:
@@ -60,6 +48,36 @@ def get_library_path(module: str, filename: str) -> list[str]:
         return [str(lib) for lib in libs]
 
 
+def _get_tensorflow_library_paths() -> tuple[list[str], list[str]]:
+    """Get TensorFlow library and preload paths when TensorFlow is installed."""
+    try:
+        tf_env = import_module("deepmd.tf.env")
+    except ModuleNotFoundError as exc:
+        if exc.name == "tensorflow":
+            return [], []
+        raise
+
+    tf_dir = tf_env.tf.sysconfig.get_lib()
+    preload_paths = []
+    if Version(tf_env.TF_VERSION) < Version("2.12"):
+        find_libpython = import_module("find_libpython").find_libpython
+        libpython = find_libpython()
+        if libpython is not None:
+            preload_paths.append(libpython)
+    return [tf_dir, os.path.join(tf_dir, "python")], preload_paths
+
+
+def _get_pytorch_library_paths() -> list[str]:
+    """Get PyTorch library paths when PyTorch is installed."""
+    try:
+        torch = import_module("torch")
+    except ModuleNotFoundError as exc:
+        if exc.name == "torch":
+            return []
+        raise
+    return [os.path.join(torch.__path__[0], "lib")]
+
+
 if platform.system() == "Linux":
     lib_env = "LD_LIBRARY_PATH"
 elif platform.system() == "Darwin":
@@ -74,54 +92,51 @@ def get_library_path(module: str, filename: str) -> list[str]:
 else:
     raise RuntimeError("Unsupported platform")
 
-tf_dir = tf.sysconfig.get_lib()
-pt_dir = os.path.join(torch.__path__[0], "lib")
 op_dir = str(SHARED_LIB_DIR)
 
-cuda_library_paths = []
-if platform.system() == "Linux":
-    cuda_library_paths.extend(
-        [
-            *get_library_path("nvidia.cuda_runtime.lib", "libcudart.so*"),
-            *get_library_path("nvidia.cublas.lib", "libcublasLt.so*"),
-            *get_library_path("nvidia.cublas.lib", "libcublas.so*"),
-            *get_library_path("nvidia.cufft.lib", "libcufft.so*"),
-            *get_library_path("nvidia.curand.lib", "libcurand.so*"),
-            *get_library_path("nvidia.cusolver.lib", "libcusolver.so*"),
-            *get_library_path("nvidia.cusparse.lib", "libcusparse.so*"),
-            *get_library_path("nvidia.cudnn.lib", "libcudnn.so*"),
-        ]
-    )
 
-os.environ[preload_env] = get_env(
-    [
-        os.environ.get(preload_env),
-        *cuda_library_paths,
-    ]
-)
+def _configure_lammps_environment() -> None:
+    """Configure library paths for the installed LAMMPS backends."""
+    cuda_library_paths = []
+    if platform.system() == "Linux":
+        cuda_library_paths.extend(
+            [
+                *get_library_path("nvidia.cuda_runtime.lib", "libcudart.so*"),
+                *get_library_path("nvidia.cublas.lib", "libcublasLt.so*"),
+                *get_library_path("nvidia.cublas.lib", "libcublas.so*"),
+                *get_library_path("nvidia.cufft.lib", "libcufft.so*"),
+                *get_library_path("nvidia.curand.lib", "libcurand.so*"),
+                *get_library_path("nvidia.cusolver.lib", "libcusolver.so*"),
+                *get_library_path("nvidia.cusparse.lib", "libcusparse.so*"),
+                *get_library_path("nvidia.cudnn.lib", "libcudnn.so*"),
+            ]
+        )
+
+    tf_library_paths, tf_preload_paths = _get_tensorflow_library_paths()
+    pt_library_paths = _get_pytorch_library_paths()
 
-# set LD_LIBRARY_PATH
-os.environ[lib_env] = get_env(
-    [
-        os.environ.get(lib_env),
-        tf_dir,
-        os.path.join(tf_dir, "python"),
-        pt_dir,
-        op_dir,
-    ]
-)
-
-# preload python library, only for TF<2.12
-if find_libpython is not None:
-    libpython = find_libpython()
     os.environ[preload_env] = get_env(
         [
             os.environ.get(preload_env),
-            libpython,
+            *cuda_library_paths,
+            *tf_preload_paths,
+        ]
+    )
+
+    # set LD_LIBRARY_PATH
+    os.environ[lib_env] = get_env(
+        [
+            os.environ.get(lib_env),
+            *tf_library_paths,
+            *pt_library_paths,
+            op_dir,
         ]
     )
 
 
+_configure_lammps_environment()
+
+
 def get_op_dir() -> str:
     """Get the directory of the deepmd-kit OP library."""
     return op_dir

doc/install/easy-install.md

@@ -199,7 +199,7 @@ The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28
 If your platform is not supported, or you want to build against the installed backends, or you want to enable ROCM support, please [build from source](install-from-source.md).
 :::
 
-[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are provided on Linux and macOS for the TensorFlow, PyTorch, and JAX backend. It requires both TensorFlow and PyTorch. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` to extras:
+[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are provided on Linux and macOS for the TensorFlow, PyTorch, and JAX backend. The LAMMPS module loads the installed TensorFlow and/or PyTorch runtime libraries dynamically, so it does not require both backends to be installed. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` to extras:
 
 ```bash
 pip install deepmd-kit[gpu,cu12,lmp,ipi]

source/tests/test_lmp.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import os
+import platform
+from types import (
+    SimpleNamespace,
+)
+
+import pytest
+
+if platform.system() not in {"Linux", "Darwin"}:
+    pytest.skip("deepmd.lmp supports Linux and Darwin", allow_module_level=True)
+
+from deepmd import (
+    lmp,
+)
+
+
+def _missing_module(name: str) -> ModuleNotFoundError:
+    return ModuleNotFoundError(f"No module named {name!r}", name=name)
+
+
+def test_tensorflow_library_paths_skip_missing_tensorflow(monkeypatch):
+    def fake_import_module(module_name):
+        if module_name == "deepmd.tf.env":
+            raise _missing_module("tensorflow")
+        raise AssertionError(module_name)
+
+    monkeypatch.setattr(lmp, "import_module", fake_import_module)
+
+    assert lmp._get_tensorflow_library_paths() == ([], [])
+
+
+def test_tensorflow_library_paths_include_libpython_for_old_tensorflow(monkeypatch):
+    tf_env = SimpleNamespace(
+        TF_VERSION="2.11.0",
+        tf=SimpleNamespace(
+            sysconfig=SimpleNamespace(get_lib=lambda: "/opt/tensorflow")
+        ),
+    )
+    find_libpython = SimpleNamespace(find_libpython=lambda: "/opt/libpython.so")
+
+    def fake_import_module(module_name):
+        if module_name == "deepmd.tf.env":
+            return tf_env
+        if module_name == "find_libpython":
+            return find_libpython
+        raise AssertionError(module_name)
+
+    monkeypatch.setattr(lmp, "import_module", fake_import_module)
+
+    assert lmp._get_tensorflow_library_paths() == (
+        ["/opt/tensorflow", "/opt/tensorflow/python"],
+        ["/opt/libpython.so"],
+    )
+
+
+def test_pytorch_library_paths_skip_missing_torch(monkeypatch):
+    def fake_import_module(module_name):
+        if module_name == "torch":
+            raise _missing_module("torch")
+        raise AssertionError(module_name)
+
+    monkeypatch.setattr(lmp, "import_module", fake_import_module)
+
+    assert lmp._get_pytorch_library_paths() == []
+
+
+def test_configure_lammps_environment_keeps_op_dir_without_backends(monkeypatch):
+    monkeypatch.delenv(lmp.lib_env, raising=False)
+    monkeypatch.delenv(lmp.preload_env, raising=False)
+    monkeypatch.setattr(lmp, "_get_tensorflow_library_paths", lambda: ([], []))
+    monkeypatch.setattr(lmp, "_get_pytorch_library_paths", lambda: [])
+    monkeypatch.setattr(lmp, "get_library_path", lambda module, filename: [])
+
+    lmp._configure_lammps_environment()
+
+    assert os.environ[lmp.lib_env] == lmp.op_dir
+    assert os.environ[lmp.preload_env] == ""
+
+
+def test_configure_lammps_environment_adds_installed_backend_paths(monkeypatch):
+    monkeypatch.setenv(lmp.lib_env, "/existing/lib")
+    monkeypatch.setenv(lmp.preload_env, "/existing/preload")
+    monkeypatch.setattr(
+        lmp,
+        "_get_tensorflow_library_paths",
+        lambda: (["/tensorflow", "/tensorflow/python"], ["/libpython.so"]),
+    )
+    monkeypatch.setattr(lmp, "_get_pytorch_library_paths", lambda: ["/torch/lib"])
+    monkeypatch.setattr(lmp, "get_library_path", lambda module, filename: [])
+
+    lmp._configure_lammps_environment()
+
+    assert os.environ[lmp.lib_env] == ":".join(
+        [
+            "/existing/lib",
+            "/tensorflow",
+            "/tensorflow/python",
+            "/torch/lib",
+            lmp.op_dir,
+        ]
+    )
+    assert os.environ[lmp.preload_env] == "/existing/preload:/libpython.so"