Handle OMP_THREAD_LIMIT. (#7390) (#7391)

tests/python/test_openmp.py

@@ -1,6 +1,12 @@
-# -*- coding: utf-8 -*-
+import os
+import tempfile
+import subprocess
+
 import xgboost as xgb
 import numpy as np
+import pytest
+
+import testing as tm
 
 
 class TestOMP:
@@ -71,3 +77,31 @@ class TestOMP:
         assert auc_1 == auc_2 == auc_3
         assert np.array_equal(auc_1, auc_2)
         assert np.array_equal(auc_1, auc_3)
+
+    @pytest.mark.skipif(**tm.no_sklearn())
+    def test_with_omp_thread_limit(self):
+        args = [
+            "python", os.path.join(
+                tm.PROJECT_ROOT, "tests", "python", "with_omp_limit.py"
+            )
+        ]
+        results = []
+        with tempfile.TemporaryDirectory() as tmpdir:
+            for i in (1, 2, 16):
+                path = os.path.join(tmpdir, str(i))
+                with open(path, "w") as fd:
+                    fd.write("\n")
+                cp = args.copy()
+                cp.append(path)
+
+                env = os.environ.copy()
+                env["OMP_THREAD_LIMIT"] = str(i)
+
+                status = subprocess.call(cp, env=env)
+                assert status == 0
+
+                with open(path, "r") as fd:
+                    results.append(float(fd.read()))
+
+        for auc in results:
+            np.testing.assert_allclose(auc, results[0])

tests/python/with_omp_limit.py

@@ -0,0 +1,26 @@
+import os
+import xgboost as xgb
+from sklearn.datasets import make_classification
+from sklearn.metrics import roc_auc_score
+import sys
+
+
+def run_omp(output_path: str):
+    X, y = make_classification(
+        n_samples=200, n_features=32, n_classes=3, n_informative=8
+    )
+    Xy = xgb.DMatrix(X, y, nthread=16)
+    booster = xgb.train(
+        {"num_class": 3, "objective": "multi:softprob", "n_jobs": 16},
+        Xy,
+        num_boost_round=8,
+    )
+    score = booster.predict(Xy)
+    auc = roc_auc_score(y, score, average="weighted", multi_class="ovr")
+    with open(output_path, "w") as fd:
+        fd.write(str(auc))
+
+
+if __name__ == "__main__":
+    out = sys.argv[1]
+    run_omp(out)