Set memory tracker backend option in build.py

triton-inference-server · May 17, 2023 · 654cbe8 · 654cbe8
1 parent 39252fd
commit 654cbe8
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/build.py b/build.py
@@ -598,6 +598,16 @@ def backend_cmake_args(images, components, be, install_dir, library_paths):
     cargs.append(
         cmake_backend_enable(be, 'TRITON_ENABLE_METRICS', FLAGS.enable_metrics))
 
+    if (target_platform() == 'windows') and (not FLAGS.no_container_build):
+        print(
+            "Warning: Detected docker build is used for Windows, backend utility 'device memory tracker' will be disabled due to missing library in CUDA Windows docker image."
+        )
+        cargs.append(
+            cmake_backend_enable(be, 'TRITON_ENABLE_MEMORY_TRACKER', False))
+    elif FLAGS.enable_gpu:
+        cargs.append(
+            cmake_backend_enable(be, 'TRITON_ENABLE_MEMORY_TRACKER', True))
+
     cargs += cmake_backend_extra_args(be)
     cargs.append('..')
     return cargs