Tag Prometheus metrics with model name and version (#560)

Co-authored-by: Paul Brena <[email protected]>
Co-authored-by: Paul Brena <[email protected]>

mlserver/handlers/dataplane.py

@@ -8,6 +8,25 @@
 )
 from ..middleware import inference_middlewares
 from ..utils import generate_uuid
+from prometheus_client import (
+    Counter,
+    Summary,
+)
+
+
+_ModelInferRequestSuccess = Counter(
+    "model_infer_request_success",
+    "Model infer request success count",
+    ["model", "version"],
+)
+_ModelInferRequestFailure = Counter(
+    "model_infer_request_failure",
+    "Model infer request failure count",
+    ["model", "version"],
+)
+_ModelInferRequestDuration = Summary(
+    "model_infer_request_duration", "Model infer request duration", ["model", "version"]
+)
 
 
 class DataPlane:
@@ -48,18 +67,27 @@ async def model_metadata(
     async def infer(
         self, payload: InferenceRequest, name: str, version: str = None
     ) -> InferenceResponse:
-        if payload.id is None:
-            payload.id = generate_uuid()
 
-        model = await self._model_registry.get_model(name, version)
+        with _ModelInferRequestDuration.labels(
+            model=name, version=version
+        ).time(), _ModelInferRequestFailure.labels(
+            model=name, version=version
+        ).count_exceptions():
 
-        # Run middlewares
-        inference_middlewares(payload, model.settings)
+            if payload.id is None:
+                payload.id = generate_uuid()
 
-        # TODO: Make await optional for sync methods
-        prediction = await model.predict(payload)
+            model = await self._model_registry.get_model(name, version)
+
+            # Run middlewares
+            inference_middlewares(payload, model.settings)
+
+            # TODO: Make await optional for sync methods
+            prediction = await model.predict(payload)
+
+            # Ensure ID matches
+            prediction.id = payload.id
 
-        # Ensure ID matches
-        prediction.id = payload.id
+            _ModelInferRequestSuccess.labels(model=name, version=version).inc()
 
-        return prediction
+            return prediction

tests/metrics/test_endpoint.py

@@ -28,7 +28,13 @@ async def test_metrics(metrics_client: MetricsClient):
         metrics = await metrics_client.metrics()
         assert metrics is not None
 
-        expected_prefixes = ("python_", "process_", "rest_server_", "grpc_server_")
+        expected_prefixes = (
+            "python_",
+            "process_",
+            "rest_server_",
+            "grpc_server_",
+            "model_infer_",
+        )
         metrics_list = list(iter(metrics))
         assert len(metrics_list) > 0
         for metric in metrics_list: