Fix cudnn.h not found error by copying to Conda env include path and setting CUDA_HOME

Dockerfile

@@ -20,27 +20,23 @@ RUN apt-get update -y && apt-get install -qqy \
     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
+    && git lfs install
 
 # 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 && \
-    # Add Conda to PATH for subsequent commands in this RUN layer
     export PATH=$CONDA_DIR/bin:$PATH && \
     conda clean --all --yes && \
     conda config --set auto_activate_base false && \
     conda config --add channels conda-forge
 
-# NEW: Set the global PATH for Conda's base environment immediately after installation.
-# This ensures 'conda' command is available for subsequent RUN commands.
+# Set the global PATH for Conda's base environment immediately after installation.
 ENV PATH=$CONDA_DIR/bin:$PATH
 
 # Accept Conda Terms of Service for default channels.
-# Now 'conda' command should be found.
 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
@@ -49,7 +45,6 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
 COPY . /app
 
 # Create the Conda environment named 'cosmos-predict1' using the provided YAML file.
-# Now 'conda' command should be found.
 RUN conda env create -f cosmos-predict1.yaml
 
 # Set the default Conda environment to be activated and update PATH (for subsequent layers and runtime)
@@ -65,15 +60,32 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
         torchaudio==2.3.1 \
         --index-url https://download.pytorch.org/whl/cu121
 
-# IMPORTANT: Symlink fix for Transformer Engine compilation.
-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
+# NEW: Ensure cudnn.h is available in the Conda environment's include path.
+# The base image has cudnn.h at /usr/local/cuda/include.
+# We explicitly copy it to the Conda environment's include directory if it's not already there.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    if [ ! -f "$CONCONDA_PREFIX/include/cudnn.h" ]; then \
+        echo "cudnn.h not found in Conda environment, copying from /usr/local/cuda/include"; \
+        mkdir -p "$CONDA_PREFIX/include" && \
+        cp /usr/local/cuda/include/cudnn.h "$CONDA_PREFIX/include/"; \
+    else \
+        echo "cudnn.h already present in Conda environment"; \
+    fi
+
+# IMPORTANT: Symlink fix for Transformer Engine compilation (from INSTALL.md).
+# These symlinks are for other NVIDIA headers that might be in Python site-packages.
+# $CONDA_PREFIX is the current activated Conda environment.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    ln -sf "$CONDA_PREFIX"/lib/python3.10/site-packages/nvidia/*/include/* "$CONDA_PREFIX"/include/ || true && \
+    ln -sf "$CONDA_PREFIX"/lib/python3.10/site-packages/nvidia/*/include/* "$CONDA_PREFIX"/include/python3.10 || true
 
-# Install Transformer Engine by attempting to compile it, relying on the robust build environment.
+# Install Transformer Engine by attempting to compile it.
+# Pass CUDA_HOME explicitly to ensure it looks in the correct Conda environment path.
 RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
     conda activate cosmos-predict1 && \
-    pip install --no-cache-dir --no-build-isolation "transformer-engine[pytorch]==1.12.0"
+    CUDA_HOME=$CONDA_PREFIX 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 && \
@@ -89,7 +101,7 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
 # Make the start.sh script executable.
 RUN chmod +x /app/start.sh
 
-# --- Verification Steps ---
+# Verification Steps
 RUN echo "Verifying Python and Conda installations..."
 RUN python --version
 RUN conda env list
@@ -104,7 +116,6 @@ 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"]