diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 75546c3..13bc671 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -11,12 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -FROM nvcr.io/nvidia/pytorch:25.06-py3 +FROM nvcr.io/nvidia/pytorch:25.09-py3 ENV PIP_CONSTRAINT="" # Install uv and python -ARG UV_VERSION=0.7.2 +ARG UV_VERSION=0.9.3 ENV PATH="/root/.local/bin:$PATH" RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh @@ -34,6 +34,7 @@ RUN --mount=type=bind,source=pyproject.toml,target=/workspace/pyproject.toml \ uv sync --link-mode symlink --locked --all-groups \ --no-install-package absl-py \ --no-install-package torch \ + --no-install-package triton \ --no-install-package nvidia-cublas-cu12 \ --no-install-package nvidia-cuda-cupti-cu12 \ --no-install-package nvidia-cuda-nvrtc-cu12 \ @@ -45,5 +46,7 @@ RUN --mount=type=bind,source=pyproject.toml,target=/workspace/pyproject.toml \ --no-install-package nvidia-cusolver-cu12 \ --no-install-package nvidia-cusparse-cu12 \ --no-install-package nvidia-cusparselt-cu12 \ - --no-install-package nvidia-nccl-cu12 + --no-install-package nvidia-nccl-cu12 \ + --no-install-package nvidia-nvjitlink-cu12 \ + --no-install-package nvidia-nvtx-cu12 EOF