docker image improvements (#1890)

<!-- .github/pull_request_template.md -->

## 📌 Description

1. Added more triggers to push a new docker image
2. Added an update-docker-tags job. This creates a PR to update
docker-tags, when new images are pushed
3. Swapped order of torch installation with requirements, to ensure the
desired torch package gets installed in the docker image


Example workflow: 
commit b3f22f2 on
https://github.com/flashinfer-ai/flashinfer/tree/refined_docker_sha
triggered the following github workflow:


https://github.com/flashinfer-ai/flashinfer/actions/runs/18382220944/job/52374357169

This succesfully resulted in generating an automated PR here:
#1896
## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 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

- [ ] 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`.
- [ ] 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

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: nvmbreughe

.github/workflows/release-ci-docker.yml

@@ -8,13 +8,14 @@ on:
     paths:
       - 'docker/**'
       - '.github/workflows/release-ci-docker.yml'
+      - 'requirements.txt'
   pull_request:
     branches:
       - main
     paths:
       - 'docker/**'
       - '.github/workflows/release-ci-docker.yml'
-
+      - 'requirements.txt'
 jobs:
   generate-tag:
     runs-on: ubuntu-latest
@@ -104,3 +105,42 @@ jobs:
             -t flashinfer/flashinfer-ci-${{ matrix.cuda }}:latest \
             flashinfer/flashinfer-ci-${{ matrix.cuda }}:amd64-${DATE_SHA} \
             flashinfer/flashinfer-ci-${{ matrix.cuda }}:arm64-${DATE_SHA}
+
+  update-docker-tags:
+    if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+    runs-on: ubuntu-latest
+    needs: [generate-tag, create-manifests]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Update docker-tags.yml
+        env:
+          DATE_SHA: ${{ needs.generate-tag.outputs.date_sha }}
+        run: |
+          cat > ci/docker-tags.yml << EOF
+          flashinfer/flashinfer-ci-cu126: ${DATE_SHA}
+          flashinfer/flashinfer-ci-cu128: ${DATE_SHA}
+          flashinfer/flashinfer-ci-cu129: ${DATE_SHA}
+          flashinfer/flashinfer-ci-cu130: ${DATE_SHA}
+          EOF
+
+      - name: Create Pull Request
+        uses: peter-evans/create-pull-request@v6
+        with:
+          commit-message: "Update Docker tags to ${{ needs.generate-tag.outputs.date_sha }}"
+          title: "Update Docker CI tags to ${{ needs.generate-tag.outputs.date_sha }}"
+          body: |
+            This PR updates the Docker CI image tags to the latest version: `${{ needs.generate-tag.outputs.date_sha }}`
+
+            Updated images:
+            - flashinfer/flashinfer-ci-cu126:${{ needs.generate-tag.outputs.date_sha }}
+            - flashinfer/flashinfer-ci-cu128:${{ needs.generate-tag.outputs.date_sha }}
+            - flashinfer/flashinfer-ci-cu129:${{ needs.generate-tag.outputs.date_sha }}
+            - flashinfer/flashinfer-ci-cu130:${{ needs.generate-tag.outputs.date_sha }}
+
+            Auto-generated by [release-ci-docker workflow](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }})
+          branch: update-docker-tags-${{ needs.generate-tag.outputs.date_sha }}
+          delete-branch: true
+          labels: |
+            docker
+            automated

docker/install/install_python_packages.sh

@@ -22,8 +22,9 @@ set -u
 # Accept CUDA version as parameter (e.g., cu126, cu128, cu129)
 CUDA_VERSION=${1:-cu128}
 
-pip3 install torch --index-url https://download.pytorch.org/whl/${CUDA_VERSION}
+# Install requirements.txt first, so we can override any versions later for docker image
 pip3 install -r /install/requirements.txt
+pip3 install torch --index-url https://download.pytorch.org/whl/${CUDA_VERSION}
 pip3 install responses pytest scipy build cuda-python nvidia-nvshmem-cu12
 
 # Install cudnn package based on CUDA version