CI: Accerate GitHub Actions using uv (#3676)

Setup [`uv`](https://github.com/astral-sh/uv) in the GitHub Actions, saving several minutes compared to pip. pip: ![image](https://github.com/deepmodeling/deepmd-kit/assets/9496702/547adb02-1bc2-47fb-953d-24d38e3e986d) uv: ![image](https://github.com/deepmodeling/deepmd-kit/assets/9496702/6ec6536b-5dcf-44c6-a4b6-c78d08b9c4f8) Using `uv` has some limitations, but it's good to use it in the CI. --------- Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Apr 16, 2024 · 219b19e · 219b19e
1 parent 25435c0
commit 219b19e
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 16 deletions.
diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml
@@ -30,9 +30,9 @@ jobs:
     - uses: actions/setup-python@v5
       with:
         python-version: '3.11'
-        cache: 'pip'
     - uses: lukka/get-cmake@latest
-    - run: python -m pip install tensorflow
+    - run: python -m pip install uv
+    - run: python -m uv pip install --system tensorflow
     - name: Download libtorch
       run: |
          wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip

diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml
@@ -26,7 +26,8 @@ jobs:
       with:
         mpi: mpich
     - uses: lukka/get-cmake@latest
-    - run: python -m pip install tensorflow
+    - run: python -m pip install uv
+    - run: python -m uv pip install --system tensorflow
     - name: Download libtorch
       run: |
          wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
@@ -49,8 +50,8 @@ jobs:
     # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843
     # TODO: remove ase version when ase has new release
     - run: |
-        python -m pip install -U pip
-        python -m pip install -e .[cpu,test,lmp] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
+        export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+        python -m uv pip install --system -e .[cpu,test,lmp] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
       env:
         DP_BUILD_TESTING: 1
       if: ${{ !matrix.check_memleak }}

diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
@@ -46,9 +46,11 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
-    - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
-    - run: python -m pip install "tensorflow>=2.15.0rc0" "torch>=2.2.0"
-    - run: python -m pip install -v -e .[gpu,test,lmp,cu12,torch] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
+    - run: python -m pip install -U uv
+    - run: python -m uv pip install --system "tensorflow>=2.15.0rc0" "torch>=2.2.0"
+    - run: |
+        export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+        python -m uv pip install --system -v -e .[gpu,test,lmp,cu12,torch] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1

diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
@@ -27,17 +27,19 @@ jobs:
     - uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python }}
-        cache: 'pip'
     - uses: mpi4py/setup-mpi@v1
       with:
         mpi: openmpi
-    # https://github.com/pypa/pip/issues/11770
-    - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
-    - run: pip install -e .[cpu,test,torch]
+    - run: python -m pip install -U uv
+    - run: uv pip install --system -e .[cpu,test,torch]
       env:
+        # Please note that uv has some issues with finding
+        # existing TensorFlow package. Currently, it uses
+        # TensorFlow in the build dependency, but if it
+        # changes, setting `TENSORFLOW_ROOT`.
         TENSORFLOW_VERSION: ${{ matrix.tf }}
         DP_BUILD_TESTING: 1
-    - run: pip install horovod mpi4py
+    - run: uv pip install --system --no-build-isolation horovod mpi4py
       env:
         HOROVOD_WITH_TENSORFLOW: 1
         HOROVOD_WITHOUT_PYTORCH: 1

diff --git a/source/install/docker/Dockerfile b/source/install/docker/Dockerfile
@@ -1,13 +1,15 @@
 FROM python:3.11 AS compile-image
 ARG VARIANT=""
 ARG CUDA_VERSION="12"
-RUN python -m venv /opt/deepmd-kit
+RUN python -m pip install uv
+RUN python -m uv venv /opt/deepmd-kit
 # Make sure we use the virtualenv
 ENV PATH="/opt/deepmd-kit/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/deepmd-kit"
 # Install package
 COPY dist /dist
-RUN if [ "${CUDA_VERSION}" = 11 ]; then pip install torch --index-url https://download.pytorch.org/whl/cu118; fi \
-    && pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \
+RUN if [ "${CUDA_VERSION}" = 11 ]; then uv pip install torch --index-url https://download.pytorch.org/whl/cu118; fi \
+    && uv pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \
     && dp -h \
     && lmp -h \
     && dp_ipi \