Spaces:
Sleeping
Sleeping
File size: 3,244 Bytes
5edbdc5 2bbee40 5edbdc5 2bbee40 5edbdc5 2bbee40 5edbdc5 2bbee40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import os
import mlflow
import pandas as pd
import uvicorn
import json
from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
from typing import Literal, List, Union
# -----------------------------------------------------------------------------
# ENV + MLflow setup
# -----------------------------------------------------------------------------
# Sur Hugging Face, ces variables sont lues depuis les "Secrets"
MLFLOW_TRACKING_URI = os.getenv("MLFLOW_TRACKING_URI")
REGISTERED_MODEL_NAME = os.getenv("MLFLOW_REGISTERED_MODEL_NAME")
MODEL_STAGE = os.getenv("MLFLOW_MODEL_STAGE")
MODEL_ALIAS = os.getenv("MLFLOW_MODEL_ALIAS")
# On force l'URI pour mlflow
if MLFLOW_TRACKING_URI:
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
def build_model_uri() -> str:
if MODEL_ALIAS:
return f"models:/{REGISTERED_MODEL_NAME}@{MODEL_ALIAS}"
return f"models:/{REGISTERED_MODEL_NAME}/{MODEL_STAGE or 'Production'}"
# models:/ibm_attrition_detector@production
MODEL_URI = build_model_uri()
MODEL = None
# -----------------------------------------------------------------------------
# FastAPI Setup
# -----------------------------------------------------------------------------
app = FastAPI(title="🌬️ Wind Turbine Prediction API")
class PredictionFeatures(BaseModel):
Hour_Index: Union[int, float]
Turbine_ID: int
Rotor_Speed_RPM: float
Wind_Speed_mps: float
Power_Output_kW: float
Gearbox_Oil_Temp_C: float
Generator_Bearing_Temp_C: float
Vibration_Level_mmps: float
Ambient_Temp_C: float
Humidity_pct: float
Maintenance_Label: int
# -----------------------------------------------------------------------------
# Startup: CHARGEMENT BLOQUANT (Solution au bug 500)
# -----------------------------------------------------------------------------
@app.on_event("startup")
def load_model_sync():
global MODEL
print(f"🚀 [INFO] Attempting to load model: {MODEL_URI}")
try:
# On attend que le chargement soit fini avant de rendre l'API disponible
MODEL = mlflow.sklearn.load_model(MODEL_URI)
# models:/wind_turbine_predictor@production
print("✅ [INFO] Model loaded successfully!")
except Exception as e:
print(f"❌ [ERROR] Failed to load model: {e}")
# En cas d'échec, on laisse MODEL à None pour que /health le signale
# -----------------------------------------------------------------------------
# Endpoints
# -----------------------------------------------------------------------------
@app.get("/health")
def health():
return {
"status": "ok",
"model_uri": MODEL_URI,
"model_loaded": MODEL is not None,
}
@app.post("/predict")
async def predict(payload: PredictionFeatures):
if MODEL is None:
raise HTTPException(
status_code=503, detail="Model is still loading or failed to load."
)
# Conversion pydantic -> dict -> DataFrame
df = pd.DataFrame([payload.dict()])
pred = MODEL.predict(df)
return {"prediction": int(pred[0])}
if __name__ == "__main__":
# Port 7860 est le standard pour Hugging Face Spaces
uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860))) |