ins-pricing 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

ins_pricing/cli/Explain_entry.py

@@ -491,54 +491,56 @@ def explain_from_config(args: argparse.Namespace) -> None:
         categorical_features = cfg.get("categorical_features")
         plot_path_style = runtime_cfg["plot_path_style"]
 
-        model = ropt.BayesOptModel(
-            train_df,
-            test_df,
-            model_name,
-            cfg["target"],
-            cfg["weight"],
-            feature_list,
-            task_type=str(cfg.get("task_type", "regression")),
-            binary_resp_nme=binary_target,
-            cate_list=categorical_features,
-            prop_test=prop_test,
-            rand_seed=rand_seed,
-            epochs=int(runtime_cfg["epochs"]),
-            use_gpu=bool(cfg.get("use_gpu", True)),
-            output_dir=output_dir,
-            xgb_max_depth_max=runtime_cfg["xgb_max_depth_max"],
-            xgb_n_estimators_max=runtime_cfg["xgb_n_estimators_max"],
-            resn_weight_decay=cfg.get("resn_weight_decay"),
-            final_ensemble=bool(cfg.get("final_ensemble", False)),
-            final_ensemble_k=int(cfg.get("final_ensemble_k", 3)),
-            final_refit=bool(cfg.get("final_refit", True)),
-            optuna_storage=runtime_cfg["optuna_storage"],
-            optuna_study_prefix=runtime_cfg["optuna_study_prefix"],
-            best_params_files=runtime_cfg["best_params_files"],
-            gnn_use_approx_knn=cfg.get("gnn_use_approx_knn", True),
-            gnn_approx_knn_threshold=cfg.get("gnn_approx_knn_threshold", 50000),
-            gnn_graph_cache=cfg.get("gnn_graph_cache"),
-            gnn_max_gpu_knn_nodes=cfg.get("gnn_max_gpu_knn_nodes", 200000),
-            gnn_knn_gpu_mem_ratio=cfg.get("gnn_knn_gpu_mem_ratio", 0.9),
-            gnn_knn_gpu_mem_overhead=cfg.get("gnn_knn_gpu_mem_overhead", 2.0),
-            region_province_col=cfg.get("region_province_col"),
-            region_city_col=cfg.get("region_city_col"),
-            region_effect_alpha=cfg.get("region_effect_alpha"),
-            geo_feature_nmes=cfg.get("geo_feature_nmes"),
-            geo_token_hidden_dim=cfg.get("geo_token_hidden_dim"),
-            geo_token_layers=cfg.get("geo_token_layers"),
-            geo_token_dropout=cfg.get("geo_token_dropout"),
-            geo_token_k_neighbors=cfg.get("geo_token_k_neighbors"),
-            geo_token_learning_rate=cfg.get("geo_token_learning_rate"),
-            geo_token_epochs=cfg.get("geo_token_epochs"),
-            ft_role=str(cfg.get("ft_role", "model")),
-            ft_feature_prefix=str(cfg.get("ft_feature_prefix", "ft_emb")),
-            ft_num_numeric_tokens=cfg.get("ft_num_numeric_tokens"),
-            infer_categorical_max_unique=int(cfg.get("infer_categorical_max_unique", 50)),
-            infer_categorical_max_ratio=float(cfg.get("infer_categorical_max_ratio", 0.05)),
-            reuse_best_params=runtime_cfg["reuse_best_params"],
-            plot_path_style=plot_path_style,
-        )
+        config_fields = getattr(ropt.BayesOptConfig, "__dataclass_fields__", {})
+        allowed_config_keys = set(config_fields.keys())
+        config_payload = {k: v for k, v in cfg.items() if k in allowed_config_keys}
+        config_payload.update({
+            "model_nme": model_name,
+            "resp_nme": cfg["target"],
+            "weight_nme": cfg["weight"],
+            "factor_nmes": feature_list,
+            "task_type": str(cfg.get("task_type", "regression")),
+            "binary_resp_nme": binary_target,
+            "cate_list": categorical_features,
+            "prop_test": prop_test,
+            "rand_seed": rand_seed,
+            "epochs": int(runtime_cfg["epochs"]),
+            "use_gpu": bool(cfg.get("use_gpu", True)),
+            "output_dir": output_dir,
+            "xgb_max_depth_max": runtime_cfg["xgb_max_depth_max"],
+            "xgb_n_estimators_max": runtime_cfg["xgb_n_estimators_max"],
+            "resn_weight_decay": cfg.get("resn_weight_decay"),
+            "final_ensemble": bool(cfg.get("final_ensemble", False)),
+            "final_ensemble_k": int(cfg.get("final_ensemble_k", 3)),
+            "final_refit": bool(cfg.get("final_refit", True)),
+            "optuna_storage": runtime_cfg["optuna_storage"],
+            "optuna_study_prefix": runtime_cfg["optuna_study_prefix"],
+            "best_params_files": runtime_cfg["best_params_files"],
+            "gnn_use_approx_knn": cfg.get("gnn_use_approx_knn", True),
+            "gnn_approx_knn_threshold": cfg.get("gnn_approx_knn_threshold", 50000),
+            "gnn_graph_cache": cfg.get("gnn_graph_cache"),
+            "gnn_max_gpu_knn_nodes": cfg.get("gnn_max_gpu_knn_nodes", 200000),
+            "gnn_knn_gpu_mem_ratio": cfg.get("gnn_knn_gpu_mem_ratio", 0.9),
+            "gnn_knn_gpu_mem_overhead": cfg.get("gnn_knn_gpu_mem_overhead", 2.0),
+            "region_province_col": cfg.get("region_province_col"),
+            "region_city_col": cfg.get("region_city_col"),
+            "region_effect_alpha": cfg.get("region_effect_alpha"),
+            "geo_feature_nmes": cfg.get("geo_feature_nmes"),
+            "geo_token_hidden_dim": cfg.get("geo_token_hidden_dim"),
+            "geo_token_layers": cfg.get("geo_token_layers"),
+            "geo_token_dropout": cfg.get("geo_token_dropout"),
+            "geo_token_k_neighbors": cfg.get("geo_token_k_neighbors"),
+            "geo_token_learning_rate": cfg.get("geo_token_learning_rate"),
+            "geo_token_epochs": cfg.get("geo_token_epochs"),
+            "ft_role": str(cfg.get("ft_role", "model")),
+            "ft_feature_prefix": str(cfg.get("ft_feature_prefix", "ft_emb")),
+            "ft_num_numeric_tokens": cfg.get("ft_num_numeric_tokens"),
+            "reuse_best_params": runtime_cfg["reuse_best_params"],
+            "plot_path_style": plot_path_style or "nested",
+        })
+        config_payload = {k: v for k, v in config_payload.items() if v is not None}
+        config = ropt.BayesOptConfig(**config_payload)
+        model = ropt.BayesOptModel(train_df, test_df, config=config)
 
         output_overrides = resolve_explain_output_overrides(
             explain_cfg,

ins_pricing/cli/bayesopt_entry_runner.py

@@ -1223,76 +1223,79 @@ def train_from_config(args: argparse.Namespace) -> None:
             cfg.get("ft_feature_prefix", args.ft_feature_prefix))
         ft_num_numeric_tokens = cfg.get("ft_num_numeric_tokens")
 
-        model = ropt.BayesOptModel(
-            train_df,
-            test_df,
-            model_name,
-            cfg["target"],
-            cfg["weight"],
-            feature_list,
-            task_type=task_type,
-            binary_resp_nme=binary_target,
-            cate_list=categorical_features,
-            prop_test=val_ratio,
-            rand_seed=rand_seed,
-            epochs=epochs,
-            use_gpu=use_gpu,
-            use_resn_data_parallel=use_resn_dp,
-            use_ft_data_parallel=use_ft_dp,
-            use_resn_ddp=use_resn_ddp,
-            use_ft_ddp=use_ft_ddp,
-            use_gnn_data_parallel=use_gnn_dp,
-            use_gnn_ddp=use_gnn_ddp,
-            output_dir=output_dir,
-            xgb_max_depth_max=xgb_max_depth_max,
-            xgb_n_estimators_max=xgb_n_estimators_max,
-            resn_weight_decay=cfg.get("resn_weight_decay"),
-            final_ensemble=bool(cfg.get("final_ensemble", False)),
-            final_ensemble_k=int(cfg.get("final_ensemble_k", 3)),
-            final_refit=bool(cfg.get("final_refit", True)),
-            optuna_storage=optuna_storage,
-            optuna_study_prefix=optuna_study_prefix,
-            best_params_files=best_params_files,
-            gnn_use_approx_knn=gnn_use_ann,
-            gnn_approx_knn_threshold=gnn_threshold,
-            gnn_graph_cache=gnn_graph_cache,
-            gnn_max_gpu_knn_nodes=gnn_max_gpu_nodes,
-            gnn_knn_gpu_mem_ratio=gnn_gpu_mem_ratio,
-            gnn_knn_gpu_mem_overhead=gnn_gpu_mem_overhead,
-            region_province_col=region_province_col,
-            region_city_col=region_city_col,
-            region_effect_alpha=region_effect_alpha,
-            geo_feature_nmes=geo_feature_nmes,
-            geo_token_hidden_dim=geo_token_hidden_dim,
-            geo_token_layers=geo_token_layers,
-            geo_token_dropout=geo_token_dropout,
-            geo_token_k_neighbors=geo_token_k_neighbors,
-            geo_token_learning_rate=geo_token_learning_rate,
-            geo_token_epochs=geo_token_epochs,
-            ft_role=ft_role,
-            ft_feature_prefix=ft_feature_prefix,
-            ft_num_numeric_tokens=ft_num_numeric_tokens,
-            infer_categorical_max_unique=int(
-                cfg.get("infer_categorical_max_unique", 50)),
-            infer_categorical_max_ratio=float(
-                cfg.get("infer_categorical_max_ratio", 0.05)),
-            reuse_best_params=reuse_best_params,
-            bo_sample_limit=bo_sample_limit,
-            cache_predictions=cache_predictions,
-            prediction_cache_dir=prediction_cache_dir,
-            prediction_cache_format=prediction_cache_format,
-            cv_strategy=cv_strategy or split_strategy,
-            cv_group_col=cv_group_col or split_group_col,
-            cv_time_col=cv_time_col or split_time_col,
-            cv_time_ascending=cv_time_ascending,
-            cv_splits=cv_splits,
-            ft_oof_folds=ft_oof_folds,
-            ft_oof_strategy=ft_oof_strategy,
-            ft_oof_shuffle=ft_oof_shuffle,
-            save_preprocess=save_preprocess,
-            preprocess_artifact_path=preprocess_artifact_path,
-            plot_path_style=plot_path_style,
-        )
+        config_fields = getattr(ropt.BayesOptConfig,
+                                "__dataclass_fields__", {})
+        allowed_config_keys = set(config_fields.keys())
+        config_payload = {k: v for k,
+                          v in cfg.items() if k in allowed_config_keys}
+        config_payload.update({
+            "model_nme": model_name,
+            "resp_nme": cfg["target"],
+            "weight_nme": cfg["weight"],
+            "factor_nmes": feature_list,
+            "task_type": task_type,
+            "binary_resp_nme": binary_target,
+            "cate_list": categorical_features,
+            "prop_test": val_ratio,
+            "rand_seed": rand_seed,
+            "epochs": epochs,
+            "use_gpu": use_gpu,
+            "use_resn_data_parallel": use_resn_dp,
+            "use_ft_data_parallel": use_ft_dp,
+            "use_gnn_data_parallel": use_gnn_dp,
+            "use_resn_ddp": use_resn_ddp,
+            "use_ft_ddp": use_ft_ddp,
+            "use_gnn_ddp": use_gnn_ddp,
+            "output_dir": output_dir,
+            "xgb_max_depth_max": xgb_max_depth_max,
+            "xgb_n_estimators_max": xgb_n_estimators_max,
+            "resn_weight_decay": cfg.get("resn_weight_decay"),
+            "final_ensemble": bool(cfg.get("final_ensemble", False)),
+            "final_ensemble_k": int(cfg.get("final_ensemble_k", 3)),
+            "final_refit": bool(cfg.get("final_refit", True)),
+            "optuna_storage": optuna_storage,
+            "optuna_study_prefix": optuna_study_prefix,
+            "best_params_files": best_params_files,
+            "gnn_use_approx_knn": gnn_use_ann,
+            "gnn_approx_knn_threshold": gnn_threshold,
+            "gnn_graph_cache": gnn_graph_cache,
+            "gnn_max_gpu_knn_nodes": gnn_max_gpu_nodes,
+            "gnn_knn_gpu_mem_ratio": gnn_gpu_mem_ratio,
+            "gnn_knn_gpu_mem_overhead": gnn_gpu_mem_overhead,
+            "region_province_col": region_province_col,
+            "region_city_col": region_city_col,
+            "region_effect_alpha": region_effect_alpha,
+            "geo_feature_nmes": geo_feature_nmes,
+            "geo_token_hidden_dim": geo_token_hidden_dim,
+            "geo_token_layers": geo_token_layers,
+            "geo_token_dropout": geo_token_dropout,
+            "geo_token_k_neighbors": geo_token_k_neighbors,
+            "geo_token_learning_rate": geo_token_learning_rate,
+            "geo_token_epochs": geo_token_epochs,
+            "ft_role": ft_role,
+            "ft_feature_prefix": ft_feature_prefix,
+            "ft_num_numeric_tokens": ft_num_numeric_tokens,
+            "reuse_best_params": reuse_best_params,
+            "bo_sample_limit": bo_sample_limit,
+            "cache_predictions": cache_predictions,
+            "prediction_cache_dir": prediction_cache_dir,
+            "prediction_cache_format": prediction_cache_format,
+            "cv_strategy": cv_strategy or split_strategy,
+            "cv_group_col": cv_group_col or split_group_col,
+            "cv_time_col": cv_time_col or split_time_col,
+            "cv_time_ascending": cv_time_ascending,
+            "cv_splits": cv_splits,
+            "ft_oof_folds": ft_oof_folds,
+            "ft_oof_strategy": ft_oof_strategy,
+            "ft_oof_shuffle": ft_oof_shuffle,
+            "save_preprocess": save_preprocess,
+            "preprocess_artifact_path": preprocess_artifact_path,
+            "plot_path_style": plot_path_style or "nested",
+        })
+        config_payload = {
+            k: v for k, v in config_payload.items() if v is not None}
+        config = ropt.BayesOptConfig(**config_payload)
+        model = ropt.BayesOptModel(train_df, test_df, config=config)
 
         if plot_requested:
             plot_cfg = cfg.get("plot", {})

ins_pricing/modelling/core/bayesopt/models/model_ft_trainer.py

@@ -626,6 +626,7 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
         best_state = None
         patience_counter = 0
         is_ddp_model = isinstance(self.ft, DDP)
+        use_collectives = dist.is_initialized() and is_ddp_model
 
         clip_fn = None
         if self.device.type == 'cuda':
@@ -669,7 +670,7 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
                             device=X_num_b.device)
                         local_bad = 0 if bool(torch.isfinite(batch_loss)) else 1
                         global_bad = local_bad
-                        if dist.is_initialized():
+                        if use_collectives:
                             bad = torch.tensor(
                                 [local_bad],
                                 device=batch_loss.device,
@@ -774,7 +775,7 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
                             total_n += float(end - start)
                     val_loss_tensor[0] = total_val / max(total_n, 1.0)
 
-                if dist.is_initialized():
+                if use_collectives:
                     dist.broadcast(val_loss_tensor, src=0)
                 val_loss_value = float(val_loss_tensor.item())
                 prune_now = False
@@ -806,7 +807,7 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
                     if trial.should_prune():
                         prune_now = True
 
-                if dist.is_initialized():
+                if use_collectives:
                     flag = torch.tensor(
                         [1 if prune_now else 0],
                         device=loss_tensor_device,

ins_pricing/modelling/core/bayesopt/models/model_gnn.py

@@ -42,6 +42,12 @@ _GNN_MPS_WARNED = False
 # Simplified GNN implementation.
 # =============================================================================
 
+def _adj_mm(adj: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """Matrix multiply that supports sparse or dense adjacency."""
+    if adj.is_sparse:
+        return torch.sparse.mm(adj, x)
+    return adj.matmul(x)
+
 
 class SimpleGraphLayer(nn.Module):
     def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.1):
@@ -52,7 +58,7 @@ class SimpleGraphLayer(nn.Module):
 
     def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
         # Message passing with normalized sparse adjacency: A_hat * X * W.
-        h = torch.sparse.mm(adj, x)
+        h = _adj_mm(adj, x)
         h = self.linear(h)
         h = self.activation(h)
         return self.dropout(h)
@@ -86,7 +92,7 @@ class SimpleGNN(nn.Module):
         h = x
         for layer in self.layers:
             h = layer(h, adj_used)
-        h = torch.sparse.mm(adj_used, h)
+        h = _adj_mm(adj_used, h)
         out = self.output(h)
         return self.output_act(out)
 
@@ -124,7 +130,11 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         self.knn_gpu_mem_ratio = max(0.0, min(1.0, knn_gpu_mem_ratio))
         self.knn_gpu_mem_overhead = max(1.0, knn_gpu_mem_overhead)
         self.knn_cpu_jobs = knn_cpu_jobs
+        self.mps_dense_max_nodes = int(
+            os.environ.get("BAYESOPT_GNN_MPS_DENSE_MAX_NODES", "5000")
+        )
         self._knn_warning_emitted = False
+        self._mps_fallback_triggered = False
         self._adj_cache_meta: Optional[Dict[str, Any]] = None
         self._adj_cache_key: Optional[Tuple[Any, ...]] = None
         self._adj_cache_tensor: Optional[torch.Tensor] = None
@@ -168,11 +178,11 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             else:
                 self.device = torch.device('cuda')
         elif torch.backends.mps.is_available():
-            self.device = torch.device('cpu')
+            self.device = torch.device('mps')
             global _GNN_MPS_WARNED
             if not _GNN_MPS_WARNED:
                 print(
-                    "[GNN] MPS backend does not support sparse ops; falling back to CPU.",
+                    "[GNN] Using MPS backend; will fall back to CPU on unsupported ops.",
                     flush=True,
                 )
                 _GNN_MPS_WARNED = True
@@ -235,6 +245,41 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         else:
             base.register_buffer("adj_buffer", adj)
 
+    @staticmethod
+    def _is_mps_unsupported_error(exc: BaseException) -> bool:
+        msg = str(exc).lower()
+        if "mps" not in msg:
+            return False
+        if any(token in msg for token in ("not supported", "not implemented", "does not support", "unimplemented", "out of memory")):
+            return True
+        return "sparse" in msg
+
+    def _fallback_to_cpu(self, reason: str) -> None:
+        if self.device.type != "mps" or self._mps_fallback_triggered:
+            return
+        self._mps_fallback_triggered = True
+        print(f"[GNN] MPS op unsupported ({reason}); falling back to CPU.", flush=True)
+        self.device = torch.device("cpu")
+        self.use_pyg_knn = False
+        self.data_parallel_enabled = False
+        self.ddp_enabled = False
+        base = self._unwrap_gnn()
+        try:
+            base = base.to(self.device)
+        except Exception:
+            pass
+        self.gnn = base
+        self.invalidate_graph_cache()
+
+    def _run_with_mps_fallback(self, fn, *args, **kwargs):
+        try:
+            return fn(*args, **kwargs)
+        except (RuntimeError, NotImplementedError) as exc:
+            if self.device.type == "mps" and self._is_mps_unsupported_error(exc):
+                self._fallback_to_cpu(str(exc))
+                return fn(*args, **kwargs)
+            raise
+
     def _graph_cache_meta(self, X_df: pd.DataFrame) -> Dict[str, Any]:
         row_hash = pd.util.hash_pandas_object(X_df, index=False).values
         idx_hash = pd.util.hash_pandas_object(X_df.index, index=False).values
@@ -255,11 +300,14 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             "knn_gpu_mem_ratio": float(self.knn_gpu_mem_ratio),
             "knn_gpu_mem_overhead": float(self.knn_gpu_mem_overhead),
         }
+        adj_format = "dense" if self.device.type == "mps" else "sparse"
         return {
             "n_samples": int(X_df.shape[0]),
             "n_features": int(X_df.shape[1]),
             "hash": hasher.hexdigest(),
             "knn_config": knn_config,
+            "adj_format": adj_format,
+            "device_type": self.device.type,
         }
 
     def _graph_cache_key(self, X_df: pd.DataFrame) -> Tuple[Any, ...]:
@@ -284,8 +332,7 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             if meta_expected is None:
                 meta_expected = self._graph_cache_meta(X_df)
             try:
-                payload = torch.load(self.graph_cache_path,
-                                     map_location=self.device)
+                payload = torch.load(self.graph_cache_path, map_location="cpu")
             except Exception as exc:
                 print(
                     f"[GNN] Failed to load cached graph from {self.graph_cache_path}: {exc}")
@@ -293,7 +340,13 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             if isinstance(payload, dict) and "adj" in payload:
                 meta_cached = payload.get("meta")
                 if meta_cached == meta_expected:
-                    return payload["adj"].to(self.device)
+                    adj = payload["adj"]
+                    if self.device.type == "mps" and getattr(adj, "is_sparse", False):
+                        print(
+                            f"[GNN] Cached sparse graph incompatible with MPS; rebuilding: {self.graph_cache_path}"
+                        )
+                        return None
+                    return adj.to(self.device)
                 print(
                     f"[GNN] Cached graph metadata mismatch; rebuilding: {self.graph_cache_path}")
                 return None
@@ -408,6 +461,11 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         return True
 
     def _normalized_adj(self, edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor:
+        if self.device.type == "mps":
+            return self._normalized_adj_dense(edge_index, num_nodes)
+        return self._normalized_adj_sparse(edge_index, num_nodes)
+
+    def _normalized_adj_sparse(self, edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor:
         values = torch.ones(edge_index.shape[1], device=self.device)
         adj = torch.sparse_coo_tensor(
             edge_index.to(self.device), values, (num_nodes, num_nodes))
@@ -421,6 +479,21 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             adj.indices(), norm_values, size=adj.shape)
         return adj_norm
 
+    def _normalized_adj_dense(self, edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor:
+        if self.mps_dense_max_nodes <= 0 or num_nodes > self.mps_dense_max_nodes:
+            raise RuntimeError(
+                f"MPS dense adjacency not supported for {num_nodes} nodes; "
+                f"max={self.mps_dense_max_nodes}. Falling back to CPU."
+            )
+        edge_index = edge_index.to(self.device)
+        adj = torch.zeros((num_nodes, num_nodes), device=self.device, dtype=torch.float32)
+        adj[edge_index[0], edge_index[1]] = 1.0
+        deg = adj.sum(dim=1)
+        deg_inv_sqrt = torch.pow(deg + 1e-8, -0.5)
+        adj = adj * deg_inv_sqrt.view(-1, 1)
+        adj = adj * deg_inv_sqrt.view(1, -1)
+        return adj
+
     def _tensorize_split(self, X, y, w, allow_none: bool = False):
         if X is None and allow_none:
             return None, None, None
@@ -462,17 +535,25 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
             if self._adj_cache_meta == meta_expected and self._adj_cache_tensor is not None:
                 cached = self._adj_cache_tensor
                 if cached.device != self.device:
-                    cached = cached.to(self.device)
-                    self._adj_cache_tensor = cached
-                return cached
+                    if self.device.type == "mps" and getattr(cached, "is_sparse", False):
+                        self._adj_cache_tensor = None
+                    else:
+                        cached = cached.to(self.device)
+                        self._adj_cache_tensor = cached
+                if self._adj_cache_tensor is not None:
+                    return self._adj_cache_tensor
         else:
             cache_key = self._graph_cache_key(X_df)
             if self._adj_cache_key == cache_key and self._adj_cache_tensor is not None:
                 cached = self._adj_cache_tensor
                 if cached.device != self.device:
-                    cached = cached.to(self.device)
-                    self._adj_cache_tensor = cached
-                return cached
+                    if self.device.type == "mps" and getattr(cached, "is_sparse", False):
+                        self._adj_cache_tensor = None
+                    else:
+                        cached = cached.to(self.device)
+                        self._adj_cache_tensor = cached
+                if self._adj_cache_tensor is not None:
+                    return self._adj_cache_tensor
         X_np = None
         if X_tensor is None:
             X_np = X_df.to_numpy(dtype=np.float32, copy=False)
@@ -511,7 +592,20 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
     def fit(self, X_train, y_train, w_train=None,
             X_val=None, y_val=None, w_val=None,
             trial: Optional[optuna.trial.Trial] = None):
+        return self._run_with_mps_fallback(
+            self._fit_impl,
+            X_train,
+            y_train,
+            w_train,
+            X_val,
+            y_val,
+            w_val,
+            trial,
+        )
 
+    def _fit_impl(self, X_train, y_train, w_train=None,
+                  X_val=None, y_val=None, w_val=None,
+                  trial: Optional[optuna.trial.Trial] = None):
         X_train_tensor, y_train_tensor, w_train_tensor = self._tensorize_split(
             X_train, y_train, w_train, allow_none=False)
         has_val = X_val is not None and y_val is not None
@@ -621,6 +715,9 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         self.best_epoch = int(best_epoch or self.epochs)
 
     def predict(self, X: pd.DataFrame) -> np.ndarray:
+        return self._run_with_mps_fallback(self._predict_impl, X)
+
+    def _predict_impl(self, X: pd.DataFrame) -> np.ndarray:
         self.gnn.eval()
         X_tensor, _, _ = self._tensorize_split(
             X, None, None, allow_none=False)
@@ -640,6 +737,9 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         return y_pred.ravel()
 
     def encode(self, X: pd.DataFrame) -> np.ndarray:
+        return self._run_with_mps_fallback(self._encode_impl, X)
+
+    def _encode_impl(self, X: pd.DataFrame) -> np.ndarray:
         """Return per-sample node embeddings (hidden representations)."""
         base = self._unwrap_gnn()
         base.eval()
@@ -655,7 +755,7 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
                 raise RuntimeError("GNN base module does not expose layers.")
             for layer in layers:
                 h = layer(h, adj)
-            h = torch.sparse.mm(adj, h)
+            h = _adj_mm(adj, h)
         return h.detach().cpu().numpy()
 
     def set_params(self, params: Dict[str, Any]):

ins_pricing/modelling/core/bayesopt/trainers/trainer_base.py

@@ -27,6 +27,7 @@ from sklearn.preprocessing import StandardScaler
 from ..config_preprocess import BayesOptConfig, OutputManager
 from ..utils import DistributedUtils, EPS, ensure_parent_dir
 from ins_pricing.utils import get_logger, GPUMemoryManager, DeviceManager
+from ins_pricing.utils.torch_compat import torch_load
 
 # Module-level logger
 _logger = get_logger("ins_pricing.trainer")
@@ -616,7 +617,7 @@ class TrainerBase:
                 pass
             else:
                 # FT-Transformer: load state_dict and reconstruct model
-                loaded = torch.load(path, map_location='cpu', weights_only=False)
+                loaded = torch_load(path, map_location='cpu', weights_only=False)
                 if isinstance(loaded, dict):
                     if "state_dict" in loaded and "model_config" in loaded:
                         # New format: state_dict + model_config
@@ -1094,7 +1095,7 @@ class TrainerBase:
                 split_iter = splitter
 
         losses: List[float] = []
-        for train_idx, val_idx in split_iter:
+        for fold_idx, (train_idx, val_idx) in enumerate(split_iter):
             X_train = X_all.iloc[train_idx]
             y_train = y_all.iloc[train_idx]
             X_val = X_all.iloc[val_idx]
@@ -1108,9 +1109,11 @@ class TrainerBase:
             model = model_builder(params)
             try:
                 if fit_predict_fn:
+                    # Avoid duplicate Optuna step reports across folds.
+                    trial_for_fold = trial if fold_idx == 0 else None
                     y_pred = fit_predict_fn(
                         model, X_train, y_train, w_train,
-                        X_val, y_val, w_val, trial
+                        X_val, y_val, w_val, trial_for_fold
                     )
                 else:
                     fit_kwargs = {}
@@ -1288,4 +1291,3 @@ class TrainerBase:
             predict_kwargs_train=predict_kwargs_train,
             predict_kwargs_test=predict_kwargs_test,
             predict_fn=predict_fn)
-

ins_pricing/modelling/core/bayesopt/trainers/trainer_gnn.py

@@ -12,6 +12,7 @@ from .trainer_base import TrainerBase
 from ..models import GraphNeuralNetSklearn
 from ..utils import EPS
 from ins_pricing.utils import get_logger
+from ins_pricing.utils.torch_compat import torch_load
 
 _logger = get_logger("ins_pricing.trainer.gnn")
 
@@ -300,7 +301,7 @@ class GNNTrainer(TrainerBase):
         if not os.path.exists(path):
             print(f"[load] Warning: Model file not found: {path}")
             return
-        payload = torch.load(path, map_location='cpu', weights_only=False)
+        payload = torch_load(path, map_location='cpu', weights_only=False)
         if not isinstance(payload, dict):
             raise ValueError(f"Invalid GNN checkpoint: {path}")
         params = payload.get("best_params") or {}
@@ -322,4 +323,3 @@ class GNNTrainer(TrainerBase):
         self.model = model
         self.best_params = dict(params) if isinstance(params, dict) else None
         self.ctx.gnn_best = self.model
-

ins_pricing/modelling/core/bayesopt/utils/torch_trainer_mixin.py

@@ -314,7 +314,10 @@ class TorchTrainerMixin:
             f">>> DataLoader config: Batch Size={batch_size}, Accum Steps={accum_steps}, "
             f"Workers={workers}, Prefetch={prefetch_factor or 'off'}, Profile={profile}")
         sampler = None
-        if dist.is_initialized():
+        use_distributed_sampler = bool(
+            dist.is_initialized() and getattr(self, "is_ddp_enabled", False)
+        )
+        if use_distributed_sampler:
             sampler = DistributedSampler(dataset, shuffle=True)
             shuffle = False
         else:
@@ -413,6 +416,7 @@ class TorchTrainerMixin:
         val_history: List[float] = []
 
         is_ddp_model = isinstance(model, DDP)
+        use_collectives = dist.is_initialized() and is_ddp_model
 
         for epoch in range(1, getattr(self, "epochs", 1) + 1):
             epoch_start_ts = time.time()
@@ -487,7 +491,7 @@ class TorchTrainerMixin:
                             val_weighted_loss = val_result
                     val_loss_tensor[0] = float(val_weighted_loss)
 
-                if dist.is_initialized():
+                if use_collectives:
                     dist.broadcast(val_loss_tensor, src=0)
                 val_weighted_loss = float(val_loss_tensor.item())
 
@@ -502,7 +506,7 @@ class TorchTrainerMixin:
                     trial.report(val_weighted_loss, epoch)
                     prune_flag = trial.should_prune()
 
-                if dist.is_initialized():
+                if use_collectives:
                     prune_device = getattr(self, "device", torch.device("cpu"))
                     if not isinstance(prune_device, torch.device):
                         prune_device = torch.device(prune_device)

ins_pricing/production/predict.py

@@ -24,6 +24,7 @@ from .scoring import batch_score
 from ..modelling.core.bayesopt.models.model_gnn import GraphNeuralNetSklearn
 from ..modelling.core.bayesopt.models.model_resn import ResNetSklearn
 from ins_pricing.utils import DeviceManager, get_logger
+from ins_pricing.utils.torch_compat import torch_load
 
 _logger = get_logger("ins_pricing.production.predict")
 
@@ -130,7 +131,7 @@ def _load_preprocess_from_model_file(
     if model_key in {"xgb", "glm"}:
         payload = joblib.load(model_path)
     else:
-        payload = torch.load(model_path, map_location="cpu")
+        payload = torch_load(model_path, map_location="cpu")
     if isinstance(payload, dict):
         return payload.get("preprocess_artifacts")
     return None
@@ -261,7 +262,7 @@ def load_saved_model(
         return payload
 
     if model_key == "ft":
-        payload = torch.load(model_path, map_location="cpu", weights_only=False)
+        payload = torch_load(model_path, map_location="cpu", weights_only=False)
         if isinstance(payload, dict):
             if "state_dict" in payload and "model_config" in payload:
                 # New format: state_dict + model_config (DDP-safe)
@@ -325,7 +326,7 @@ def load_saved_model(
     if model_key == "resn":
         if input_dim is None:
             raise ValueError("input_dim is required for ResNet loading")
-        payload = torch.load(model_path, map_location="cpu")
+        payload = torch_load(model_path, map_location="cpu")
         if isinstance(payload, dict) and "state_dict" in payload:
             state_dict = payload.get("state_dict")
             params = payload.get("best_params") or load_best_params(
@@ -351,7 +352,7 @@ def load_saved_model(
     if model_key == "gnn":
         if input_dim is None:
             raise ValueError("input_dim is required for GNN loading")
-        payload = torch.load(model_path, map_location="cpu")
+        payload = torch_load(model_path, map_location="cpu")
         if not isinstance(payload, dict):
             raise ValueError(f"Invalid GNN checkpoint: {model_path}")
         params = payload.get("best_params") or {}

ins_pricing/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="ins_pricing",
-    version="0.3.0",
+    version="0.3.2",
     description="Reusable modelling, pricing, governance, and reporting utilities.",
     author="meishi125478",
     license="Proprietary",

ins_pricing/utils/torch_compat.py

@@ -0,0 +1,45 @@
+"""Torch compatibility helpers for 1.x and 2.x support."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any, Optional
+
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except ImportError:  # pragma: no cover - handled by callers
+    TORCH_AVAILABLE = False
+    torch = None
+
+_SUPPORTS_WEIGHTS_ONLY: Optional[bool] = None
+
+
+def _supports_weights_only() -> bool:
+    """Check whether torch.load supports the weights_only argument."""
+    global _SUPPORTS_WEIGHTS_ONLY
+    if _SUPPORTS_WEIGHTS_ONLY is None:
+        if not TORCH_AVAILABLE:
+            _SUPPORTS_WEIGHTS_ONLY = False
+        else:
+            try:
+                sig = inspect.signature(torch.load)
+                _SUPPORTS_WEIGHTS_ONLY = "weights_only" in sig.parameters
+            except (TypeError, ValueError):
+                _SUPPORTS_WEIGHTS_ONLY = False
+    return bool(_SUPPORTS_WEIGHTS_ONLY)
+
+
+def torch_load(
+    path: Any,
+    *args: Any,
+    weights_only: Optional[bool] = None,
+    **kwargs: Any,
+) -> Any:
+    """Load a torch artifact while handling 1.x/2.x API differences."""
+    if not TORCH_AVAILABLE:
+        raise RuntimeError("torch is required to load model files.")
+    if weights_only is not None and _supports_weights_only():
+        return torch.load(path, *args, weights_only=weights_only, **kwargs)
+    return torch.load(path, *args, **kwargs)

{ins_pricing-0.3.0.dist-info → ins_pricing-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ins_pricing
-Version: 0.3.0
+Version: 0.3.2
 Summary: Reusable modelling, pricing, governance, and reporting utilities.
 Author: meishi125478
 License: Proprietary

{ins_pricing-0.3.0.dist-info → ins_pricing-0.3.2.dist-info}/RECORD

@@ -3,14 +3,14 @@ ins_pricing/README.md,sha256=W4V2xtzM6pyQzwJPvWP7cNn-We9rxM8xrxRlBVQwoY8,3399
 ins_pricing/RELEASE_NOTES_0.2.8.md,sha256=KIJzk1jbZbZPKjwnkPSDHO_2Ipv3SP3CzCNDdf07jI0,9331
 ins_pricing/__init__.py,sha256=46j1wCdLVrgrofeBwKl-3NXTxzjbTv-w3KjW-dyKGiY,2622
 ins_pricing/exceptions.py,sha256=5fZavPV4zNJ7wPC75L215KkHXX9pRrfDAYZOdSKJMGo,4778
-ins_pricing/setup.py,sha256=Rwx8jUs1m2BbK0xw8aaHvUpr6HOln42XCmHb7QE7BmU,1702
+ins_pricing/setup.py,sha256=ORpXP7CtmQeJRu2-QbRNzo96MuPT1vWY39-yUMydqd0,1702
 ins_pricing/cli/BayesOpt_entry.py,sha256=X3AiNQQh5ARcjVMM2vOKWPYPDIId40n_RPZA76pTGl4,558
 ins_pricing/cli/BayesOpt_incremental.py,sha256=_Klr5vvNoq_TbgwrH_T3f0a6cHmA9iVJMViiji6ahJY,35927
 ins_pricing/cli/Explain_Run.py,sha256=gEPQjqHiXyXlCTKjUzwSvbAn5_h74ABgb_sEGs-YHVE,664
-ins_pricing/cli/Explain_entry.py,sha256=PP6j4uQspmettCk21lhh9L9jurDwZfWXGxS2EBUZpMs,23191
+ins_pricing/cli/Explain_entry.py,sha256=xS23x31qRvXqy2Wvo21wbkxmPRLHy77aZbtjY-J59NA,23570
 ins_pricing/cli/Pricing_Run.py,sha256=qZribQ_ku4NK4oIvlrLJdM-jyyKtIUoCbbvo8Wh_RQ4,711
 ins_pricing/cli/__init__.py,sha256=F296f1J_tBPv33lDJQ6LaN_CPwMJTMtOuTsMof0dr2o,50
-ins_pricing/cli/bayesopt_entry_runner.py,sha256=TKqaupVDoG0pkVBrsU5_IQidJ2dx4tms52tzptMOp_0,53365
+ins_pricing/cli/bayesopt_entry_runner.py,sha256=Cl5_uw_nHxkRESUW_OLjSJaHQMfiXdbitttN7ZS-tVM,53841
 ins_pricing/cli/watchdog_run.py,sha256=ehUkN9VqsQkxc6YC_WLanU6Pu-ers-nvPEtCaJ9UsgY,6188
 ins_pricing/cli/utils/__init__.py,sha256=u3kt1B27OiuOEgw6PQN-fNs9vNiAjdPyybsRQsZkM_I,54
 ins_pricing/cli/utils/cli_common.py,sha256=CgMnN_0NQQt7Bh5RjClOydz0LzU2TBmIOsFa3KxYQOQ,8254
@@ -44,14 +44,14 @@ ins_pricing/modelling/core/bayesopt/utils.py,sha256=fTDqBHCxsOVte0QhPqnMw8vkefIv
 ins_pricing/modelling/core/bayesopt/utils_backup.py,sha256=5RKizpR3j6KwR87WqqaXPtgjQXWPW4vM75sIkx38SSM,57924
 ins_pricing/modelling/core/bayesopt/models/__init__.py,sha256=vFFCkGnO6rm50TbxR6QekKKQjq-NW4UFwog6fng8-p8,700
 ins_pricing/modelling/core/bayesopt/models/model_ft_components.py,sha256=0I0NiDf1D3cOhTRQwatsNTw9Julmxv5v3HZV8fTrvcQ,10989
-ins_pricing/modelling/core/bayesopt/models/model_ft_trainer.py,sha256=2rOKekH20_9Ccb9X-S6adW-TZcSN9Aj1Bmym8c3Sgrc,39206
-ins_pricing/modelling/core/bayesopt/models/model_gnn.py,sha256=l_oIgLeTJndfYR8lpZbqUq0MKH6wReE2z1B8n1E0P8k,28095
+ins_pricing/modelling/core/bayesopt/models/model_ft_trainer.py,sha256=srBsI4KmijRC6s4G-XlVI9ahzjc2yquQxgdeICP1_4A,39253
+ins_pricing/modelling/core/bayesopt/models/model_gnn.py,sha256=1LdcuMwCRXsQcdj1Xjm8NlLjpX-cmunJlelkc2pfGB4,32400
 ins_pricing/modelling/core/bayesopt/models/model_resn.py,sha256=hAU77GcGC1mHbOOLfZ9vC5nhUhHlaPZmXjrkH3BrjKc,17128
 ins_pricing/modelling/core/bayesopt/trainers/__init__.py,sha256=ODYKjT-v4IDxu4ohGLCXY8r1-pMME9LAaNx6pmj5_38,481
-ins_pricing/modelling/core/bayesopt/trainers/trainer_base.py,sha256=-AaGt9pIeynG5ZcCZkemXsyjUzNbAzzWN_-hoh2Yo5Q,54788
+ins_pricing/modelling/core/bayesopt/trainers/trainer_base.py,sha256=101Til7y9yuRQuFiYQHz9UGWmPghNiHHqmnNtQZxVok,55015
 ins_pricing/modelling/core/bayesopt/trainers/trainer_ft.py,sha256=jkafhvfEqIV_PYJ90e8kkOXVGvFpCKS0CRSKLvQ_elQ,34730
 ins_pricing/modelling/core/bayesopt/trainers/trainer_glm.py,sha256=wVU6F2Ubyu3IxP6K-epjkUTGOE8gKPCdpVxGW-JP9rM,7806
-ins_pricing/modelling/core/bayesopt/trainers/trainer_gnn.py,sha256=OChMeBELl_XWYZmZcpKq5BLY-srtTgempk4L5D0I7ys,13453
+ins_pricing/modelling/core/bayesopt/trainers/trainer_gnn.py,sha256=G7an1PSRoJBWIfqZ_iIV5IOzeM03Z1amW4CPmFTklMU,13506
 ins_pricing/modelling/core/bayesopt/trainers/trainer_resn.py,sha256=y9KyKoUHLiKoXiX6XP4QGu0MKX5LyvlmSAT5tgn826M,11140
 ins_pricing/modelling/core/bayesopt/trainers/trainer_xgb.py,sha256=Ha3PtkB7AedlnRwWEfIBTEAPO69LY5lnjUrdGNL-yas,13783
 ins_pricing/modelling/core/bayesopt/utils/__init__.py,sha256=dbf4DrWOH4rABOuaZdBF7drYOBH5prjvM0TexT6DYyg,1911
@@ -59,7 +59,7 @@ ins_pricing/modelling/core/bayesopt/utils/constants.py,sha256=0ihYxGlJ8tIElYvkhI
 ins_pricing/modelling/core/bayesopt/utils/distributed_utils.py,sha256=cu01dHyYE5EREbmtJgCHSH6z5mQIqajz8_-oWZV6zVc,5787
 ins_pricing/modelling/core/bayesopt/utils/io_utils.py,sha256=vXDlAc_taCG2joxnC6wu0jVYA76UhRbX9OT_5z_im-E,3857
 ins_pricing/modelling/core/bayesopt/utils/metrics_and_devices.py,sha256=irrsE-oDDpOd82QyW9uxfDueF8RFKkUHME4wyWr7vRI,20533
-ins_pricing/modelling/core/bayesopt/utils/torch_trainer_mixin.py,sha256=q5aEy9Lp8Vemf3iCyWi-OGsbe9roVEbKYkN3qRJlVUo,24283
+ins_pricing/modelling/core/bayesopt/utils/torch_trainer_mixin.py,sha256=_ELpu5dxucq5luO_cqe_kxT7oMIVXd63pb4EZ7fWlUQ,24465
 ins_pricing/modelling/explain/__init__.py,sha256=CPoGzGu8TTO3FOXjxoXC13VkuIDCf3YTH6L3BqJq3Ok,1171
 ins_pricing/modelling/explain/gradients.py,sha256=9TqCws_p49nFxVMcjVxe4KCZ7frezeL0uV_LCdoM5yo,11088
 ins_pricing/modelling/explain/metrics.py,sha256=K_xOY7ZrHWhbJ79RNB7eXN3VXeTe8vq68ZLH2BlZufA,5389
@@ -82,7 +82,7 @@ ins_pricing/pricing/rate_table.py,sha256=llDW95i7gR6cCtGFwcGqgpgFvOOPCURaJWmuQw1
 ins_pricing/production/__init__.py,sha256=plUjyiwxrzHDDgXKezyGp9UHOg7Mav4f0ryXYtNmbfs,885
 ins_pricing/production/drift.py,sha256=q_oE_h2NbVETTBkh9QUu8Y68ERuFFcrfKpOb3zBcvsA,383
 ins_pricing/production/monitoring.py,sha256=A6Hyc5WSKhFkDZOIrqmFteuDee75CdcwdTq644vrk-U,4836
-ins_pricing/production/predict.py,sha256=mJog-RGLHxIJxx4oh0D1gbhJwBZPQs1z1P6YTDfWtfg,21791
+ins_pricing/production/predict.py,sha256=fJoROPxW3474TrD4tnG6aEXXzoYVSZJd-kwOsUSEgQ8,21845
 ins_pricing/production/preprocess.py,sha256=cl20X0rVcKNCjVJswB8SdHffMgox6Qga4Ac29L6pW5g,9404
 ins_pricing/production/scoring.py,sha256=yFmMmbYb7w_RC4uZOCMnAjLMRcjXQWIuT1nsfu-bwuc,1379
 ins_pricing/reporting/README.md,sha256=kTVdB6pNewwh1HlCHrI2SzWTgprtQoQprLRQ2qLdgNA,486
@@ -118,8 +118,9 @@ ins_pricing/utils/logging.py,sha256=_AKB4ErmvygwGLtu7Ai7ESemj6Hh8FTgh4cs8j_gVW4,
 ins_pricing/utils/metrics.py,sha256=zhKYgXgki8RDscjP_GO2lEgzrtMIZCqOX_aLpQzdw6k,8390
 ins_pricing/utils/paths.py,sha256=o_tBiclFvBci4cYg9WANwKPxrMcglEdOjDP-EZgGjdQ,8749
 ins_pricing/utils/profiling.py,sha256=kmbykHLcYywlZxAf_aVU8HXID3zOvUcBoO5Q58AijhA,11132
+ins_pricing/utils/torch_compat.py,sha256=ztTMmVxkQiNGvC5Ksz1VrzEA4QNZzRy72kW1MKIyA6M,1366
 ins_pricing/utils/validation.py,sha256=4Tw9VUJPk0N-WO3YUqZP-xXRl1Xpubkm0vi3WzzZrv4,13348
-ins_pricing-0.3.0.dist-info/METADATA,sha256=dCrwU2NZzA2I-VyTq85Pr5ucyoQA8gV4Un8bOMCejkg,6101
-ins_pricing-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ins_pricing-0.3.0.dist-info/top_level.txt,sha256=haZuNQpHKNBEPZx3NjLnHp8pV3I_J9QG8-HyJn00FA0,12
-ins_pricing-0.3.0.dist-info/RECORD,,
+ins_pricing-0.3.2.dist-info/METADATA,sha256=K-Gae8gNhr87oAcVXOColnccx7J1CAH0SsfZeHd75vo,6101
+ins_pricing-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ins_pricing-0.3.2.dist-info/top_level.txt,sha256=haZuNQpHKNBEPZx3NjLnHp8pV3I_J9QG8-HyJn00FA0,12
+ins_pricing-0.3.2.dist-info/RECORD,,