Fix num_threads param, check for n_jobs param (#118)

Author: Yard1

.github/workflows/test.yaml

@@ -30,7 +30,7 @@ jobs:
 
   test_linux_ray_master:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 40
     strategy:
       matrix:
         python-version: [3.6.9, 3.7, 3.8]
@@ -68,7 +68,7 @@ jobs:
 
   test_linux_ray_release:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 40
     strategy:
       matrix:
         python-version: [3.6.9, 3.7, 3.8]
@@ -101,7 +101,7 @@ jobs:
     # Test compatibility when some optional libraries are missing
     # Test runs on latest ray release
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 40
     strategy:
       matrix:
         python-version: [3.6.9, 3.7, 3.8]
@@ -138,7 +138,7 @@ jobs:
   test_linux_cutting_edge:
     # Tests on cutting edge, i.e. latest Ray master, latest XGBoost master
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 40
     strategy:
       matrix:
         python-version: [3.6.9, 3.7, 3.8]
@@ -194,7 +194,7 @@ jobs:
   test_linux_xgboost_legacy:
     # Tests on XGBoost 0.90 and latest Ray release
     runs-on: ubuntu-latest
-    timeout-minutes: 28
+    timeout-minutes: 38
     strategy:
       matrix:
         python-version: [3.6.9]

xgboost_ray/main.py

@@ -501,7 +501,8 @@ def train(self, rabit_args: List[str], return_bst: bool,
 
         if "nthread" not in local_params and "n_jobs" not in local_params:
             if num_threads > 0:
-                local_params["num_threads"] = num_threads
+                local_params["nthread"] = num_threads
+                local_params["n_jobs"] = num_threads
             else:
                 local_params["nthread"] = sum(
                     num
@@ -838,14 +839,16 @@ def _train(params: Dict,
     # Un-schedule possible scheduled restarts
     _training_state.restart_training_at = None
 
-    if "nthread" in params:
-        if params["nthread"] > cpus_per_actor:
+    if "nthread" in params or "n_jobs" in params:
+        if ("nthread" in params and params["nthread"] > cpus_per_actor) or (
+                "n_jobs" in params and params["n_jobs"] > cpus_per_actor):
             raise ValueError(
                 "Specified number of threads greater than number of CPUs. "
                 "\nFIX THIS by passing a lower value for the `nthread` "
                 "parameter or a higher number for `cpus_per_actor`.")
     else:
         params["nthread"] = cpus_per_actor
+        params["n_jobs"] = cpus_per_actor
 
     # This is a callback that handles actor failures.
     # We identify the rank of the failed actor, add this to a set of