561a7f527c
Merge in RR/image-prediction from RED-4206-fix-unwanted-restart-bug to master
Squashed commit of the following:
commit 3dfe7b861816ef9019103e16a23efd97a08fb617
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date: Thu Sep 22 13:53:32 2022 +0200
RED-4206 wrap queue callback in process to manage memory allocation with the operating system and force deallocation after processing.
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
|