[R-package] [c++] add tighter multithreading control, avoid global Op…

…enMP side effects (fixes #4705, fixes #5102) (#6226)
microsoft · Dec 7, 2023 · 1548b42 · 1548b42
1 parent e797985
commit 1548b42
Show file tree

Hide file tree

Showing 58 changed files with 429 additions and 21 deletions.
diff --git a/.ci/lint-cpp.sh b/.ci/lint-cpp.sh
@@ -30,8 +30,7 @@ get_omp_pragmas_without_num_threads() {
         --include='*.h' \
         --include='*.hpp' \
         'pragma omp parallel' \
-    | grep -v ' num_threads' \
-    | grep -v 'openmp_wrapper.h'
+    | grep -v ' num_threads'
 }
 PROBLEMATIC_LINES=$(
     get_omp_pragmas_without_num_threads

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -432,6 +432,7 @@ file(
       src/objective/*.cpp
       src/network/*.cpp
       src/treelearner/*.cpp
+      src/utils/*.cpp
 if(USE_CUDA)
       src/treelearner/*.cu
       src/boosting/cuda/*.cpp

diff --git a/R-package/NAMESPACE b/R-package/NAMESPACE
@@ -9,6 +9,7 @@ S3method(print,lgb.Booster)
 S3method(set_field,lgb.Dataset)
 S3method(slice,lgb.Dataset)
 S3method(summary,lgb.Booster)
+export(getLGBMthreads)
 export(get_field)
 export(lgb.Dataset)
 export(lgb.Dataset.construct)
@@ -35,6 +36,7 @@ export(lgb.train)
 export(lightgbm)
 export(readRDS.lgb.Booster)
 export(saveRDS.lgb.Booster)
+export(setLGBMthreads)
 export(set_field)
 export(slice)
 import(methods)

diff --git a/R-package/R/lgb.Booster.R b/R-package/R/lgb.Booster.R
@@ -917,6 +917,8 @@ NULL
 #'         the factor levels not being present in the output.
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1082,6 +1084,8 @@ predict.lgb.Booster <- function(object,
 #'         \link{predict.lgb.Booster}.
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' library(lightgbm)
 #' data(mtcars)
 #' X <- as.matrix(mtcars[, -1L])
@@ -1224,6 +1228,8 @@ summary.lgb.Booster <- function(object, ...) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1289,6 +1295,8 @@ lgb.load <- function(filename = NULL, model_str = NULL) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' library(lightgbm)
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
@@ -1346,6 +1354,8 @@ lgb.save <- function(booster, filename, num_iteration = NULL) {
 #' @examples
 #' \donttest{
 #' library(lightgbm)
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1396,6 +1406,8 @@ lgb.dump <- function(booster, num_iteration = NULL) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' # train a regression model
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train

diff --git a/R-package/R/lgb.Dataset.R b/R-package/R/lgb.Dataset.R
@@ -780,6 +780,8 @@ Dataset <- R6::R6Class(
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -837,6 +839,8 @@ lgb.Dataset <- function(data,
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -913,6 +917,8 @@ lgb.Dataset.create.valid <- function(dataset,
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -942,6 +948,8 @@ lgb.Dataset.construct <- function(dataset) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -975,6 +983,8 @@ dim.lgb.Dataset <- function(x) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1045,6 +1055,8 @@ dimnames.lgb.Dataset <- function(x) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1089,6 +1101,8 @@ slice.lgb.Dataset <- function(dataset, idxset) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1138,6 +1152,8 @@ get_field.lgb.Dataset <- function(dataset, field_name) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1177,6 +1193,8 @@ set_field.lgb.Dataset <- function(dataset, field_name, data) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)
@@ -1207,6 +1225,8 @@ lgb.Dataset.set.categorical <- function(dataset, categorical_feature) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' # create training Dataset
 #' data(agaricus.train, package ="lightgbm")
 #' train <- agaricus.train
@@ -1240,6 +1260,8 @@ lgb.Dataset.set.reference <- function(dataset, reference) {
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/lgb.cv.R b/R-package/R/lgb.cv.R
@@ -51,6 +51,8 @@ CVBooster <- R6::R6Class(
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/lgb.importance.R b/R-package/R/lgb.importance.R
@@ -14,6 +14,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/lgb.interprete.R b/R-package/R/lgb.interprete.R
@@ -17,6 +17,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' Logit <- function(x) log(x / (1.0 - x))
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train

diff --git a/R-package/R/lgb.model.dt.tree.R b/R-package/R/lgb.model.dt.tree.R
@@ -29,6 +29,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/lgb.plot.importance.R b/R-package/R/lgb.plot.importance.R
@@ -19,6 +19,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/lgb.plot.interpretation.R b/R-package/R/lgb.plot.interpretation.R
@@ -16,6 +16,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' Logit <- function(x) {
 #'   log(x / (1.0 - x))
 #' }

diff --git a/R-package/R/lgb.restore_handle.R b/R-package/R/lgb.restore_handle.R
@@ -16,7 +16,10 @@
 #' @return \code{lgb.Booster} (the same `model` object that was passed as input, invisibly).
 #' @seealso \link{lgb.make_serializable}, \link{lgb.drop_serialized}.
 #' @examples
+#' \donttest{
 #' library(lightgbm)
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data("agaricus.train")
 #' model <- lightgbm(
 #'   agaricus.train$data
@@ -33,6 +36,7 @@
 #' model_new$check_null_handle()
 #' lgb.restore_handle(model_new)
 #' model_new$check_null_handle()
+#' }
 #' @export
 lgb.restore_handle <- function(model) {
   if (!.is_Booster(x = model)) {

diff --git a/R-package/R/lgb.train.R b/R-package/R/lgb.train.R
@@ -19,6 +19,8 @@
 #'
 #' @examples
 #' \donttest{
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/multithreading.R b/R-package/R/multithreading.R
@@ -0,0 +1,51 @@
+#' @name setLGBMThreads
+#' @title Set maximum number of threads used by LightGBM
+#' @description LightGBM attempts to speed up many operations by using multi-threading.
+#'              The number of threads used in those operations can be controlled via the
+#'              \code{num_threads} parameter passed through \code{params} to functions like
+#'              \link{lgb.train} and \link{lgb.Dataset}. However, some operations (like materializing
+#'              a model from a text file) are done via code paths that don't explicitly accept thread-control
+#'              configuration.
+#'
+#'              Use this function to set the maximum number of threads LightGBM will use for such operations.
+#'
+#'              This function affects all LightGBM operations in the same process.
+#'
+#'              So, for example, if you call \code{setLGBMthreads(4)}, no other multi-threaded LightGBM
+#'              operation in the same process will use more than 4 threads.
+#'
+#'              Call \code{setLGBMthreads(-1)} to remove this limitation.
+#' @param num_threads maximum number of threads to be used by LightGBM in multi-threaded operations
+#' @return NULL
+#' @seealso \link{getLGBMthreads}
+#' @export
+setLGBMthreads <- function(num_threads) {
+    .Call(
+        LGBM_SetMaxThreads_R,
+        num_threads
+    )
+    return(invisible(NULL))
+}
+
+#' @name getLGBMThreads
+#' @title Get default number of threads used by LightGBM
+#' @description LightGBM attempts to speed up many operations by using multi-threading.
+#'              The number of threads used in those operations can be controlled via the
+#'              \code{num_threads} parameter passed through \code{params} to functions like
+#'              \link{lgb.train} and \link{lgb.Dataset}. However, some operations (like materializing
+#'              a model from a text file) are done via code paths that don't explicitly accept thread-control
+#'              configuration.
+#'
+#'              Use this function to see the default number of threads LightGBM will use for such operations.
+#' @return number of threads as an integer. \code{-1} means that in situations where parameter \code{num_threads} is
+#'         not explicitly supplied, LightGBM will choose a number of threads to use automatically.
+#' @seealso \link{setLGBMthreads}
+#' @export
+getLGBMthreads <- function() {
+    out <- 0L
+    .Call(
+        LGBM_GetMaxThreads_R,
+        out
+    )
+    return(out)
+}
diff --git a/R-package/R/readRDS.lgb.Booster.R b/R-package/R/readRDS.lgb.Booster.R
@@ -12,6 +12,8 @@
 #' @examples
 #' \donttest{
 #' library(lightgbm)
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/R/saveRDS.lgb.Booster.R b/R-package/R/saveRDS.lgb.Booster.R
@@ -22,6 +22,8 @@
 #' @examples
 #' \donttest{
 #' library(lightgbm)
+#' \dontshow{setLGBMthreads(2L)}
+#' \dontshow{data.table::setDTthreads(1L)}
 #' data(agaricus.train, package = "lightgbm")
 #' train <- agaricus.train
 #' dtrain <- lgb.Dataset(train$data, label = train$label)

diff --git a/R-package/man/dim.Rd b/R-package/man/dim.Rd
diff --git a/R-package/man/dimnames.lgb.Dataset.Rd b/R-package/man/dimnames.lgb.Dataset.Rd
diff --git a/R-package/man/getLGBMThreads.Rd b/R-package/man/getLGBMThreads.Rd
diff --git a/R-package/man/get_field.Rd b/R-package/man/get_field.Rd