flask_model_serving_container_with_health_and_predict_endpoints.py

1import os
2from flask import Flask, request, jsonify
3
4app = Flask(__name__)
5
6# The standard requires health checks to return 200 OK.
7# Default path for health checks is often /health or defined via environment variables.
8@app.route(os.environ.get('AIP_HEALTH_ROUTE', '/health'), methods=['GET'])
9def health_check():
10    return {"status": "healthy"}, 200
11
12# The standard requires a prediction endpoint.
13# Default path is /predict or defined via environment variables.
14@app.route(os.environ.get('AIP_PREDICT_ROUTE', '/predict'), methods=['POST'])
15def predict():
16    # Requests are typically JSON with an "instances" key.
17    data = request.get_json()
18    if not data or "instances" not in data:
19        return {"error": "Invalid input, 'instances' key required"}, 400
20
21    instances = data["instances"]
22    
23    # Placeholder for model inference logic
24    predictions = [f"prediction_for_{instance}" for instance in instances]
25
26    return jsonify({"predictions": predictions})
27
28if __name__ == "__main__":
29    # The port is typically provided by the AIP_HTTP_PORT environment variable (default 8080).
30    port = int(os.environ.get('AIP_HTTP_PORT', 8080))
31    app.run(host='0.0.0.0', port=port)