flashinfer/docker/Dockerfile.cu126 at 3404aacbbdd0e2321123ad5d61ffc0fcd6c6cf82 · bkryu/flashinfer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
FROM nvidia/cuda:12.6.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"

# Ensure pip-installed nvidia-cublas takes precedence over system libraries
ENV LD_LIBRARY_PATH="/opt/conda/envs/py312/lib/python3.12/site-packages/nvidia/cublas/lib/:$LD_LIBRARY_PATH"

# Install torch and other python packages
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 cu126

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