ci(cuda): overlap AOTI .pt2 compiles with GPU unit tests #11202
Workflow file for this run
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 | |
| jobs: | |
| test_cuda: | |
| name: Test Python and C++ on CUDA | |
| runs-on: gpu | |
| # The full CUDA suite (serial pytest ~3h43m + C++ + LAMMPS) exceeds the | |
| # default 360-min job limit; the self-hosted GPU runner has no hard cap. | |
| timeout-minutes: 480 | |
| # Share ONE AOTInductor on-disk compile cache across every step: the two | |
| # pytest lanes AND the C++ fixture builds (gen_*.py in test_cc_local.sh) | |
| # all drive the same inductor backend on the same architectures, so the C++ | |
| # fixture compiles reuse the kernels the Python lane just built. actions/cache | |
| # (below) persists it across runs so unchanged models never recompile. | |
| env: | |
| TORCHINDUCTOR_CACHE_DIR: ${{ github.workspace }}/.inductor-cache | |
| # 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 | |
| - name: Cache AOTInductor compile artifacts | |
| # Restored AFTER checkout so checkout's git-clean can't wipe it. The repo | |
| # Actions cache is capped at 10 GB; if the inductor cache outgrows it the | |
| # oldest entries are evicted -- correctness is unaffected (inductor | |
| # re-validates each graph; a miss just recompiles). | |
| uses: actions/cache@v6 | |
| with: | |
| path: ${{ github.workspace }}/.inductor-cache | |
| key: inductor-cuda-${{ github.sha }} | |
| restore-keys: | | |
| inductor-cuda- | |
| - uses: actions/setup-python@v6 | |
| 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)') | |
| pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310" | |
| 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: Run Python tests + C++ build (overlap CPU compile work with the GPU lane) | |
| # Two lanes on one GPU box, sized so only TWO heavy workstreams ever run | |
| # at once -- the shape that was green before the C++ prep was (wrongly) | |
| # added as a THIRD concurrent stream (that 3-way overlap SIGFAULTed both | |
| # Python lanes via TORCHINDUCTOR_CACHE_DIR contention; see #5882): | |
| # Lane 1 (CPU compile stream): the .pt2-freezing tests (-m aoti_compile) | |
| # -- CPU-bound, GPU ~98% idle while inductor/g++/ptxas run -- THEN the | |
| # C++ build + gen_*.py .pt2 fixtures (also CPU-bound, GPU idle). They | |
| # run SEQUENTIALLY in one lane, so the compile work never races the | |
| # shared inductor cache and cc-prep reuses the kernels the aoti tests | |
| # just built. | |
| # Lane 2 (GPU stream): the GPU-bound tests (-m "not aoti_compile") -- the | |
| # ~19k-test bulk and the wall-clock bottleneck. It does not freeze | |
| # .pt2, so it only reads the inductor cache, never competing to write. | |
| # The whole CPU stream is hidden under lane 2; ctest + LAMMPS run after | |
| # BOTH lanes join. | |
| run: | | |
| # Cap the compile lane's inductor threads to ~half the cores so lane 2 | |
| # keeps CPU for its host-side GPU dispatch; nice the compile work too. | |
| nc=$(( $(nproc) / 2 )); [ "$nc" -lt 1 ] && nc=1 | |
| # rc=0 + `|| rc=$?` so a lane's failure does not let the step's default | |
| # `set -e` skip the wall-clock echo (the timing we most want on failure). | |
| # Lane 1: aoti_compile pytest, THEN the C++ build + fixture generation. | |
| ( ta=$(date +%s); rc=0 | |
| TORCHINDUCTOR_COMPILE_THREADS=$nc nice -n 15 \ | |
| python -m pytest source/tests -m "aoti_compile" \ | |
| -p no:cacheprovider --junit-xml=junit-aoti.xml || rc=$? | |
| # pytest exit 5 == "no tests collected"; harmless for the compile lane. | |
| [ "$rc" = 5 ] && rc=0 | |
| echo "[lane aoti_compile] wall=$(( $(date +%s) - ta ))s exit=$rc" | |
| tc=$(date +%s); rc_cc=0 | |
| ( export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH | |
| export 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 | |
| source/tests/infer/convert-models.sh | |
| DP_CC_SKIP_CTEST=1 nice -n 19 source/install/test_cc_local.sh ) || rc_cc=$? | |
| echo "[lane cc-prep] wall=$(( $(date +%s) - tc ))s exit=$rc_cc" | |
| exit $(( rc != 0 || rc_cc != 0 )) ) & | |
| pid_a=$! | |
| # Lane 2: the GPU-bound unit tests (the bulk). | |
| ( tb=$(date +%s); rc=0 | |
| python -m pytest source/tests -m "not aoti_compile" \ | |
| -p no:cacheprovider --junit-xml=junit-gpu.xml || rc=$? | |
| echo "[lane gpu] wall=$(( $(date +%s) - tb ))s exit=$rc" | |
| exit $rc ) & | |
| pid_b=$! | |
| rc_a=0; wait $pid_a || rc_a=$? | |
| rc_b=0; wait $pid_b || rc_b=$? | |
| echo "lane exit codes: aoti_compile+cc-prep=$rc_a gpu=$rc_b" | |
| exit $(( rc_a != 0 || rc_b != 0 )) | |
| 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 | |
| - name: Upload per-lane JUnit reports | |
| # junit-aoti.xml = compile lane (T_A), junit-gpu.xml = GPU lane (T_B); | |
| # used to compare the two lanes against the pre-split serial baseline. | |
| if: always() | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: cuda-pytest-junit | |
| path: | | |
| junit-aoti.xml | |
| junit-gpu.xml | |
| if-no-files-found: warn | |
| - name: C++ tests (ctest; build + fixtures done in lane 1 above) | |
| run: | | |
| export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH | |
| DP_CC_SKIP_BUILD=1 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 | |
| - 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 | |
| cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/paddle/lib/* $GITHUB_WORKSPACE/dp_test/lib/ | |
| cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/third_party/install/onednn/lib/* $GITHUB_WORKSPACE/dp_test/lib/ | |
| cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/third_party/install/mklml/lib/* $GITHUB_WORKSPACE/dp_test/lib/ | |
| 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 | |
| pass: | |
| name: Pass testing on CUDA | |
| needs: [test_cuda] | |
| 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_cuda |