Dockerfile

# Start from a clean NVIDIA CUDA base image.
FROM nvidia/cuda:12.4.0-devel-ubuntu22.04

# Set environment variables for non-interactive installations to prevent prompts during apt-get.
ENV DEBIAN_FRONTEND=noninteractive
ENV CONDA_DIR=/opt/conda
ENV PATH=$CONDA_DIR/bin:$PATH

WORKDIR /app

# Install essential system dependencies required for Miniconda and general build tools.
RUN apt-get update && apt-get install -y --no-install-recommends \
    wget \
    git \
    build-essential \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Install Miniconda
RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh && \
    /bin/bash miniconda.sh -b -p $CONDA_DIR && \
    rm miniconda.sh && \
    conda clean --all --yes && \
    conda config --set auto_activate_base false && \
    conda config --add channels conda-forge

# Accept Conda Terms of Service for default channels.
RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r

# Copy all local project files into the container's working directory (/app).
COPY . /app

# Create the Conda environment named 'cosmos-predict1' using the provided YAML file.
RUN conda env create -f cosmos-predict1.yaml

# Set the default Conda environment to be activated.
ENV CONDA_DEFAULT_ENV=cosmos-predict1
ENV PATH=$CONDA_DIR/envs/cosmos-predict1/bin:$PATH

# Install PyTorch and TorchVision via pip with specific CUDA index.
RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
    conda activate cosmos-predict1 && \
    pip install --no-cache-dir \
        torch==2.3.1 \
        torchvision==0.18.1 \
        torchaudio==2.3.1 \
        --index-url https://download.pytorch.org/whl/cu121

# Install Transformer Engine using the pre-built wheel
# Ensure the filename matches your actual wheel file.
COPY ./transformer_engine.whl /tmp/
RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
    conda activate cosmos-predict1 && \
    pip install --no-cache-dir /tmp/transformer_engine.whl

# Make the start.sh script executable.
# THIS MUST BE A SEPARATE RUN COMMAND.
RUN chmod +x /app/start.sh

# Set the default command to run when the container starts.
CMD ["/app/start.sh"]