From 2b93cb652c934de183439883effccedd8964e892 Mon Sep 17 00:00:00 2001 From: Thierry RAMORASOAVINA Date: Fri, 28 Aug 2026 17:15:33 +0200 Subject: [PATCH] Add `max_cores` to the sklearn estimators to limit the number of CPU cores pre-allocated for the training --- CHANGELOG.md | 3 ++ khiops/sklearn/estimators.py | 57 ++++++++++++++++++++++++++++++++++-- tests/test_sklearn.py | 15 +++++++++- 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f15c72dc..0bb80348 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ ### Fixed - (`core`) Fix a RuntimeException error when automatically correcting deprecated data paths in task args +### Changed +- (`sklearn`) Add `max_cores` to the sklearn estimators to limit the number of CPU cores pre-allocated and used by Khiops. + ## 11.0.1.0 - 2026-07-02 ### Added diff --git a/khiops/sklearn/estimators.py b/khiops/sklearn/estimators.py index d2b28d94..d04ba9eb 100644 --- a/khiops/sklearn/estimators.py +++ b/khiops/sklearn/estimators.py @@ -235,6 +235,8 @@ class KhiopsEstimator(ABC, BaseEstimator): Parameters ---------- + n_cores : int, optional + Maximum number of CPU cores allocated and used for all estimator operations. verbose : bool, default `False` If `True` it prints debug information and it does not erase temporary files when fitting, predicting or transforming. @@ -248,12 +250,14 @@ class KhiopsEstimator(ABC, BaseEstimator): def __init__( self, + n_cores=None, verbose=False, output_dir=None, auto_sort=True, ): # Set the estimator parameters and internal variables self._khiops_model_prefix = None + self.n_cores = n_cores self.output_dir = output_dir self.verbose = verbose self.auto_sort = auto_sort @@ -351,6 +355,16 @@ def _cleanup_computation_dir(self, computation_dir): def fit(self, X, y=None, **kwargs): """Fit the estimator + Parameters + ---------- + X : array-like of shape (n_samples, n_features_in) or dict + Training dataset. Either an [array-like][] or a [dict][] + specification for multi-table datasets + (see [Multi-Table Learning Primer](../multi_table_primer.md)). + + y : array-like of shape (n_samples,) + The target values. + Returns ------- self : KhiopsEstimator @@ -637,6 +651,8 @@ class KhiopsCoclustering(ClusterMixin, KhiopsEstimator): to speed up the processing. This affects the [predict][] method. *Note* The sort by key is performed in a left-to-right, hierarchical, lexicographic manner. + n_cores : int, optional + Maximum number of CPU cores allocated and used for all estimator operations. Attributes ---------- @@ -664,11 +680,13 @@ def __init__( build_name_var=True, build_distance_vars=False, build_frequency_vars=False, + n_cores=None, ): super().__init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) self._khiops_model_prefix = "CC_" self.build_name_var = build_name_var @@ -770,6 +788,7 @@ def _fit_train_model(self, ds, computation_dir, **kwargs): main_table_path, variables, coclustering_file_path, + max_cores=self.n_cores, log_file_path=train_log_file_path, trace=self.verbose, ) @@ -1186,7 +1205,23 @@ def _transform_prepare_deployment_for_predict(self, _): return self.model_.copy(), None def fit_predict(self, X, y=None, **kwargs): - """Performs clustering on X and returns result (instead of labels)""" + """Performs clustering on X and returns result (instead of labels) + + Parameters + ---------- + X : array-like of shape (n_samples, n_features_in) or dict + Training dataset. Either an [array-like][] or a [dict][] + specification for multi-table datasets + (see [Multi-Table Learning Primer](../multi_table_primer.md)). + + y : array-like of shape (n_samples,) + The target values. + + Returns + ------- + results : `numpy.array` + """ + return self.fit(X, y, **kwargs).predict(X) @@ -1207,11 +1242,13 @@ def __init__( verbose=False, output_dir=None, auto_sort=True, + n_cores=None, ): super().__init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) self.n_features = n_features self.n_trees = n_trees @@ -1382,7 +1419,8 @@ def _fit_prepare_training_function_inputs(self, ds, computation_dir): report_file_path, ] - # Build the optional parameters from a copy of the estimator parameters + # Build the optional parameters from a copy + # of the estimator initializer parameters kwargs = self.get_params() # Remove non core.api params @@ -1403,6 +1441,7 @@ def _fit_prepare_training_function_inputs(self, ds, computation_dir): kwargs["max_text_features"] = kwargs.pop("n_text_features") kwargs["text_features"] = kwargs.pop("type_text_features") kwargs["max_parts"] = kwargs.pop("n_feature_parts") + kwargs["max_cores"] = kwargs.pop("n_cores") # Add the additional_data_tables parameter kwargs["additional_data_tables"] = additional_data_tables @@ -1537,6 +1576,7 @@ def __init__( verbose=False, output_dir=None, auto_sort=True, + n_cores=None, ): super().__init__( n_features=n_features, @@ -1551,6 +1591,7 @@ def __init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) # Data to be specified by inherited classes self._predicted_target_meta_data_tag = None @@ -1731,6 +1772,8 @@ class KhiopsClassifier(ClassifierMixin, KhiopsPredictor): affects the [fit][], [predict][] and [predict_proba][] methods. *Note* The sort by key is performed in a left-to-right, hierarchical, lexicographic manner. + n_cores : int, optional + Maximum number of CPU cores allocated and used for all estimator operations. Attributes ---------- @@ -1779,6 +1822,7 @@ def __init__( verbose=False, output_dir=None, auto_sort=True, + n_cores=None, ): super().__init__( n_features=n_features, @@ -1793,6 +1837,7 @@ def __init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) self.n_pairs = n_pairs self.specific_pairs = specific_pairs @@ -2138,6 +2183,8 @@ class KhiopsRegressor(RegressorMixin, KhiopsPredictor): affects the [fit][] and [predict][] methods. *Note* The sort by key is performed in a left-to-right, hierarchical, lexicographic manner. + n_cores : int, optional + Maximum number of CPU cores allocated and used for all estimator operations. Attributes ---------- @@ -2173,6 +2220,7 @@ def __init__( verbose=False, output_dir=None, auto_sort=True, + n_cores=None, ): super().__init__( n_features=n_features, @@ -2187,6 +2235,7 @@ def __init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) self._khiops_model_prefix = "SNB_" self._khiops_baseline_model_prefix = "B_" @@ -2390,6 +2439,8 @@ class KhiopsEncoder(TransformerMixin, KhiopsSupervisedEstimator): affects the [fit][] and [transform][] methods. *Note* The sort by key is performed in a left-to-right, hierarchical, lexicographic manner. + n_cores : int, optional + Maximum number of CPU cores allocated and used for all estimator operations. Attributes ---------- @@ -2431,6 +2482,7 @@ def __init__( verbose=False, output_dir=None, auto_sort=True, + n_cores=None, ): super().__init__( n_features=n_features, @@ -2442,6 +2494,7 @@ def __init__( verbose=verbose, output_dir=output_dir, auto_sort=auto_sort, + n_cores=n_cores, ) self.n_pairs = n_pairs self.specific_pairs = specific_pairs diff --git a/tests/test_sklearn.py b/tests/test_sklearn.py index 9eb2f117..84bb1928 100644 --- a/tests/test_sklearn.py +++ b/tests/test_sklearn.py @@ -709,7 +709,8 @@ def setUpClass(cls): ("khiops.core", "train_coclustering"): { "log_file_path": os.path.join( cls.output_dir, "khiops_train_cc.log" - ) + ), + "max_cores": 62, }, ("khiops.core", "simplify_coclustering"): { "max_part_numbers": {"SampleId": 2}, @@ -767,6 +768,7 @@ def setUpClass(cls): "group_target_value": False, "additional_data_tables": {}, "keep_selected_variables_only": False, + "max_cores": 63, } }, "predict": { @@ -797,6 +799,7 @@ def setUpClass(cls): "max_parts": 5, "keep_selected_variables_only": False, "additional_data_tables": {}, + "max_cores": 65, } }, "predict": { @@ -834,6 +837,7 @@ def setUpClass(cls): "numerical_recoding_method": "part Id", "pairs_recoding_method": "part Id", "additional_data_tables": {}, + "max_cores": 67, } }, "predict": { @@ -1452,6 +1456,7 @@ def test_parameter_transfer_classifier_fit_from_monotable_dataframe(self): "n_feature_parts": 3, "group_target_value": False, "keep_selected_variables_only": False, + "n_cores": 63, }, ) @@ -1478,6 +1483,7 @@ def test_parameter_transfer_classifier_fit_from_monotable_dataframe_with_df_y( "n_feature_parts": 3, "group_target_value": False, "keep_selected_variables_only": False, + "n_cores": 63, }, ) @@ -1545,6 +1551,7 @@ def test_parameter_transfer_encoder_fit_from_monotable_dataframe(self): "transform_type_categorical": "part_id", "transform_type_numerical": "part_id", "transform_type_pairs": "part_id", + "n_cores": 67, }, ) @@ -1573,6 +1580,7 @@ def test_parameter_transfer_encoder_fit_from_monotable_dataframe_with_df_y( "transform_type_categorical": "part_id", "transform_type_numerical": "part_id", "transform_type_pairs": "part_id", + "n_cores": 67, }, ) @@ -1636,6 +1644,7 @@ def test_parameter_transfer_regressor_fit_from_monotable_dataframe(self): "construction_rules": ["TableMode", "TableSelection"], "n_feature_parts": 5, "keep_selected_variables_only": False, + "n_cores": 65, }, ) @@ -1657,6 +1666,7 @@ def test_parameter_transfer_regressor_fit_from_monotable_dataframe_with_df_y( "construction_rules": ["TableMode", "TableSelection"], "n_feature_parts": 5, "keep_selected_variables_only": False, + "n_cores": 65, }, ) @@ -1704,6 +1714,9 @@ def test_parameter_transfer_coclustering_fit_from_dataframe(self): estimator_method="fit", schema_type="not_applicable", source_type="dataframe", + extra_estimator_kwargs={ + "n_cores": 62, + }, custom_kwargs={ "fit": { "columns": ("SampleId", "Pos", "Char"),