Reduce time for some multi-gpu tests (#8288)

* Faster dask tests

* Reuse AllReducer objects in tests.

* Faster boost from prediction tests.

* Use rmm dask fixture.

* Speed up dask demo.

* mypy

* Format with black.

* mypy

* Clang-tidy

Co-authored-by: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Rory Mitchell
2022-10-04 12:49:33 +02:00
committed by GitHub
parent ca0547bb65
commit d686bf52a6
8 changed files with 337 additions and 336 deletions

View File

@@ -22,8 +22,8 @@ def setup_rmm_pool(request, pytestconfig):
rmm.reinitialize(pool_allocator=True, initial_pool_size=1024*1024*1024,
devices=list(range(get_n_gpus())))
@pytest.fixture(scope='function')
def local_cuda_cluster(request, pytestconfig):
@pytest.fixture(scope='class')
def local_cuda_client(request, pytestconfig):
kwargs = {}
if hasattr(request, 'param'):
kwargs.update(request.param)
@@ -31,13 +31,12 @@ def local_cuda_cluster(request, pytestconfig):
if not has_rmm():
raise ImportError('The --use-rmm-pool option requires the RMM package')
import rmm
from dask_cuda.utils import get_n_gpus
kwargs['rmm_pool_size'] = '2GB'
if tm.no_dask_cuda()['condition']:
raise ImportError('The local_cuda_cluster fixture requires dask_cuda package')
from dask_cuda import LocalCUDACluster
with LocalCUDACluster(**kwargs) as cluster:
yield cluster
from dask.distributed import Client
yield Client(LocalCUDACluster(**kwargs))
def pytest_addoption(parser):
parser.addoption('--use-rmm-pool', action='store_true', default=False, help='Use RMM pool')

View File

@@ -32,8 +32,5 @@ def test_categorical_demo():
@pytest.mark.mgpu
def test_dask_training():
script = os.path.join(tm.PROJECT_ROOT, 'demo', 'dask', 'gpu_training.py')
cmd = ['python', script, '--ddqdm=1']
subprocess.check_call(cmd)
cmd = ['python', script, '--ddqdm=0']
subprocess.check_call(cmd)
cmd = ['python', script]
subprocess.check_call(cmd)

View File

@@ -17,26 +17,26 @@ if sys.platform.startswith("win"):
pytest.skip("Skipping dask tests on Windows", allow_module_level=True)
sys.path.append("tests/python")
import testing as tm # noqa
import testing as tm # noqa
if tm.no_dask_cuda()["condition"]:
pytest.skip(tm.no_dask_cuda()["reason"], allow_module_level=True)
from test_with_dask import run_empty_dmatrix_reg # noqa
from test_with_dask import run_empty_dmatrix_auc # noqa
from test_with_dask import run_auc # noqa
from test_with_dask import run_empty_dmatrix_reg # noqa
from test_with_dask import run_empty_dmatrix_auc # noqa
from test_with_dask import run_auc # noqa
from test_with_dask import run_boost_from_prediction # noqa
from test_with_dask import run_boost_from_prediction_multi_class # noqa
from test_with_dask import run_dask_classifier # noqa
from test_with_dask import run_empty_dmatrix_cls # noqa
from test_with_dask import _get_client_workers # noqa
from test_with_dask import generate_array # noqa
from test_with_dask import kCols as random_cols # noqa
from test_with_dask import suppress # noqa
from test_with_dask import run_tree_stats # noqa
from test_with_dask import run_categorical # noqa
from test_with_dask import make_categorical # noqa
from test_with_dask import run_dask_classifier # noqa
from test_with_dask import run_empty_dmatrix_cls # noqa
from test_with_dask import _get_client_workers # noqa
from test_with_dask import generate_array # noqa
from test_with_dask import kCols as random_cols # noqa
from test_with_dask import suppress # noqa
from test_with_dask import run_tree_stats # noqa
from test_with_dask import run_categorical # noqa
from test_with_dask import make_categorical # noqa
try:
@@ -45,7 +45,7 @@ try:
import xgboost as xgb
from dask.distributed import Client
from dask import array as da
from dask_cuda import LocalCUDACluster
from dask_cuda import LocalCUDACluster, utils
import cudf
except ImportError:
pass
@@ -53,6 +53,7 @@ except ImportError:
def run_with_dask_dataframe(DMatrixT: Type, client: Client) -> None:
import cupy as cp
cp.cuda.runtime.setDevice(0)
X, y, _ = generate_array()
@@ -63,14 +64,16 @@ def run_with_dask_dataframe(DMatrixT: Type, client: Client) -> None:
y = y.map_partitions(cudf.from_pandas)
dtrain = DMatrixT(client, X, y)
out = dxgb.train(client, {'tree_method': 'gpu_hist',
'debug_synchronize': True},
dtrain=dtrain,
evals=[(dtrain, 'X')],
num_boost_round=4)
out = dxgb.train(
client,
{"tree_method": "gpu_hist", "debug_synchronize": True},
dtrain=dtrain,
evals=[(dtrain, "X")],
num_boost_round=4,
)
assert isinstance(out['booster'], dxgb.Booster)
assert len(out['history']['X']['rmse']) == 4
assert isinstance(out["booster"], dxgb.Booster)
assert len(out["history"]["X"]["rmse"]) == 4
predictions = dxgb.predict(client, out, dtrain)
assert isinstance(predictions.compute(), np.ndarray)
@@ -78,27 +81,23 @@ def run_with_dask_dataframe(DMatrixT: Type, client: Client) -> None:
series_predictions = dxgb.inplace_predict(client, out, X)
assert isinstance(series_predictions, dd.Series)
single_node = out['booster'].predict(xgboost.DMatrix(X.compute()))
single_node = out["booster"].predict(xgboost.DMatrix(X.compute()))
cp.testing.assert_allclose(single_node, predictions.compute())
np.testing.assert_allclose(single_node,
series_predictions.compute().to_numpy())
np.testing.assert_allclose(single_node, series_predictions.compute().to_numpy())
predt = dxgb.predict(client, out, X)
assert isinstance(predt, dd.Series)
T = TypeVar('T')
T = TypeVar("T")
def is_df(part: T) -> T:
assert isinstance(part, cudf.DataFrame), part
return part
predt.map_partitions(
is_df,
meta=dd.utils.make_meta({'prediction': 'f4'}))
predt.map_partitions(is_df, meta=dd.utils.make_meta({"prediction": "f4"}))
cp.testing.assert_allclose(
predt.values.compute(), single_node)
cp.testing.assert_allclose(predt.values.compute(), single_node)
# Make sure the output can be integrated back to original dataframe
X["predict"] = predictions
@@ -110,49 +109,35 @@ def run_with_dask_dataframe(DMatrixT: Type, client: Client) -> None:
def run_with_dask_array(DMatrixT: Type, client: Client) -> None:
import cupy as cp
cp.cuda.runtime.setDevice(0)
X, y, _ = generate_array()
X = X.map_blocks(cp.asarray)
y = y.map_blocks(cp.asarray)
dtrain = DMatrixT(client, X, y)
out = dxgb.train(client, {'tree_method': 'gpu_hist',
'debug_synchronize': True},
dtrain=dtrain,
evals=[(dtrain, 'X')],
num_boost_round=2)
out = dxgb.train(
client,
{"tree_method": "gpu_hist", "debug_synchronize": True},
dtrain=dtrain,
evals=[(dtrain, "X")],
num_boost_round=2,
)
from_dmatrix = dxgb.predict(client, out, dtrain).compute()
inplace_predictions = dxgb.inplace_predict(
client, out, X).compute()
single_node = out['booster'].predict(
xgboost.DMatrix(X.compute()))
inplace_predictions = dxgb.inplace_predict(client, out, X).compute()
single_node = out["booster"].predict(xgboost.DMatrix(X.compute()))
np.testing.assert_allclose(single_node, from_dmatrix)
device = cp.cuda.runtime.getDevice()
assert device == inplace_predictions.device.id
single_node = cp.array(single_node)
assert device == single_node.device.id
cp.testing.assert_allclose(
single_node,
inplace_predictions)
@pytest.mark.skipif(**tm.no_dask_cudf())
def test_categorical(local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
import dask_cudf
X, y = make_categorical(client, 10000, 30, 13)
X = dask_cudf.from_dask_dataframe(X)
X_onehot, _ = make_categorical(client, 10000, 30, 13, True)
X_onehot = dask_cudf.from_dask_dataframe(X_onehot)
run_categorical(client, "gpu_hist", X, X_onehot, y)
cp.testing.assert_allclose(single_node, inplace_predictions)
def to_cp(x: Any, DMatrixT: Type) -> Any:
import cupy
if isinstance(x, np.ndarray) and \
DMatrixT is dxgb.DaskDeviceQuantileDMatrix:
if isinstance(x, np.ndarray) and DMatrixT is dxgb.DaskDeviceQuantileDMatrix:
X = cupy.array(x)
else:
X = x
@@ -213,217 +198,250 @@ def run_gpu_hist(
assert tm.non_increasing(history)
@pytest.mark.skipif(**tm.no_cudf())
def test_boost_from_prediction(local_cuda_cluster: LocalCUDACluster) -> None:
import cudf
from sklearn.datasets import load_breast_cancer, load_digits
with Client(local_cuda_cluster) as client:
X_, y_ = load_breast_cancer(return_X_y=True)
X = dd.from_array(X_, chunksize=100).map_partitions(cudf.from_pandas)
y = dd.from_array(y_, chunksize=100).map_partitions(cudf.from_pandas)
run_boost_from_prediction(X, y, "gpu_hist", client)
def test_tree_stats() -> None:
with LocalCUDACluster(n_workers=1) as cluster:
with Client(cluster) as client:
local = run_tree_stats(client, "gpu_hist")
X_, y_ = load_digits(return_X_y=True)
X = dd.from_array(X_, chunksize=100).map_partitions(cudf.from_pandas)
y = dd.from_array(y_, chunksize=100).map_partitions(cudf.from_pandas)
run_boost_from_prediction_multi_class(X, y, "gpu_hist", client)
with LocalCUDACluster(n_workers=2) as cluster:
with Client(cluster) as client:
distributed = run_tree_stats(client, "gpu_hist")
assert local == distributed
class TestDistributedGPU:
@pytest.mark.skipif(**tm.no_cudf())
def test_boost_from_prediction(self, local_cuda_client: Client) -> None:
import cudf
from sklearn.datasets import load_breast_cancer, load_iris
X_, y_ = load_breast_cancer(return_X_y=True)
X = dd.from_array(X_, chunksize=100).map_partitions(cudf.from_pandas)
y = dd.from_array(y_, chunksize=100).map_partitions(cudf.from_pandas)
run_boost_from_prediction(X, y, "gpu_hist", local_cuda_client)
X_, y_ = load_iris(return_X_y=True)
X = dd.from_array(X_, chunksize=50).map_partitions(cudf.from_pandas)
y = dd.from_array(y_, chunksize=50).map_partitions(cudf.from_pandas)
run_boost_from_prediction_multi_class(X, y, "gpu_hist", local_cuda_client)
@pytest.mark.skipif(**tm.no_dask_cudf())
def test_dask_dataframe(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
run_with_dask_dataframe(dxgb.DaskDMatrix, client)
run_with_dask_dataframe(dxgb.DaskDeviceQuantileDMatrix, client)
def test_dask_dataframe(self, local_cuda_client: Client) -> None:
run_with_dask_dataframe(dxgb.DaskDMatrix, local_cuda_client)
run_with_dask_dataframe(dxgb.DaskDeviceQuantileDMatrix, local_cuda_client)
@pytest.mark.skipif(**tm.no_dask_cudf())
def test_categorical(self, local_cuda_client: Client) -> None:
import dask_cudf
X, y = make_categorical(local_cuda_client, 10000, 30, 13)
X = dask_cudf.from_dask_dataframe(X)
X_onehot, _ = make_categorical(local_cuda_client, 10000, 30, 13, True)
X_onehot = dask_cudf.from_dask_dataframe(X_onehot)
run_categorical(local_cuda_client, "gpu_hist", X, X_onehot, y)
@given(
params=parameter_strategy,
num_rounds=strategies.integers(1, 20),
dataset=tm.dataset_strategy,
dmatrix_type=strategies.sampled_from(
[dxgb.DaskDMatrix, dxgb.DaskDeviceQuantileDMatrix]
),
)
@settings(
deadline=duration(seconds=120),
max_examples=20,
suppress_health_check=suppress,
print_blob=True,
)
@settings(deadline=duration(seconds=120), suppress_health_check=suppress, print_blob=True)
@pytest.mark.skipif(**tm.no_cupy())
@pytest.mark.parametrize(
"local_cuda_cluster", [{"n_workers": 2}], indirect=["local_cuda_cluster"]
)
def test_gpu_hist(
self,
params: Dict,
num_rounds: int,
dataset: tm.TestDataset,
local_cuda_cluster: LocalCUDACluster,
dmatrix_type: type,
local_cuda_client: Client,
) -> None:
with Client(local_cuda_cluster) as client:
run_gpu_hist(params, num_rounds, dataset, dxgb.DaskDMatrix, client)
run_gpu_hist(
params, num_rounds, dataset, dxgb.DaskDeviceQuantileDMatrix, client
)
run_gpu_hist(params, num_rounds, dataset, dmatrix_type, local_cuda_client)
@pytest.mark.skipif(**tm.no_cupy())
def test_dask_array(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
run_with_dask_array(dxgb.DaskDMatrix, client)
run_with_dask_array(dxgb.DaskDeviceQuantileDMatrix, client)
def test_dask_array(self, local_cuda_client: Client) -> None:
run_with_dask_array(dxgb.DaskDMatrix, local_cuda_client)
run_with_dask_array(dxgb.DaskDeviceQuantileDMatrix, local_cuda_client)
@pytest.mark.skipif(**tm.no_cupy())
def test_early_stopping(self, local_cuda_cluster: LocalCUDACluster) -> None:
def test_early_stopping(self, local_cuda_client: Client) -> None:
from sklearn.datasets import load_breast_cancer
with Client(local_cuda_cluster) as client:
X, y = load_breast_cancer(return_X_y=True)
X, y = da.from_array(X), da.from_array(y)
m = dxgb.DaskDMatrix(client, X, y)
X, y = load_breast_cancer(return_X_y=True)
X, y = da.from_array(X), da.from_array(y)
valid = dxgb.DaskDMatrix(client, X, y)
early_stopping_rounds = 5
booster = dxgb.train(client, {'objective': 'binary:logistic',
'eval_metric': 'error',
'tree_method': 'gpu_hist'}, m,
evals=[(valid, 'Valid')],
num_boost_round=1000,
early_stopping_rounds=early_stopping_rounds)[
'booster']
assert hasattr(booster, 'best_score')
dump = booster.get_dump(dump_format='json')
print(booster.best_iteration)
assert len(dump) - booster.best_iteration == early_stopping_rounds + 1
m = dxgb.DaskDMatrix(local_cuda_client, X, y)
valid_X = X
valid_y = y
cls = dxgb.DaskXGBClassifier(objective='binary:logistic',
tree_method='gpu_hist',
eval_metric='error',
n_estimators=100)
cls.client = client
cls.fit(X, y, early_stopping_rounds=early_stopping_rounds,
eval_set=[(valid_X, valid_y)])
booster = cls.get_booster()
dump = booster.get_dump(dump_format='json')
assert len(dump) - booster.best_iteration == early_stopping_rounds + 1
valid = dxgb.DaskDMatrix(local_cuda_client, X, y)
early_stopping_rounds = 5
booster = dxgb.train(
local_cuda_client,
{
"objective": "binary:logistic",
"eval_metric": "error",
"tree_method": "gpu_hist",
},
m,
evals=[(valid, "Valid")],
num_boost_round=1000,
early_stopping_rounds=early_stopping_rounds,
)["booster"]
assert hasattr(booster, "best_score")
dump = booster.get_dump(dump_format="json")
assert len(dump) - booster.best_iteration == early_stopping_rounds + 1
valid_X = X
valid_y = y
cls = dxgb.DaskXGBClassifier(
objective="binary:logistic",
tree_method="gpu_hist",
eval_metric="error",
n_estimators=100,
)
cls.client = local_cuda_client
cls.fit(
X,
y,
early_stopping_rounds=early_stopping_rounds,
eval_set=[(valid_X, valid_y)],
)
booster = cls.get_booster()
dump = booster.get_dump(dump_format="json")
assert len(dump) - booster.best_iteration == early_stopping_rounds + 1
@pytest.mark.skipif(**tm.no_cudf())
@pytest.mark.parametrize("model", ["boosting"])
def test_dask_classifier(
self, model: str, local_cuda_cluster: LocalCUDACluster
) -> None:
def test_dask_classifier(self, model: str, local_cuda_client: Client) -> None:
import dask_cudf
with Client(local_cuda_cluster) as client:
X_, y_, w_ = generate_array(with_weights=True)
y_ = (y_ * 10).astype(np.int32)
X = dask_cudf.from_dask_dataframe(dd.from_dask_array(X_))
y = dask_cudf.from_dask_dataframe(dd.from_dask_array(y_))
w = dask_cudf.from_dask_dataframe(dd.from_dask_array(w_))
run_dask_classifier(X, y, w, model, "gpu_hist", client, 10)
def test_empty_dmatrix(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
parameters = {'tree_method': 'gpu_hist', 'debug_synchronize': True}
run_empty_dmatrix_reg(client, parameters)
run_empty_dmatrix_cls(client, parameters)
X_, y_, w_ = generate_array(with_weights=True)
y_ = (y_ * 10).astype(np.int32)
X = dask_cudf.from_dask_dataframe(dd.from_dask_array(X_))
y = dask_cudf.from_dask_dataframe(dd.from_dask_array(y_))
w = dask_cudf.from_dask_dataframe(dd.from_dask_array(w_))
run_dask_classifier(X, y, w, model, "gpu_hist", local_cuda_client, 10)
def test_empty_dmatrix(self, local_cuda_client: Client) -> None:
parameters = {"tree_method": "gpu_hist", "debug_synchronize": True}
run_empty_dmatrix_reg(local_cuda_client, parameters)
run_empty_dmatrix_cls(local_cuda_client, parameters)
@pytest.mark.skipif(**tm.no_dask_cudf())
def test_empty_partition(self, local_cuda_cluster: LocalCUDACluster) -> None:
def test_empty_partition(self, local_cuda_client: Client) -> None:
import dask_cudf
import cudf
import cupy
with Client(local_cuda_cluster) as client:
mult = 100
df = cudf.DataFrame(
{
"a": [1, 2, 3, 4, 5.1] * mult,
"b": [10, 15, 29.3, 30, 31] * mult,
"y": [10, 20, 30, 40., 50] * mult,
}
)
parameters = {"tree_method": "gpu_hist", "debug_synchronize": True}
empty = df.iloc[:0]
ddf = dask_cudf.concat(
[dask_cudf.from_cudf(empty, npartitions=1)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
dtrain = dxgb.DaskDeviceQuantileDMatrix(client, X, y)
bst_empty = xgb.dask.train(
client, parameters, dtrain, evals=[(dtrain, "train")]
)
predt_empty = dxgb.predict(client, bst_empty, X).compute().values
mult = 100
df = cudf.DataFrame(
{
"a": [1, 2, 3, 4, 5.1] * mult,
"b": [10, 15, 29.3, 30, 31] * mult,
"y": [10, 20, 30, 40.0, 50] * mult,
}
)
parameters = {"tree_method": "gpu_hist", "debug_synchronize": True}
ddf = dask_cudf.concat(
[dask_cudf.from_cudf(df, npartitions=3)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
dtrain = dxgb.DaskDeviceQuantileDMatrix(client, X, y)
bst = xgb.dask.train(client, parameters, dtrain, evals=[(dtrain, "train")])
empty = df.iloc[:0]
ddf = dask_cudf.concat(
[dask_cudf.from_cudf(empty, npartitions=1)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
dtrain = dxgb.DaskDeviceQuantileDMatrix(local_cuda_client, X, y)
bst_empty = xgb.dask.train(
local_cuda_client, parameters, dtrain, evals=[(dtrain, "train")]
)
predt_empty = dxgb.predict(local_cuda_client, bst_empty, X).compute().values
predt = dxgb.predict(client, bst, X).compute().values
cupy.testing.assert_allclose(predt, predt_empty)
ddf = dask_cudf.concat(
[dask_cudf.from_cudf(df, npartitions=3)]
+ [dask_cudf.from_cudf(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
dtrain = dxgb.DaskDeviceQuantileDMatrix(local_cuda_client, X, y)
bst = xgb.dask.train(
local_cuda_client, parameters, dtrain, evals=[(dtrain, "train")]
)
predt = dxgb.predict(client, bst, dtrain).compute()
cupy.testing.assert_allclose(predt, predt_empty)
predt = dxgb.predict(local_cuda_client, bst, X).compute().values
cupy.testing.assert_allclose(predt, predt_empty)
predt = dxgb.inplace_predict(client, bst, X).compute().values
cupy.testing.assert_allclose(predt, predt_empty)
predt = dxgb.predict(local_cuda_client, bst, dtrain).compute()
cupy.testing.assert_allclose(predt, predt_empty)
df = df.to_pandas()
empty = df.iloc[:0]
ddf = dd.concat(
[dd.from_pandas(empty, npartitions=1)]
+ [dd.from_pandas(df, npartitions=3)]
+ [dd.from_pandas(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
predt = dxgb.inplace_predict(local_cuda_client, bst, X).compute().values
cupy.testing.assert_allclose(predt, predt_empty)
predt_empty = cupy.asnumpy(predt_empty)
df = df.to_pandas()
empty = df.iloc[:0]
ddf = dd.concat(
[dd.from_pandas(empty, npartitions=1)]
+ [dd.from_pandas(df, npartitions=3)]
+ [dd.from_pandas(df, npartitions=3)]
)
X = ddf[ddf.columns.difference(["y"])]
y = ddf[["y"]]
predt = dxgb.predict(client, bst_empty, X).compute().values
np.testing.assert_allclose(predt, predt_empty)
predt_empty = cupy.asnumpy(predt_empty)
in_predt = dxgb.inplace_predict(client, bst_empty, X).compute().values
np.testing.assert_allclose(predt, in_predt)
predt = dxgb.predict(local_cuda_client, bst_empty, X).compute().values
np.testing.assert_allclose(predt, predt_empty)
def test_empty_dmatrix_auc(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
n_workers = len(_get_client_workers(client))
run_empty_dmatrix_auc(client, "gpu_hist", n_workers)
in_predt = (
dxgb.inplace_predict(local_cuda_client, bst_empty, X).compute().values
)
np.testing.assert_allclose(predt, in_predt)
def test_auc(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
run_auc(client, "gpu_hist")
def test_empty_dmatrix_auc(self, local_cuda_client: Client) -> None:
n_workers = len(_get_client_workers(local_cuda_client))
run_empty_dmatrix_auc(local_cuda_client, "gpu_hist", n_workers)
def test_data_initialization(self, local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
X, y, _ = generate_array()
fw = da.random.random((random_cols, ))
fw = fw - fw.min()
m = dxgb.DaskDMatrix(client, X, y, feature_weights=fw)
def test_auc(self, local_cuda_client: Client) -> None:
run_auc(local_cuda_client, "gpu_hist")
workers = _get_client_workers(client)
rabit_args = client.sync(dxgb._get_rabit_args, len(workers), None, client)
def test_data_initialization(self, local_cuda_client: Client) -> None:
def worker_fn(worker_addr: str, data_ref: Dict) -> None:
with dxgb.RabitContext(rabit_args):
local_dtrain = dxgb._dmatrix_from_list_of_parts(**data_ref, nthread=7)
fw_rows = local_dtrain.get_float_info("feature_weights").shape[0]
assert fw_rows == local_dtrain.num_col()
X, y, _ = generate_array()
fw = da.random.random((random_cols,))
fw = fw - fw.min()
m = dxgb.DaskDMatrix(local_cuda_client, X, y, feature_weights=fw)
futures = []
for i in range(len(workers)):
futures.append(
client.submit(
worker_fn,
workers[i],
m._create_fn_args(workers[i]),
pure=False,
workers=[workers[i]]
)
workers = _get_client_workers(local_cuda_client)
rabit_args = local_cuda_client.sync(
dxgb._get_rabit_args, len(workers), None, local_cuda_client
)
def worker_fn(worker_addr: str, data_ref: Dict) -> None:
with dxgb.RabitContext(rabit_args):
local_dtrain = dxgb._dmatrix_from_list_of_parts(**data_ref, nthread=7)
fw_rows = local_dtrain.get_float_info("feature_weights").shape[0]
assert fw_rows == local_dtrain.num_col()
futures = []
for i in range(len(workers)):
futures.append(
local_cuda_client.submit(
worker_fn,
workers[i],
m._create_fn_args(workers[i]),
pure=False,
workers=[workers[i]],
)
client.gather(futures)
)
local_cuda_client.gather(futures)
def test_interface_consistency(self) -> None:
sig = OrderedDict(signature(dxgb.DaskDMatrix).parameters)
@@ -441,7 +459,7 @@ class TestDistributedGPU:
assert ddm_names[i] == ddqdm_names[i]
sig = OrderedDict(signature(xgb.DMatrix).parameters)
del sig["nthread"] # no nthread in dask
del sig["nthread"] # no nthread in dask
dm_names = list(sig.keys())
sig = OrderedDict(signature(xgb.QuantileDMatrix).parameters)
del sig["nthread"]
@@ -470,81 +488,79 @@ class TestDistributedGPU:
for rn, drn in zip(ranker_names, dranker_names):
assert rn == drn
def test_tree_stats(self) -> None:
with LocalCUDACluster(n_workers=1) as cluster:
with Client(cluster) as client:
local = run_tree_stats(client, "gpu_hist")
with LocalCUDACluster(n_workers=2) as cluster:
with Client(cluster) as client:
distributed = run_tree_stats(client, "gpu_hist")
assert local == distributed
def run_quantile(self, name: str, local_cuda_cluster: LocalCUDACluster) -> None:
def run_quantile(self, name: str, local_cuda_client: Client) -> None:
if sys.platform.startswith("win"):
pytest.skip("Skipping dask tests on Windows")
exe = None
for possible_path in {'./testxgboost', './build/testxgboost',
'../build/testxgboost', '../gpu-build/testxgboost'}:
for possible_path in {
"./testxgboost",
"./build/testxgboost",
"../build/testxgboost",
"../gpu-build/testxgboost",
}:
if os.path.exists(possible_path):
exe = possible_path
assert exe, 'No testxgboost executable found.'
assert exe, "No testxgboost executable found."
test = "--gtest_filter=GPUQuantile." + name
def runit(
worker_addr: str, rabit_args: List[bytes]
) -> subprocess.CompletedProcess:
port_env = ''
port_env = ""
# setup environment for running the c++ part.
for arg in rabit_args:
if arg.decode('utf-8').startswith('DMLC_TRACKER_PORT'):
port_env = arg.decode('utf-8')
if arg.decode("utf-8").startswith("DMLC_TRACKER_PORT"):
port_env = arg.decode("utf-8")
if arg.decode("utf-8").startswith("DMLC_TRACKER_URI"):
uri_env = arg.decode("utf-8")
port = port_env.split('=')
port = port_env.split("=")
env = os.environ.copy()
env[port[0]] = port[1]
uri = uri_env.split("=")
env[uri[0]] = uri[1]
return subprocess.run([str(exe), test], env=env, stdout=subprocess.PIPE)
with Client(local_cuda_cluster) as client:
workers = _get_client_workers(client)
rabit_args = client.sync(dxgb._get_rabit_args, len(workers), None, client)
futures = client.map(runit,
workers,
pure=False,
workers=workers,
rabit_args=rabit_args)
results = client.gather(futures)
for ret in results:
msg = ret.stdout.decode('utf-8')
assert msg.find('1 test from GPUQuantile') != -1, msg
assert ret.returncode == 0, msg
workers = _get_client_workers(local_cuda_client)
rabit_args = local_cuda_client.sync(
dxgb._get_rabit_args, len(workers), None, local_cuda_client
)
futures = local_cuda_client.map(
runit, workers, pure=False, workers=workers, rabit_args=rabit_args
)
results = local_cuda_client.gather(futures)
for ret in results:
msg = ret.stdout.decode("utf-8")
assert msg.find("1 test from GPUQuantile") != -1, msg
assert ret.returncode == 0, msg
@pytest.mark.gtest
def test_quantile_basic(self, local_cuda_cluster: LocalCUDACluster) -> None:
self.run_quantile('AllReduceBasic', local_cuda_cluster)
def test_quantile_basic(self, local_cuda_client: Client) -> None:
self.run_quantile("AllReduceBasic", local_cuda_client)
@pytest.mark.gtest
def test_quantile_same_on_all_workers(
self, local_cuda_cluster: LocalCUDACluster
) -> None:
self.run_quantile('SameOnAllWorkers', local_cuda_cluster)
def test_quantile_same_on_all_workers(self, local_cuda_client: Client) -> None:
self.run_quantile("SameOnAllWorkers", local_cuda_client)
@pytest.mark.skipif(**tm.no_cupy())
def test_with_asyncio(local_cuda_client: Client) -> None:
address = local_cuda_client.scheduler.address
output = asyncio.run(run_from_dask_array_asyncio(address))
assert isinstance(output["booster"], xgboost.Booster)
assert isinstance(output["history"], dict)
async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainReturnT:
async with Client(scheduler_address, asynchronous=True) as client:
import cupy as cp
X, y, _ = generate_array()
X = X.map_blocks(cp.array)
y = y.map_blocks(cp.array)
m = await xgboost.dask.DaskDeviceQuantileDMatrix(client, X, y)
output = await xgboost.dask.train(client, {'tree_method': 'gpu_hist'},
dtrain=m)
output = await xgboost.dask.train(client, {"tree_method": "gpu_hist"}, dtrain=m)
with_m = await xgboost.dask.predict(client, output, m)
with_X = await xgboost.dask.predict(client, output, X)
@@ -553,19 +569,12 @@ async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainRetur
assert isinstance(with_X, da.Array)
assert isinstance(inplace, da.Array)
cp.testing.assert_allclose(await client.compute(with_m),
await client.compute(with_X))
cp.testing.assert_allclose(await client.compute(with_m),
await client.compute(inplace))
cp.testing.assert_allclose(
await client.compute(with_m), await client.compute(with_X)
)
cp.testing.assert_allclose(
await client.compute(with_m), await client.compute(inplace)
)
client.shutdown()
return output
@pytest.mark.skipif(**tm.no_cupy())
def test_with_asyncio(local_cuda_cluster: LocalCUDACluster) -> None:
with Client(local_cuda_cluster) as client:
address = client.scheduler.address
output = asyncio.run(run_from_dask_array_asyncio(address))
assert isinstance(output['booster'], xgboost.Booster)
assert isinstance(output['history'], dict)