From 6404f14e0b62816048b6e1272a2521b1780810fc Mon Sep 17 00:00:00 2001
From: Yingge He <yinggeh@nvidia.com>
Date: Wed, 7 Aug 2024 17:06:15 -0700
Subject: [PATCH] Fix "metrics not supported error" when building with
 TRITON_ENABLE_METRICS=OFF flag

---
 src/pb_stub.cc | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/pb_stub.cc b/src/pb_stub.cc
index 2a6be556..416b58ae 100644
--- a/src/pb_stub.cc
+++ b/src/pb_stub.cc
@@ -564,6 +564,24 @@ Stub::Initialize(bi::managed_external_buffer::handle_t map_handle)
     model_config_params[pair.first.c_str()] = pair.second;
   }
 
+  // Set metrics_mode value based on the build flag.
+#ifdef TRITON_ENABLE_METRICS
+#if defined(TRITON_ENABLE_METRICS_CPU) && defined(TRITON_ENABLE_METRICS_GPU)
+  model_config_params["metrics_mode"] = "all";
+#else  // defined(TRITON_ENABLE_METRICS_CPU) &&
+       // defined(TRITON_ENABLE_METRICS_GPU)
+#ifdef TRITON_ENABLE_METRICS_CPU
+  model_config_params["metrics_mode"] = "cpu";
+#endif  // TRITON_ENABLE_METRICS_CPU
+#ifdef TRITON_ENABLE_METRICS_GPU
+  model_config_params["metrics_mode"] = "gpu";
+#endif  // TRITON_ENABLE_METRICS_GPU
+#endif  // defined(TRITON_ENABLE_METRICS_CPU) &&
+        // defined(TRITON_ENABLE_METRICS_GPU)
+#else   // TRITON_ENABLE_METRICS
+  model_config_params["metrics_mode"] = "off";
+#endif  // TRITON_ENABLE_METRICS
+
   LaunchStubToParentQueueMonitor();
   LaunchParentToStubQueueMonitor();