From bc8aaa3d3e1d70915930669d7e4754b498207827 Mon Sep 17 00:00:00 2001
From: Jirka <jirka.borovec@seznam.cz>
Date: Sat, 16 Mar 2024 21:56:47 +0100
Subject: [PATCH] docker: pip install nvFuser for Torch wheels

---
 dockers/ubuntu-cuda/Dockerfile | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/dockers/ubuntu-cuda/Dockerfile b/dockers/ubuntu-cuda/Dockerfile
index e815d827f6..2213921897 100644
--- a/dockers/ubuntu-cuda/Dockerfile
+++ b/dockers/ubuntu-cuda/Dockerfile
@@ -24,6 +24,7 @@ ARG CUDNN_FRONTEND_CHECKOUT="v1.1.0"
 ARG PYTHON_VERSION="3.10"
 ARG TORCH_VERSION="2.2.1"
 ARG TRITON_VERSION="2.2.0"
+ARG TORCH_INSTALL="stable"
 
 SHELL ["/bin/bash", "-c"]
 # https://techoverflow.net/2019/05/18/how-to-fix-configuring-tzdata-interactive-input-when-building-docker-images/
@@ -96,7 +97,7 @@ ENV \
     TORCH_CUDA_ARCH_LIST="8.0" \
     CUDA_SELECT_NVCC_ARCH_FLAGS="8.0"
 
-ARG TORCH_INSTALL="wheel"
+ARG TORCH_INSTALL
 
 RUN \
     if [ "${TORCH_INSTALL}" == "source" ]; then \
@@ -122,15 +123,26 @@ RUN \
           --index-url="https://download.pytorch.org/whl/cu${CUDA_VERSION_MM//'.'/''}"; \
     fi
 
+ARG TORCH_INSTALL
+
 RUN \
-    # building nvFuser from source
-    git clone https://github.com/NVIDIA/Fuser.git && \
-    cd Fuser && \
-    git submodule update --init --recursive && \
-    pip install -r requirements.txt && \
-    python setup.py install --no-test --no-benchmark && \
-    cd .. && \
-    rm -rf Fuser
+    if [ "${TORCH_INSTALL}" == "source" ]; then \
+        # building nvFuser from source
+        git clone https://github.com/NVIDIA/Fuser.git && \
+        cd Fuser && \
+        git submodule update --init --recursive && \
+        pip install -r requirements.txt && \
+        python setup.py install --no-test --no-benchmark && \
+        cd .. && \
+        rm -rf Fuser ; \
+    elif [ "${TORCH_INSTALL}" == "test" ]; then \
+        echo "Not supported option" ; \
+    else \
+        # installing pytorch from wheels \
+        CUDA_VERSION_MM=${CUDA_VERSION%.*} && \
+        TORCH_VERSION_MM=${TORCH_VERSION%.*} && \
+        pip install -U "nvfuser-cu${CUDA_VERSION_MM/./}-torch${TORCH_VERSION_MM/./}" ; \
+    fi
 
 RUN \
     ls -lh requirements/ && \