Fix cudnn.h not found error by dynamically locating and symlinking it from pip installed package

Dockerfile

@@ -60,22 +60,23 @@ RUN . $CONDA_DIR/etc/profile.d/conda.sh && \
         torchaudio==2.3.1 \
         --index-url https://download.pytorch.org/whl/cu121
 
-# 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.
+# NEW: Dynamically find cudnn.h and symlink it to the Conda environment's include path.
+# This ensures Transformer Engine can find it during compilation.
 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"; \
+    CUDNN_HEADER_PATH=$(find "$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/" -name "cudnn.h" | head -n 1) && \
+    if [ -f "$CUDNN_HEADER_PATH" ]; then \
+        echo "Found cudnn.h at: $CUDNN_HEADER_PATH"; \
         mkdir -p "$CONDA_PREFIX/include" && \
-        cp /usr/local/cuda/include/cudnn.h "$CONDA_PREFIX/include/"; \
+        ln -sf "$CUDNN_HEADER_PATH" "$CONDA_PREFIX/include/cudnn.h" || \
+        cp "$CUDNN_HEADER_PATH" "$CONDA_PREFIX/include/cudnn.h"; \
     else \
-        echo "cudnn.h already present in Conda environment"; \
+        echo "Error: cudnn.h not found in any expected location within Conda environment. This will likely cause compilation failures."; \
+        exit 1; \
     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 && \