61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
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
|
|
from image_prediction.utils.process_wrapping import wrap_in_process
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
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():
|
|
resp = jsonify("OK")
|
|
resp.status_code = 200
|
|
return resp
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def healthy():
|
|
resp = jsonify("OK")
|
|
resp.status_code = 200
|
|
return resp
|
|
|
|
def __failure():
|
|
response = jsonify("Analysis failed")
|
|
response.status_code = 500
|
|
return response
|
|
|
|
@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.
|
|
# See: https://stackoverflow.com/questions/39758094/clearing-tensorflow-gpu-memory-after-model-execution
|
|
predict_fn_wrapped = wrap_in_process(predict_fn)
|
|
|
|
logger.info("Analysing...")
|
|
predictions = predict_fn_wrapped(request.data)
|
|
|
|
if predictions is not None:
|
|
response = jsonify(predictions)
|
|
logger.info("Analysis completed.")
|
|
return response
|
|
else:
|
|
logger.error("Analysis failed.")
|
|
return __failure()
|
|
|
|
@app.route("/prometheus", methods=["GET"])
|
|
def prometheus():
|
|
return generate_latest(registry=registry)
|
|
|
|
return app
|