"""
Custom images must provide their own webserver, with port 8080 and GET ping and POST invocations endpoints

This webserver calls inference.py which follows standard sagemaker entry_point conventions, but as a custom webserver you could define your own convention
"""

from flask import Flask, request, Response

from inference import model_fn, input_fn, predict_fn, output_fn

app = Flask(__name__)

model = None

@app.before_request
def load_model():
    global model
    if not model:
        model = model_fn('.')

@app.route('/ping', methods=['GET'])
def ping():
    return '', 200

@app.route('/invocations', methods=['POST'])
def invocations():
    if not model:
        return Response(response='Model not loaded', status=503)
    input_data = input_fn(request.data, request.content_type)
    predictions = predict_fn(input_data, model)
    output = output_fn(predictions, request.accept_mimetypes)
    return Response(output,
                    status=200,
                    mimetype='application/json')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080)