139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
"""
|
|
Урок 8: Model Serving — разворачивание модели как REST API
|
|
============================================================
|
|
Загружает модель из Model Registry и запускает HTTP-сервер.
|
|
Модель принимает JSON-запросы и возвращает предсказания.
|
|
|
|
Два способа:
|
|
1. Встроенный MLflow server (простой):
|
|
mlflow models serve -m models:/digits_rf_model/Production -p 5001
|
|
|
|
2. Этот скрипт — кастомный сервер с дополнительной логикой:
|
|
|
|
Запуск:
|
|
python src/serve_model.py --port 5001
|
|
|
|
Запрос:
|
|
curl -X POST http://localhost:5001/predict \
|
|
-H "Content-Type: application/json" \
|
|
-d '{"instances": [[0,0,0,...,0]]}'
|
|
"""
|
|
import argparse
|
|
import json
|
|
import numpy as np
|
|
from sklearn.datasets import load_digits
|
|
|
|
import mlflow
|
|
from flask import Flask, request, jsonify
|
|
|
|
|
|
def create_app(model_name, stage):
|
|
app = Flask(__name__)
|
|
|
|
# Загружаем модель один раз при старте
|
|
print(f"📥 Загружаем модель '{model_name}' ({stage})...")
|
|
model = mlflow.sklearn.load_model(f"models:/{model_name}/{stage}")
|
|
print(f"✅ Модель загружена!")
|
|
|
|
# Загружаем датасет для демо-генерации примеров
|
|
digits = load_digits()
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def health():
|
|
return jsonify({"status": "ok", "model": model_name, "stage": stage})
|
|
|
|
@app.route("/info", methods=["GET"])
|
|
def info():
|
|
return jsonify({
|
|
"model": model_name,
|
|
"stage": stage,
|
|
"n_features": 64,
|
|
"classes": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
|
"endpoints": ["/health", "/info", "/predict", "/predict_random"],
|
|
})
|
|
|
|
@app.route("/predict", methods=["POST"])
|
|
def predict():
|
|
"""Предсказание для переданных образцов.
|
|
Формат запроса:
|
|
{"instances": [[...64 числа...], [...64 числа...]]}
|
|
"""
|
|
data = request.get_json(force=True)
|
|
|
|
if "instances" not in data:
|
|
return jsonify({"error": "поле 'instances' обязательно"}), 400
|
|
|
|
instances = np.array(data["instances"])
|
|
|
|
if instances.shape[-1] != 64:
|
|
return jsonify({
|
|
"error": f"ожидалось 64 признака, получено {instances.shape[-1]}"
|
|
}), 400
|
|
|
|
predictions = model.predict(instances)
|
|
|
|
# Если модель поддерживает predict_proba — добавим вероятности
|
|
try:
|
|
probabilities = model.predict_proba(instances)
|
|
return jsonify({
|
|
"predictions": predictions.tolist(),
|
|
"probabilities": probabilities.tolist(),
|
|
"n_samples": len(predictions),
|
|
})
|
|
except AttributeError:
|
|
return jsonify({
|
|
"predictions": predictions.tolist(),
|
|
"n_samples": len(predictions),
|
|
})
|
|
|
|
@app.route("/predict_random", methods=["GET"])
|
|
def predict_random():
|
|
"""Демо-эндпоинт: берёт случайный образец из digits и предсказывает."""
|
|
idx = np.random.randint(0, len(digits.data))
|
|
sample = digits.data[idx:idx+1]
|
|
true_label = int(digits.target[idx])
|
|
prediction = int(model.predict(sample)[0])
|
|
|
|
try:
|
|
probs = model.predict_proba(sample)[0]
|
|
confidence = float(max(probs))
|
|
except AttributeError:
|
|
confidence = None
|
|
|
|
return jsonify({
|
|
"true_label": true_label,
|
|
"prediction": prediction,
|
|
"correct": true_label == prediction,
|
|
"confidence": confidence,
|
|
"sample_image": sample[0].reshape(8, 8).tolist(),
|
|
})
|
|
|
|
return app
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="MLflow model serving")
|
|
parser.add_argument("--model-name", default="digits_rf_model")
|
|
parser.add_argument("--stage", default="Production")
|
|
parser.add_argument("--port", type=int, default=5001)
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
args = parser.parse_args()
|
|
|
|
app = create_app(args.model_name, args.stage)
|
|
|
|
print(f"\n🚀 Сервер запускается на http://{args.host}:{args.port}")
|
|
print(f" /health — проверка состояния")
|
|
print(f" /info — информация о модели")
|
|
print(f" /predict — предсказание (POST JSON)")
|
|
print(f" /predict_random — случайный образец (GET)")
|
|
print(f"\n📋 Примеры запросов:")
|
|
print(f" curl http://localhost:{args.port}/health")
|
|
print(f" curl http://localhost:{args.port}/predict_random")
|
|
print()
|
|
|
|
app.run(host=args.host, port=args.port, debug=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|