chore: manage CI pinnings in pyproject.toml

.github/workflows/build_cc.yml

@@ -35,8 +35,7 @@ jobs:
         python-version: '3.11'
     - uses: lukka/get-cmake@latest
     - run: python -m pip install uv
-    - run: source/install/uv_with_retry.sh pip install --system tensorflow
-    - run: source/install/uv_with_retry.sh pip install --system 'torch==2.8.*' --index-url https://download.pytorch.org/whl/cpu
+    - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \

.github/workflows/codeql.yml

@@ -43,8 +43,8 @@ jobs:
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
          && sudo apt-get update \
          && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2
-        python -m pip install tensorflow
-        python -m pip install 'torch==2.8.*' --index-url https://download.pytorch.org/whl/cpu
+        python -m pip install uv
+        uv pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
       env:
         DEBIAN_FRONTEND: noninteractive
     # Initializes the CodeQL tools for scanning.

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

@@ -49,10 +49,7 @@ jobs:
         run: echo "VIRTUAL_ENV=$PWD/venv" >> $GITHUB_ENV && echo "$PWD/venv/bin" >> $GITHUB_PATH
 
       - name: Install base dependencies
-        run: uv pip install tensorflow-cpu
-
-      - name: Install PyTorch
-        run: uv pip install torch --index-url https://download.pytorch.org/whl/cpu
+        run: uv pip install --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
 
       - name: Build Python package
         run: uv pip install -e .[cpu,test]

.github/workflows/test_cc.yml

@@ -28,10 +28,9 @@ jobs:
     - run: python -m pip install uv
     - name: Install Python dependencies
       run: |
-        source/install/uv_with_retry.sh pip install --system tensorflow-cpu~=2.18.0 jax==0.5.0
+        source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax --torch-backend cpu
         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 -e .[cpu,test,lmp,jax] mpi4py mpich
-        source/install/uv_with_retry.sh pip install --system 'torch==2.8.*' --index-url https://download.pytorch.org/whl/cpu
     - name: Convert models
       run: source/tests/infer/convert-models.sh
     # https://github.com/actions/runner-images/issues/9491

.github/workflows/test_cuda.yml

@@ -43,7 +43,7 @@ jobs:
          && 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 uv
-    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.7.0" "jax[cuda12]==0.5.0"
+    - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax "jax[cuda12]"
     - 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)')

.github/workflows/test_python.yml

@@ -28,11 +28,10 @@ jobs:
         python-version: ${{ matrix.python }}
     - run: python -m pip install -U uv
     - run: |
-        source/install/uv_with_retry.sh pip install --system openmpi tensorflow-cpu~=2.18.0
-        source/install/uv_with_retry.sh pip install --system torch -i https://download.pytorch.org/whl/cpu
+        source/install/uv_with_retry.sh pip install --system openmpi --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
         export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.5.0;python_version>='3.10'"
+        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py --group pin_jax
         source/install/uv_with_retry.sh pip install --system -U setuptools
         source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
         source/install/uv_with_retry.sh pip install --system --pre "paddlepaddle==3.0.0" -i https://www.paddlepaddle.org.cn/packages/stable/cpu/

pyproject.toml

@@ -158,6 +158,21 @@ dev = [
   "cmake",
   "mpich",
 ]
+pin_tensorflow_cpu = [
+  "tensorflow-cpu~=2.18.0",
+]
+pin_tensorflow_gpu = [
+  "tensorflow~=2.18.0",
+]
+pin_pytorch_cpu = [
+  "torch~=2.8.0",
+]
+pin_pytorch_gpu = [
+  "torch~=2.7.0",
+]
+pin_jax = [
+  "jax==0.5.0;python_version>='3.10'",
+]
 
 [tool.setuptools_scm]