diff --git a/deepmd/lmp.py b/deepmd/lmp.py index 3daf269c92..fb5a406710 100644 --- a/deepmd/lmp.py +++ b/deepmd/lmp.py @@ -10,7 +10,6 @@ Path, ) -import torch # noqa: TID253 from packaging.version import ( Version, ) @@ -18,17 +17,6 @@ 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 diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 6488d0f31e..7c6e47a67a 100644 --- a/doc/install/easy-install.md +++ b/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] diff --git a/source/tests/test_lmp.py b/source/tests/test_lmp.py new file mode 100644 index 0000000000..0fc2fed002 --- /dev/null +++ b/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"