ins-pricing 0.5.0__py3-none-any.whl → 0.5.3__py3-none-any.whl

ins_pricing/modelling/bayesopt/trainers/trainer_glm.py

@@ -9,9 +9,15 @@ import statsmodels.api as sm
 from sklearn.metrics import log_loss
 
 from ins_pricing.modelling.bayesopt.trainers.trainer_base import TrainerBase
-from ins_pricing.utils import EPS
+from ins_pricing.utils import EPS, get_logger, log_print
 from ins_pricing.utils.losses import regression_loss
 
+_logger = get_logger("ins_pricing.trainer.glm")
+
+
+def _log(*args, **kwargs) -> None:
+    log_print(_logger, *args, **kwargs)
+
 class GLMTrainer(TrainerBase):
     def __init__(self, context: "BayesOptModel") -> None:
         super().__init__(context, 'GLM', 'GLM')
@@ -160,7 +166,7 @@ class GLMTrainer(TrainerBase):
 
         split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
         if split_iter is None:
-            print(
+            _log(
                 f"[GLM Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
                 flush=True,
             )
@@ -187,7 +193,7 @@ class GLMTrainer(TrainerBase):
             split_count += 1
 
         if split_count < 1:
-            print(
+            _log(
                 f"[GLM Ensemble] no CV splits generated; skip ensemble.",
                 flush=True,
             )

ins_pricing/modelling/bayesopt/trainers/trainer_gnn.py

@@ -10,18 +10,30 @@ from sklearn.metrics import log_loss
 
 from ins_pricing.modelling.bayesopt.trainers.trainer_base import TrainerBase
 from ins_pricing.modelling.bayesopt.models import GraphNeuralNetSklearn
-from ins_pricing.utils import EPS
+from ins_pricing.utils import EPS, get_logger, log_print
 from ins_pricing.utils.losses import regression_loss
-from ins_pricing.utils import get_logger
 from ins_pricing.utils.torch_compat import torch_load
 
 _logger = get_logger("ins_pricing.trainer.gnn")
 
-class GNNTrainer(TrainerBase):
-    def __init__(self, context: "BayesOptModel") -> None:
-        super().__init__(context, 'GNN', 'GNN')
-        self.model: Optional[GraphNeuralNetSklearn] = None
-        self.enable_distributed_optuna = bool(context.config.use_gnn_ddp)
+
+def _log(*args, **kwargs) -> None:
+    log_print(_logger, *args, **kwargs)
+
+class GNNTrainer(TrainerBase):
+    def __init__(self, context: "BayesOptModel") -> None:
+        super().__init__(context, 'GNN', 'GNN')
+        self.model: Optional[GraphNeuralNetSklearn] = None
+        self.enable_distributed_optuna = bool(context.config.use_gnn_ddp)
+
+    def _maybe_cleanup_gpu(self, model: Optional[GraphNeuralNetSklearn]) -> None:
+        if not bool(getattr(self.ctx.config, "gnn_cleanup_per_fold", False)):
+            return
+        if model is not None:
+            getattr(getattr(model, "gnn", None), "to",
+                    lambda *_args, **_kwargs: None)("cpu")
+        synchronize = bool(getattr(self.ctx.config, "gnn_cleanup_synchronize", False))
+        self._clean_gpu(synchronize=synchronize)
 
     def _build_model(self, params: Optional[Dict[str, Any]] = None) -> GraphNeuralNetSklearn:
         params = params or {}
@@ -167,19 +179,17 @@ class GNNTrainer(TrainerBase):
 
             if use_refit:
                 tmp_model = self._build_model(self.best_params)
-                tmp_model.fit(
-                    X_train,
-                    y_train,
-                    w_train=w_train,
-                    X_val=X_val,
-                    y_val=y_val,
-                    w_val=w_val,
-                    trial=None,
-                )
-                refit_epochs = int(getattr(tmp_model, "best_epoch", None) or self.ctx.epochs)
-                getattr(getattr(tmp_model, "gnn", None), "to",
-                        lambda *_args, **_kwargs: None)("cpu")
-                self._clean_gpu()
+                tmp_model.fit(
+                    X_train,
+                    y_train,
+                    w_train=w_train,
+                    X_val=X_val,
+                    y_val=y_val,
+                    w_val=w_val,
+                    trial=None,
+                )
+                refit_epochs = int(getattr(tmp_model, "best_epoch", None) or self.ctx.epochs)
+                self._maybe_cleanup_gpu(tmp_model)
             else:
                 self.model = self._build_model(self.best_params)
                 self.model.fit(
@@ -242,7 +252,7 @@ class GNNTrainer(TrainerBase):
         n_samples = len(X_all)
         split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
         if split_iter is None:
-            print(
+            _log(
                 f"[GNN Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
                 flush=True,
             )
@@ -264,15 +274,13 @@ class GNNTrainer(TrainerBase):
             )
             pred_train = model.predict(X_all)
             pred_test = model.predict(X_test)
-            preds_train_sum += np.asarray(pred_train, dtype=np.float64)
-            preds_test_sum += np.asarray(pred_test, dtype=np.float64)
-            getattr(getattr(model, "gnn", None), "to",
-                    lambda *_args, **_kwargs: None)("cpu")
-            self._clean_gpu()
-            split_count += 1
+            preds_train_sum += np.asarray(pred_train, dtype=np.float64)
+            preds_test_sum += np.asarray(pred_test, dtype=np.float64)
+            self._maybe_cleanup_gpu(model)
+            split_count += 1
 
         if split_count < 1:
-            print(
+            _log(
                 f"[GNN Ensemble] no CV splits generated; skip ensemble.",
                 flush=True,
             )
@@ -297,11 +305,11 @@ class GNNTrainer(TrainerBase):
         self.ctx.test_geo_tokens = test_tokens
         self.ctx.geo_token_cols = cols
         self.ctx.geo_gnn_model = geo_gnn
-        print(f"[GeoToken][GNNTrainer] Generated {len(cols)} dims and injected into FT.", flush=True)
+        _log(f"[GeoToken][GNNTrainer] Generated {len(cols)} dims and injected into FT.", flush=True)
 
     def save(self) -> None:
         if self.model is None:
-            print(f"[save] Warning: No model to save for {self.label}")
+            _log(f"[save] Warning: No model to save for {self.label}")
             return
         path = self.output.model_path(self._get_model_filename())
         base_gnn = getattr(self.model, "_unwrap_gnn", lambda: None)()
@@ -318,7 +326,7 @@ class GNNTrainer(TrainerBase):
     def load(self) -> None:
         path = self.output.model_path(self._get_model_filename())
         if not os.path.exists(path):
-            print(f"[load] Warning: Model file not found: {path}")
+            _log(f"[load] Warning: Model file not found: {path}")
             return
         payload = torch_load(path, map_location='cpu', weights_only=False)
         if not isinstance(payload, dict):
@@ -335,7 +343,7 @@ class GNNTrainer(TrainerBase):
                 base_gnn.load_state_dict(state_dict, strict=True)
             except RuntimeError as e:
                 if "Missing key" in str(e) or "Unexpected key" in str(e):
-                    print(f"[GNN load] Warning: State dict mismatch, loading with strict=False: {e}")
+                    _log(f"[GNN load] Warning: State dict mismatch, loading with strict=False: {e}")
                     base_gnn.load_state_dict(state_dict, strict=False)
                 else:
                     raise

ins_pricing/modelling/bayesopt/trainers/trainer_resn.py

@@ -11,15 +11,31 @@ from sklearn.metrics import log_loss
 from ins_pricing.modelling.bayesopt.trainers.trainer_base import TrainerBase
 from ins_pricing.modelling.bayesopt.models import ResNetSklearn
 from ins_pricing.utils.losses import regression_loss
+from ins_pricing.utils import get_logger, log_print
 
-class ResNetTrainer(TrainerBase):
-    def __init__(self, context: "BayesOptModel") -> None:
-        if context.task_type == 'classification':
-            super().__init__(context, 'ResNetClassifier', 'ResNet')
-        else:
-            super().__init__(context, 'ResNet', 'ResNet')
-        self.model: Optional[ResNetSklearn] = None
-        self.enable_distributed_optuna = bool(context.config.use_resn_ddp)
+_logger = get_logger("ins_pricing.trainer.resn")
+
+
+def _log(*args, **kwargs) -> None:
+    log_print(_logger, *args, **kwargs)
+
+class ResNetTrainer(TrainerBase):
+    def __init__(self, context: "BayesOptModel") -> None:
+        if context.task_type == 'classification':
+            super().__init__(context, 'ResNetClassifier', 'ResNet')
+        else:
+            super().__init__(context, 'ResNet', 'ResNet')
+        self.model: Optional[ResNetSklearn] = None
+        self.enable_distributed_optuna = bool(context.config.use_resn_ddp)
+
+    def _maybe_cleanup_gpu(self, model: Optional[ResNetSklearn]) -> None:
+        if not bool(getattr(self.ctx.config, "resn_cleanup_per_fold", False)):
+            return
+        if model is not None:
+            getattr(getattr(model, "resnet", None), "to",
+                    lambda *_args, **_kwargs: None)("cpu")
+        synchronize = bool(getattr(self.ctx.config, "resn_cleanup_synchronize", False))
+        self._clean_gpu(synchronize=synchronize)
 
     def _resolve_input_dim(self) -> int:
         data = getattr(self.ctx, "train_oht_scl_data", None)
@@ -174,13 +190,11 @@ class ResNetTrainer(TrainerBase):
                 w_all.iloc[val_idx],
                 trial=None,
             )
-            refit_epochs = self._resolve_best_epoch(
-                getattr(tmp_model, "training_history", None),
-                default_epochs=int(self.ctx.epochs),
-            )
-            getattr(getattr(tmp_model, "resnet", None), "to",
-                    lambda *_args, **_kwargs: None)("cpu")
-            self._clean_gpu()
+            refit_epochs = self._resolve_best_epoch(
+                getattr(tmp_model, "training_history", None),
+                default_epochs=int(self.ctx.epochs),
+            )
+            self._maybe_cleanup_gpu(tmp_model)
 
         self.model = self._build_model(params)
         if refit_epochs is not None:
@@ -219,7 +233,7 @@ class ResNetTrainer(TrainerBase):
         n_samples = len(X_all)
         split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
         if split_iter is None:
-            print(
+            _log(
                 f"[ResNet Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
                 flush=True,
             )
@@ -241,15 +255,13 @@ class ResNetTrainer(TrainerBase):
             )
             pred_train = model.predict(X_all)
             pred_test = model.predict(X_test)
-            preds_train_sum += np.asarray(pred_train, dtype=np.float64)
-            preds_test_sum += np.asarray(pred_test, dtype=np.float64)
-            getattr(getattr(model, "resnet", None), "to",
-                    lambda *_args, **_kwargs: None)("cpu")
-            self._clean_gpu()
-            split_count += 1
+            preds_train_sum += np.asarray(pred_train, dtype=np.float64)
+            preds_test_sum += np.asarray(pred_test, dtype=np.float64)
+            self._maybe_cleanup_gpu(model)
+            split_count += 1
 
         if split_count < 1:
-            print(
+            _log(
                 f"[ResNet Ensemble] no CV splits generated; skip ensemble.",
                 flush=True,
             )
@@ -280,4 +292,4 @@ class ResNetTrainer(TrainerBase):
             self.model = resn_loaded
             self.ctx.resn_best = self.model
         else:
-            print(f"[ResNetTrainer.load] Model file not found: {path}")
+            _log(f"[ResNetTrainer.load] Model file not found: {path}")

ins_pricing/modelling/bayesopt/trainers/trainer_xgb.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-import inspect
-from typing import Any, Dict, List, Optional, Tuple
+import inspect
+import os
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import optuna
@@ -10,14 +11,190 @@ import xgboost as xgb
 from sklearn.metrics import log_loss
 
 from ins_pricing.modelling.bayesopt.trainers.trainer_base import TrainerBase
-from ins_pricing.utils import EPS
+from ins_pricing.utils import EPS, get_logger, log_print
 from ins_pricing.utils.losses import regression_loss
 
-_XGB_CUDA_CHECKED = False
-_XGB_HAS_CUDA = False
+_logger = get_logger("ins_pricing.trainer.xgb")
 
-_XGB_CUDA_CHECKED = False
-_XGB_HAS_CUDA = False
+
+def _log(*args, **kwargs) -> None:
+    log_print(_logger, *args, **kwargs)
+
+_XGB_CUDA_CHECKED = False
+_XGB_HAS_CUDA = False
+
+
+def _is_oom_error(exc: Exception) -> bool:
+    msg = str(exc).lower()
+    return "out of memory" in msg or ("cuda" in msg and "memory" in msg)
+
+
+class _XGBDMatrixWrapper:
+    """Sklearn-like wrapper that uses xgb.train + (Quantile)DMatrix internally."""
+
+    def __init__(
+        self,
+        params: Dict[str, Any],
+        *,
+        task_type: str,
+        use_gpu: bool,
+        allow_cpu_fallback: bool = True,
+    ) -> None:
+        self.params = dict(params)
+        self.task_type = task_type
+        self.use_gpu = bool(use_gpu)
+        self.allow_cpu_fallback = allow_cpu_fallback
+        self._booster: Optional[xgb.Booster] = None
+        self.best_iteration: Optional[int] = None
+
+    def set_params(self, **params: Any) -> "_XGBDMatrixWrapper":
+        self.params.update(params)
+        return self
+
+    def get_params(self, deep: bool = True) -> Dict[str, Any]:
+        _ = deep
+        return dict(self.params)
+
+    def _select_dmatrix_class(self) -> Any:
+        if self.use_gpu and hasattr(xgb, "DeviceQuantileDMatrix"):
+            return xgb.DeviceQuantileDMatrix
+        if hasattr(xgb, "QuantileDMatrix"):
+            return xgb.QuantileDMatrix
+        return xgb.DMatrix
+
+    def _build_dmatrix(self, X, y=None, weight=None) -> xgb.DMatrix:
+        if isinstance(X, (str, os.PathLike)):
+            raise ValueError(
+                "External-memory DMatrix is disabled; pass in-memory data instead."
+            )
+        if isinstance(X, xgb.DMatrix):
+            raise ValueError(
+                "DMatrix inputs are disabled; pass raw in-memory data instead."
+            )
+        dmatrix_cls = self._select_dmatrix_class()
+        kwargs: Dict[str, Any] = {}
+        if y is not None:
+            kwargs["label"] = y
+        if weight is not None:
+            kwargs["weight"] = weight
+        if bool(self.params.get("enable_categorical", False)):
+            kwargs["enable_categorical"] = True
+        try:
+            return dmatrix_cls(X, **kwargs)
+        except TypeError:
+            kwargs.pop("enable_categorical", None)
+            return dmatrix_cls(X, **kwargs)
+        except Exception:
+            if dmatrix_cls is not xgb.DMatrix:
+                return xgb.DMatrix(X, **kwargs)
+            raise
+
+    def _resolve_train_params(self) -> Dict[str, Any]:
+        params = dict(self.params)
+        if not self.use_gpu:
+            params["tree_method"] = "hist"
+            params["predictor"] = "cpu_predictor"
+            params.pop("gpu_id", None)
+        return params
+
+    def _train_booster(
+        self,
+        X,
+        y,
+        *,
+        sample_weight=None,
+        eval_set=None,
+        sample_weight_eval_set=None,
+        early_stopping_rounds: Optional[int] = None,
+        verbose: bool = False,
+    ) -> None:
+        params = self._resolve_train_params()
+        num_boost_round = int(params.pop("n_estimators", 100))
+        dtrain = self._build_dmatrix(X, y, sample_weight)
+        evals = []
+        if eval_set:
+            weights = sample_weight_eval_set or []
+            for idx, (X_val, y_val) in enumerate(eval_set):
+                w_val = weights[idx] if idx < len(weights) else None
+                dval = self._build_dmatrix(X_val, y_val, w_val)
+                evals.append((dval, f"val{idx}"))
+        self._booster = xgb.train(
+            params,
+            dtrain,
+            num_boost_round=num_boost_round,
+            evals=evals,
+            early_stopping_rounds=early_stopping_rounds,
+            verbose_eval=verbose,
+        )
+        self.best_iteration = getattr(self._booster, "best_iteration", None)
+
+    def fit(self, X, y, **fit_kwargs) -> "_XGBDMatrixWrapper":
+        sample_weight = fit_kwargs.pop("sample_weight", None)
+        eval_set = fit_kwargs.pop("eval_set", None)
+        sample_weight_eval_set = fit_kwargs.pop("sample_weight_eval_set", None)
+        early_stopping_rounds = fit_kwargs.pop("early_stopping_rounds", None)
+        verbose = bool(fit_kwargs.pop("verbose", False))
+        fit_kwargs.pop("eval_metric", None)
+        try:
+            self._train_booster(
+                X,
+                y,
+                sample_weight=sample_weight,
+                eval_set=eval_set,
+                sample_weight_eval_set=sample_weight_eval_set,
+                early_stopping_rounds=early_stopping_rounds,
+                verbose=verbose,
+            )
+        except Exception as exc:
+            if self.use_gpu and self.allow_cpu_fallback and _is_oom_error(exc):
+                _log("[XGBoost] GPU OOM detected; retrying with CPU.", flush=True)
+                self.use_gpu = False
+                self._train_booster(
+                    X,
+                    y,
+                    sample_weight=sample_weight,
+                    eval_set=eval_set,
+                    sample_weight_eval_set=sample_weight_eval_set,
+                    early_stopping_rounds=early_stopping_rounds,
+                    verbose=verbose,
+                )
+            else:
+                raise
+        return self
+
+    def _resolve_iteration_range(self) -> Optional[Tuple[int, int]]:
+        if self.best_iteration is None:
+            return None
+        return (0, int(self.best_iteration) + 1)
+
+    def _predict_raw(self, X) -> np.ndarray:
+        if self._booster is None:
+            raise RuntimeError("Booster not trained.")
+        dtest = self._build_dmatrix(X)
+        iteration_range = self._resolve_iteration_range()
+        if iteration_range is None:
+            return self._booster.predict(dtest)
+        try:
+            return self._booster.predict(dtest, iteration_range=iteration_range)
+        except TypeError:
+            return self._booster.predict(dtest, ntree_limit=iteration_range[1])
+
+    def predict(self, X, **_kwargs) -> np.ndarray:
+        pred = self._predict_raw(X)
+        if self.task_type == "classification":
+            if pred.ndim == 1:
+                return (pred > 0.5).astype(int)
+            return np.argmax(pred, axis=1)
+        return pred
+
+    def predict_proba(self, X, **_kwargs) -> np.ndarray:
+        pred = self._predict_raw(X)
+        if pred.ndim == 1:
+            return np.column_stack([1 - pred, pred])
+        return pred
+
+    def get_booster(self) -> Optional[xgb.Booster]:
+        return self._booster
 
 
 def _xgb_cuda_available() -> bool:
@@ -54,39 +231,65 @@ def _xgb_cuda_available() -> bool:
     _XGB_HAS_CUDA = False
     return False
 
-class XGBTrainer(TrainerBase):
+class XGBTrainer(TrainerBase):
     def __init__(self, context: "BayesOptModel") -> None:
         super().__init__(context, 'Xgboost', 'Xgboost')
         self.model: Optional[xgb.XGBModel] = None
         self._xgb_use_gpu = False
         self._xgb_gpu_warned = False
 
-    def _build_estimator(self) -> xgb.XGBModel:
-        use_gpu = bool(self.ctx.use_gpu and _xgb_cuda_available())
-        self._xgb_use_gpu = use_gpu
-        params = dict(
-            objective=self.ctx.obj,
-            random_state=self.ctx.rand_seed,
-            subsample=0.9,
-            tree_method='gpu_hist' if use_gpu else 'hist',
-            enable_categorical=True,
-            predictor='gpu_predictor' if use_gpu else 'cpu_predictor'
-        )
+    def _build_sklearn_estimator(self, params: Dict[str, Any]) -> xgb.XGBModel:
+        if self.ctx.task_type == 'classification':
+            return xgb.XGBClassifier(**params)
+        return xgb.XGBRegressor(**params)
+
+    def _build_estimator(self) -> xgb.XGBModel:
+        use_gpu = bool(self.ctx.use_gpu and _xgb_cuda_available())
+        self._xgb_use_gpu = use_gpu
+        params = dict(
+            objective=self.ctx.obj,
+            random_state=self.ctx.rand_seed,
+            subsample=0.9,
+            tree_method='gpu_hist' if use_gpu else 'hist',
+            enable_categorical=True,
+            predictor='gpu_predictor' if use_gpu else 'cpu_predictor'
+        )
         if self.ctx.use_gpu and not use_gpu and not self._xgb_gpu_warned:
-            print(
+            _log(
                 "[XGBoost] CUDA requested but not available; falling back to CPU.",
                 flush=True,
             )
             self._xgb_gpu_warned = True
-        if use_gpu:
-            params['gpu_id'] = 0
-            print(f">>> XGBoost using GPU ID: 0 (Single GPU Mode)")
-        eval_metric = self._resolve_eval_metric()
-        if eval_metric is not None:
-            params.setdefault("eval_metric", eval_metric)
-        if self.ctx.task_type == 'classification':
-            return xgb.XGBClassifier(**params)
-        return xgb.XGBRegressor(**params)
+        if use_gpu:
+            gpu_id = self._resolve_gpu_id()
+            params['gpu_id'] = gpu_id
+            _log(f">>> XGBoost using GPU ID: {gpu_id}")
+        eval_metric = self._resolve_eval_metric()
+        if eval_metric is not None:
+            params.setdefault("eval_metric", eval_metric)
+        use_dmatrix = bool(getattr(self.config, "xgb_use_dmatrix", True))
+        if use_dmatrix:
+            return _XGBDMatrixWrapper(
+                params,
+                task_type=self.ctx.task_type,
+                use_gpu=use_gpu,
+            )
+        return self._build_sklearn_estimator(params)
+
+    def _resolve_gpu_id(self) -> int:
+        gpu_id = getattr(self.config, "xgb_gpu_id", None)
+        if gpu_id is None:
+            return 0
+        try:
+            return int(gpu_id)
+        except (TypeError, ValueError):
+            return 0
+
+    def _maybe_cleanup_gpu(self) -> None:
+        if not bool(getattr(self.config, "xgb_cleanup_per_fold", False)):
+            return
+        synchronize = bool(getattr(self.config, "xgb_cleanup_synchronize", False))
+        self._clean_gpu(synchronize=synchronize)
 
     def _resolve_eval_metric(self) -> Optional[Any]:
         fit_params = self.ctx.fit_params or {}
@@ -148,7 +351,7 @@ class XGBTrainer(TrainerBase):
         n_samples = len(X_all)
         split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
         if split_iter is None:
-            print(
+            _log(
                 f"[XGB Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
                 flush=True,
             )
@@ -184,11 +387,11 @@ class XGBTrainer(TrainerBase):
                 pred_test = clf.predict(X_test)
             preds_train_sum += np.asarray(pred_train, dtype=np.float64)
             preds_test_sum += np.asarray(pred_test, dtype=np.float64)
-            self._clean_gpu()
-            split_count += 1
+            self._maybe_cleanup_gpu()
+            split_count += 1
 
         if split_count < 1:
-            print(
+            _log(
                 f"[XGB Ensemble] no CV splits generated; skip ensemble.",
                 flush=True,
             )
@@ -213,7 +416,7 @@ class XGBTrainer(TrainerBase):
         reg_alpha = trial.suggest_float('reg_alpha', 1e-10, 1, log=True)
         reg_lambda = trial.suggest_float('reg_lambda', 1e-10, 1, log=True)
         if trial is not None:
-            print(
+            _log(
                 f"[Optuna][Xgboost] trial_id={trial.number} max_depth={max_depth} "
                 f"n_estimators={n_estimators}",
                 flush=True,
@@ -280,9 +483,9 @@ class XGBTrainer(TrainerBase):
                     tweedie_power=tweedie_variance_power,
                 )
             losses.append(float(loss))
-            self._clean_gpu()
-
-        return float(np.mean(losses))
+            self._maybe_cleanup_gpu()
+
+        return float(np.mean(losses))
 
     def train(self) -> None:
         if not self.best_params: