Create a queue requests count metric

frontend/server/src/main/java/org/pytorch/serve/grpcimpl/InferenceImpl.java

@@ -92,7 +92,6 @@ public void predictions(
                     new InputParameter(entry.getKey(), entry.getValue().toByteArray()));
         }
 
-        MetricAggregator.handleInferenceMetric(modelName, modelVersion);
         Job job =
                 new GRPCJob(
                         responseObserver,
@@ -109,6 +108,8 @@ public void predictions(
                 InternalServerException e = new InternalServerException(responseMessage);
                 sendErrorResponse(
                         responseObserver, Status.INTERNAL, e, "InternalServerException.()");
+            } else {
+                MetricAggregator.handleInferenceMetric(modelName, modelVersion, job.getPriority());
             }
         } catch (ModelNotFoundException | ModelVersionNotFoundException e) {
             sendErrorResponse(responseObserver, Status.INTERNAL, e, null);

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/InferenceRequestHandler.java

@@ -21,7 +21,6 @@
 import org.pytorch.serve.http.HttpRequestHandlerChain;
 import org.pytorch.serve.http.ResourceNotFoundException;
 import org.pytorch.serve.http.StatusResponse;
-import org.pytorch.serve.metrics.api.MetricAggregator;
 import org.pytorch.serve.openapi.OpenApiUtils;
 import org.pytorch.serve.servingsdk.ModelServerEndpoint;
 import org.pytorch.serve.util.ApiUtils;
@@ -254,8 +253,6 @@ private void predict(
             NettyUtils.sendJsonResponse(ctx, resp);
             return;
         }
-
-        MetricAggregator.handleInferenceMetric(modelName, modelVersion);
         ApiUtils.addRESTInferenceJob(ctx, modelName, modelVersion, input);
     }
 

frontend/server/src/main/java/org/pytorch/serve/job/GRPCJob.java

@@ -14,6 +14,7 @@
 import org.pytorch.serve.http.messages.DescribeModelResponse;
 import org.pytorch.serve.metrics.Dimension;
 import org.pytorch.serve.metrics.Metric;
+import org.pytorch.serve.metrics.api.MetricAggregator;
 import org.pytorch.serve.util.ApiUtils;
 import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.GRPCUtils;
@@ -66,6 +67,10 @@ public void response(
             predictionResponseObserver.onNext(reply);
             predictionResponseObserver.onCompleted();
 
+            long inferTime = System.nanoTime() - getBegin();
+            MetricAggregator.handleInferenceMetric(
+                    getModelName(), getModelVersion(), getPriority(), getScheduled() - getBegin(), inferTime);
+
             logger.debug(
                     "Waiting time ns: {}, Backend time ns: {}",
                     getScheduled() - getBegin(),

frontend/server/src/main/java/org/pytorch/serve/job/RestJob.java

@@ -136,13 +136,14 @@ private void responseInference(
          * by external clients.
          */
         if (ctx != null) {
-            MetricAggregator.handleInferenceMetric(
-                    getModelName(), getModelVersion(), getScheduled() - getBegin(), inferTime);
             NettyUtils.sendHttpResponse(ctx, resp, true);
         } else if (responsePromise != null) {
             responsePromise.complete(body);
         }
 
+        MetricAggregator.handleInferenceMetric(
+                getModelName(), getModelVersion(), getPriority(), getScheduled() - getBegin(), inferTime);
+
         logger.debug(
                 "Waiting time ns: {}, Backend time ns: {}",
                 getScheduled() - getBegin(),

frontend/server/src/main/java/org/pytorch/serve/metrics/api/MetricAggregator.java

@@ -2,27 +2,33 @@
 
 import org.pytorch.serve.metrics.format.prometheous.PrometheusMetricManager;
 import org.pytorch.serve.util.ConfigManager;
+import org.pytorch.serve.util.Priority;
 
 public final class MetricAggregator {
 
     private MetricAggregator() {}
 
-    public static void handleInferenceMetric(final String modelName, final String modelVersion) {
+    // Is executed upon successful Job insertion in queue
+    public static void handleInferenceMetric(final String modelName, final String modelVersion, Priority priority) {
         ConfigManager configMgr = ConfigManager.getInstance();
         if (configMgr.isMetricApiEnable()
                 && configMgr.getMetricsFormat().equals(ConfigManager.METRIC_FORMAT_PROMETHEUS)) {
-            PrometheusMetricManager.getInstance().incInferCount(modelName, modelVersion);
+            PrometheusMetricManager metrics = PrometheusMetricManager.getInstance();
+            metrics.incInferCount(modelName, modelVersion);
+            metrics.incQueueCount(modelName, modelVersion, priority);
         }
     }
 
+    // Is executed upon successful Job completion
     public static void handleInferenceMetric(
-            final String modelName, final String modelVersion, long timeInQueue, long inferTime) {
+            final String modelName, final String modelVersion, Priority priority, long timeInQueue, long inferTime) {
         ConfigManager configMgr = ConfigManager.getInstance();
         if (configMgr.isMetricApiEnable()
                 && configMgr.getMetricsFormat().equals(ConfigManager.METRIC_FORMAT_PROMETHEUS)) {
             PrometheusMetricManager metrics = PrometheusMetricManager.getInstance();
             metrics.incInferLatency(inferTime, modelName, modelVersion);
-            metrics.incQueueLatency(timeInQueue, modelName, modelVersion);
+            metrics.incQueueLatency(timeInQueue, modelName, modelVersion, priority);
+            metrics.decQueueCount(modelName, modelVersion, priority);
         }
     }
 }

frontend/server/src/main/java/org/pytorch/serve/metrics/format/prometheous/PrometheusMetricManager.java

@@ -1,35 +1,47 @@
 package org.pytorch.serve.metrics.format.prometheous;
 
+import org.pytorch.serve.util.ConfigManager;
+import org.pytorch.serve.util.Priority;
 import io.prometheus.client.Counter;
+import io.prometheus.client.Gauge;
 import java.util.UUID;
 
 public final class PrometheusMetricManager {
 
     private static final PrometheusMetricManager METRIC_MANAGER = new PrometheusMetricManager();
     private static final String METRICS_UUID = UUID.randomUUID().toString();
-    private Counter inferRequestCount;
-    private Counter inferLatency;
-    private Counter queueLatency;
+    private final Counter inferRequestCount;
+    private final Counter inferLatency;
+    private final Counter queueLatency;
+    private final Gauge queueRequestCount;
 
     private PrometheusMetricManager() {
-        String[] metricsLabels = {"uuid", "model_name", "model_version"};
+        String[] metricsLabelsNonQueue = {"uuid", "model_name", "model_version"};
         inferRequestCount =
                 Counter.build()
                         .name("ts_inference_requests_total")
-                        .labelNames(metricsLabels)
+                        .labelNames(metricsLabelsNonQueue)
                         .help("Total number of inference requests.")
                         .register();
         inferLatency =
                 Counter.build()
                         .name("ts_inference_latency_microseconds")
-                        .labelNames(metricsLabels)
-                        .help("Cumulative inference duration in microseconds")
+                        .labelNames(metricsLabelsNonQueue)
+                        .help("Cumulative inference duration in microseconds.")
                         .register();
+
+        String[] metricsLabelsQueue = {"uuid", "model_name", "model_version", "priority", "max_queue_size"};
         queueLatency =
                 Counter.build()
                         .name("ts_queue_latency_microseconds")
-                        .labelNames(metricsLabels)
-                        .help("Cumulative queue duration in microseconds")
+                        .labelNames(metricsLabelsQueue)
+                        .help("Cumulative queue duration in microseconds.")
+                        .register();
+        queueRequestCount =
+                Gauge.build()
+                        .name("ts_queue_requests_total")
+                        .labelNames(metricsLabelsQueue)
+                        .help("Current queue inference request count.")
                         .register();
     }
 
@@ -61,9 +73,10 @@ public void incInferLatency(long inferTime, String modelName, String modelVersio
      * @param modelName name of the model
      * @param modelVersion version of the model
      */
-    public void incQueueLatency(long queueTime, String modelName, String modelVersion) {
+    public void incQueueLatency(long queueTime, String modelName, String modelVersion, Priority priority) {
+        int queueSize = ConfigManager.getInstance().getJobQueueSize();
         queueLatency
-                .labels(METRICS_UUID, modelName, getOrDefaultModelVersion(modelVersion))
+                .labels(METRICS_UUID, modelName, getOrDefaultModelVersion(modelVersion), priority.toString(), String.valueOf(queueSize))
                 .inc(queueTime / 1000.0);
     }
 
@@ -78,4 +91,31 @@ public void incInferCount(String modelName, String modelVersion) {
                 .labels(METRICS_UUID, modelName, getOrDefaultModelVersion(modelVersion))
                 .inc();
     }
+
+
+    /**
+     * Counts a valid inference request that has been added to a queue
+     *
+     * @param modelName name of the model
+     * @param modelVersion version of the model
+     */
+    public void incQueueCount(String modelName, String modelVersion, Priority priority) {
+        int queueSize = ConfigManager.getInstance().getJobQueueSize();
+        queueRequestCount
+                .labels(METRICS_UUID, modelName, getOrDefaultModelVersion(modelVersion), priority.toString(), String.valueOf(queueSize))
+                .inc();
+    }
+
+    /**
+     * Counts a valid inference request that has been removed from a queue
+     *
+     * @param modelName name of the model
+     * @param modelVersion version of the model
+     */
+    public void decQueueCount(String modelName, String modelVersion, Priority priority) {
+        int queueSize = ConfigManager.getInstance().getJobQueueSize();
+        queueRequestCount
+                .labels(METRICS_UUID, modelName, getOrDefaultModelVersion(modelVersion), priority.toString(), String.valueOf(queueSize))
+                .dec();
+    }
 }

frontend/server/src/main/java/org/pytorch/serve/util/ApiUtils.java

@@ -30,6 +30,7 @@
 import org.pytorch.serve.http.messages.RegisterModelRequest;
 import org.pytorch.serve.job.RestJob;
 import org.pytorch.serve.snapshot.SnapshotManager;
+import org.pytorch.serve.metrics.api.MetricAggregator;
 import org.pytorch.serve.util.messages.RequestInput;
 import org.pytorch.serve.util.messages.WorkerCommands;
 import org.pytorch.serve.wlm.Model;
@@ -397,6 +398,8 @@ public static RestJob addRESTInferenceJob(
             String priority = job.getPriority().toString();
             String responseMessage = getInferenceErrorResponseMessage(modelName, version, priority);
             throw new ServiceUnavailableException(responseMessage);
+        } else {
+            MetricAggregator.handleInferenceMetric(modelName, version, job.getPriority());
         }
         return job;
     }