Update cudnn home

triton-inference-server · Jan 10, 2025 · 9cf0c01 · 9cf0c01
1 parent 8f8ea27
commit 9cf0c01
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/tools/gen_ort_dockerfile.py b/tools/gen_ort_dockerfile.py
@@ -259,10 +259,6 @@ def dockerfile_for_linux(output_file):
             ep_flags += ' --cuda_version "{}"'.format(FLAGS.cuda_version)
         if FLAGS.cuda_home is not None:
             ep_flags += ' --cuda_home "{}"'.format(FLAGS.cuda_home)
-        if FLAGS.cudnn_home is not None:
-            ep_flags += ' --cudnn_home "{}"'.format(FLAGS.cudnn_home)
-        elif target_platform() == "igpu":
-            ep_flags += ' --cudnn_home "/usr/lib/aarch64-linux-gnu"'
         if FLAGS.ort_tensorrt:
             ep_flags += " --use_tensorrt"
             if FLAGS.ort_version >= "1.12.1":
@@ -292,6 +288,16 @@ def dockerfile_for_linux(output_file):
 """.format(
         cuda_archs
     )
+    if FLAGS.enable_gpu : #and target_platform() != "igpu"
+    # For GPU build, include the cudnn_home flag
+        df += """
+        RUN _CUDNN_VERSION=$(echo $CUDNN_VERSION | cut -d. -f1-2) && ./build.sh ${{COMMON_BUILD_ARGS}} --update --build {} --cudnn_home /usr/local/cudnn-$_CUDNN_VERSION/cuda
+        """.format(ep_flags)
+    else:
+    # For non-GPU
+        df += """
+        RUN ./build.sh ${{COMMON_BUILD_ARGS}} --update --build {}
+        """.format(ep_flags)
 
     df += """
 RUN ./build.sh ${{COMMON_BUILD_ARGS}} --update --build {}