Spaces:

yashb18
/

deepreads

Running

yash-b18

fix: resolve SIGSEGV and OpenMP deadlock for LightGBM + PyTorch coexistence

09fcfe1 7 months ago

8.52 kB

"""
Model loading utilities for the inference server.
Provides a single ModelRegistry class that loads all trained artifacts
once at startup and exposes them to the API layer.
Usage:
    from scripts.serve_utils import ModelRegistry
    registry = ModelRegistry.load(cfg)
"""

from __future__ import annotations

import pickle
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

from src.config import Config, MODELS_DIR, PROCESSED_DIR
from src.explainability.language_explainer import LanguageExplainer
from src.models.classical import LightGBMReranker
from src.models.naive import GlobalPopularityRecommender
from src.utils.logging import get_logger

logger = get_logger(__name__)

ITEM_META_COLS = [
    "item_avg_rating", "item_num_ratings", "item_popularity",
    "item_positive_rate", "price_norm", "category_idx", "brand_idx",
]


@dataclass
class ModelRegistry:
    """
    Holds all loaded model artifacts for the inference API.
    All attributes are populated by ModelRegistry.load(cfg).
    """
    naive_model: Optional[GlobalPopularityRecommender] = None
    lgbm_model: Optional[LightGBMReranker] = None
    deep_trainer: Optional[object] = None  # TwoTowerTrainer (lazy-imported)
    item_embeddings: Optional[np.ndarray] = None
    item_features: Optional[pd.DataFrame] = None
    user_features: Optional[pd.DataFrame] = None
    train_df: Optional[pd.DataFrame] = None
    feature_artifacts: dict = field(default_factory=dict)
    item_index: dict = field(default_factory=dict)
    language_explainer: Optional[LanguageExplainer] = None
    user_seen: dict[int, set[int]] = field(default_factory=dict)
    n_users: int = 0
    n_items: int = 0

    # Which models are available
    naive_available: bool = False
    classical_available: bool = False
    deep_available: bool = False

    @classmethod
    def load(cls, cfg: Config) -> "ModelRegistry":
        """
        Load all available model artifacts from disk.
        Missing artifacts are skipped with a warning (so partial deployments work).
        Args:
            cfg: Full Config object.
        Returns:
            Populated ModelRegistry.
        """
        registry = cls()

        # --- Shared data ---
        try:
            registry.item_features = pd.read_parquet(cfg.data.item_features_path)
            registry.n_items = int(registry.item_features["item_idx"].max() + 1)
        except Exception as e:
            logger.warning("Could not load item_features: %s", e)

        try:
            registry.user_features = pd.read_parquet(cfg.data.user_features_path)
        except Exception as e:
            logger.warning("Could not load user_features: %s", e)

        try:
            registry.train_df = pd.read_parquet(cfg.data.train_path)
            registry.n_users = int(registry.train_df["user_idx"].max() + 1)
            registry.user_seen = (
                registry.train_df.groupby("user_idx")["item_idx"].apply(set).to_dict()
            )
        except Exception as e:
            logger.warning("Could not load train_df: %s", e)

        try:
            fa_path = PROCESSED_DIR / "feature_artifacts.pkl"
            with open(fa_path, "rb") as f:
                registry.feature_artifacts = pickle.load(f)
        except Exception as e:
            logger.warning("Could not load feature_artifacts: %s", e)

        try:
            with open(cfg.deep.item_index_path, "rb") as f:
                registry.item_index = pickle.load(f)
        except Exception as e:
            logger.warning("Could not load item_index: %s", e)

        # --- Naive model ---
        naive_path = MODELS_DIR / "naive_baseline.pkl"
        if naive_path.exists():
            try:
                registry.naive_model = GlobalPopularityRecommender.load(naive_path)
                registry.naive_available = True
                logger.info("Naive model loaded.")
            except Exception as e:
                logger.warning("Failed to load naive model: %s", e)

        # --- Classical model ---
        lgbm_path = cfg.classical.model_path
        if lgbm_path.exists():
            try:
                registry.lgbm_model = LightGBMReranker.load(lgbm_path)
                registry.classical_available = True
                logger.info("LightGBM model loaded.")
            except Exception as e:
                logger.warning("Failed to load LightGBM model: %s", e)

        # --- Deep model ---
        if cfg.deep.model_path.exists() and cfg.deep.item_embeddings_path.exists():
            try:
                # Lazy import: torch must be imported AFTER LightGBM pickle load on Apple Silicon
                # to avoid a SIGSEGV caused by OpenMP/allocator conflicts.
                # Set single-threaded mode IMMEDIATELY after import to prevent OpenMP deadlock
                # when LightGBM's thread pool and PyTorch's OpenMP pool co-exist.
                logger.info("Deep model: importing torch…")
                import torch  # noqa: PLC0415
                torch.set_num_threads(1)
                torch.set_num_interop_threads(1)
                logger.info("Deep model: torch imported (single-thread mode), importing deep module…")
                from src.models.deep import TwoTowerModel, TwoTowerTrainer, build_item_feature_matrix  # noqa: PLC0415
                logger.info("Deep model: module imported, detecting device…")

                device = "cuda" if torch.cuda.is_available() else "cpu"
                logger.info("Deep model: device=%s", device)
                n_meta = 0
                if registry.item_features is not None:
                    logger.info("Deep model: building item feature matrix…")
                    item_feat_matrix = build_item_feature_matrix(
                        registry.item_features, ITEM_META_COLS
                    )
                    n_meta = item_feat_matrix.shape[1]
                    logger.info("Deep model: n_meta=%d", n_meta)
                else:
                    item_feat_matrix = None

                logger.info("Deep model: loading state dict from %s…", cfg.deep.model_path)
                state_dict = torch.load(cfg.deep.model_path, map_location=device, weights_only=True)
                logger.info("Deep model: state dict loaded (%d keys), creating model…", len(state_dict))
                model = TwoTowerModel(
                    n_users=registry.n_users,
                    n_items=registry.n_items,
                    embedding_dim=cfg.deep.embedding_dim,
                    output_dim=cfg.deep.embedding_dim,
                    hidden_dims=cfg.deep.hidden_dims,
                    dropout=0.0,          # no dropout at inference
                    n_meta_features=n_meta,
                )
                logger.info("Deep model: applying state dict…")
                model.load_state_dict(state_dict)
                del state_dict
                logger.info("Deep model: parameters loaded.")
                logger.info("Deep model: state dict loaded, creating trainer…")
                registry.deep_trainer = TwoTowerTrainer(
                    model=model, device=device, model_path=cfg.deep.model_path
                )
                logger.info("Deep model: loading embeddings…")
                registry.item_embeddings = np.load(cfg.deep.item_embeddings_path)
                registry.deep_available = True
                logger.info("Two-Tower model loaded. Item embeddings shape: %s",
                            registry.item_embeddings.shape)
            except Exception as e:
                logger.warning("Failed to load deep model: %s", e)

        # --- Language explainer (for deep model) ---
        if registry.train_df is not None and registry.item_features is not None:
            try:
                registry.language_explainer = LanguageExplainer(
                    registry.train_df,
                    registry.item_features,
                    positive_threshold=cfg.data.positive_rating_threshold,
                )
                logger.info("Language explainer initialised.")
            except Exception as e:
                logger.warning("Failed to initialise language explainer: %s", e)

        logger.info(
            "ModelRegistry ready — naive: %s | classical: %s | deep: %s",
            registry.naive_available,
            registry.classical_available,
            registry.deep_available,
        )
        return registry

	"""
	Model loading utilities for the inference server.

	Provides a single ModelRegistry class that loads all trained artifacts
	once at startup and exposes them to the API layer.

	Usage:
	from scripts.serve_utils import ModelRegistry
	registry = ModelRegistry.load(cfg)
	"""

	from __future__ import annotations

	import pickle
	from dataclasses import dataclass, field
	from pathlib import Path
	from typing import Optional

	import numpy as np
	import pandas as pd

	from src.config import Config, MODELS_DIR, PROCESSED_DIR
	from src.explainability.language_explainer import LanguageExplainer
	from src.models.classical import LightGBMReranker
	from src.models.naive import GlobalPopularityRecommender
	from src.utils.logging import get_logger

	logger = get_logger(__name__)

	ITEM_META_COLS = [
	"item_avg_rating", "item_num_ratings", "item_popularity",
	"item_positive_rate", "price_norm", "category_idx", "brand_idx",
	]


	@dataclass
	class ModelRegistry:
	"""
	Holds all loaded model artifacts for the inference API.

	All attributes are populated by ModelRegistry.load(cfg).
	"""
	naive_model: Optional[GlobalPopularityRecommender] = None
	lgbm_model: Optional[LightGBMReranker] = None
	deep_trainer: Optional[object] = None # TwoTowerTrainer (lazy-imported)
	item_embeddings: Optional[np.ndarray] = None
	item_features: Optional[pd.DataFrame] = None
	user_features: Optional[pd.DataFrame] = None
	train_df: Optional[pd.DataFrame] = None
	feature_artifacts: dict = field(default_factory=dict)
	item_index: dict = field(default_factory=dict)
	language_explainer: Optional[LanguageExplainer] = None
	user_seen: dict[int, set[int]] = field(default_factory=dict)
	n_users: int = 0
	n_items: int = 0

	# Which models are available
	naive_available: bool = False
	classical_available: bool = False
	deep_available: bool = False

	@classmethod
	def load(cls, cfg: Config) -> "ModelRegistry":
	"""
	Load all available model artifacts from disk.

	Missing artifacts are skipped with a warning (so partial deployments work).

	Args:
	cfg: Full Config object.

	Returns:
	Populated ModelRegistry.
	"""
	registry = cls()

	# --- Shared data ---
	try:
	registry.item_features = pd.read_parquet(cfg.data.item_features_path)
	registry.n_items = int(registry.item_features["item_idx"].max() + 1)
	except Exception as e:
	logger.warning("Could not load item_features: %s", e)

	try:
	registry.user_features = pd.read_parquet(cfg.data.user_features_path)
	except Exception as e:
	logger.warning("Could not load user_features: %s", e)

	try:
	registry.train_df = pd.read_parquet(cfg.data.train_path)
	registry.n_users = int(registry.train_df["user_idx"].max() + 1)
	registry.user_seen = (
	registry.train_df.groupby("user_idx")["item_idx"].apply(set).to_dict()
	)
	except Exception as e:
	logger.warning("Could not load train_df: %s", e)

	try:
	fa_path = PROCESSED_DIR / "feature_artifacts.pkl"
	with open(fa_path, "rb") as f:
	registry.feature_artifacts = pickle.load(f)
	except Exception as e:
	logger.warning("Could not load feature_artifacts: %s", e)

	try:
	with open(cfg.deep.item_index_path, "rb") as f:
	registry.item_index = pickle.load(f)
	except Exception as e:
	logger.warning("Could not load item_index: %s", e)

	# --- Naive model ---
	naive_path = MODELS_DIR / "naive_baseline.pkl"
	if naive_path.exists():
	try:
	registry.naive_model = GlobalPopularityRecommender.load(naive_path)
	registry.naive_available = True
	logger.info("Naive model loaded.")
	except Exception as e:
	logger.warning("Failed to load naive model: %s", e)

	# --- Classical model ---
	lgbm_path = cfg.classical.model_path
	if lgbm_path.exists():
	try:
	registry.lgbm_model = LightGBMReranker.load(lgbm_path)
	registry.classical_available = True
	logger.info("LightGBM model loaded.")
	except Exception as e:
	logger.warning("Failed to load LightGBM model: %s", e)

	# --- Deep model ---
	if cfg.deep.model_path.exists() and cfg.deep.item_embeddings_path.exists():
	try:
	# Lazy import: torch must be imported AFTER LightGBM pickle load on Apple Silicon
	# to avoid a SIGSEGV caused by OpenMP/allocator conflicts.
	# Set single-threaded mode IMMEDIATELY after import to prevent OpenMP deadlock
	# when LightGBM's thread pool and PyTorch's OpenMP pool co-exist.
	logger.info("Deep model: importing torch…")
	import torch # noqa: PLC0415
	torch.set_num_threads(1)
	torch.set_num_interop_threads(1)
	logger.info("Deep model: torch imported (single-thread mode), importing deep module…")
	from src.models.deep import TwoTowerModel, TwoTowerTrainer, build_item_feature_matrix # noqa: PLC0415
	logger.info("Deep model: module imported, detecting device…")

	device = "cuda" if torch.cuda.is_available() else "cpu"
	logger.info("Deep model: device=%s", device)
	n_meta = 0
	if registry.item_features is not None:
	logger.info("Deep model: building item feature matrix…")
	item_feat_matrix = build_item_feature_matrix(
	registry.item_features, ITEM_META_COLS
	)
	n_meta = item_feat_matrix.shape[1]
	logger.info("Deep model: n_meta=%d", n_meta)
	else:
	item_feat_matrix = None

	logger.info("Deep model: loading state dict from %s…", cfg.deep.model_path)
	state_dict = torch.load(cfg.deep.model_path, map_location=device, weights_only=True)
	logger.info("Deep model: state dict loaded (%d keys), creating model…", len(state_dict))
	model = TwoTowerModel(
	n_users=registry.n_users,
	n_items=registry.n_items,
	embedding_dim=cfg.deep.embedding_dim,
	output_dim=cfg.deep.embedding_dim,
	hidden_dims=cfg.deep.hidden_dims,
	dropout=0.0, # no dropout at inference
	n_meta_features=n_meta,
	)
	logger.info("Deep model: applying state dict…")
	model.load_state_dict(state_dict)
	del state_dict
	logger.info("Deep model: parameters loaded.")
	logger.info("Deep model: state dict loaded, creating trainer…")
	registry.deep_trainer = TwoTowerTrainer(
	model=model, device=device, model_path=cfg.deep.model_path
	)
	logger.info("Deep model: loading embeddings…")
	registry.item_embeddings = np.load(cfg.deep.item_embeddings_path)
	registry.deep_available = True
	logger.info("Two-Tower model loaded. Item embeddings shape: %s",
	registry.item_embeddings.shape)
	except Exception as e:
	logger.warning("Failed to load deep model: %s", e)

	# --- Language explainer (for deep model) ---
	if registry.train_df is not None and registry.item_features is not None:
	try:
	registry.language_explainer = LanguageExplainer(
	registry.train_df,
	registry.item_features,
	positive_threshold=cfg.data.positive_rating_threshold,
	)
	logger.info("Language explainer initialised.")
	except Exception as e:
	logger.warning("Failed to initialise language explainer: %s", e)

	logger.info(
	"ModelRegistry ready — naive: %s \| classical: %s \| deep: %s",
	registry.naive_available,
	registry.classical_available,
	registry.deep_available,
	)
	return registry