Building ML model serving APIs with Flask, request handling, and deployment.
Flask is a lightweight web framework perfect for serving ML models via HTTP endpoints. It's easy to learn and deploy.
Why Flask for ML?
Advantages:
- Simple: Minimal boilerplate, easy to learn
- Flexible: Highly customizable
- Lightweight: Minimal dependencies
- Testing: Easy to test
- Mature: Well-established, many tutorials
When to use:
- Simple to moderate complexity models
- Learning/prototyping
- Microservices
- When you need full control
When NOT to use:
- High-traffic APIs (use FastAPI, async)
- Complex request pipelines
- Need automatic documentation
Basic ML API with Flask
from flask import Flask, request, jsonify
import pickle
import numpy as np
import logging
app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Load model at startup
try:
with open('model.pkl', 'rb') as f:
model = pickle.load(f)
logger.info("Model loaded successfully")
except Exception as e:
logger.error(f"Failed to load model: {e}")
model = None
@app.route('/health', methods=['GET'])
def health():
"""Health check endpoint"""
return jsonify({'status': 'healthy', 'model_loaded': model is not None}), 200
@app.route('/predict', methods=['POST'])
def predict():
"""Single prediction endpoint"""
try:
if not model:
return jsonify({'error': 'Model not loaded'}), 500
data = request.get_json()
if not data or 'features' not in data:
return jsonify({'error': 'No features provided'}), 400
# Convert to array
features = np.array(data['features']).reshape(1, -1)
# Predict
prediction = model.predict(features)[0]
return jsonify({
'prediction': float(prediction),
'status': 'success'
}), 200
except Exception as e:
logger.error(f"Prediction error: {str(e)}")
return jsonify({'error': str(e)}), 500
@app.route('/batch_predict', methods=['POST'])
def batch_predict():
"""Batch prediction endpoint"""
try:
data = request.get_json()
features = np.array(data['features'])
predictions = model.predict(features)
return jsonify({
'predictions': predictions.tolist(),
'count': len(predictions)
}), 200
except Exception as e:
return jsonify({'error': str(e)}), 500
@app.route('/model_info', methods=['GET'])
def model_info():
"""Get model information"""
return jsonify({
'model_type': str(type(model).__name__),
'n_features': getattr(model, 'n_features_in_', 'unknown'),
'version': '1.0'
}), 200
if __name__ == '__main__':
# Development server
app.run(debug=False, host='0.0.0.0', port=8000)
Testing Flask API
import pytest
from main import app
@pytest.fixture
def client():
"""Create test client"""
app.config['TESTING'] = True
with app.test_client() as client:
yield client
class TestMLAPI:
def test_health(self, client):
"""Test health check endpoint"""
response = client.get('/health')
assert response.status_code == 200
assert response.json['status'] == 'healthy'
def test_predict_valid(self, client):
"""Test prediction with valid input"""
response = client.post('/predict', json={
'features': [1.0, 2.0, 3.0, 4.0]
})
assert response.status_code == 200
assert 'prediction' in response.json
assert 'status' in response.json
def test_predict_missing_features(self, client):
"""Test prediction with missing features"""
response = client.post('/predict', json={})
assert response.status_code == 400
assert 'error' in response.json
def test_batch_predict(self, client):
"""Test batch prediction"""
response = client.post('/batch_predict', json={
'features': [
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
[7.0, 8.0, 9.0]
]
})
assert response.status_code == 200
assert response.json['count'] == 3
def test_model_info(self, client):
"""Test model info endpoint"""
response = client.get('/model_info')
assert response.status_code == 200
assert 'model_type' in response.json
# Run tests: pytest test_api.py
Error Handling & Validation
from flask import Flask, request, jsonify
from functools import wraps
app = Flask(__name__)
def validate_input(required_fields):
"""Decorator for input validation"""
def decorator(f):
@wraps(f)
def decorated_function(*args, **kwargs):
data = request.get_json()
if not data:
return jsonify({'error': 'Empty request body'}), 400
for field in required_fields:
if field not in data:
return jsonify({'error': f'Missing field: {field}'}), 400
return f(*args, **kwargs)
return decorated_function
return decorator
@app.route('/predict', methods=['POST'])
@validate_input(['features'])
def predict():
data = request.get_json()
# features guaranteed to exist here
return jsonify({'prediction': 1.0})
@app.errorhandler(400)
def bad_request(error):
return jsonify({'error': 'Bad request', 'details': str(error)}), 400
@app.errorhandler(500)
def internal_error(error):
return jsonify({'error': 'Internal server error'}), 500
Production Deployment with Gunicorn
# Install gunicorn
pip install gunicorn
# Basic deployment (4 workers)
gunicorn -w 4 -b 0.0.0.0:8000 main:app
# With logging
gunicorn -w 4 \
-b 0.0.0.0:8000 \
--log-level info \
--access-logfile - \
--error-logfile - \
main:app
# With timeout and restart settings
gunicorn -w 4 \
-b 0.0.0.0:8000 \
--timeout 60 \
--max-requests 1000 \
--max-requests-jitter 100 \
main:app
Docker Deployment
FROM python:3.9-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:8000", "main:app"]
docker build -t ml-api:1.0 .
docker run -p 8000:8000 ml-api:1.0
API Client Example
import requests
API_URL = 'http://localhost:8000'
# Health check
response = requests.get(f'{API_URL}/health')
print(response.json())
# Single prediction
response = requests.post(f'{API_URL}/predict', json={
'features': [1.0, 2.0, 3.0, 4.0]
})
print(f"Prediction: {response.json()['prediction']}")
# Batch prediction
response = requests.post(f'{API_URL}/batch_predict', json={
'features': [
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0]
]
})
print(f"Predictions: {response.json()['predictions']}")
Comparison: Flask vs FastAPI
| Feature | Flask | FastAPI |
|---|
| Learning Curve | Easy | Moderate |
| Performance | Good | Excellent |
| Async Support | Limited | Native |
| Documentation | Manual | Automatic |
| Type Hints | Optional | Required |
| Deployment | Gunicorn | Uvicorn/Gunicorn |
Choose Flask for:
- Simple models
- Learning
- Moderate traffic
Choose FastAPI for:
- High-traffic APIs
- Complex features
- Auto documentation needed
Summary
Flask for ML:
- Simple to implement
- Easy to test
- Lightweight
- Good for learning
- Production-ready with gunicorn