Fix xgboost.dask import

Author: jrbourbeau

tests/runtime/test_xgboost.py

@@ -2,7 +2,7 @@
 import pytest
 
 dask_ml = pytest.importorskip("dask_ml")
-xgb = pytest.importorskip("xgboost")
+dxgb = pytest.importorskip("xgboost.dask")
 
 
 def test_xgboost_distributed_training(small_client):
@@ -20,8 +20,8 @@ def test_xgboost_distributed_training(small_client):
     )
 
     # Create the XGBoost DMatrix for our training and testing splits
-    dtrain = xgb.dask.DaskDMatrix(small_client, X_train, y_train)
-    dtest = xgb.dask.DaskDMatrix(small_client, X_test, y_test)
+    dtrain = dxgb.DaskDMatrix(small_client, X_train, y_train)
+    dtest = dxgb.DaskDMatrix(small_client, X_test, y_test)
 
     # Set model parameters (XGBoost defaults)
     params = {
@@ -32,10 +32,10 @@ def test_xgboost_distributed_training(small_client):
         "objective": "reg:squarederror",
         "grow_policy": "depthwise",
     }
-    output = xgb.dask.train(
+    output = dxgb.train(
         small_client, params, dtrain, num_boost_round=5, evals=[(dtrain, "train")]
     )
 
     # make predictions
-    y_pred = xgb.dask.predict(small_client, output, dtest)
+    y_pred = dxgb.predict(small_client, output, dtest)
     assert y_pred.shape[0] == y_test.shape[0].compute()