Check specified n_gpus.

src/common/common.h

@@ -155,24 +155,31 @@ class GPUSet {
 
   static GPUSet Empty() { return GPUSet(); }
 
-  static GPUSet Range(GpuIndex start, GpuIndex n_gpu) {
-    return n_gpu <=0 ? Empty() : GPUSet{start, n_gpu};
+  static GPUSet Range(GpuIndex start, GpuIndex n_gpus) {
+    return n_gpus <=0 ? Empty() : GPUSet{start, n_gpus};
   }
   /*! \brief ndevices and num_rows both are upper bounds. */
-  static GPUSet All(GpuIndex gpu_id, GpuIndex n_gpu,
+  static GPUSet All(GpuIndex gpu_id, GpuIndex n_gpus,
                     GpuIndex num_rows = std::numeric_limits<GpuIndex>::max()) {
+    CHECK_GE(gpu_id, 0) << "gpu_id must be >= 0.";
+    CHECK_GE(n_gpus, -1) << "n_gpus must be >= -1.";
     GpuIndex n_devices_visible = AllVisible().Size();
-    if (n_gpu == kAll) {  // Use all devices starting from `gpu_id'.
+    if (n_gpus == kAll) {  // Use all devices starting from `gpu_id'.
       CHECK(gpu_id < n_devices_visible || gpu_id == 0)
           << "\ngpu_id should be less than available devices.\ngpu_id: "
           << gpu_id
           << ", number of available devices: "
-          << n_devices_visible << std::endl;
+          << n_devices_visible;
+      GpuIndex n_devices = n_devices_visible - gpu_id;
+      n_devices = n_devices < num_rows ? n_devices : num_rows;
       return Range(gpu_id, n_devices_visible - gpu_id);
-    } else {
+    } else {  // Use devices in ( gpu_id, gpu_id + n_gpus ).
       GpuIndex n_available_devices = n_devices_visible - gpu_id;
-      GpuIndex n_devices =
-          n_available_devices < n_gpu ? n_devices_visible : n_gpu;
+      CHECK_LE(n_gpus, n_available_devices)
+          << "Starting from gpu id: " << gpu_id << ", there are only "
+          << n_available_devices << " available devices, while n_gpus is set to: "
+          << n_gpus;
+      GpuIndex n_devices = n_gpus < num_rows ? n_gpus : num_rows;
       return Range(gpu_id, n_devices);
     }
   }

src/predictor/gpu_predictor.cu

@@ -1,5 +1,5 @@
 /*!
- * Copyright by Contributors 2017
+ * Copyright 2017-2018 by Contributors
  */
 #include <dmlc/parameter.h>
 #include <thrust/copy.h>