Fix cudnn.h not found during Transformer Engine build by adding symlinks as per INSTALL.md

Dockerfile

@@ -15,8 +15,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
     libgl1-mesa-glx \
     libglib2.0-0 \
-    # Add any other system libraries that might be needed by transformer_engine or other deep learning libs
-    # For example, libstdc++6, libgomp1, etc., though usually covered by base image or build-essential
     && rm -rf /var/lib/apt/lists/*
 
 # Install Miniconda
@@ -51,13 +49,19 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
         torchaudio==2.3.1 \
         --index-url https://download.pytorch.org/whl/cu121
 
-# --- NEW: Install Transformer Engine separately after PyTorch and cuDNN are in place ---
-# This ensures it can find the necessary CUDA/cuDNN libraries.
-# Also add git to apt-get install above if not already there.
+# --- NEW SECTION: Patch Transformer Engine linking issues by creating symlinks for headers ---
+# These commands must run AFTER conda env creation and BEFORE transformer-engine pip install.
+RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
+    conda activate cosmos-predict1 && \
+    echo "Creating symlinks for NVIDIA headers..." && \
+    ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/ && \
+    ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10/
+# --- END NEW SECTION ---
+
+# Install Transformer Engine separately after PyTorch and cuDNN are in place and headers are linked.
 RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
     conda activate cosmos-predict1 && \
     pip install --no-cache-dir transformer-engine[pytorch]==1.12.0
-# --- END NEW ---
 
 # --- Verification Steps ---
 RUN echo "Verifying Python and Conda installations..."