Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[dask] Disable broadcast in the scatter call. (#10632)

Browse Source
This commit is contained in:
Jiaming Yuan 2024-07-25 04:16:34 +08:00 committed by GitHub
parent 411c8466bd
commit fcae6301ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
python-package/xgboost/dask/__init__.py
View File

@ -1223,12 +1223,14 @@ def _infer_predict_output(
async def _get_model_future(
client: "distributed.Client", model: Union[Booster, Dict, "distributed.Future"]
) -> "distributed.Future":
# See https://github.com/dask/dask/issues/11179#issuecomment-2168094529 for
# the use of hash.
# See https://github.com/dask/dask/issues/11179#issuecomment-2168094529 for the use
# of hash.
# https://github.com/dask/distributed/pull/8796 Don't use broadcast in the `scatter`
# call, otherwise, the predict function might hang.
if isinstance(model, Booster):
booster = await client.scatter(model, broadcast=True, hash=False)
booster = await client.scatter(model, hash=False)
elif isinstance(model, dict):
booster = await client.scatter(model["booster"], broadcast=True, hash=False)
booster = await client.scatter(model["booster"], hash=False)
elif isinstance(model, distributed.Future):
booster = model
t = booster.type