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 aa6a0f0..da99202 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,4 +21,5 @@ Pillow==9.1.0 PDFNetPython3==9.1.0 pdf2image==1.16.0 frozendict==2.3.0 -protobuf<=3.20.* \ No newline at end of file +protobuf<=3.20.* +prometheus-client==0.13.1