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

[dask] Fix asyncio. (#7508) (#7561)

Browse Source
This commit is contained in:
Jiaming Yuan 2022-01-13 21:49:11 +08:00 committed by GitHub
parent afb9dfd421
commit 3e2d7519a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python-package/xgboost/dask.py
View File

@ -1606,8 +1606,9 @@ class DaskScikitLearnBase(XGBModel):
should use `worker_client' instead of default client. should use `worker_client' instead of default client.
""" """
asynchronous = getattr(self, "_asynchronous", False)
if self._client is None: if self._client is None:
asynchronous = getattr(self, "_asynchronous", False)
try: try:
distributed.get_worker() distributed.get_worker()
in_worker = True in_worker = True
@ -1620,7 +1621,7 @@ class DaskScikitLearnBase(XGBModel):
return ret return ret
return ret return ret
return self.client.sync(func, **kwargs, asynchronous=asynchronous) return self.client.sync(func, **kwargs, asynchronous=self.client.asynchronous)
@xgboost_model_doc( @xgboost_model_doc(

3
tests/python/test_with_dask.py
View File

@ -705,8 +705,7 @@ async def run_from_dask_array_asyncio(scheduler_address: str) -> xgb.dask.TrainR
async def run_dask_regressor_asyncio(scheduler_address: str) -> None: async def run_dask_regressor_asyncio(scheduler_address: str) -> None:
async with Client(scheduler_address, asynchronous=True) as client: async with Client(scheduler_address, asynchronous=True) as client:
X, y, _ = generate_array() X, y, _ = generate_array()
regressor = await xgb.dask.DaskXGBRegressor(verbosity=1, regressor = await xgb.dask.DaskXGBRegressor(verbosity=1, n_estimators=2)
n_estimators=2)
regressor.set_params(tree_method='hist') regressor.set_params(tree_method='hist')
regressor.client = client regressor.client = client
await regressor.fit(X, y, eval_set=[(X, y)]) await regressor.fit(X, y, eval_set=[(X, y)])