Merged Dockerfile with robust build environment for transformer-engine compilation

Dockerfile

@@ -1,23 +1,31 @@
-# Start from a clean NVIDIA CUDA base image.
-FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
+# Adopt new base image with cuDNN pre-installed
+FROM nvidia/cuda:12.4.1-cudnn-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 \
+# Install essential system dependencies from both Dockerfiles
+RUN apt-get update -y && apt-get install -qqy \
     wget \
     git \
     build-essential \
     libgl1-mesa-glx \
     libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/*
+    rsync \
+    make \
+    libssl-dev zlib1g-dev \
+    libbz2-dev libreadline-dev libsqlite3-dev curl llvm \
+    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev \
+    ffmpeg libsm6 libxext6 cmake libmagickwand-dev \
+    # Ensure git-lfs is installed and initialized
+    git-lfs \
+    && rm -rf /var/lib/apt/lists/* \
+    && git lfs install # Initialize LFS system-wide
 
-# Install Miniconda
+# Install Miniconda (retain our existing approach)
 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 && \
@@ -36,7 +44,7 @@ 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.
+# Set the default Conda environment to be activated and update PATH
 ENV CONDA_DEFAULT_ENV=cosmos-predict1
 ENV PATH=$CONDA_DIR/envs/cosmos-predict1/bin:$PATH
 
@@ -49,35 +57,31 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
         torchaudio==2.3.1 \
         --index-url https://download.pytorch.org/whl/cu121
 
-# Copy the pre-built Transformer Engine wheel into the container
-# Ensure the filename matches your actual wheel file.
-COPY ./transformer_engine.whl /tmp/
+# IMPORTANT: Symlink fix for Transformer Engine compilation.
+# The `nvidia/cuda` base images place CUDA libraries and headers in /usr/local/cuda.
+# We need to ensure that the build system can find cuDNN headers.
+ENV CONDA_PREFIX_FIX=/usr/local/cuda
+RUN ln -sf $CONDA_PREFIX_FIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX_FIX/include/ || true && \
+    ln -sf $CONDA_PREFIX_FIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX_FIX/include/python3.10 || true
 
-# Install Transformer Engine using the pre-built wheel
+# Install Transformer Engine by attempting to compile it, relying on the robust build environment.
 RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
     conda activate cosmos-predict1 && \
-    pip install --no-cache-dir /tmp/transformer_engine.whl
+    pip install --no-cache-dir --no-build-isolation "transformer-engine[pytorch]==1.12.0"
+
+# Install Apex for inference.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    git clone https://github.com/NVIDIA/apex /app/apex && \
+    CUDA_HOME=$CONDA_PREFIX pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" /app/apex
+
+# Install MoGe for inference.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    pip install --no-cache-dir git+https://github.com/microsoft/MoGe.git
 
 # Make the start.sh script executable.
-# THIS IS A STANDALONE RUN COMMAND.
 RUN chmod +x /app/start.sh
 
-# --- Verification Steps ---
-RUN echo "Verifying Python and Conda installations..."
-RUN python --version
-RUN conda env list
-RUN echo "Verifying PyTorch and CUDA availability..."
-RUN conda run -n cosmos-predict1 python <<EOF
-import torch
-print('PyTorch Version: ' + torch.__version__)
-print('CUDA Available: ' + str(torch.cuda.is_available()))
-if torch.cuda.is_available():
-    print('CUDA Device Name: ' + torch.cuda.get_device_name(0))
-else:
-    print('CUDA Device Name: N/A')
-EOF
-RUN [ $? -eq 0 ] || echo "PyTorch verification failed. Check dependencies in cosmos-predict1.yaml."
-# --- End Verification Steps ---
-
 # Set the default command to run when the container starts.
 CMD ["/app/start.sh"]