Dockerfile.cu132

FROM nvidia/cuda:13.2.0-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive

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

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

# Set home directory
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-cublas takes precedence over system libraries
ENV LD_LIBRARY_PATH="/opt/conda/envs/py312/lib/python3.12/site-packages/nvidia/cu13/lib/:$LD_LIBRARY_PATH"

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

# Install torch and other python packages
# use nightly/cu132 temporarily and change to cu132 when torch releases stable version
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 nightly/cu132

# Install tilelang and cuda-tile
RUN pip install tilelang cuda-tile

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