From 92ae3abc972c100a11d7e75fb6dbeebdb7cefd1e Mon Sep 17 00:00:00 2001 From: Gil Forsyth Date: Tue, 27 Jul 2021 14:07:23 -0400 Subject: [PATCH] [dask] Disallow importing non-dask estimators from xgboost.dask (#7133) * Disallow importing non-dask estimators from xgboost.dask This is mostly a style change, but also avoids a user error (that I have committed on a few occasions). Since `XGBRegressor` and `XGBClassifier` are imported as parent classes for the `dask` estimators, without defining an `__all__`, autocomplete (or muscle) memory will produce the following with little prompting: ``` from xgboost.dask import XGBClassifier ``` There's nothing inherently wrong with that, but given that `XGBClassifier` is not `dask` enabled, it can lead to confusing behavior until you figure out you should've typed ``` from xgboost.dask import DaskXGBClassifier ``` Another option is to alias import the existing non-dask estimators. * Remove base/iter class, add train predict funcs --- python-package/xgboost/dask.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py index a73dc2638..580897ece 100644 --- a/python-package/xgboost/dask.py +++ b/python-package/xgboost/dask.py @@ -48,7 +48,6 @@ from .sklearn import xgboost_model_doc from .sklearn import _cls_predict_proba from .sklearn import XGBRanker - if TYPE_CHECKING: from dask import dataframe as dd from dask import array as da @@ -71,6 +70,20 @@ try: except ImportError: TrainReturnT = Dict[str, Any] # type:ignore +__all__ = [ + "RabitContext", + "DaskDMatrix", + "DaskDeviceQuantileDMatrix", + "DaskXGBRegressor", + "DaskXGBClassifier", + "DaskXGBRanker", + "DaskXGBRFRegressor", + "DaskXGBRFClassifier", + "train", + "predict", + "inplace_predict", +] + # TODOs: # - CV #