Fix MLflow deprecation error and fallback model loading

Dockerfile

@@ -41,5 +41,6 @@ EXPOSE 7860
 
 # Default model run ID
 ENV KRONECTOR_MODEL_RUN_ID=8d8d20f14dce44d991c5fccfdc090a68
+ENV MLFLOW_ALLOW_FILE_STORE=true
 
 CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "7860"]

Dockerfile.hf

@@ -58,4 +58,6 @@ EOF
 # HF Spaces expects port 7860
 EXPOSE 7860
 
+ENV MLFLOW_ALLOW_FILE_STORE=true
+
 CMD ["supervisord", "-c", "/etc/supervisor/supervisord.conf"]

api/main.py

@@ -103,6 +103,9 @@ async def lifespan(app: FastAPI):
         
         load_dotenv()  # Load variables from .env file
         
+        # Allow MLflow to use the local filesystem store (required in MLflow 2.13+)
+        os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true"
+        
         # Ensure MLFLOW_TRACKING_URI is set correctly to local mlruns if empty or unset
         if not os.getenv("MLFLOW_TRACKING_URI"):
             os.environ["MLFLOW_TRACKING_URI"] = "./mlruns"

ml/predict.py

@@ -19,8 +19,19 @@ def load_model_and_encoders(run_id: str):
     """Load a logged MLflow model and its fitted categorical encoders."""
     import mlflow.lightgbm
     from mlflow.artifacts import download_artifacts
-
-    model = mlflow.lightgbm.load_model(f"runs:/{run_id}/model")
+    import logging
+    
+    logger = logging.getLogger(__name__)
+
+    try:
+        model = mlflow.lightgbm.load_model(f"runs:/{run_id}/model")
+    except Exception as e:
+        logger.warning(f"Could not load from runs:/ URI: {e}. Trying registered model...")
+        try:
+            model = mlflow.lightgbm.load_model("models:/kronector-f1-lgbm/latest")
+        except Exception as e2:
+            logger.error(f"Failed to load registered model: {e2}")
+            raise
     encoder_path = download_artifacts(
         run_id=run_id, artifact_path="encoders/label_encoders.pkl"
     )