model_server/app.py - master - NeonBite/big-data-5

model_server/app.py
92 строки3 KB
NeonBite
feat: model-server added
20 дек 2025, 18:08
20 дек 2025, 18:088f01feb
import datetime as dt
import logging
import os
from typing import List, Optional

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)


class PredictRequest(BaseModel):
    features: List[float]


app = FastAPI(title="Model Server", version="1.0")
_session: Optional[ort.InferenceSession] = None
_input_name: Optional[str] = None


def setup_logging():
    log_dir = os.environ.get("LOG_DIR", ".logs")
    os.makedirs(log_dir, exist_ok=True)
    script_name = os.path.splitext(os.path.basename(__file__))[0]
    timestamp = dt.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    log_path = os.path.join(log_dir, f"{script_name}-{timestamp}.log")
    logging.basicConfig(
        level=os.environ.get("LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(name)s - %(message)s",
        handlers=[
            logging.FileHandler(log_path),
            logging.StreamHandler(),
        ],
    )
    logger.info("Logging to %s", log_path)


def load_model() -> None:
    global _session, _input_name
    model_path = os.environ.get("MODEL_PATH", "/models/model.onnx")
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model not found: {model_path}")
    _session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    _input_name = _session.get_inputs()[0].name
    logger.info("Model loaded from %s", model_path)


@app.on_event("startup")
def on_startup() -> None:
    setup_logging()
    try:
        load_model()
    except Exception:
        logger.exception("Failed to load model")
        raise


@app.get("/health")
def health() -> dict:
    status = "ok" if _session is not None else "model_missing"
    return {"status": status}


@app.post("/predict")
def predict(request: PredictRequest) -> dict:
    if _session is None or _input_name is None:
        raise HTTPException(status_code=503, detail="Model is not loaded")
    if not request.features:
        raise HTTPException(status_code=400, detail="Empty features")

    features = np.asarray(request.features, dtype=np.float32).reshape(1, -1)
    try:
        outputs = _session.run(None, {_input_name: features})
    except Exception as exc:
        logger.exception("Inference failed")
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    prob = None
    for output in outputs:
        arr = np.asarray(output)
        if arr.ndim == 2 and arr.shape[1] >= 2:
            prob = float(arr[0, 1])
            break
        if arr.ndim == 1 and arr.size == 1:
            prob = float(arr[0])

    if prob is None:
        raise HTTPException(status_code=500, detail="Unexpected model output format")

    return {"probability_like": prob, "label": int(prob >= 0.5)}