Skip to content

refactor(docker): consolidate Dockerfiles into one multi-stage file #2806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dierksen wants to merge 6 commits into
base: main
Choose a base branch
from
Open

refactor(docker): consolidate Dockerfiles into one multi-stage file #2806

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
29e990e
refactor(docker): consolidate 8 per-version Dockerfiles into 2 parame…
dierksen Mar 12, 2026
bcba450
refactor(docker): merge Dockerfile.dev into Dockerfile as a multi-sta…
dierksen Mar 17, 2026
05e8ebc
refactor(docker): rename prod stage to test
dierksen Mar 17, 2026
b272f37
Merge main, resolve conflict, drop Dockerfile.cu131[.dev]
dierksen Mar 17, 2026
7e990f6
fix(docker): use PYTORCH_INDEX arg to decouple PyTorch wheel index fr…
dierksen Mar 17, 2026
f13c1b7
fix(devcontainer): update devcontainer.json files to use consolidated…
dierksen Mar 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions .devcontainer/cu126/devcontainer.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
{
"name": "CUDA Development Container",
"build": {
"dockerfile": "../../docker/Dockerfile.cu126.dev",
"context": "../../"
"dockerfile": "../../docker/Dockerfile",
"context": "../../",
"target": "dev",
"args": {
"CUDA_BASE_IMAGE": "nvidia/cuda:12.6.0-devel-ubuntu24.04",
"CUDA_VERSION": "cu126",
"PYTORCH_INDEX": "cu126"
}
},
"runArgs": [
"--gpus=all"
Expand Down
10 changes: 8 additions & 2 deletions .devcontainer/cu128/devcontainer.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
{
"name": "CUDA Development Container",
"build": {
"dockerfile": "../../docker/Dockerfile.cu128.dev",
"context": "../../"
"dockerfile": "../../docker/Dockerfile",
"context": "../../",
"target": "dev",
"args": {
"CUDA_BASE_IMAGE": "nvidia/cuda:12.8.0-devel-ubuntu24.04",
"CUDA_VERSION": "cu128",
"PYTORCH_INDEX": "cu128"
}
},
"runArgs": [
"--gpus=all"
Expand Down
10 changes: 8 additions & 2 deletions .devcontainer/cu129/devcontainer.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
{
"name": "CUDA Development Container",
"build": {
"dockerfile": "../../docker/Dockerfile.cu129.dev",
"context": "../../"
"dockerfile": "../../docker/Dockerfile",
"context": "../../",
"target": "dev",
"args": {
"CUDA_BASE_IMAGE": "nvidia/cuda:12.9.0-devel-ubuntu24.04",
"CUDA_VERSION": "cu129",
"PYTORCH_INDEX": "cu129"
}
},
"runArgs": [
"--gpus=all"
Expand Down
10 changes: 8 additions & 2 deletions .devcontainer/cu130/devcontainer.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
{
"name": "CUDA Development Container",
"build": {
"dockerfile": "../../docker/Dockerfile.cu130.dev",
"context": "../../"
"dockerfile": "../../docker/Dockerfile",
"context": "../../",
"target": "dev",
"args": {
"CUDA_BASE_IMAGE": "nvidia/cuda:13.0.1-devel-ubuntu24.04",
"CUDA_VERSION": "cu130",
"PYTORCH_INDEX": "cu130"
}
},
"runArgs": [
"--gpus=all"
Expand Down
36 changes: 34 additions & 2 deletions .github/workflows/release-ci-docker.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,32 @@ jobs:
needs: generate-tag
strategy:
matrix:
cuda: [cu126, cu128, cu129, cu130, cu131]
include:
- cuda: cu126
base_image: nvidia/cuda:12.6.0-devel-ubuntu24.04
nvidia_lib_path: nvidia/cublas/lib
pytorch_index: cu126
install_tilelang: "false"
- cuda: cu128
base_image: nvidia/cuda:12.8.0-devel-ubuntu24.04
nvidia_lib_path: nvidia/cublas/lib
pytorch_index: cu128
install_tilelang: "false"
- cuda: cu129
base_image: nvidia/cuda:12.9.0-devel-ubuntu24.04
nvidia_lib_path: nvidia/cublas/lib
pytorch_index: cu129
install_tilelang: "false"
- cuda: cu130
base_image: nvidia/cuda:13.0.1-devel-ubuntu24.04
nvidia_lib_path: nvidia/cu13/lib
pytorch_index: cu130
install_tilelang: "false"
- cuda: cu131
base_image: nvidia/cuda:13.1.1-cudnn-devel-ubuntu24.04
nvidia_lib_path: nvidia/cu13/lib
pytorch_index: cu130 # TODO: update to cu131 when PyTorch publishes cu131 wheels
install_tilelang: "true"
arch: [amd64, arm64]
steps:
- uses: actions/checkout@v4
Expand All @@ -58,7 +83,14 @@ jobs:
uses: docker/build-push-action@v5
with:
context: .
file: docker/Dockerfile.${{ matrix.cuda }}
file: docker/Dockerfile
target: test
build-args: |
CUDA_BASE_IMAGE=${{ matrix.base_image }}
CUDA_VERSION=${{ matrix.cuda }}
PYTORCH_INDEX=${{ matrix.pytorch_index }}
NVIDIA_LIB_PATH=${{ matrix.nvidia_lib_path }}
INSTALL_TILELANG=${{ matrix.install_tilelang }}
platforms: linux/${{ matrix.arch }}
push: ${{ github.event_name != 'pull_request' }}
pull: true # Always pull the latest base image
Expand Down
60 changes: 52 additions & 8 deletions docker/Dockerfile.cu131.dev → docker/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,56 @@
FROM nvidia/cuda:13.1.1-cudnn-devel-ubuntu24.04
ARG CUDA_BASE_IMAGE
FROM ${CUDA_BASE_IMAGE} AS base

ENV DEBIAN_FRONTEND=noninteractive

# Update package lists and install system dependencies
RUN apt-get update && apt-get install -y \
curl \
git \
wget \
wget
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To keep the Docker image size minimal, it's a good practice to clean up the apt cache after installing packages. You can do this by chaining && rm -rf /var/lib/apt/lists/* to your apt-get install command. This will reduce the size of this base layer and all subsequent images built from it.

    wget && rm -rf /var/lib/apt/lists/*


# ---- test target ----
FROM base AS test

# Install python
COPY docker/install/install_python.sh /install/install_python.sh
RUN bash /install/install_python.sh /opt/conda py312

WORKDIR /workspace

RUN echo "source activate py312" >> ~/.bashrc
ENV PATH="/opt/conda/bin:$PATH"
ENV PATH="/opt/conda/envs/py312/bin:$PATH"

# Set LD_LIBRARY_PATH to ensure pip-installed nvidia libs take precedence over system libraries
ARG NVIDIA_LIB_PATH
ENV LD_LIBRARY_PATH="/opt/conda/envs/py312/lib/python3.12/site-packages/${NVIDIA_LIB_PATH}/:$LD_LIBRARY_PATH"

# Triton
ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"

# Install torch and other python packages
COPY requirements.txt /install/requirements.txt
COPY docker/install/install_python_packages.sh /install/install_python_packages.sh
ARG CUDA_VERSION
ARG PYTORCH_INDEX=${CUDA_VERSION}
RUN bash /install/install_python_packages.sh ${PYTORCH_INDEX}

ARG INSTALL_TILELANG=false
RUN if [ "$INSTALL_TILELANG" = "true" ]; then pip install tilelang cuda-tile; fi

# Install mpi4py in the conda environment
RUN conda install -n py312 -y mpi4py mpich

# Configure pip for user-site installations (allows arbitrary users to install packages)
# This enables 'pip install --user' and 'pip install -e .' to work for any user
RUN mkdir -p /opt/pip-user && chmod 1777 /opt/pip-user
ENV PYTHONUSERBASE=/opt/pip-user
ENV PATH="/opt/pip-user/bin:$PATH"

# ---- dev target ----
FROM base AS dev

RUN apt-get update && apt-get install -y \
clang-format \
clangd-19 \
vim \
Expand Down Expand Up @@ -45,13 +89,13 @@ ENV PATH="/home/$USERNAME/conda/bin:$PATH"
ENV PATH="/home/$USERNAME/conda/envs/py312/bin:$PATH"

# Install torch and other python packages
# TODO: update cu130 -> cu131 when PyTorch starts publishing cu131 wheels
COPY requirements.txt /install/requirements.txt
COPY docker/install/install_python_packages.sh /install/install_python_packages.sh
RUN bash /install/install_python_packages.sh cu130 && pip3 install pre-commit

# Install tilelang and cuda-tile
RUN pip install tilelang cuda-tile
ARG CUDA_VERSION
ARG PYTORCH_INDEX=${CUDA_VERSION}
ARG INSTALL_TILELANG=false
RUN bash /install/install_python_packages.sh ${PYTORCH_INDEX} && pip3 install pre-commit && \
if [ "$INSTALL_TILELANG" = "true" ]; then pip install tilelang cuda-tile; fi

# Install mpi4py in the conda environment
RUN conda install -n py312 -y mpi4py mpich
Expand Down
37 changes: 0 additions & 37 deletions docker/Dockerfile.cu126
View file
Open in desktop

This file was deleted.

73 changes: 0 additions & 73 deletions docker/Dockerfile.cu126.dev
View file
Open in desktop

This file was deleted.

37 changes: 0 additions & 37 deletions docker/Dockerfile.cu128
View file
Open in desktop

This file was deleted.

Loading
Loading