From a78fe4bd0bf81ecb8a26604a6530f466bf2c6329 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 4 Jan 2023 01:50:57 -0500 Subject: [PATCH] add cuda toolkit to optional dependencies In this way, one can install CUDA Toolkit and cuDNN via pip if one does not want to install them manually. These packages are provided officially by NVIDIA and can be dlopen before tensorflow. Signed-off-by: Jinzhe Zeng --- deepmd/env.py | 36 +++++++++++++++++++++++++++++++++++- doc/install/easy-install.md | 6 ++++-- setup.py | 18 ++++++++++++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/deepmd/env.py b/deepmd/env.py index 9fcde6fd05..620a87008d 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -4,8 +4,9 @@ import os import re import platform +import ctypes from configparser import ConfigParser -from importlib import reload +from importlib import reload, import_module from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from packaging.version import Version @@ -15,6 +16,39 @@ if TYPE_CHECKING: from types import ModuleType + +def dlopen_library(module: str, filename: str): + """Dlopen a library from a module. + + Parameters + ---------- + module : str + The module name. + filename : str + The library filename pattern. + """ + try: + m = import_module(module) + except ModuleNotFoundError: + pass + else: + libs = sorted(Path(m.__file__).parent.glob(filename)) + # hope that there is only one version installed... + if len(libs): + ctypes.CDLL(str(libs[0].absolute())) + +# dlopen pip cuda library before tensorflow +if platform.system() == "Linux": + dlopen_library("nvidia.cuda_runtime.lib", "libcudart.so*") + dlopen_library("nvidia.cublas.lib", "libcublas.so*") + dlopen_library("nvidia.cublas.lib", "libcublasLt.so*") + dlopen_library("nvidia.cufft.lib", "libcufft.so*") + dlopen_library("nvidia.curand.lib", "libcurand.so*") + dlopen_library("nvidia.cusolver.lib", "libcusolver.so*") + dlopen_library("nvidia.cusparse.lib", "libcusparse.so*") + dlopen_library("nvidia.cudnn.lib", "libcudnn.so*") + + # import tensorflow v1 compatability try: import tensorflow.compat.v1 as tf diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 9e5e955332..011c5869b2 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -87,9 +87,11 @@ docker pull deepmodeling/dpmdkit-rocm:dp2.0.3-rocm4.5.2-tf2.6-lmp29Sep2021 If you have no existing TensorFlow installed, you can use `pip` to install the pre-built package of the Python interface with CUDA 11 supported: ```bash -pip install deepmd-kit[gpu] +pip install deepmd-kit[gpu,cu11] ``` +`cu11` is required only when CUDA Toolkit and cuDNN were not installed. + Or install the CPU version without CUDA supported: ```bash pip install deepmd-kit[cpu] @@ -97,7 +99,7 @@ pip install deepmd-kit[cpu] [LAMMPS module](../third-party/lammps-command.md) is only provided on Linux and macOS. To enable it, add `lmp` to extras: ```bash -pip install deepmd-kit[gpu,lmp] +pip install deepmd-kit[gpu,cu11,lmp] ``` MPICH is required for parallel running. diff --git a/setup.py b/setup.py index 76b2cbafd5..e89f373080 100644 --- a/setup.py +++ b/setup.py @@ -107,6 +107,24 @@ def get_tag(self): "find_libpython", ], **get_tf_requirement(tf_version), + "cu11": [ + "nvidia-cuda-runtime-cu11", + "nvidia-cublas-cu11", + "nvidia-cufft-cu11", + "nvidia-curand-cu11", + "nvidia-cusolver-cu11", + "nvidia-cusparse-cu11", + "nvidia-cudnn-cu11", + ], + "cu12": [ + "nvidia-cuda-runtime-cu12", + "nvidia-cublas-cu12", + "nvidia-cufft-cu12", + "nvidia-curand-cu12", + "nvidia-cusolver-cu12", + "nvidia-cusparse-cu12", + "nvidia-cudnn-cu12", + ], }, entry_points={ "console_scripts": ["dp = deepmd.entrypoints.main:main"],