ins-pricing 0.4.5__py3-none-any.whl → 0.5.1__py3-none-any.whl

ins_pricing/modelling/{core/bayesopt → bayesopt}/trainers/trainer_xgb.py

@@ -1,65 +1,247 @@
-from __future__ import annotations
-
+from __future__ import annotations
+
 import inspect
+import os
 from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import optuna
+import torch
+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, get_logger, log_print
+from ins_pricing.utils.losses import regression_loss
+
+_logger = get_logger("ins_pricing.trainer.xgb")
+
+
+def _log(*args, **kwargs) -> None:
+    log_print(_logger, *args, **kwargs)
+
+_XGB_CUDA_CHECKED = False
+_XGB_HAS_CUDA = False
 
-import numpy as np
-import optuna
-import torch
-import xgboost as xgb
-from sklearn.metrics import log_loss
 
-from .trainer_base import TrainerBase
-from ..utils import EPS
-from ..utils.losses import regression_loss
+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)
 
-_XGB_CUDA_CHECKED = False
-_XGB_HAS_CUDA = False
 
-_XGB_CUDA_CHECKED = False
-_XGB_HAS_CUDA = False
+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 _xgb_cuda_available() -> bool:
-    # Best-effort check for XGBoost CUDA build; cached to avoid repeated checks.
-    global _XGB_CUDA_CHECKED, _XGB_HAS_CUDA
-    if _XGB_CUDA_CHECKED:
-        return _XGB_HAS_CUDA
-    _XGB_CUDA_CHECKED = True
-    if not torch.cuda.is_available():
-        _XGB_HAS_CUDA = False
-        return False
-    try:
-        build_info = getattr(xgb, "build_info", None)
-        if callable(build_info):
-            info = build_info()
-            for key in ("USE_CUDA", "use_cuda", "cuda"):
-                if key in info:
-                    val = info[key]
-                    if isinstance(val, str):
-                        _XGB_HAS_CUDA = val.strip().upper() in (
-                            "ON", "YES", "TRUE", "1")
-                    else:
-                        _XGB_HAS_CUDA = bool(val)
-                    return _XGB_HAS_CUDA
-    except Exception:
-        pass
-    try:
-        has_cuda = getattr(getattr(xgb, "core", None), "_has_cuda_support", None)
-        if callable(has_cuda):
-            _XGB_HAS_CUDA = bool(has_cuda())
-            return _XGB_HAS_CUDA
-    except Exception:
-        pass
-    _XGB_HAS_CUDA = False
-    return False
+    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:
+    # Best-effort check for XGBoost CUDA build; cached to avoid repeated checks.
+    global _XGB_CUDA_CHECKED, _XGB_HAS_CUDA
+    if _XGB_CUDA_CHECKED:
+        return _XGB_HAS_CUDA
+    _XGB_CUDA_CHECKED = True
+    if not torch.cuda.is_available():
+        _XGB_HAS_CUDA = False
+        return False
+    try:
+        build_info = getattr(xgb, "build_info", None)
+        if callable(build_info):
+            info = build_info()
+            for key in ("USE_CUDA", "use_cuda", "cuda"):
+                if key in info:
+                    val = info[key]
+                    if isinstance(val, str):
+                        _XGB_HAS_CUDA = val.strip().upper() in (
+                            "ON", "YES", "TRUE", "1")
+                    else:
+                        _XGB_HAS_CUDA = bool(val)
+                    return _XGB_HAS_CUDA
+    except Exception:
+        pass
+    try:
+        has_cuda = getattr(getattr(xgb, "core", None), "_has_cuda_support", None)
+        if callable(has_cuda):
+            _XGB_HAS_CUDA = bool(has_cuda())
+            return _XGB_HAS_CUDA
+    except Exception:
+        pass
+    _XGB_HAS_CUDA = False
+    return False
+
 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 __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_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())
@@ -72,276 +254,296 @@ class XGBTrainer(TrainerBase):
             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(
-                "[XGBoost] CUDA requested but not available; falling back to CPU.",
-                flush=True,
-            )
-            self._xgb_gpu_warned = True
+        if self.ctx.use_gpu and not use_gpu and not self._xgb_gpu_warned:
+            _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)")
+            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)
-        if self.ctx.task_type == 'classification':
-            return xgb.XGBClassifier(**params)
-        return xgb.XGBRegressor(**params)
-
-    def _resolve_eval_metric(self) -> Optional[Any]:
-        fit_params = self.ctx.fit_params or {}
-        eval_metric = fit_params.get("eval_metric")
-        if eval_metric is None:
-            return "logloss" if self.ctx.task_type == 'classification' else "rmse"
-        return 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 _fit_supports_param(self, name: str) -> bool:
+    def _resolve_gpu_id(self) -> int:
+        gpu_id = getattr(self.config, "xgb_gpu_id", None)
+        if gpu_id is None:
+            return 0
         try:
-            fit = xgb.XGBClassifier.fit if self.ctx.task_type == 'classification' else xgb.XGBRegressor.fit
-            return name in inspect.signature(fit).parameters
+            return int(gpu_id)
         except (TypeError, ValueError):
-            return True
-
-    def _resolve_early_stopping_rounds(self, n_estimators: int) -> int:
-        n_estimators = max(1, int(n_estimators))
-        base = max(5, n_estimators // 10)
-        return min(50, base)
-
-    def _build_fit_kwargs(self,
-                          w_train,
-                          X_val=None,
-                          y_val=None,
-                          w_val=None,
-                          n_estimators: Optional[int] = None) -> Dict[str, Any]:
-        supports_early = self._fit_supports_param("early_stopping_rounds")
-        fit_kwargs = dict(self.ctx.fit_params or {})
-        fit_kwargs.pop("sample_weight", None)
-        fit_kwargs.pop("eval_metric", None)
-        fit_kwargs["sample_weight"] = w_train
-
-        if "eval_set" not in fit_kwargs and X_val is not None and y_val is not None:
-            fit_kwargs["eval_set"] = [(X_val, y_val)]
-            if w_val is not None:
-                fit_kwargs["sample_weight_eval_set"] = [w_val]
-
-        if (
-            supports_early
-            and "early_stopping_rounds" not in fit_kwargs
-            and "eval_set" in fit_kwargs
-        ):
-            rounds = self._resolve_early_stopping_rounds(n_estimators or 100)
-            fit_kwargs["early_stopping_rounds"] = rounds
-        if not supports_early:
-            fit_kwargs.pop("early_stopping_rounds", None)
-
-        fit_kwargs.setdefault("verbose", False)
-        return fit_kwargs
+            return 0
 
-    def ensemble_predict(self, k: int) -> None:
-        if not self.best_params:
-            raise RuntimeError("Run tune() first to obtain best XGB parameters.")
-        k = max(2, int(k))
-        X_all = self.ctx.train_data[self.ctx.factor_nmes]
-        y_all = self.ctx.train_data[self.ctx.resp_nme].values
-        w_all = self.ctx.train_data[self.ctx.weight_nme].values
-        X_test = self.ctx.test_data[self.ctx.factor_nmes]
-        n_samples = len(X_all)
-        split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
-        if split_iter is None:
-            print(
-                f"[XGB Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
-                flush=True,
-            )
+    def _maybe_cleanup_gpu(self) -> None:
+        if not bool(getattr(self.config, "xgb_cleanup_per_fold", False)):
             return
-        preds_train_sum = np.zeros(n_samples, dtype=np.float64)
-        preds_test_sum = np.zeros(len(X_test), dtype=np.float64)
-
-        split_count = 0
-        for train_idx, val_idx in split_iter:
-            X_train = X_all.iloc[train_idx]
-            y_train = y_all[train_idx]
-            w_train = w_all[train_idx]
-            X_val = X_all.iloc[val_idx]
-            y_val = y_all[val_idx]
-            w_val = w_all[val_idx]
-
-            clf = self._build_estimator()
-            clf.set_params(**self.best_params)
-            fit_kwargs = self._build_fit_kwargs(
-                w_train=w_train,
-                X_val=X_val,
-                y_val=y_val,
-                w_val=w_val,
-                n_estimators=self.best_params.get("n_estimators", 100),
-            )
-            clf.fit(X_train, y_train, **fit_kwargs)
-
-            if self.ctx.task_type == 'classification':
-                pred_train = clf.predict_proba(X_all)[:, 1]
-                pred_test = clf.predict_proba(X_test)[:, 1]
-            else:
-                pred_train = clf.predict(X_all)
-                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()
+        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 {}
+        eval_metric = fit_params.get("eval_metric")
+        if eval_metric is None:
+            return "logloss" if self.ctx.task_type == 'classification' else "rmse"
+        return eval_metric
+
+    def _fit_supports_param(self, name: str) -> bool:
+        try:
+            fit = xgb.XGBClassifier.fit if self.ctx.task_type == 'classification' else xgb.XGBRegressor.fit
+            return name in inspect.signature(fit).parameters
+        except (TypeError, ValueError):
+            return True
+
+    def _resolve_early_stopping_rounds(self, n_estimators: int) -> int:
+        n_estimators = max(1, int(n_estimators))
+        base = max(5, n_estimators // 10)
+        return min(50, base)
+
+    def _build_fit_kwargs(self,
+                          w_train,
+                          X_val=None,
+                          y_val=None,
+                          w_val=None,
+                          n_estimators: Optional[int] = None) -> Dict[str, Any]:
+        supports_early = self._fit_supports_param("early_stopping_rounds")
+        fit_kwargs = dict(self.ctx.fit_params or {})
+        fit_kwargs.pop("sample_weight", None)
+        fit_kwargs.pop("eval_metric", None)
+        fit_kwargs["sample_weight"] = w_train
+
+        if "eval_set" not in fit_kwargs and X_val is not None and y_val is not None:
+            fit_kwargs["eval_set"] = [(X_val, y_val)]
+            if w_val is not None:
+                fit_kwargs["sample_weight_eval_set"] = [w_val]
+
+        if (
+            supports_early
+            and "early_stopping_rounds" not in fit_kwargs
+            and "eval_set" in fit_kwargs
+        ):
+            rounds = self._resolve_early_stopping_rounds(n_estimators or 100)
+            fit_kwargs["early_stopping_rounds"] = rounds
+        if not supports_early:
+            fit_kwargs.pop("early_stopping_rounds", None)
+
+        fit_kwargs.setdefault("verbose", False)
+        return fit_kwargs
+
+    def ensemble_predict(self, k: int) -> None:
+        if not self.best_params:
+            raise RuntimeError("Run tune() first to obtain best XGB parameters.")
+        k = max(2, int(k))
+        X_all = self.ctx.train_data[self.ctx.factor_nmes]
+        y_all = self.ctx.train_data[self.ctx.resp_nme].values
+        w_all = self.ctx.train_data[self.ctx.weight_nme].values
+        X_test = self.ctx.test_data[self.ctx.factor_nmes]
+        n_samples = len(X_all)
+        split_iter, _ = self._resolve_ensemble_splits(X_all, k=k)
+        if split_iter is None:
+            _log(
+                f"[XGB Ensemble] unable to build CV split (n_samples={n_samples}); skip ensemble.",
+                flush=True,
+            )
+            return
+        preds_train_sum = np.zeros(n_samples, dtype=np.float64)
+        preds_test_sum = np.zeros(len(X_test), dtype=np.float64)
+
+        split_count = 0
+        for train_idx, val_idx in split_iter:
+            X_train = X_all.iloc[train_idx]
+            y_train = y_all[train_idx]
+            w_train = w_all[train_idx]
+            X_val = X_all.iloc[val_idx]
+            y_val = y_all[val_idx]
+            w_val = w_all[val_idx]
+
+            clf = self._build_estimator()
+            clf.set_params(**self.best_params)
+            fit_kwargs = self._build_fit_kwargs(
+                w_train=w_train,
+                X_val=X_val,
+                y_val=y_val,
+                w_val=w_val,
+                n_estimators=self.best_params.get("n_estimators", 100),
+            )
+            clf.fit(X_train, y_train, **fit_kwargs)
+
+            if self.ctx.task_type == 'classification':
+                pred_train = clf.predict_proba(X_all)[:, 1]
+                pred_test = clf.predict_proba(X_test)[:, 1]
+            else:
+                pred_train = clf.predict(X_all)
+                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._maybe_cleanup_gpu()
             split_count += 1
-
-        if split_count < 1:
-            print(
-                f"[XGB Ensemble] no CV splits generated; skip ensemble.",
-                flush=True,
-            )
-            return
-        preds_train = preds_train_sum / float(split_count)
-        preds_test = preds_test_sum / float(split_count)
-        self._cache_predictions("xgb", preds_train, preds_test)
-
-    def cross_val(self, trial: optuna.trial.Trial) -> float:
-        learning_rate = trial.suggest_float(
-            'learning_rate', 1e-5, 1e-1, log=True)
-        gamma = trial.suggest_float('gamma', 0, 10000)
-        max_depth_max = max(
-            3, int(getattr(self.config, "xgb_max_depth_max", 25)))
-        n_estimators_max = max(
-            10, int(getattr(self.config, "xgb_n_estimators_max", 500)))
-        max_depth = trial.suggest_int('max_depth', 3, max_depth_max)
-        n_estimators = trial.suggest_int(
-            'n_estimators', 10, n_estimators_max, step=10)
-        min_child_weight = trial.suggest_int(
-            'min_child_weight', 100, 10000, step=100)
-        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(
-                f"[Optuna][Xgboost] trial_id={trial.number} max_depth={max_depth} "
-                f"n_estimators={n_estimators}",
-                flush=True,
-            )
-        if max_depth >= 20 and n_estimators >= 300:
-            raise optuna.TrialPruned(
-                "XGB config is likely too slow (max_depth>=20 & n_estimators>=300)")
-        clf = self._build_estimator()
-        params = {
-            'learning_rate': learning_rate,
-            'gamma': gamma,
-            'max_depth': max_depth,
-            'n_estimators': n_estimators,
-            'min_child_weight': min_child_weight,
-            'reg_alpha': reg_alpha,
-            'reg_lambda': reg_lambda
-        }
-        loss_name = getattr(self.ctx, "loss_name", "tweedie")
-        tweedie_variance_power = None
-        if self.ctx.task_type != 'classification':
-            if loss_name == "tweedie":
-                tweedie_variance_power = trial.suggest_float(
-                    'tweedie_variance_power', 1, 2)
-                params['tweedie_variance_power'] = tweedie_variance_power
-            elif loss_name == "poisson":
-                tweedie_variance_power = 1.0
-            elif loss_name == "gamma":
-                tweedie_variance_power = 2.0
-        X_all = self.ctx.train_data[self.ctx.factor_nmes]
-        y_all = self.ctx.train_data[self.ctx.resp_nme].values
-        w_all = self.ctx.train_data[self.ctx.weight_nme].values
-
-        losses: List[float] = []
-        for train_idx, val_idx in self.ctx.cv.split(X_all):
-            X_train = X_all.iloc[train_idx]
-            y_train = y_all[train_idx]
-            w_train = w_all[train_idx]
-            X_val = X_all.iloc[val_idx]
-            y_val = y_all[val_idx]
-            w_val = w_all[val_idx]
-
-            clf = self._build_estimator()
-            clf.set_params(**params)
-            fit_kwargs = self._build_fit_kwargs(
-                w_train=w_train,
-                X_val=X_val,
-                y_val=y_val,
-                w_val=w_val,
-                n_estimators=n_estimators,
-            )
-            clf.fit(X_train, y_train, **fit_kwargs)
-
-            if self.ctx.task_type == 'classification':
-                y_pred = clf.predict_proba(X_val)[:, 1]
-                y_pred = np.clip(y_pred, EPS, 1 - EPS)
-                loss = log_loss(y_val, y_pred, sample_weight=w_val)
-            else:
-                y_pred = clf.predict(X_val)
-                loss = regression_loss(
-                    y_val,
-                    y_pred,
-                    w_val,
-                    loss_name=loss_name,
-                    tweedie_power=tweedie_variance_power,
-                )
-            losses.append(float(loss))
-            self._clean_gpu()
+
+        if split_count < 1:
+            _log(
+                f"[XGB Ensemble] no CV splits generated; skip ensemble.",
+                flush=True,
+            )
+            return
+        preds_train = preds_train_sum / float(split_count)
+        preds_test = preds_test_sum / float(split_count)
+        self._cache_predictions("xgb", preds_train, preds_test)
+
+    def cross_val(self, trial: optuna.trial.Trial) -> float:
+        learning_rate = trial.suggest_float(
+            'learning_rate', 1e-5, 1e-1, log=True)
+        gamma = trial.suggest_float('gamma', 0, 10000)
+        max_depth_max = max(
+            3, int(getattr(self.config, "xgb_max_depth_max", 25)))
+        n_estimators_max = max(
+            10, int(getattr(self.config, "xgb_n_estimators_max", 500)))
+        max_depth = trial.suggest_int('max_depth', 3, max_depth_max)
+        n_estimators = trial.suggest_int(
+            'n_estimators', 10, n_estimators_max, step=10)
+        min_child_weight = trial.suggest_int(
+            'min_child_weight', 100, 10000, step=100)
+        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:
+            _log(
+                f"[Optuna][Xgboost] trial_id={trial.number} max_depth={max_depth} "
+                f"n_estimators={n_estimators}",
+                flush=True,
+            )
+        if max_depth >= 20 and n_estimators >= 300:
+            raise optuna.TrialPruned(
+                "XGB config is likely too slow (max_depth>=20 & n_estimators>=300)")
+        clf = self._build_estimator()
+        params = {
+            'learning_rate': learning_rate,
+            'gamma': gamma,
+            'max_depth': max_depth,
+            'n_estimators': n_estimators,
+            'min_child_weight': min_child_weight,
+            'reg_alpha': reg_alpha,
+            'reg_lambda': reg_lambda
+        }
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
+        tweedie_variance_power = None
+        if self.ctx.task_type != 'classification':
+            if loss_name == "tweedie":
+                tweedie_variance_power = trial.suggest_float(
+                    'tweedie_variance_power', 1, 2)
+                params['tweedie_variance_power'] = tweedie_variance_power
+            elif loss_name == "poisson":
+                tweedie_variance_power = 1.0
+            elif loss_name == "gamma":
+                tweedie_variance_power = 2.0
+        X_all = self.ctx.train_data[self.ctx.factor_nmes]
+        y_all = self.ctx.train_data[self.ctx.resp_nme].values
+        w_all = self.ctx.train_data[self.ctx.weight_nme].values
+
+        losses: List[float] = []
+        for train_idx, val_idx in self.ctx.cv.split(X_all):
+            X_train = X_all.iloc[train_idx]
+            y_train = y_all[train_idx]
+            w_train = w_all[train_idx]
+            X_val = X_all.iloc[val_idx]
+            y_val = y_all[val_idx]
+            w_val = w_all[val_idx]
+
+            clf = self._build_estimator()
+            clf.set_params(**params)
+            fit_kwargs = self._build_fit_kwargs(
+                w_train=w_train,
+                X_val=X_val,
+                y_val=y_val,
+                w_val=w_val,
+                n_estimators=n_estimators,
+            )
+            clf.fit(X_train, y_train, **fit_kwargs)
+
+            if self.ctx.task_type == 'classification':
+                y_pred = clf.predict_proba(X_val)[:, 1]
+                y_pred = np.clip(y_pred, EPS, 1 - EPS)
+                loss = log_loss(y_val, y_pred, sample_weight=w_val)
+            else:
+                y_pred = clf.predict(X_val)
+                loss = regression_loss(
+                    y_val,
+                    y_pred,
+                    w_val,
+                    loss_name=loss_name,
+                    tweedie_power=tweedie_variance_power,
+                )
+            losses.append(float(loss))
+            self._maybe_cleanup_gpu()
 
         return float(np.mean(losses))
-
-    def train(self) -> None:
-        if not self.best_params:
-            raise RuntimeError("Run tune() first to obtain best XGB parameters.")
-        self.model = self._build_estimator()
-        self.model.set_params(**self.best_params)
-        use_refit = bool(getattr(self.ctx.config, "final_refit", True))
-        predict_fn = None
-        if self.ctx.task_type == 'classification':
-            def _predict_proba(X, **_kwargs):
-                return self.model.predict_proba(X)[:, 1]
-            predict_fn = _predict_proba
-        X_all = self.ctx.train_data[self.ctx.factor_nmes]
-        y_all = self.ctx.train_data[self.ctx.resp_nme].values
-        w_all = self.ctx.train_data[self.ctx.weight_nme].values
-
-        split = self._resolve_train_val_indices(X_all)
-        if split is not None:
-            train_idx, val_idx = split
-            X_train = X_all.iloc[train_idx]
-            y_train = y_all[train_idx]
-            w_train = w_all[train_idx]
-            X_val = X_all.iloc[val_idx]
-            y_val = y_all[val_idx]
-            w_val = w_all[val_idx]
-            fit_kwargs = self._build_fit_kwargs(
-                w_train=w_train,
-                X_val=X_val,
-                y_val=y_val,
-                w_val=w_val,
-                n_estimators=self.best_params.get("n_estimators", 100),
-            )
-            self.model.fit(X_train, y_train, **fit_kwargs)
-            best_iter = getattr(self.model, "best_iteration", None)
-            if use_refit and best_iter is not None:
-                refit_model = self._build_estimator()
-                refit_params = dict(self.best_params)
-                refit_params["n_estimators"] = int(best_iter) + 1
-                refit_model.set_params(**refit_params)
-                refit_kwargs = dict(self.ctx.fit_params or {})
-                refit_kwargs.setdefault("sample_weight", w_all)
-                refit_kwargs.pop("eval_set", None)
-                refit_kwargs.pop("sample_weight_eval_set", None)
-                refit_kwargs.pop("early_stopping_rounds", None)
-                refit_kwargs.pop("eval_metric", None)
-                refit_kwargs.setdefault("verbose", False)
-                refit_model.fit(X_all, y_all, **refit_kwargs)
-                self.model = refit_model
-        else:
-            fit_kwargs = dict(self.ctx.fit_params or {})
-            fit_kwargs.setdefault("sample_weight", w_all)
-            fit_kwargs.pop("eval_metric", None)
-            self.model.fit(X_all, y_all, **fit_kwargs)
-
-        self.ctx.model_label.append(self.label)
-        self._predict_and_cache(
-            self.model,
-            pred_prefix='xgb',
-            predict_fn=predict_fn
-        )
-        self.ctx.xgb_best = self.model
-
+
+    def train(self) -> None:
+        if not self.best_params:
+            raise RuntimeError("Run tune() first to obtain best XGB parameters.")
+        self.model = self._build_estimator()
+        self.model.set_params(**self.best_params)
+        use_refit = bool(getattr(self.ctx.config, "final_refit", True))
+        predict_fn = None
+        if self.ctx.task_type == 'classification':
+            def _predict_proba(X, **_kwargs):
+                return self.model.predict_proba(X)[:, 1]
+            predict_fn = _predict_proba
+        X_all = self.ctx.train_data[self.ctx.factor_nmes]
+        y_all = self.ctx.train_data[self.ctx.resp_nme].values
+        w_all = self.ctx.train_data[self.ctx.weight_nme].values
+
+        split = self._resolve_train_val_indices(X_all)
+        if split is not None:
+            train_idx, val_idx = split
+            X_train = X_all.iloc[train_idx]
+            y_train = y_all[train_idx]
+            w_train = w_all[train_idx]
+            X_val = X_all.iloc[val_idx]
+            y_val = y_all[val_idx]
+            w_val = w_all[val_idx]
+            fit_kwargs = self._build_fit_kwargs(
+                w_train=w_train,
+                X_val=X_val,
+                y_val=y_val,
+                w_val=w_val,
+                n_estimators=self.best_params.get("n_estimators", 100),
+            )
+            self.model.fit(X_train, y_train, **fit_kwargs)
+            best_iter = getattr(self.model, "best_iteration", None)
+            if use_refit and best_iter is not None:
+                refit_model = self._build_estimator()
+                refit_params = dict(self.best_params)
+                refit_params["n_estimators"] = int(best_iter) + 1
+                refit_model.set_params(**refit_params)
+                refit_kwargs = dict(self.ctx.fit_params or {})
+                refit_kwargs.setdefault("sample_weight", w_all)
+                refit_kwargs.pop("eval_set", None)
+                refit_kwargs.pop("sample_weight_eval_set", None)
+                refit_kwargs.pop("early_stopping_rounds", None)
+                refit_kwargs.pop("eval_metric", None)
+                refit_kwargs.setdefault("verbose", False)
+                refit_model.fit(X_all, y_all, **refit_kwargs)
+                self.model = refit_model
+        else:
+            fit_kwargs = dict(self.ctx.fit_params or {})
+            fit_kwargs.setdefault("sample_weight", w_all)
+            fit_kwargs.pop("eval_metric", None)
+            self.model.fit(X_all, y_all, **fit_kwargs)
+
+        self.ctx.model_label.append(self.label)
+        self._predict_and_cache(
+            self.model,
+            pred_prefix='xgb',
+            predict_fn=predict_fn
+        )
+        self.ctx.xgb_best = self.model