FIX Missing num_threads in some HGBT pranges (#22566)

jeremiedbb · lorentzenchr · web-flow · commit 5c7dac0131eb · 2022-02-21T17:04:14.000+01:00
Co-authored-by: Christian Lorentzen &lt;lorentzen.ch@gmail.com&gt;
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
@@ -330,6 +330,12 @@ Changelog
   :class:`ensemble.ExtraTreesClassifier`.
   :pr:`20803` by :user:`Brian Sun <bsun94>`.
 
+- |Fix| Removed a potential source of CPU oversubscription in
+  :class:`ensemble.HistGradientBoostingClassifier` and
+  :class:`ensemble.HistGradientBoostingRegressor` when CPU resource usage is limited,
+  for instance using cgroups quota in a docker container. :pr:`22566` by
+  :user:`Jérémie du Boisberranger <jeremiedbb>`.
+
 :mod:`sklearn.feature_extraction`
 .................................
 
diff --git a/sklearn/ensemble/_hist_gradient_boosting/splitting.pyx b/sklearn/ensemble/_hist_gradient_boosting/splitting.pyx
@@ -335,7 +335,7 @@ cdef class Splitter:
 
             # map indices from sample_indices to left/right_indices_buffer
             for thread_idx in prange(n_threads, schedule='static',
-                                     chunksize=1):
+                                     chunksize=1, num_threads=n_threads):
                 left_count = 0
                 right_count = 0
 
@@ -377,7 +377,7 @@ cdef class Splitter:
             # sample_indices. This also updates self.partition since
             # sample_indices is a view.
             for thread_idx in prange(n_threads, schedule='static',
-                                     chunksize=1):
+                                     chunksize=1, num_threads=n_threads):
                 memcpy(
                     &sample_indices[left_offset[thread_idx]],
                     &left_indices_buffer[offset_in_buffers[thread_idx]],
