ci: Install mpi4py (#1149)

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

## 📌 Description

Install the python packages for CI docker: mpi4py, pynvml. They will be
used for the comm ops.

## 🔍 Related Issues

#1145,
#1134

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

- [ ] 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: yongwww

docker/Dockerfile.ci_gpu

@@ -19,6 +19,9 @@ RUN echo "source activate py312" >> ~/.bashrc
 ENV PATH="/opt/conda/bin:$PATH"
 ENV PATH="/opt/conda/envs/py312/bin:$PATH"
 
-# Install torch
+# Install torch and other python packages
 COPY install/install_python_packages.sh /install/install_python_packages.sh
 RUN bash /install/install_python_packages.sh
+
+# Install mpi4py in the conda environment
+RUN conda install -n py312 -y mpi4py

docker/install/install_python_packages.sh

@@ -20,4 +20,4 @@ set -e
 set -u
 
 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
-pip3 install ninja pytest numpy scipy build
+pip3 install ninja pytest numpy scipy build pynvml