Update product_return_prediction/monitoring.py

product_return_prediction/monitoring.py

@@ -1,13 +1,13 @@
 import os
 from typing import Callable
 
-from prometheus_client import Histogram
+from prometheus_client import Histogram, Counter
 from prometheus_fastapi_instrumentator import Instrumentator, metrics
 from prometheus_fastapi_instrumentator.metrics import Info
 
 # Let's use environment variables METRICS_NAMESPACE and METRICS_SUBSYSTEM for logical grouping of metrics
 NAMESPACE = os.environ.get("METRICS_NAMESPACE", "fastapi")
-SUBSYSTEM = os.environ.get("METRICS_SUBSYSTEM", "model")
+SUBSYSTEM = os.environ.get("METRICS_SUBSYSTEM", "svm_model")
 
 instrumentator = Instrumentator(
     should_group_status_codes=True,
@@ -18,7 +18,7 @@ instrumentator = Instrumentator(
     inprogress_labels=True,
 )
 
-# ----- add metrics -----
+# ----- add general metrics -----
 
 instrumentator.add(
     metrics.request_size(
@@ -57,14 +57,13 @@ instrumentator.add(
     )
 )
 
-
-# ----- custom metrics -----
+# ----- custom metrics for SVM -----
 def model_output(
     metric_name: str = "model_output",
-    metric_doc: str = "Output value of model",
+    metric_doc: str = "Output value of SVM model",
     metric_namespace: str = "",
     metric_subsystem: str = "",
-    buckets=(0, 1, 2),
+    buckets=(0, 1),  # Adjust buckets based on the SVM model output range
 ) -> Callable[[Info], None]:
     METRIC = Histogram(
         metric_name,
@@ -74,17 +73,40 @@ def model_output(
         namespace=metric_namespace,
         subsystem=metric_subsystem,
     )
-    METRIC.labels("SVC")
-    METRIC.labels("LogisticRegression")
+    METRIC.labels("SVM")
 
     def instrumentation(info: Info) -> None:
-        if info.modified_handler == "/models/{type}":
-            predicted_flower_type = info.response.headers.get("X-model-prediction") # type: ignore
-            model_type = info.response.headers.get("X-model-type") # type: ignore
-            if predicted_flower_type:
-                METRIC.labels(model_type).observe(float(predicted_flower_type))
+        if info.modified_handler == "/predict":  # Adjust based on your endpoint
+            predicted_class = info.response.headers.get("X-model-prediction")  # type: ignore
+            model_type = info.response.headers.get("X-model-type")  # type: ignore
+            if predicted_class:
+                METRIC.labels(model_type or "SVM").observe(float(predicted_class))
 
     return instrumentation
 
+# Counter for model inference calls
+def inference_counter(
+    metric_name: str = "model_inferences_total",
+    metric_doc: str = "Total number of inferences made by the SVM model",
+    metric_namespace: str = "",
+    metric_subsystem: str = "",
+) -> Callable[[Info], None]:
+    METRIC = Counter(
+        metric_name,
+        metric_doc,
+        labelnames=["model_type"],
+        namespace=metric_namespace,
+        subsystem=metric_subsystem,
+    )
+    METRIC.labels("SVM")
+
+    def instrumentation(info: Info) -> None:
+        if info.modified_handler == "/predict":  # Adjust based on your endpoint
+            model_type = info.response.headers.get("X-model-type")  # type: ignore
+            METRIC.labels(model_type or "SVM").inc()
+
+    return instrumentation
 
+# Add metrics to the instrumentator
 instrumentator.add(model_output(metric_namespace=NAMESPACE, metric_subsystem=SUBSYSTEM))
+instrumentator.add(inference_counter(metric_namespace=NAMESPACE, metric_subsystem=SUBSYSTEM))