import os import warnings # Suppress deprecation and future warnings from third-party libraries (like MLflow) warnings.filterwarnings("ignore", category=FutureWarning) os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true" from fastapi import FastAPI, HTTPException from pydantic import BaseModel import mlflow import pandas as pd import logging from contextlib import asynccontextmanager logging.basicConfig(level=logging.INFO) # Configuration # Force API to use local models bundled via Git LFS instead of downloading from DagsHub MLFLOW_TRACKING_URI = "file:///app/mlruns" if os.path.exists("/app/mlruns") else "file:./mlruns" EXPERIMENT_NAME = "SL_Commodity_Forecasting" # Global model variable model = None COMMODITY_MAP = { "Samba": 0, "Kekulu": 1, "Big Onion": 2, "Potato": 3, "Dried Chilli": 4, "Coconut": 5 } class PredictionRequest(BaseModel): commodity: str lag_prices: list[float] model_config = { "json_schema_extra": { "example": { "commodity": "Samba", "lag_prices": [220.5, 219.0, 218.5, 221.0, 222.5, 220.0, 219.5] } } } def load_best_model(): """Loads the best XGBoost model from the local MLflow registry on startup.""" global model mlflow.set_tracking_uri(MLFLOW_TRACKING_URI) # Configure DagsHub credentials automatically if URI is provided dagshub_token = os.environ.get("DAGSHUB_USER_TOKEN") if dagshub_token and "dagshub.com" in MLFLOW_TRACKING_URI: os.environ["MLFLOW_TRACKING_USERNAME"] = MLFLOW_TRACKING_URI.split('/')[-2] os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token logging.info("Configured DagsHub MLflow remote for API") try: experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME) if not experiment: logging.error(f"Experiment '{EXPERIMENT_NAME}' not found. Make sure you trained the model.") return # Fetch the best XGBoost run by RMSE runs = mlflow.search_runs( experiment_ids=[experiment.experiment_id], filter_string="tags.mlflow.runName = 'XGBoost_Candidate'", order_by=["metrics.rmse ASC"], max_results=1 ) is_empty = False if runs is None: is_empty = True elif hasattr(runs, "empty"): is_empty = runs.empty else: is_empty = not runs if is_empty: logging.error("No XGBoost runs found in MLflow.") return if hasattr(runs, "iloc"): best_run_id = runs.iloc[0].run_id else: # Handle list/PagedList of Run objects robustly best_run = runs[0] best_run_id = None # 1. Try direct attribute/key 'run_id' if hasattr(best_run, "run_id"): best_run_id = best_run.run_id elif isinstance(best_run, dict) and "run_id" in best_run: best_run_id = best_run["run_id"] # 2. Try info attribute or callable info if not best_run_id and hasattr(best_run, "info"): info_attr = best_run.info if callable(info_attr): try: info_obj = info_attr() best_run_id = getattr(info_obj, "run_id", None) except Exception: pass if not best_run_id: best_run_id = getattr(info_attr, "run_id", None) # 3. Try info key/dict if not best_run_id and isinstance(best_run, dict) and "info" in best_run: info_val = best_run["info"] if isinstance(info_val, dict): best_run_id = info_val.get("run_id") else: best_run_id = getattr(info_val, "run_id", None) # 4. Last resort fallback if not best_run_id: best_run_id = getattr(best_run, "run_id", None) # Try finding the model in the git-tracked models folder model_found = False model_uri = None models_dir = os.path.join("mlruns", experiment.experiment_id, "models") if os.path.exists(models_dir): for model_id in os.listdir(models_dir): meta_path = os.path.join(models_dir, model_id, "meta.yaml") if os.path.exists(meta_path): with open(meta_path, "r") as f: if f"source_run_id: {best_run_id}" in f.read(): model_uri = os.path.join(models_dir, model_id, "artifacts") model_found = True break if not model_found or model_uri is None: model_uri = f"runs:/{best_run_id}/xgboost_model" logging.info(f"Loading latest XGBoost model (Run ID: {best_run_id}) from {model_uri}") model = mlflow.xgboost.load_model(model_uri) logging.info("✅ Model loaded successfully and ready for inference.") except Exception as e: logging.error(f"❌ Failed to load model: {e}") @asynccontextmanager async def lifespan(app: FastAPI): load_best_model() yield app = FastAPI( title="SL Commodity Forecasting API", description="MLOps API serving the XGBoost model for Sri Lankan commodity prices", version="1.0", lifespan=lifespan ) @app.get("/") def health_check(): return { "service": "MLOps Forecasting API", "status": "healthy", "model_loaded": model is not None } @app.post("/predict") def predict_price(request: PredictionRequest): if not model: raise HTTPException(status_code=503, detail="Model is currently unavailable or failed to load.") if len(request.lag_prices) != 7: raise HTTPException(status_code=400, detail="Exactly 7 lag prices must be provided (chronological order: oldest to newest).") if request.commodity not in COMMODITY_MAP: raise HTTPException(status_code=400, detail=f"Unknown commodity. Must be one of {list(COMMODITY_MAP.keys())}") try: current_lags = request.lag_prices.copy() forecast = [] # 7-day recursive forecasting for day in range(7): # Construct the feature DataFrame based on what the model expects expected_features = model.get_booster().feature_names features: dict[str, list] = {} if "Commodity_ID" in expected_features: features["Commodity_ID"] = [COMMODITY_MAP[request.commodity]] for i in range(1, 8): if f"Lag_{i}" in expected_features: features[f"Lag_{i}"] = [current_lags[-i]] # Ensure columns are in the exact order the model expects df_features = pd.DataFrame(features)[expected_features] # Predict prediction = float(model.predict(df_features)[0]) forecast.append(round(prediction, 2)) # Append new prediction to the end of the lags (it becomes the new t-1 for the next iteration) current_lags.append(prediction) return { "predicted_price_lkr": forecast[0], "7_day_forecast": forecast, "currency": "LKR" } except Exception as e: raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")