build Docker images for each push (#2324)

and push the image to the GitHub container.

Also, provide the pypi CUDA support for pypi LAMMPS.

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

.github/workflows/build_wheel.yml

@@ -36,6 +36,8 @@ jobs:
       - uses: actions/checkout@v2
         with:
           submodules: true
+          # https://github.com/pypa/setuptools_scm/issues/480
+          fetch-depth: 0
       - uses: docker/setup-qemu-action@v2
         name: Setup QEMU
         if: matrix.platform_id == 'manylinux_aarch64'
@@ -82,6 +84,37 @@ jobs:
           user: __token__
           password: ${{ secrets.pypi_password }}
 
+  build_docker:
+    # use the already built wheels to build docker
+    needs: [build_wheels]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/download-artifact@v3
+        with:
+          name: artifact
+          path: source/install/docker/dist
+      - name: Log in to the Container registry
+        uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract metadata (tags, labels) for Docker
+        id: meta
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ghcr.io/deepmodeling/deepmd-kit
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@ad44023a93711e3deb337508980b4b5e9bcdc5dc
+        with:
+          context: source/install/docker
+          push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' }}
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+
   pass:
     name: Pass testing build wheels
     needs: [build_wheels, build_sdist]

deepmd/lmp.py

@@ -1,6 +1,9 @@
 """Register entry points for lammps-wheel."""
 import os
 import platform
+from importlib import (
+    import_module,
+)
 from pathlib import (
     Path,
 )
@@ -23,6 +26,27 @@ def get_env(paths: List[Optional[str]]) -> str:
     return ":".join((p for p in paths if p is not None))
 
 
+def get_library_path(module: str) -> List[str]:
+    """Get library path from a module.
+
+    Parameters
+    ----------
+    module : str
+        The module name.
+
+    Returns
+    -------
+    list[str]
+        The library path.
+    """
+    try:
+        m = import_module(module)
+    except ModuleNotFoundError:
+        return []
+    else:
+        return [str(Path(m.__file__).parent)]
+
+
 if platform.system() == "Linux":
     lib_env = "LD_LIBRARY_PATH"
 elif platform.system() == "Darwin":
@@ -32,13 +56,31 @@ def get_env(paths: List[Optional[str]]) -> str:
 
 tf_dir = tf.sysconfig.get_lib()
 op_dir = str((Path(__file__).parent / "op").absolute())
+
+
+cuda_library_paths = []
+if platform.system() == "Linux":
+    cuda_library_paths.extend(
+        [
+            *get_library_path("nvidia.cuda_runtime.lib"),
+            *get_library_path("nvidia.cublas.lib"),
+            *get_library_path("nvidia.cublas.lib"),
+            *get_library_path("nvidia.cufft.lib"),
+            *get_library_path("nvidia.curand.lib"),
+            *get_library_path("nvidia.cusolver.lib"),
+            *get_library_path("nvidia.cusparse.lib"),
+            *get_library_path("nvidia.cudnn.lib"),
+        ]
+    )
+
 # set LD_LIBRARY_PATH
 os.environ[lib_env] = get_env(
     [
         os.environ.get(lib_env),
         tf_dir,
         os.path.join(tf_dir, "python"),
         op_dir,
+        *cuda_library_paths,
     ]
 )
 

source/install/docker/Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.10 AS compile-image
+RUN python -m venv /opt/deepmd-kit
+# Make sure we use the virtualenv
+ENV PATH="/opt/deepmd-kit/bin:$PATH"
+# Install package
+COPY dist /dist
+RUN pip install "$(ls /dist/deepmd_kit-*manylinux*_x86_64.whl)[gpu,cu11,lmp,ipi]" \
+    && dp -h \
+    && lmp -h \
+    && dp_ipi
+
+FROM python:3.10 AS build-image
+COPY --from=compile-image /opt/deepmd-kit /opt/deepmd-kit
+ENV PATH="/opt/deepmd-kit/bin:$PATH"
+CMD ["/bin/bash"]