[dask] Use client to persist collections (#6722)

Co-authored-by: fis <jm.yuan@outlook.com>
This commit is contained in:
capybara
2021-02-25 02:40:38 -06:00
committed by GitHub
parent 9b530e5697
commit b6167cd2ff
2 changed files with 22 additions and 5 deletions

View File

@@ -303,11 +303,11 @@ class DaskDMatrix:
' of columns for your dask Array explicitly. e.g.' \
' chunks=(partition_size, X.shape[1])'
data = data.persist()
data = client.persist(data)
for meta in [label, weights, base_margin, label_lower_bound,
label_upper_bound]:
if meta is not None:
meta = meta.persist()
meta = client.persist(meta)
# Breaking data into partitions, a trick borrowed from dask_xgboost.
# `to_delayed` downgrades high-level objects into numpy or pandas