flask_model_hosting_container_health_and_predict_endpoints.py

1import os
2from flask import Flask, request, jsonify
3
4app = Flask(__name__)
5
6# The environment variable AIP_HEALTH_ROUTE is injected by the platform, 
7# but defaults to /health if not specified.
8HEALTH_ROUTE = os.environ.get('AIP_HEALTH_ROUTE', '/health')
9
10# The environment variable AIP_PREDICT_ROUTE is injected by the platform, 
11# but defaults to /predict if not specified.
12PREDICT_ROUTE = os.environ.get('AIP_PREDICT_ROUTE', '/predict')
13
14@app.route(HEALTH_ROUTE, methods=['GET'])
15def health_check():
16    """Status check to ensure the container is ready to handle requests."""
17    return {"status": "healthy"}, 200
18
19@app.route(PREDICT_ROUTE, methods=['POST'])
20def predict():
21    """Handler for prediction requests."""
22    # Retrieve the instances from the request body
23    body = request.get_json()
24    instances = body.get("instances", [])
25
26    # Perform inference (Example logic)
27    predictions = [f"Processed: {instance}" for instance in instances]
28
29    return jsonify({"predictions": predictions})
30
31if __name__ == "__main__":
32    # The platform expects the server to listen on port 8080 by default,
33    # or the port specified in the AIP_HTTP_PORT environment variable.
34    port = int(os.environ.get('AIP_HTTP_PORT', 8080))
35    app.run(host='0.0.0.0', port=port)