deepmd-kit/.github/workflows/test_cuda.yml at devel · Johnsonjz/deepmd-kit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
on:
  # manually trigger
  workflow_dispatch:
  pull_request:
    types:
      - "labeled"
      # to let the PR pass the test
      - "opened"
      - "reopened"
      - "synchronize"
  merge_group:
concurrency:
  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
  cancel-in-progress: true
name: Test CUDA
env:
  # Temporarily disable Paddle in CUDA CI while its external downloads are
  # unreliable. Keep this switch explicit so the coverage can be restored.
  DP_CI_ALLOW_MISSING_PADDLE: "1"
  DP_ENABLE_PADDLE: "0"
jobs:
  test_python:
    name: Test Python on CUDA
    runs-on: gpu
    # The serial Python test suite can take several hours on the GPU runner.
    timeout-minutes: 480
    # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
    # container:
    #   image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
    #   options: --gpus all
    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
    steps:
      - uses: actions/checkout@v7
      - uses: actions/setup-python@v7
        with:
          python-version: "3.11"
          # cache: 'pip'
      - name: Install wget and unzip
        run: sudo apt-get update && sudo apt-get install -y wget unzip
      - uses: lukka/get-cmake@latest
        with:
          useLocalCache: true
          useCloudCache: false
      - run: |
          UBUNTU_VERSION=$(lsb_release -rs | tr -d '.')
          OS_NAME="ubuntu${UBUNTU_VERSION}"
          echo "Current OS: ${OS_NAME}"

          wget https://developer.download.nvidia.com/compute/cuda/repos/${OS_NAME}/x86_64/cuda-keyring_1.1-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.1-1_all.deb \
          && sudo apt-get update \
          && sudo apt-get -y install cuda-toolkit-12-9 cudnn9-cuda-12
          echo "CUDA_PATH=/usr/local/cuda-12.9" >> $GITHUB_ENV
          echo "/usr/local/cuda-12.9/bin" >> $GITHUB_PATH
        # if: false # skip as we use nvidia image
      - run: python -m pip install -U uv
      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax_gpu
      - run: |
          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
          export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
          # See https://github.com/jax-ml/jax/issues/29042
          source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
        env:
          DP_VARIANT: cuda
          DP_ENABLE_NATIVE_OPTIMIZATION: 1
          DP_ENABLE_PYTORCH: 1
      - run: dp --version
      - run: python -m pytest source/tests --ignore=source/tests/pd
        env:
          NUM_WORKERS: 0
          CUDA_VISIBLE_DEVICES: 0
          # See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
          XLA_PYTHON_CLIENT_PREALLOCATE: false
          XLA_PYTHON_CLIENT_ALLOCATOR: platform
          FLAGS_use_stride_compute_kernel: 0

  test_cc:
    name: Test C++ on CUDA
    runs-on: gpu
    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
    steps:
      # Jobs run on separate runners, so the C++ job needs its own complete
      # CUDA toolchain and Python dependency installation.
      - uses: actions/checkout@v7
      - uses: actions/setup-python@v7
        with:
          python-version: "3.11"
          # cache: 'pip'
      - name: Install wget and unzip
        run: sudo apt-get update && sudo apt-get install -y wget unzip
      - uses: lukka/get-cmake@latest
        with:
          useLocalCache: true
          useCloudCache: false
      - run: |
          UBUNTU_VERSION=$(lsb_release -rs | tr -d '.')
          OS_NAME="ubuntu${UBUNTU_VERSION}"
          echo "Current OS: ${OS_NAME}"

          wget https://developer.download.nvidia.com/compute/cuda/repos/${OS_NAME}/x86_64/cuda-keyring_1.1-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.1-1_all.deb \
          && sudo apt-get update \
          && sudo apt-get -y install cuda-toolkit-12-9 cudnn9-cuda-12
          echo "CUDA_PATH=/usr/local/cuda-12.9" >> $GITHUB_ENV
          echo "/usr/local/cuda-12.9/bin" >> $GITHUB_PATH
        # if: false # skip as we use nvidia image
      - run: python -m pip install -U uv
      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax_gpu
      - run: |
          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
          export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
          # See https://github.com/jax-ml/jax/issues/29042
          source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
        env:
          DP_VARIANT: cuda
          DP_ENABLE_NATIVE_OPTIMIZATION: 1
          DP_ENABLE_PYTORCH: 1
      - run: dp --version
      - name: Convert models
        run: source/tests/infer/convert-models.sh
      - run: |
          export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
          source/install/test_cc_local.sh
        env:
          OMP_NUM_THREADS: 1
          TF_INTRA_OP_PARALLELISM_THREADS: 1
          TF_INTER_OP_PARALLELISM_THREADS: 1
          CMAKE_GENERATOR: Ninja
          DP_VARIANT: cuda
          DP_USE_MPICH2: 1
          ENABLE_PADDLE: FALSE
      - run: |
          export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH
          export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
          python -m pytest -s source/lmp/tests || (cat log.lammps && exit 1)
          python -m pytest source/ipi/tests
        env:
          OMP_NUM_THREADS: 1
          TF_INTRA_OP_PARALLELISM_THREADS: 1
          TF_INTER_OP_PARALLELISM_THREADS: 1
          LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
          CUDA_VISIBLE_DEVICES: 0
          ENABLE_PADDLE: 0
  pass:
    name: Pass testing on CUDA
    needs: [test_python, test_cc]
    runs-on: ubuntu-slim
    if: always()
    steps:
      - name: Decide whether the needed jobs succeeded or failed
        uses: re-actors/alls-green@release/v1
        with:
          jobs: ${{ toJSON(needs) }}
          allowed-skips: test_python, test_cc