Spaces:
Sleeping
Sleeping
| """Model Registry for loading, caching, and serving trained model artifacts. | |
| Supports eager loading at startup and lazy loading on first access. | |
| Reports per-model health status. | |
| """ | |
| import json | |
| import logging | |
| from pathlib import Path | |
| from typing import Any | |
| import joblib | |
| from app.core.exceptions import ModelNotLoadedError | |
| logger = logging.getLogger(__name__) | |
| # All registered model names and their artifact subdirectory names | |
| REGISTERED_MODELS: list[str] = [ | |
| "lo_tagger", | |
| "bloom_classifier", | |
| "mastery_model", | |
| "risk_model", | |
| "answer_scorer", | |
| "recommender", | |
| ] | |
| class ModelRegistry: | |
| """Loads and manages trained model artifacts. | |
| Supports eager loading at startup and lazy loading on first access. | |
| Reports per-model health status. | |
| """ | |
| def __init__(self, artifact_dir: str | Path) -> None: | |
| self._artifact_dir = Path(artifact_dir) | |
| self._models: dict[str, dict[str, Any]] = {} | |
| self._status: dict[str, str] = {} # "loaded" | "not_loaded" | "error" | |
| self._metadata: dict[str, dict] = {} # metrics.json content per model | |
| # Initialize all registered models as not_loaded | |
| for model_name in REGISTERED_MODELS: | |
| self._status[model_name] = "not_loaded" | |
| def load_all(self) -> None: | |
| """Eagerly load all model artifacts from artifact_dir subdirectories.""" | |
| for model_name in REGISTERED_MODELS: | |
| self._load_model(model_name) | |
| def _load_model(self, model_name: str) -> None: | |
| """Load a single model's artifacts from its subdirectory. | |
| On failure, sets status to 'error' and logs the exception without crashing. | |
| """ | |
| model_dir = self._artifact_dir / model_name | |
| if not model_dir.exists(): | |
| logger.warning( | |
| "Model directory not found for '%s': %s", model_name, model_dir | |
| ) | |
| self._status[model_name] = "not_loaded" | |
| return | |
| try: | |
| model_data: dict[str, Any] = {} | |
| # Required: model.joblib | |
| model_path = model_dir / "model.joblib" | |
| if not model_path.exists(): | |
| logger.warning( | |
| "model.joblib not found for '%s' at %s", model_name, model_path | |
| ) | |
| self._status[model_name] = "not_loaded" | |
| return | |
| model_data["model"] = joblib.load(model_path) | |
| # Optional: vectorizer.joblib | |
| vectorizer_path = model_dir / "vectorizer.joblib" | |
| if vectorizer_path.exists(): | |
| model_data["vectorizer"] = joblib.load(vectorizer_path) | |
| # Optional: label_encoder.joblib | |
| label_encoder_path = model_dir / "label_encoder.joblib" | |
| if label_encoder_path.exists(): | |
| model_data["label_encoder"] = joblib.load(label_encoder_path) | |
| # Optional: feature_columns.json | |
| feature_columns_path = model_dir / "feature_columns.json" | |
| if feature_columns_path.exists(): | |
| with open(feature_columns_path, "r", encoding="utf-8") as f: | |
| model_data["feature_columns"] = json.load(f) | |
| # Optional: metrics.json (stored as metadata) | |
| metrics_path = model_dir / "metrics.json" | |
| if metrics_path.exists(): | |
| with open(metrics_path, "r", encoding="utf-8") as f: | |
| self._metadata[model_name] = json.load(f) | |
| self._models[model_name] = model_data | |
| self._status[model_name] = "loaded" | |
| logger.info("Model '%s' loaded successfully from %s", model_name, model_dir) | |
| except Exception: | |
| logger.exception("Failed to load model '%s' from %s", model_name, model_dir) | |
| self._status[model_name] = "error" | |
| def get_model(self, model_name: str) -> dict[str, Any]: | |
| """Retrieve a loaded model by name. Triggers lazy load if not yet loaded. | |
| Returns a dict containing 'model' and optional 'vectorizer', | |
| 'label_encoder', 'feature_columns' keys. | |
| Raises ModelNotLoadedError if artifact is missing/corrupted after load attempt. | |
| """ | |
| if model_name not in self._status: | |
| raise ModelNotLoadedError( | |
| f"Model '{model_name}' is not registered in the registry." | |
| ) | |
| # Lazy loading: attempt to load if not yet loaded | |
| if self._status[model_name] == "not_loaded": | |
| self._load_model(model_name) | |
| if self._status[model_name] != "loaded": | |
| raise ModelNotLoadedError( | |
| f"Model '{model_name}' is not available (status: {self._status[model_name]})." | |
| ) | |
| return self._models[model_name] | |
| def get_status(self, model_name: str) -> str: | |
| """Return 'loaded', 'not_loaded', or 'error' for a given model.""" | |
| return self._status.get(model_name, "not_loaded") | |
| def get_metadata(self, model_name: str) -> dict | None: | |
| """Return metrics.json content for a model, or None if not available.""" | |
| return self._metadata.get(model_name) | |
| def get_all_status(self) -> dict[str, dict]: | |
| """Return status + metadata for all registered models.""" | |
| result: dict[str, dict] = {} | |
| for model_name in REGISTERED_MODELS: | |
| entry: dict[str, Any] = { | |
| "status": self._status.get(model_name, "not_loaded"), | |
| } | |
| metadata = self._metadata.get(model_name) | |
| if metadata: | |
| entry["metadata"] = metadata | |
| result[model_name] = entry | |
| return result | |
| def is_loaded(self, model_name: str) -> bool: | |
| """Check if a model is loaded and ready for inference.""" | |
| return self._status.get(model_name) == "loaded" | |