Fix wrapping GPU ID and prevent data copying. (#5160)

* Removed some data copying.

* Make sure gpu_id is valid before any configuration is carried out.

tests/cpp/data/test_sparse_page_dmatrix.cc

@@ -51,14 +51,14 @@ TEST(SparsePageDMatrix, ColAccess) {
   EXPECT_EQ(dmat->GetColDensity(1), 0.5);
 
   // Loop over the batches and assert the data is as expected
-  for (auto col_batch : dmat->GetBatches<xgboost::SortedCSCPage>()) {
+  for (auto const& col_batch : dmat->GetBatches<xgboost::SortedCSCPage>()) {
     EXPECT_EQ(col_batch.Size(), dmat->Info().num_col_);
     EXPECT_EQ(col_batch[1][0].fvalue, 10.0f);
     EXPECT_EQ(col_batch[1].size(), 1);
   }
 
   // Loop over the batches and assert the data is as expected
-  for (auto col_batch : dmat->GetBatches<xgboost::CSCPage>()) {
+  for (auto const& col_batch : dmat->GetBatches<xgboost::CSCPage>()) {
     EXPECT_EQ(col_batch.Size(), dmat->Info().num_col_);
     EXPECT_EQ(col_batch[1][0].fvalue, 10.0f);
     EXPECT_EQ(col_batch[1].size(), 1);
@@ -223,7 +223,7 @@ TEST(SparsePageDMatrix, FromFile) {
   const std::string tmp_file = tempdir.path + "/simple.libsvm";
   data::SparsePageDMatrix dmat(
       &adapter, std::numeric_limits<float>::quiet_NaN(), -1, tmp_file, 1);
-  
+
   for (auto &batch : dmat.GetBatches<SparsePage>()) {
     std::vector<bst_row_t> expected_offset(batch.Size() + 1);
     int n = -3;

tests/cpp/helpers.cc

@@ -46,7 +46,7 @@ void CheckObjFunctionImpl(std::unique_ptr<xgboost::ObjFunction> const& obj,
                           std::vector<xgboost::bst_float> preds,
                           std::vector<xgboost::bst_float> labels,
                           std::vector<xgboost::bst_float> weights,
-                          xgboost::MetaInfo info,
+                          xgboost::MetaInfo const& info,
                           std::vector<xgboost::bst_float> out_grad,
                           std::vector<xgboost::bst_float> out_hess) {
   xgboost::HostDeviceVector<xgboost::bst_float> in_preds(preds);

tests/python-gpu/load_pickle.py

@@ -37,3 +37,15 @@ class TestLoadPickle(unittest.TestCase):
         config = json.loads(config)
         assert config['learner']['gradient_booster']['gbtree_train_param'][
             'predictor'] == 'gpu_predictor'
+
+    def test_wrap_gpu_id(self):
+        assert os.environ['CUDA_VISIBLE_DEVICES'] == '0'
+        bst = load_pickle(model_path)
+        config = bst.save_config()
+        config = json.loads(config)
+        assert config['learner']['generic_param']['gpu_id'] == '0'
+
+        x, y = build_dataset()
+        test_x = xgb.DMatrix(x)
+        res = bst.predict(test_x)
+        assert len(res) == 10

tests/python-gpu/test_gpu_pickling.py

@@ -5,6 +5,9 @@ import numpy as np
 import subprocess
 import os
 import json
+import pytest
+import copy
+
 import xgboost as xgb
 from xgboost import XGBClassifier
 
@@ -31,6 +34,12 @@ def load_pickle(path):
 
 
 class TestPickling(unittest.TestCase):
+    args_template = [
+        "pytest",
+        "--verbose",
+        "-s",
+        "--fulltrace"]
+
     def test_pickling(self):
         x, y = build_dataset()
         train_x = xgb.DMatrix(x, label=y)
@@ -61,13 +70,29 @@ class TestPickling(unittest.TestCase):
         assert status == 0
         os.remove(model_path)
 
-    def test_pickled_predictor(self):
-        args_templae = [
-            "pytest",
-            "--verbose",
-            "-s",
-            "--fulltrace"]
+    @pytest.mark.mgpu
+    def test_wrap_gpu_id(self):
+        X, y = build_dataset()
+        dtrain = xgb.DMatrix(X, y)
 
+        bst = xgb.train({'tree_method': 'gpu_hist',
+                         'gpu_id': 1},
+                        dtrain, num_boost_round=6)
+
+        model_path = 'model.pkl'
+        save_pickle(bst, model_path)
+        cuda_environment = {'CUDA_VISIBLE_DEVICES': '0'}
+        env = os.environ.copy()
+        env.update(cuda_environment)
+        args = self.args_template.copy()
+        args.append(
+            "./tests/python-gpu/"
+            "load_pickle.py::TestLoadPickle::test_wrap_gpu_id"
+        )
+        status = subprocess.call(args, env=env)
+        assert status == 0
+
+    def test_pickled_predictor(self):
         x, y = build_dataset()
         train_x = xgb.DMatrix(x, label=y)
 
@@ -80,7 +105,7 @@ class TestPickling(unittest.TestCase):
 
         save_pickle(bst, model_path)
 
-        args = args_templae.copy()
+        args = self.args_template.copy()
         args.append(
             "./tests/python-gpu/"
             "load_pickle.py::TestLoadPickle::test_predictor_type_is_auto")
@@ -93,7 +118,7 @@ class TestPickling(unittest.TestCase):
         status = subprocess.call(args, env=env)
         assert status == 0
 
-        args = args_templae.copy()
+        args = self.args_template.copy()
         args.append(
             "./tests/python-gpu/"
             "load_pickle.py::TestLoadPickle::test_predictor_type_is_gpu")
@@ -109,7 +134,6 @@ class TestPickling(unittest.TestCase):
 
         kwargs = {'tree_method': 'gpu_hist',
                   'predictor': 'gpu_predictor',
-                  'verbosity': 1,
                   'objective': 'binary:logistic',
                   'n_estimators': 10}