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

model_server/app.py
98 строк3 KB
NeonBite
feat: report done
20 дек 2025, 19:28
20 дек 2025, 19:28d12a6ed
import datetime as dt
import logging
import os
from typing import List, Optional

import numpy as np
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[object] = 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
    try:
        import onnxruntime as ort
    except ImportError:
        logger.exception("onnxruntime import failed; model server will run without inference")
        _session = None
        _input_name = None
        return
    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)}