Add max_cores to the sklearn fit methods in estimators to limit the number of CPU cores allocated to the training

CHANGELOG.md

@@ -6,6 +6,12 @@
   - Example: 10.2.1.4 is the 5th version that supports khiops 10.2.1.
 - Internals: Changes in *Internals* sections are unlikely to be of interest for data scientists.
 
+## Unreleased
+
+### Changed
+- (`sklearn`) The `fit` methods in estimators accept `max_cores` to limit the number of CPU cores allocated to the training
+- (General) When `max_cores` is set, use also its value to limit the number of pre-allocated CPU cores (instead of using all the available ones).
+
 ## 11.0.1.0 - 2026-07-02
 
 ### Added

khiops/core/api.py

@@ -193,7 +193,12 @@ def _preprocess_arguments(args):
         if arg == "max_cores":
             max_cores = args[arg]
             if max_cores is not None:
+                # This `max_cores` system setting will be used in the khiops scenario
+                # to limit the CPU cores to use for the training
                 system_settings.max_cores = int(max_cores)
+                # An additional environment variable (local to this specific run)
+                # MUST also be set to avoid allocating all the available CPU cores.
+                # Thus, allocated CPU cores = max number of CPU cores used
         elif arg == "memory_limit_mb":
             memory_limit_mb = args[arg]
             if memory_limit_mb is not None:

khiops/core/internals/runner.py

@@ -332,18 +332,30 @@ def _get_current_library_installer():
         return "unknown"
 
 
-def _build_khiops_process_environment():
+def _build_khiops_process_environment(system_settings=None):
     """Build a specific environment used for the execution of khiops in a process
 
     This environment can be modified freely without interfering
     with the global one.
+
+    Parameters
+    ----------
+
+    system_settings: `SystemSettings`
+            Set of settings that must be taken into account
+            for this specific run
     """
     khiops_env = os.environ.copy()
 
     # Ensure HOME is always set for OpenMPI 5+
     # (using KHIOPS_MPI_HOME if it exists)
     if "HOME" not in khiops_env:
         khiops_env["HOME"] = khiops_env.get("KHIOPS_MPI_HOME", "")
+    if system_settings is not None and system_settings.max_cores is not None:
+        # An additional environment variable (local to this specific run)
+        # must also be set to avoid allocating all the available CPU cores.
+        # Thus, allocated CPU cores = max number of CPU cores used
+        khiops_env["KHIOPS_PROC_NUMBER"] = system_settings.max_cores
     return khiops_env
 
 
@@ -690,6 +702,7 @@ class for more information.
                 scenario_path,
                 command_line_options,
                 trace,
+                system_settings,
             )
             # pylint: enable=assignment-from-no-return
         # Catch an OS level error if any
@@ -895,6 +908,7 @@ def _run(
         scenario_path,
         command_line_options,
         trace,
+        system_settings,
     ):
         """Abstract run method to be implemented in child classes
 
@@ -1449,7 +1463,14 @@ def _get_samples_dir(self):
             self._samples_dir_checked = True
         return self._samples_dir
 
-    def raw_run(self, tool_name, command_line_args=None, use_mpi=True, trace=False):
+    def raw_run(
+        self,
+        tool_name,
+        command_line_args=None,
+        use_mpi=True,
+        trace=False,
+        system_settings=None,
+    ):
         """Execute a Khiops tool with given command line arguments
 
         Parameters
@@ -1462,6 +1483,9 @@ def raw_run(self, tool_name, command_line_args=None, use_mpi=True, trace=False):
             Whether to execute the application with MPI
         trace : bool, default False
             If ``True`` print the trace of the process.
+        system_settings: `SystemSettings`
+            Set of settings that must be taken into account
+            for this specific run
 
         Examples
         --------
@@ -1499,9 +1523,9 @@ def raw_run(self, tool_name, command_line_args=None, use_mpi=True, trace=False):
             print(f"Khiops execution call: {khiops_call}")
 
         # Build custom Khiops process environment
-        # which makes sure HOME is defined and set
+        # which makes sure for example HOME is defined and set
         # according to khiops_env's KHIOPS_MPI_HOME
-        khiops_env = _build_khiops_process_environment()
+        khiops_env = _build_khiops_process_environment(system_settings)
 
         # Execute the process
         with subprocess.Popen(
@@ -1524,11 +1548,15 @@ def _run(
         scenario_path,
         command_line_options,
         trace,
+        system_settings,
     ):
         # Execute the tool
         khiops_args = command_line_options.build_command_line_options(scenario_path)
         stdout, stderr, return_code = self.raw_run(
-            tool_name, command_line_args=khiops_args, trace=trace
+            tool_name,
+            command_line_args=khiops_args,
+            trace=trace,
+            system_settings=system_settings,
         )
 
         return return_code, stdout, stderr

khiops/extras/docker.py

@@ -103,6 +103,7 @@ def _run(
         scenario_path,
         command_line_options,
         trace,
+        system_settings,
     ):
         # Check arguments
         if command_line_options.output_scenario_path:

khiops/sklearn/estimators.py

@@ -348,6 +348,18 @@ def _cleanup_computation_dir(self, computation_dir):
     def fit(self, X, y=None, **kwargs):
         """Fit the estimator
 
+        Parameters
+        ----------
+        X : :external:term:`array-like` of shape (n_samples, n_features_in) or dict
+            Training dataset. Either an :external:term:`array-like` or a ``dict``
+            specification for multi-table datasets (see :doc:`/multi_table_primer`).
+
+        y : :external:term:`array-like` of shape (n_samples,)
+            The target values.
+
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsEstimator`
@@ -689,6 +701,8 @@ def fit(self, X, y=None, **kwargs):
             The column that contains the id of the instance.
         columns : list, optional
             The columns to be co-clustered. If not specified it uses all columns.
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
 
         Returns
         -------
@@ -764,6 +778,7 @@ def _fit_train_model(self, ds, computation_dir, **kwargs):
             main_table_path,
             variables,
             coclustering_file_path,
+            max_cores=kwargs.get("max_cores"),
             log_file_path=train_log_file_path,
             trace=self.verbose,
         )
@@ -1179,7 +1194,25 @@ 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 : :external:term:`array-like` of shape (n_samples, n_features_in) or dict
+            Training dataset. Either an :external:term:`array-like` or a ``dict``
+            specification for multi-table datasets (see :doc:`/multi_table_primer`).
+
+        y : :external:term:`array-like` of shape (n_samples,)
+            The target values.
+
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
+        Returns
+        -------
+        results : `numpy.array`
+        """
+
         return self.fit(X, y, **kwargs).predict(X)
 
 
@@ -1253,6 +1286,9 @@ def fit(self, X, y=None, **kwargs):
         y : :external:term:`array-like` of shape (n_samples,)
             The target values.
 
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsSupervisedEstimator`
@@ -1314,7 +1350,7 @@ def _fit_check_params(self, ds, **kwargs):
     def _fit_train_model(self, ds, computation_dir, **kwargs):
         # Train the model with Khiops
         train_args, train_kwargs = self._fit_prepare_training_function_inputs(
-            ds, computation_dir
+            ds, computation_dir, **kwargs
         )
         report_file_path, model_kdic_file_path = self._fit_core_training_function(
             *train_args, **train_kwargs
@@ -1335,7 +1371,7 @@ def _fit_train_model(self, ds, computation_dir, **kwargs):
     def _fit_core_training_function(self, *args, **kwargs):
         """A wrapper to the khiops.core training function for the estimator"""
 
-    def _fit_prepare_training_function_inputs(self, ds, computation_dir):
+    def _fit_prepare_training_function_inputs(self, ds, computation_dir, **fit_kwargs):
         # Set output path files
         output_dir = self._get_output_dir(computation_dir)
         report_file_path = fs.get_child_path(
@@ -1374,7 +1410,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
@@ -1404,6 +1441,10 @@ def _fit_prepare_training_function_inputs(self, ds, computation_dir):
         kwargs["trace"] = kwargs["verbose"]
         del kwargs["verbose"]
 
+        # Set the technical parameters
+        if "max_cores" in fit_kwargs:
+            kwargs["max_cores"] = fit_kwargs["max_cores"]
+
         return args, kwargs
 
     def _fit_training_post_process(self, ds):
@@ -1584,10 +1625,10 @@ def predict(self, X):
         assert isinstance(y_pred, (str, pd.DataFrame)), "Expected str or DataFrame"
         return y_pred
 
-    def _fit_prepare_training_function_inputs(self, ds, computation_dir):
+    def _fit_prepare_training_function_inputs(self, ds, computation_dir, **fit_kwargs):
         # Call the parent method
         args, kwargs = super()._fit_prepare_training_function_inputs(
-            ds, computation_dir
+            ds, computation_dir, **fit_kwargs
         )
 
         # Rename parameters to be compatible with khiops.core
@@ -1855,10 +1896,10 @@ def _fit_check_params(self, ds, **kwargs):
         # Check the pair related parameters
         _check_pair_parameters(self)
 
-    def _fit_prepare_training_function_inputs(self, ds, computation_dir):
+    def _fit_prepare_training_function_inputs(self, ds, computation_dir, **fit_kwargs):
         # Call the parent method
         args, kwargs = super()._fit_prepare_training_function_inputs(
-            ds, computation_dir
+            ds, computation_dir, **fit_kwargs
         )
 
         # Rename parameters to be compatible with khiops.core
@@ -1878,6 +1919,9 @@ def fit(self, X, y, **kwargs):
         y : :external:term:`array-like` of shape (n_samples,)
             The target values.
 
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsClassifier`
@@ -2200,6 +2244,9 @@ def fit(self, X, y=None, **kwargs):
         y : :external:term:`array-like` of shape (n_samples,)
             The target values.
 
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsRegressor`
@@ -2606,6 +2653,9 @@ def fit(self, X, y=None, **kwargs):
         y : :external:term:`array-like` of shape (n_samples,)
             The target values.
 
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsEncoder`
@@ -2616,10 +2666,10 @@ def fit(self, X, y=None, **kwargs):
 
     # pylint: enable=useless-super-delegation
 
-    def _fit_prepare_training_function_inputs(self, ds, computation_dir):
+    def _fit_prepare_training_function_inputs(self, ds, computation_dir, **fit_kwargs):
         # Call the parent method
         args, kwargs = super()._fit_prepare_training_function_inputs(
-            ds, computation_dir
+            ds, computation_dir, **fit_kwargs
         )
         # Rename encoder parameters, delete unused ones
         # to be compatible with khiops.core
@@ -2725,6 +2775,9 @@ def fit_transform(self, X, y=None, **kwargs):
         y : :external:term:`array-like` of shape (n_samples,)
             The target values.
 
+        max_cores : int, optional
+            Maximum number of CPU cores allocated and used for the training.
+
         Returns
         -------
         self : `KhiopsEncoder`

tests/test_core.py

@@ -2639,9 +2639,11 @@ def run(
         self,
         task,
         task_args,
-        command_line_options,
+        command_line_options=None,
         trace=False,
         system_settings=None,
+        stdout_file_path="",
+        stderr_file_path="",
         force_ansi_scenario=False,
         **kwargs,
     ):
@@ -2677,6 +2679,7 @@ def _run(
         scenario_path,
         command_line_options,
         trace,
+        system_settings,
     ):
         return 0, "", ""