Implantacao

from importlib import import_module
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

from nyc_taxi_demand_forecasting.features.hourly_demand import FEATURE_COLUMNS


class BatchDemandPredictor:
    """Score processed demand rows with an MLflow model URI.

    Example:
        BatchDemandPredictor().predict("models:/name/latest", input_path, output_path)
    """

    def predict(
        self, model_uri: str, input_path: Path, output_path: Path
    ) -> Path:
        mlflow = import_module("mlflow")
        model = mlflow.pyfunc.load_model(model_uri)
        input_frame = pd.read_parquet(input_path)
        # pyfunc emits an ndarray, so normalize the output shape before the
        # frame assignment instead of trusting a Series.
        raw_predictions: Any = model.predict(
            input_frame.loc[:, list(FEATURE_COLUMNS)]
        )
        predictions = np.asarray(raw_predictions).ravel()
        output_frame = input_frame.assign(prediction=predictions)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_frame.to_parquet(output_path, index=False)
        return output_path

import logging
from dataclasses import dataclass
from importlib import import_module
from pathlib import Path
from typing import TYPE_CHECKING, Any

import mlflow
import numpy as np
import pandas as pd
from fastapi import FastAPI, HTTPException
from mlflow.exceptions import MlflowException
from mlflow.pyfunc import PyFuncModel

from nyc_taxi_demand_forecasting.configuration import ProjectConfigLoader
from nyc_taxi_demand_forecasting.features.hourly_demand import (
    FEATURE_COLUMNS,
    FEATURE_REFS,
)

if TYPE_CHECKING:
    from feast import FeatureStore

logger = logging.getLogger(__name__)

app = FastAPI(title="NYC Taxi Demand Serving API")

# api.py sits at src/nyc_taxi_demand_forecasting/inference/api.py — the
# project config lives at <project root>/configs/project.yaml, three parents up.
_CONFIG_PATH = Path(__file__).resolve().parents[3] / "configs" / "project.yaml"


@dataclass(frozen=True)
class _ServingResources:
    store: "FeatureStore"
    model: PyFuncModel | None


_resources: _ServingResources | None = None


def _load_serving_resources() -> _ServingResources:
    """Build the store and champion model; a registry-level load failure logs
    the cause and degrades to the 503 contract — anything else propagates."""
    config = ProjectConfigLoader().load(_CONFIG_PATH)
    feature_store_type = import_module("feast").FeatureStore
    store = feature_store_type(repo_path=str(config.feast.repo_path))
    mlflow.set_tracking_uri(config.mlflow.tracking_uri)
    model_uri = f"models:/{config.mlflow.registered_model_name}@champion"
    try:
        model: PyFuncModel | None = mlflow.pyfunc.load_model(model_uri)
    except MlflowException:
        logger.exception(
            "champion_model_load_failed", extra={"model_uri": model_uri}
        )
        model = None
    return _ServingResources(store=store, model=model)


def _serving_resources() -> _ServingResources:
    """Build serving resources on first request so importing this module
    stays side-effect free; a cached ``model=None`` keeps a missing champion
    a fast 503 instead of a per-request retry."""
    global _resources
    if _resources is None:
        _resources = _load_serving_resources()
    return _resources


def fetch_online_features(location_id: int) -> pd.DataFrame:
    """Fetch recent demand features from Feast online SQLite store.

    Example:
        features = fetch_online_features(142)
    """
    entity_rows = [{"pickup_location_id": location_id}]
    feature_refs = list(FEATURE_REFS)
    features_df = (
        _serving_resources()
        .store.get_online_features(
            features=feature_refs,
            entity_rows=entity_rows,
        )
        .to_df()
    )
    return features_df


def _features_missing(features: pd.DataFrame) -> bool:
    if features.empty or not set(FEATURE_COLUMNS) <= set(features.columns):
        return True
    count_na: bool = features["pickup_count"].isna().iloc[0]
    return count_na


@app.get("/predict/{location_id}")
def predict_demand(location_id: int) -> dict[str, Any]:
    """Retrieve online features and run prediction with the champion model.

    Example:
        GET /predict/142
    """
    model = _serving_resources().model
    if model is None:
        raise HTTPException(
            status_code=503,
            detail="Model is not loaded. Ensure champion model is promoted in MLflow.",
        )

    try:
        features = fetch_online_features(location_id)
    except Exception as exc:
        logger.exception(
            "online_feature_fetch_failed",
            extra={"pickup_location_id": location_id},
        )
        raise HTTPException(
            status_code=500,
            detail="Failed to retrieve features from the feature store",
        ) from exc

    if _features_missing(features):
        raise HTTPException(
            status_code=404,
            detail=f"No online features found for location ID {location_id}",
        )

    pred_features = features[list(FEATURE_COLUMNS)]

    raw_prediction: Any = model.predict(  # pyright: ignore[reportUnknownMemberType]
        pred_features
    )
    pred_val = float(np.asarray(raw_prediction).ravel()[0])

    return {
        "pickup_location_id": location_id,
        "features": features.to_dict(orient="records")[0],
        "predicted_demand": pred_val,
    }

import logging
from pathlib import Path

import mlflow
from mlflow.client import MlflowClient
from mlflow.exceptions import MlflowException

from nyc_taxi_demand_forecasting.configuration import ProjectConfigLoader


def run(config_path: Path) -> None:
    config = ProjectConfigLoader().load(config_path)

    # Connect to MLflow Model Registry
    mlflow.set_tracking_uri(config.mlflow.tracking_uri)
    client = MlflowClient()

    model_name = config.mlflow.registered_model_name

    # Promote only the gated candidate: the evaluate pipeline binds this
    # alias after the quality gate passes, so an unevaluated version can
    # never reach champion.
    try:
        candidate_version = client.get_model_version_by_alias(
            model_name, "candidate"
        )
    except MlflowException as err:
        raise ValueError(
            f"No evaluated candidate for model {model_name}: expected the "
            "evaluate pipeline to set the 'candidate' alias"
        ) from err

    # Promote model by setting the alias "champion"
    client.set_registered_model_alias(
        name=model_name,
        alias="champion",
        version=candidate_version.version,
    )

    logging.getLogger(__name__).info(
        "model_deployment_completed",
        extra={
            "model_name": model_name,
            "version": candidate_version.version,
            "alias": "champion",
        },
    )