ci: Install CUDA version specified torch first during container building. (#2167)

bkryu · web-flow · commit 1e15fede3bd9 · 2025-12-03T15:12:49.000-08:00
## 📌 Description In the existing `install_python_packages.sh`, we install requirements in requirements.txt first followed by a torch installation with cuda version specified. This means `pip install torch` is done with no CUDA version specified. The result is [this example](https://github.com/flashinfer-ai/flashinfer/actions/runs/19879133946/job/56973105583#step:6:413) where cu12 dependencies are installed. Then, torch cu130 is installed, which updates the torch itself, but not the dependencies, resulting in duplicated torch dependencies -- in cu12 and cu13. Current PR moves the torch installation earlier so that the torch requirement in requirements.txt is already satisfied. Resultant pipeline: * [torch cu13 installs first, with cu13 dependencies](https://github.com/flashinfer-ai/flashinfer/actions/runs/19905525901/job/57060644811?pr=2167#step:6:332) * Then the [torch in requirements.txt is skipped](https://github.com/flashinfer-ai/flashinfer/actions/runs/19905525901/job/57060644811?pr=2167#step:6:475)  ## 🔍 Related Issues  ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes
diff --git a/docker/install/install_python_packages.sh b/docker/install/install_python_packages.sh
@@ -22,9 +22,10 @@ set -u
 # Accept CUDA version as parameter (e.g., cu126, cu128, cu129)
 CUDA_VERSION=${1:-cu128}
 
-# Install requirements.txt first, so we can override any versions later for docker image
-pip3 install -r /install/requirements.txt
+# Install torch with specific CUDA version first, followed by others in requirements.txt, and then others.
+# This is to ensure that the torch version is compatible with the CUDA version.
 pip3 install --force-reinstall torch --index-url https://download.pytorch.org/whl/${CUDA_VERSION}
+pip3 install -r /install/requirements.txt
 pip3 install responses pytest scipy build cuda-python nvidia-nvshmem-cu12
 
 # Install cudnn package based on CUDA version
