Merge branch 'master' into test/reuse-pt2-aot-fixtures

Author: njzjz

.github/workflows/build_cc.yml

@@ -28,7 +28,7 @@ jobs:
             dp_variant: clang
     steps:
       - uses: actions/checkout@v7
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.11"
       - uses: lukka/get-cmake@latest

.github/workflows/build_wheel.yml

@@ -162,7 +162,7 @@ jobs:
           path: dist/packages
           pattern: cibw-*
           merge-multiple: true
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@v7
         name: Install Python
         with:
           python-version: "3.11"

.github/workflows/codeql.yml

@@ -31,7 +31,7 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v7
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.11"
           cache: "pip"

.github/workflows/copilot-setup-steps.yml

@@ -33,7 +33,7 @@ jobs:
         uses: actions/checkout@v7
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: "3.10"
 

.github/workflows/labeler.yml

@@ -9,6 +9,6 @@ jobs:
       pull-requests: write
     runs-on: ubuntu-slim
     steps:
-      - uses: actions/labeler@v6
+      - uses: actions/labeler@v7
         with:
           repo-token: "${{ secrets.GITHUB_TOKEN }}"

.github/workflows/test_cc.yml

@@ -14,9 +14,11 @@ name: Test C++
 jobs:
   testcc:
     name: Test C++
-    runs-on: ubuntu-22.04
+    runs-on: ${{ matrix.check_memleak == false && github.repository_owner == 'deepmodeling' && 'cpu' || 'ubuntu-22.04' }}
     strategy:
       matrix:
+        # Only regular jobs run the LAMMPS tests. Use self-hosted CPU runners
+        # for those jobs in deepmodeling, with a hosted fallback for forks.
         include:
           - check_memleak: true
             enable_tensorflow: true
@@ -36,7 +38,12 @@ jobs:
             enable_paddle: true
     steps:
       - uses: actions/checkout@v7
-      - uses: actions/setup-python@v6
+      - name: Install GCC
+        run: |
+          sudo apt-get update
+          sudo apt-get install --yes build-essential
+        if: ${{ matrix.check_memleak == false && github.repository_owner == 'deepmodeling' }}
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.11"
           cache: "pip"

.github/workflows/test_cuda.yml

@@ -14,17 +14,19 @@ concurrency:
   cancel-in-progress: true
 name: Test CUDA
 jobs:
-  test_cuda:
-    name: Test Python and C++ on CUDA
+  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@v6
+      - uses: actions/setup-python@v7
         with:
           python-version: "3.11"
           # cache: 'pip'
@@ -68,6 +70,51 @@ jobs:
           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)')
+          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: Convert models
         run: source/tests/infer/convert-models.sh
       - run: |
@@ -96,12 +143,12 @@ jobs:
           CUDA_VISIBLE_DEVICES: 0
   pass:
     name: Pass testing on CUDA
-    needs: [test_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_cuda
+          allowed-skips: test_python, test_cc

.github/workflows/test_python.yml

@@ -25,7 +25,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v7
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@v7
         with:
           python-version: ${{ matrix.python }}
       - run: python -m pip install -U uv
@@ -56,41 +56,59 @@ jobs:
           # the key must never match, even when restarting workflows, as that
           # will cause durations to get out of sync between groups, the
           # combined durations will be loaded if available
-          key: test2-durations-split-${{ github.run_id }}-${{ github.run_number}}-${{ matrix.python }}-${{ matrix.group }}
+          key: test4-durations-split-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{ matrix.python }}-${{ matrix.group }}
           restore-keys: |
-            test2-durations-combined-${{ matrix.python }}-${{ github.sha }}
-            test2-durations-combined-${{ matrix.python }}
-      - run: pytest --cov=deepmd source/tests --ignore=source/tests/tf2 --splits 12 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration
+            test4-durations-combined-${{ matrix.python }}-${{ github.sha }}
+            test4-durations-combined-${{ matrix.python }}
+      - name: Prepare independent duration inputs
+        run: |
+          for output in main tf2 jax2tf consistent_tf2; do
+            if [ -f .test_durations ]; then
+              cp .test_durations ".test_durations_${output}"
+            fi
+          done
+      # Keep fixture-sharing classes/modules together, then balance those
+      # indivisible units by their recorded durations.
+      - run: python -m pytest -p source.tests.ci_split --cov=deepmd source/tests --ignore=source/tests/tf2 --ci-splits 12 --ci-group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations_main
         env:
           NUM_WORKERS: 0
           DP_CI_IMPORT_PADDLE_BEFORE_TF: 1
           FLAGS_use_stride_compute_kernel: 0
       - name: Test TF2 eager mode
         run: |
           run_pytest_allow_no_tests() {
+            local durations_path=$1
+            shift
             set +e
-            pytest "$@"
+            pytest "$@" \
+              --store-durations \
+              --clean-durations \
+              --durations-path="$durations_path"
             local status=$?
             set -e
             if [ "$status" -eq 5 ]; then
               # pytest-split may leave an individual shard with no selected
               # tests after path/-k filtering. Other shards still cover the
               # selected tests, so do not fail the whole matrix for exit 5.
+              printf '{}\n' > "$durations_path"
               return 0
             fi
             return "$status"
           }
 
-          run_pytest_allow_no_tests --cov=deepmd --cov-append \
+          run_pytest_allow_no_tests .test_durations_tf2 \
+            --cov=deepmd --cov-append \
             source/tests/tf2 \
             --splits 12 \
             --group ${{ matrix.group }}
-          run_pytest_allow_no_tests --cov=deepmd --cov-append \
+          run_pytest_allow_no_tests .test_durations_jax2tf \
+            --cov=deepmd --cov-append \
             source/tests/consistent/io/test_io.py \
             source/jax2tf_tests \
             --splits 12 \
             --group ${{ matrix.group }}
-          run_pytest_allow_no_tests --cov=deepmd --cov-append \
+          run_pytest_allow_no_tests .test_durations_consistent_tf2 \
+            --cov=deepmd --cov-append \
             source/tests/consistent \
             -k tf2 \
             --splits 12 \
@@ -100,7 +118,8 @@ jobs:
           DP_TEST_TF2_ONLY: 1
           DP_DTYPE_PROMOTION_STRICT: 1
           DP_CI_IMPORT_PADDLE_BEFORE_TF: 1
-      - run: mv .test_durations .test_durations_${{ matrix.group }}
+      - name: Combine shard durations
+        run: jq -s 'reduce .[] as $durations ({}; . * $durations)' .test_durations_main .test_durations_tf2 .test_durations_jax2tf .test_durations_consistent_tf2 > .test_durations_${{ matrix.group }}
       - name: Upload partial durations
         uses: actions/upload-artifact@v7
         with:
@@ -132,15 +151,15 @@ jobs:
           # key won't match during the first run for the given commit, but
           # restore-key will if there's a previous stored durations file,
           # so cache will both be loaded and stored
-          key: test2-durations-combined-${{ matrix.python }}-${{ github.sha }}
-          restore-keys: test2-durations-combined-${{ matrix.python }}
+          key: test4-durations-combined-${{ matrix.python }}-${{ github.sha }}
+          restore-keys: test4-durations-combined-${{ matrix.python }}
       - name: Download artifacts
         uses: actions/download-artifact@v8
         with:
           pattern: split-${{ matrix.python }}-*
           merge-multiple: true
       - name: Combine test durations
-        run: jq -s add .test_durations_* > .test_durations
+        run: jq -s 'reduce .[] as $durations ({}; . * $durations)' .test_durations_* > .test_durations
   pass:
     name: Pass testing Python
     needs: [testpython, update_durations]

.pre-commit-config.yaml

@@ -23,14 +23,14 @@ repos:
       - id: check-toml
   # Python
   - repo: https://github.com/PyCQA/isort
-    rev: 9.0.0a3
+    rev: 9.0.0b1
     hooks:
       - id: isort
         files: \.py$
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.18
+    rev: v0.16.0
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -62,7 +62,7 @@ repos:
           - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.5
+    rev: v22.1.8
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)

deepmd/_vendors/ndtensorflow/_namespace.py

@@ -1650,7 +1650,11 @@ def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array:
         axis = 0
     axis = _normalize_axis(axis, tensor.shape.rank)
     indices_ = tf.cast(_unwrap(indices), tf.int64)
-    dim = tensor.shape[axis]
+    # ``TensorShape.__getitem__`` returns either ``None`` or
+    # ``Dimension(None)`` depending on TensorFlow's global v2-shape setting.
+    # ``as_list`` normalizes both representations, which keeps dynamic-axis
+    # indexing valid when another test or caller enables legacy TensorShape.
+    dim = tensor.shape.as_list()[axis]
     dim = (
         tf.shape(tensor, out_type=tf.int64)[axis]
         if dim is None
@@ -1664,7 +1668,7 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     tensor = _unwrap(x)
     indices_ = tf.cast(_unwrap(indices), tf.int64)
     axis = _normalize_axis(axis, tensor.shape.rank)
-    dim = tensor.shape[axis]
+    dim = tensor.shape.as_list()[axis]
     dim = (
         tf.shape(tensor, out_type=tf.int64)[axis]
         if dim is None