Added statistics retrival

Signed-off-by: pierantoniomerlino <[email protected]>
eclipse-kura · Feb 14, 2025 · 8cfde6b · 8cfde6b
1 parent 8d0d496
commit 8cfde6b
Show file tree

Hide file tree

Showing 5 changed files with 702 additions and 45 deletions.
diff --git a/kura/org.eclipse.kura.ai.triton.server/META-INF/MANIFEST.MF b/kura/org.eclipse.kura.ai.triton.server/META-INF/MANIFEST.MF
@@ -22,6 +22,7 @@ Import-Package: com.google.common.base;version="25.0.0",
  com.google.common.util.concurrent;version="25.0.0",
  com.google.gson;version="2.7.0",
  com.google.gson.annotations;version="2.9.0",
+ com.google.gson.reflect;version="[2.9.0,3.0.0)",
  com.google.protobuf;version="3.19.3",
  org.apache.commons.io;version="2.4.0",
  org.apache.commons.io.filefilter;version="2.11.0",

diff --git a/...ra.ai.triton.server/src/main/java/org/eclipse/kura/ai/triton/server/GpuMetricsParser.java b/...ra.ai.triton.server/src/main/java/org/eclipse/kura/ai/triton/server/GpuMetricsParser.java
@@ -64,7 +64,7 @@ public Map<String, String> parse() {
                     Optional<String> uuid = parseUuid(line);
                     Optional<String> name = parseName(line);
                     String value = parseValue(line);
-                    if (uuid.isEmpty() || name.isEmpty()) {
+                    if (!uuid.isPresent() || !name.isPresent()) {
                         return;
                     }
                     if (gpuMetricsMap.containsKey(uuid.get())) {

diff --git a/...triton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceAbs.java b/...triton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceAbs.java
@@ -58,6 +58,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.ProtocolStringList;
 
@@ -99,6 +101,9 @@ public abstract class TritonServerServiceAbs implements InferenceEngineMetricsSe
     private String decryptionFolderPath = "";
     private boolean decryptionFolderNeedsCleanup = false;
 
+    private final GsonBuilder gsonBuilder = new GsonBuilder();
+    private final Gson gson = gsonBuilder.create();
+
     public void setCommandExecutorService(CommandExecutorService executorService) {
         this.commandExecutorService = executorService;
     }
@@ -400,36 +405,66 @@ public List<Tensor> infer(ModelInfo modelInfo, List<Tensor> inputData) throws Ku
 
     @Override
     public Map<String, String> getMetrics() throws KuraException {
+        if (!this.options.areMetricsEnabled()) {
+            logger.debug("Triton Server Metrics not enabled.");
+            return new HashMap<>();
+        }
+
         Map<String, String> metrics = new HashMap<>();
         metrics.putAll(getModelStatistics());
         metrics.putAll(getGpuMetrics());
 
         return metrics;
     }
 
-    private Map<String, String> getModelStatistics() {
-        Map<String, String> statistics = new HashMap<>();
-        return statistics;
+    private Map<String, String> getModelStatistics() throws KuraException {
+        Map<String, String> modelStatistics = new HashMap<>();
+        List<String> response = getListMetrics(
+                "http://" + getServerAddress() + ":" + this.options.getHttpPort() + "/v2/models/stats");
+
+        response.forEach(content -> {
+            // Remove first '"model_stats": [' string and split on model 'name'
+            String[] statistics = content.substring(16, content.length() - 2).split("\\{\\\"name\\\":");
+            for (String modelStats : statistics) {
+                if (!modelStats.isEmpty()) {
+                    String name = getModelName(modelStats);
+                    String version = getModelVersion(modelStats);
+                    modelStatistics.put(name + "." + version, "{\"name\":" + sanitizeModelStatistics(modelStats));
+                }
+            }
+        });
+
+        return modelStatistics;
+    }
+
+    private String getModelName(String statistic) {
+        String[] name = statistic.split(",");
+        return name[0].substring(1, name[0].length() - 1);
+    }
+
+    private String getModelVersion(String statistics) {
+        String[] version = statistics.split(",");
+        return version[1].substring(11, version[1].length() - 1);
+    }
+
+    private String sanitizeModelStatistics(String statistics) {
+        // Remove last comma if present
+        return statistics.endsWith(",") ? statistics.substring(0, statistics.length() - 1) : statistics;
     }
 
     private Map<String, String> getGpuMetrics() throws KuraException {
-        List<String> response = getListMetrics();
+        List<String> response = getListMetrics(
+                "http://" + getServerAddress() + ":" + this.options.getMetricsPort() + "/metrics");
 
         GpuMetricsParser metricsParser = new GpuMetricsParser(response);
         return metricsParser.parse();
     }
 
-    private List<String> getListMetrics() throws KuraException {
+    private List<String> getListMetrics(String resourceURL) throws KuraException {
         List<String> metrics = new ArrayList<>();
-        if (!this.options.areMetricsEnabled()) {
-            logger.debug("Triton Server Metrics not enabled.");
-            return metrics;
-        }
-
         HttpURLConnection connection = null;
         try {
-            URL url = new URI("http://" + getServerAddress() + ":" + this.options.getMetricsPort() + "/metrics")
-                    .toURL();
+            URL url = new URI(resourceURL).toURL();
             connection = (HttpURLConnection) url.openConnection();
             connection.setRequestMethod("GET");
             connection.setReadTimeout(30000);