From fd0e4dc3cf437c54f7f75c97f602e48dfa39f6fe Mon Sep 17 00:00:00 2001
From: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date: Tue, 21 Jun 2022 14:57:39 +0200
Subject: [PATCH] Pull request #14: RED-4329 add prometheus endpoint with
 summary metric

Merge in RR/image-prediction from RED-4329-add-prometheus to master

Squashed commit of the following:

commit 2b8446e703617c6897b6149846f2548ec292a9a1
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Tue Jun 21 14:40:44 2022 +0200

    RED-4329 add prometheus endpoint with summary metric
---
 image_prediction/flask.py | 10 ++++++++++
 requirements.txt          |  1 +
 2 files changed, 11 insertions(+)

diff --git a/image_prediction/flask.py b/image_prediction/flask.py
index b1473da..4297a6f 100644
--- a/image_prediction/flask.py
+++ b/image_prediction/flask.py
@@ -3,6 +3,7 @@ import traceback
 from typing import Callable
 
 from flask import Flask, request, jsonify
+from prometheus_client import generate_latest, CollectorRegistry, Summary
 
 from image_prediction.utils import get_logger
 
@@ -34,6 +35,10 @@ def wrap_in_process(func_to_wrap):
 
 def make_prediction_server(predict_fn: Callable):
     app = Flask(__name__)
+    registry = CollectorRegistry(auto_describe=True)
+    metric = Summary(
+        f"redactmanager_imageClassification_seconds", f"Time spent on image-service classification.", registry=registry
+    )
 
     @app.route("/ready", methods=["GET"])
     def ready():
@@ -54,6 +59,7 @@ def make_prediction_server(predict_fn: Callable):
 
     @app.route("/predict", methods=["POST"])
     @app.route("/", methods=["POST"])
+    @metric.time()
     def predict():
 
         # Tensorflow does not free RAM. Workaround: Run prediction function (which instantiates a model) in sub-process.
@@ -71,4 +77,8 @@ def make_prediction_server(predict_fn: Callable):
             logger.error("Analysis failed.")
             return __failure()
 
+    @app.route("/prometheus", methods=["GET"])
+    def prometheus():
+        return generate_latest(registry=registry)
+
     return app
diff --git a/requirements.txt b/requirements.txt
index 8b33526..56fd8aa 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -21,3 +21,4 @@ Pillow==9.1.0
 PDFNetPython3==9.1.0
 pdf2image==1.16.0
 frozendict==2.3.0
+prometheus-client==0.13.1
\ No newline at end of file