test_dask.py 69 KB
Newer Older
1
# coding: utf-8
2
3
"""Tests for lightgbm.dask module"""

4
import inspect
5
import re
6
import socket
7
8
9
from itertools import groupby
from os import getenv
from sys import platform
10
from urllib.parse import urlparse
11
12

import pytest
13
from sklearn.metrics import accuracy_score, r2_score
14
15
16

import lightgbm as lgb

17
18
from .utils import sklearn_multiclass_custom_objective

19
20
if platform in {"cygwin", "win32"}:
    pytest.skip("lightgbm.dask is not currently supported on Windows", allow_module_level=True)
21
if not lgb.compat.DASK_INSTALLED:
22
    pytest.skip("Dask is not installed", allow_module_level=True)
23
24
25
26
27

import dask.array as da
import dask.dataframe as dd
import numpy as np
import pandas as pd
28
import sklearn.utils.estimator_checks as sklearn_checks
29
from dask.array.utils import assert_eq
30
from dask.distributed import Client, LocalCluster, default_client, wait
31
from scipy.sparse import csc_matrix, csr_matrix
32
from scipy.stats import spearmanr
33
34
from sklearn.datasets import make_blobs, make_regression

35
from .utils import make_ranking, pickle_obj, unpickle_obj
36

37
38
39
40
tasks = ["binary-classification", "multiclass-classification", "regression", "ranking"]
distributed_training_algorithms = ["data", "voting"]
data_output = ["array", "scipy_csr_matrix", "dataframe", "dataframe-with-categorical"]
boosting_types = ["gbdt", "dart", "goss", "rf"]
41
group_sizes = [5, 5, 5, 10, 10, 10, 20, 20, 20, 50, 50]
42
task_to_dask_factory = {
43
44
45
46
    "regression": lgb.DaskLGBMRegressor,
    "binary-classification": lgb.DaskLGBMClassifier,
    "multiclass-classification": lgb.DaskLGBMClassifier,
    "ranking": lgb.DaskLGBMRanker,
47
48
}
task_to_local_factory = {
49
50
51
52
    "regression": lgb.LGBMRegressor,
    "binary-classification": lgb.LGBMClassifier,
    "multiclass-classification": lgb.LGBMClassifier,
    "ranking": lgb.LGBMRanker,
53
}
54
55

pytestmark = [
56
57
58
    pytest.mark.skipif(getenv("TASK", "") == "mpi", reason="Fails to run with MPI interface"),
    pytest.mark.skipif(getenv("TASK", "") == "gpu", reason="Fails to run with GPU interface"),
    pytest.mark.skipif(getenv("TASK", "") == "cuda", reason="Fails to run with CUDA interface"),
59
60
61
]


62
@pytest.fixture(scope="module")
63
64
65
66
67
68
def cluster():
    dask_cluster = LocalCluster(n_workers=2, threads_per_worker=2, dashboard_address=None)
    yield dask_cluster
    dask_cluster.close()


69
@pytest.fixture(scope="module")
70
71
72
73
74
75
def cluster2():
    dask_cluster = LocalCluster(n_workers=2, threads_per_worker=2, dashboard_address=None)
    yield dask_cluster
    dask_cluster.close()


76
@pytest.fixture(scope="module")
77
78
79
80
81
82
def cluster_three_workers():
    dask_cluster = LocalCluster(n_workers=3, threads_per_worker=1, dashboard_address=None)
    yield dask_cluster
    dask_cluster.close()


83
@pytest.fixture
84
85
86
87
88
89
90
91
def listen_port():
    listen_port.port += 10
    return listen_port.port


listen_port.port = 13000


92
def _get_workers_hostname(cluster: LocalCluster) -> str:
93
    one_worker_address = next(iter(cluster.scheduler_info["workers"]))
94
95
96
    return urlparse(one_worker_address).hostname


97
def _create_ranking_data(n_samples=100, output="array", chunk_size=50, **kwargs):
98
    X, y, g = make_ranking(n_samples=n_samples, random_state=42, **kwargs)
99
100
    rnd = np.random.RandomState(42)
    w = rnd.rand(X.shape[0]) * 0.01
101
    g_rle = np.array([len(list(grp)) for _, grp in groupby(g)])
102

103
    if output.startswith("dataframe"):
104
        # add target, weight, and group to DataFrame so that partitions abide by group boundaries.
105
106
        X_df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(X.shape[1])])
        if output == "dataframe-with-categorical":
107
            for i in range(5):
108
                col_name = f"cat_col{i}"
109
110
                cat_values = rnd.choice(["a", "b"], X.shape[0])
                cat_series = pd.Series(cat_values, dtype="category")
111
                X_df[col_name] = cat_series
112
113
114
115
116
        X = X_df.copy()
        X_df = X_df.assign(y=y, g=g, w=w)

        # set_index ensures partitions are based on group id.
        # See https://stackoverflow.com/questions/49532824/dask-dataframe-split-partitions-based-on-a-column-or-function.
117
        X_df.set_index("g", inplace=True)
118
119
120
        dX = dd.from_pandas(X_df, chunksize=chunk_size)

        # separate target, weight from features.
121
122
123
        dy = dX["y"]
        dw = dX["w"]
        dX = dX.drop(columns=["y", "w"])
124
125
126
127
        dg = dX.index.to_series()

        # encode group identifiers into run-length encoding, the format LightGBMRanker is expecting
        # so that within each partition, sum(g) = n_samples.
128
129
        dg = dg.map_partitions(lambda p: p.groupby("g", sort=False).apply(lambda z: z.shape[0]))
    elif output == "array":
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        # ranking arrays: one chunk per group. Each chunk must include all columns.
        p = X.shape[1]
        dX, dy, dw, dg = [], [], [], []
        for g_idx, rhs in enumerate(np.cumsum(g_rle)):
            lhs = rhs - g_rle[g_idx]
            dX.append(da.from_array(X[lhs:rhs, :], chunks=(rhs - lhs, p)))
            dy.append(da.from_array(y[lhs:rhs]))
            dw.append(da.from_array(w[lhs:rhs]))
            dg.append(da.from_array(np.array([g_rle[g_idx]])))

        dX = da.concatenate(dX, axis=0)
        dy = da.concatenate(dy, axis=0)
        dw = da.concatenate(dw, axis=0)
        dg = da.concatenate(dg, axis=0)
    else:
145
        raise ValueError("Ranking data creation only supported for Dask arrays and dataframes")
146
147
148
149

    return X, y, w, g_rle, dX, dy, dw, dg


150
151
152
def _create_data(objective, n_samples=1_000, output="array", chunk_size=500, **kwargs):
    if objective.endswith("classification"):
        if objective == "binary-classification":
153
            centers = [[-4, -4], [4, 4]]
154
        elif objective == "multiclass-classification":
155
156
157
            centers = [[-4, -4], [4, 4], [-4, 4]]
        else:
            raise ValueError(f"Unknown classification task '{objective}'")
158
        X, y = make_blobs(n_samples=n_samples, centers=centers, random_state=42)
159
    elif objective == "regression":
160
        X, y = make_regression(n_samples=n_samples, n_features=4, n_informative=2, random_state=42)
161
162
    elif objective == "ranking":
        return _create_ranking_data(n_samples=n_samples, output=output, chunk_size=chunk_size, **kwargs)
163
    else:
164
        raise ValueError(f"Unknown objective '{objective}'")
165
166
167
    rnd = np.random.RandomState(42)
    weights = rnd.random(X.shape[0]) * 0.01

168
    if output == "array":
169
170
171
        dX = da.from_array(X, (chunk_size, X.shape[1]))
        dy = da.from_array(y, chunk_size)
        dw = da.from_array(weights, chunk_size)
172
173
174
    elif output.startswith("dataframe"):
        X_df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(X.shape[1])])
        if output == "dataframe-with-categorical":
175
            num_cat_cols = 2
176
            for i in range(num_cat_cols):
177
                col_name = f"cat_col{i}"
178
179
                cat_values = rnd.choice(["a", "b"], X.shape[0])
                cat_series = pd.Series(cat_values, dtype="category")
180
181
182
                X_df[col_name] = cat_series
                X = np.hstack((X, cat_series.cat.codes.values.reshape(-1, 1)))

183
            # make one categorical feature relevant to the target
184
185
            cat_col_is_a = X_df["cat_col0"] == "a"
            if objective == "regression":
186
                y = np.where(cat_col_is_a, y, 2 * y)
187
            elif objective == "binary-classification":
188
                y = np.where(cat_col_is_a, y, 1 - y)
189
            elif objective == "multiclass-classification":
190
191
                n_classes = 3
                y = np.where(cat_col_is_a, y, (1 + y) % n_classes)
192
        y_df = pd.Series(y, name="target")
193
194
195
        dX = dd.from_pandas(X_df, chunksize=chunk_size)
        dy = dd.from_pandas(y_df, chunksize=chunk_size)
        dw = dd.from_array(weights, chunksize=chunk_size)
196
    elif output == "scipy_csr_matrix":
197
        dX = da.from_array(X, chunks=(chunk_size, X.shape[1])).map_blocks(csr_matrix)
198
199
        dy = da.from_array(y, chunks=chunk_size)
        dw = da.from_array(weights, chunk_size)
200
        X = csr_matrix(X)
201
    elif output == "scipy_csc_matrix":
202
203
204
205
        dX = da.from_array(X, chunks=(chunk_size, X.shape[1])).map_blocks(csc_matrix)
        dy = da.from_array(y, chunks=chunk_size)
        dw = da.from_array(weights, chunk_size)
        X = csc_matrix(X)
206
    else:
207
        raise ValueError(f"Unknown output type '{output}'")
208

209
    return X, y, weights, None, dX, dy, dw, None
210
211


212
def _r2_score(dy_true, dy_pred):
213
214
215
216
217
    y_true = dy_true.compute()
    y_pred = dy_pred.compute()
    numerator = ((y_true - y_pred) ** 2).sum(axis=0)
    denominator = ((y_true - y_true.mean(axis=0)) ** 2).sum(axis=0)
    return 1 - numerator / denominator
218
219
220


def _accuracy_score(dy_true, dy_pred):
221
222
223
    y_true = dy_true.compute()
    y_pred = dy_pred.compute()
    return (y_true == y_pred).mean()
224
225


226
def _constant_metric(y_true, y_pred):
227
    metric_name = "constant_metric"
228
229
230
231
232
    value = 0.708
    is_higher_better = False
    return metric_name, value, is_higher_better


233
234
235
236
237
238
239
240
241
242
243
244
245
def _objective_least_squares(y_true, y_pred):
    grad = y_pred - y_true
    hess = np.ones(len(y_true))
    return grad, hess


def _objective_logistic_regression(y_true, y_pred):
    y_pred = 1.0 / (1.0 + np.exp(-y_pred))
    grad = y_pred - y_true
    hess = y_pred * (1.0 - y_pred)
    return grad, hess


246
247
248
249
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("task", ["binary-classification", "multiclass-classification"])
@pytest.mark.parametrize("boosting_type", boosting_types)
@pytest.mark.parametrize("tree_learner", distributed_training_algorithms)
250
251
def test_classifier(output, task, boosting_type, tree_learner, cluster):
    with Client(cluster) as client:
252
253
254
255
256
257
258
259
260
261
262
263
        X, y, w, _, dX, dy, dw, _ = _create_data(objective=task, output=output)

        params = {"boosting_type": boosting_type, "tree_learner": tree_learner, "n_estimators": 50, "num_leaves": 31}
        if boosting_type == "rf":
            params.update(
                {
                    "bagging_freq": 1,
                    "bagging_fraction": 0.9,
                }
            )
        elif boosting_type == "goss":
            params["top_rate"] = 0.5
264

265
        dask_classifier = lgb.DaskLGBMClassifier(client=client, time_out=5, **params)
266
267
        dask_classifier = dask_classifier.fit(dX, dy, sample_weight=dw)
        p1 = dask_classifier.predict(dX)
268
269
270
        p1_raw = dask_classifier.predict(dX, raw_score=True).compute()
        p1_first_iter_raw = dask_classifier.predict(dX, start_iteration=0, num_iteration=1, raw_score=True).compute()
        p1_early_stop_raw = dask_classifier.predict(
271
            dX, pred_early_stop=True, pred_early_stop_margin=1.0, pred_early_stop_freq=2, raw_score=True
272
        ).compute()
273
274
275
276
277
278
279
280
281
282
283
284
        p1_proba = dask_classifier.predict_proba(dX).compute()
        p1_pred_leaf = dask_classifier.predict(dX, pred_leaf=True)
        p1_local = dask_classifier.to_local().predict(X)
        s1 = _accuracy_score(dy, p1)
        p1 = p1.compute()

        local_classifier = lgb.LGBMClassifier(**params)
        local_classifier.fit(X, y, sample_weight=w)
        p2 = local_classifier.predict(X)
        p2_proba = local_classifier.predict_proba(X)
        s2 = local_classifier.score(X, y)

285
        if boosting_type == "rf":
286
287
288
289
290
291
292
293
294
295
296
297
            # https://github.com/microsoft/LightGBM/issues/4118
            assert_eq(s1, s2, atol=0.01)
            assert_eq(p1_proba, p2_proba, atol=0.8)
        else:
            assert_eq(s1, s2)
            assert_eq(p1, p2)
            assert_eq(p1, y)
            assert_eq(p2, y)
            assert_eq(p1_proba, p2_proba, atol=0.03)
            assert_eq(p1_local, p2)
            assert_eq(p1_local, y)

298
        # extra predict() parameters should be passed through correctly
299
        with pytest.raises(AssertionError):  # noqa: PT011
300
301
            assert_eq(p1_raw, p1_first_iter_raw)

302
        with pytest.raises(AssertionError):  # noqa: PT011
303
304
            assert_eq(p1_raw, p1_early_stop_raw)

305
306
307
        # pref_leaf values should have the right shape
        # and values that look like valid tree nodes
        pred_leaf_vals = p1_pred_leaf.compute()
308
309
        assert pred_leaf_vals.shape == (X.shape[0], dask_classifier.booster_.num_trees())
        assert np.max(pred_leaf_vals) <= params["num_leaves"]
310
        assert np.min(pred_leaf_vals) >= 0
311
        assert len(np.unique(pred_leaf_vals)) <= params["num_leaves"]
312
313
314

        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
315
316
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
317
            tree_df = dask_classifier.booster_.trees_to_dataframe()
318
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
319
            assert node_uses_cat_col.sum() > 0
320
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
321

322

323
324
@pytest.mark.parametrize("output", data_output + ["scipy_csc_matrix"])
@pytest.mark.parametrize("task", ["binary-classification", "multiclass-classification"])
325
326
def test_classifier_pred_contrib(output, task, cluster):
    with Client(cluster) as client:
327
        X, y, w, _, dX, dy, dw, _ = _create_data(objective=task, output=output)
328

329
        params = {"n_estimators": 10, "num_leaves": 10}
330

331
        dask_classifier = lgb.DaskLGBMClassifier(client=client, time_out=5, tree_learner="data", **params)
332
        dask_classifier = dask_classifier.fit(dX, dy, sample_weight=dw)
333
        preds_with_contrib = dask_classifier.predict(dX, pred_contrib=True)
334
335
336
337
338

        local_classifier = lgb.LGBMClassifier(**params)
        local_classifier.fit(X, y, sample_weight=w)
        local_preds_with_contrib = local_classifier.predict(X, pred_contrib=True)

339
340
341
342
343
344
345
346
347
348
349
350
351
        # shape depends on whether it is binary or multiclass classification
        num_features = dask_classifier.n_features_
        num_classes = dask_classifier.n_classes_
        if num_classes == 2:
            expected_num_cols = num_features + 1
        else:
            expected_num_cols = (num_features + 1) * num_classes

        # in the special case of multi-class classification using scipy sparse matrices,
        # the output of `.predict(..., pred_contrib=True)` is a list of sparse matrices (one per class)
        #
        # since that case is so different than all other cases, check the relevant things here
        # and then return early
352
353
        if output.startswith("scipy") and task == "multiclass-classification":
            if output == "scipy_csr_matrix":
354
                expected_type = csr_matrix
355
            elif output == "scipy_csc_matrix":
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
                expected_type = csc_matrix
            else:
                raise ValueError(f"Unrecognized output type: {output}")
            assert isinstance(preds_with_contrib, list)
            assert all(isinstance(arr, da.Array) for arr in preds_with_contrib)
            assert all(isinstance(arr._meta, expected_type) for arr in preds_with_contrib)
            assert len(preds_with_contrib) == num_classes
            assert len(preds_with_contrib) == len(local_preds_with_contrib)
            for i in range(num_classes):
                computed_preds = preds_with_contrib[i].compute()
                assert isinstance(computed_preds, expected_type)
                assert computed_preds.shape[1] == num_classes
                assert computed_preds.shape == local_preds_with_contrib[i].shape
                assert len(np.unique(computed_preds[:, -1])) == 1
                # raw scores will probably be different, but at least check that all predicted classes are the same
                pred_classes = np.argmax(computed_preds.toarray(), axis=1)
                local_pred_classes = np.argmax(local_preds_with_contrib[i].toarray(), axis=1)
                np.testing.assert_array_equal(pred_classes, local_pred_classes)
            return

        preds_with_contrib = preds_with_contrib.compute()
377
        if output.startswith("scipy"):
378
            preds_with_contrib = preds_with_contrib.toarray()
379
380
381

        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
382
383
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
384
            tree_df = dask_classifier.booster_.trees_to_dataframe()
385
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
386
            assert node_uses_cat_col.sum() > 0
387
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
388
389
390
391
392
393
394
395
396
397

        # * shape depends on whether it is binary or multiclass classification
        # * matrix for binary classification is of the form [feature_contrib, base_value],
        #   for multi-class it's [feat_contrib_class1, base_value_class1, feat_contrib_class2, base_value_class2, etc.]
        # * contrib outputs for distributed training are different than from local training, so we can just test
        #   that the output has the right shape and base values are in the right position
        assert preds_with_contrib.shape[1] == expected_num_cols
        assert preds_with_contrib.shape == local_preds_with_contrib.shape

        if num_classes == 2:
398
            assert len(np.unique(preds_with_contrib[:, num_features])) == 1
399
400
401
402
403
404
        else:
            for i in range(num_classes):
                base_value_col = num_features * (i + 1) + i
                assert len(np.unique(preds_with_contrib[:, base_value_col]) == 1)


405
406
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("task", ["binary-classification", "multiclass-classification"])
407
408
409
410
411
412
413
414
415
416
417
418
419
def test_classifier_custom_objective(output, task, cluster):
    with Client(cluster) as client:
        X, y, w, _, dX, dy, dw, _ = _create_data(
            objective=task,
            output=output,
        )

        params = {
            "n_estimators": 50,
            "num_leaves": 31,
            "verbose": -1,
            "seed": 708,
            "deterministic": True,
420
            "force_col_wise": True,
421
422
        }

423
424
425
426
427
428
429
430
431
432
        if task == "binary-classification":
            params.update(
                {
                    "objective": _objective_logistic_regression,
                }
            )
        elif task == "multiclass-classification":
            params.update({"objective": sklearn_multiclass_custom_objective, "num_classes": 3})

        dask_classifier = lgb.DaskLGBMClassifier(client=client, time_out=5, tree_learner="data", **params)
433
434
435
436
437
438
439
440
441
442
        dask_classifier = dask_classifier.fit(dX, dy, sample_weight=dw)
        dask_classifier_local = dask_classifier.to_local()
        p1_raw = dask_classifier.predict(dX, raw_score=True).compute()
        p1_raw_local = dask_classifier_local.predict(X, raw_score=True)

        local_classifier = lgb.LGBMClassifier(**params)
        local_classifier.fit(X, y, sample_weight=w)
        p2_raw = local_classifier.predict(X, raw_score=True)

        # with a custom objective, prediction result is a raw score instead of predicted class
443
        if task == "binary-classification":
444
445
446
447
448
449
            p1_proba = 1.0 / (1.0 + np.exp(-p1_raw))
            p1_class = (p1_proba > 0.5).astype(np.int64)
            p1_proba_local = 1.0 / (1.0 + np.exp(-p1_raw_local))
            p1_class_local = (p1_proba_local > 0.5).astype(np.int64)
            p2_proba = 1.0 / (1.0 + np.exp(-p2_raw))
            p2_class = (p2_proba > 0.5).astype(np.int64)
450
        elif task == "multiclass-classification":
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
            p1_proba = np.exp(p1_raw) / np.sum(np.exp(p1_raw), axis=1).reshape(-1, 1)
            p1_class = p1_proba.argmax(axis=1)
            p1_proba_local = np.exp(p1_raw_local) / np.sum(np.exp(p1_raw_local), axis=1).reshape(-1, 1)
            p1_class_local = p1_proba_local.argmax(axis=1)
            p2_proba = np.exp(p2_raw) / np.sum(np.exp(p2_raw), axis=1).reshape(-1, 1)
            p2_class = p2_proba.argmax(axis=1)

        # function should have been preserved
        assert callable(dask_classifier.objective_)
        assert callable(dask_classifier_local.objective_)

        # should correctly classify every sample
        assert_eq(p1_class, y)
        assert_eq(p1_class_local, y)
        assert_eq(p2_class, y)

        # probability estimates should be similar
        assert_eq(p1_proba, p2_proba, atol=0.03)
        assert_eq(p1_proba, p1_proba_local)


472
def test_machines_to_worker_map_unparsable_host_names():
473
    workers = {"0.0.0.1:80": {}, "0.0.0.2:80": {}}
474
475
476
477
478
    machines = "0.0.0.1:80,0.0.0.2:80"
    with pytest.raises(ValueError, match="Could not parse host name from worker address '0.0.0.1:80'"):
        lgb.dask._machines_to_worker_map(machines=machines, worker_addresses=workers.keys())


479
480
def test_training_does_not_fail_on_port_conflicts(cluster):
    with Client(cluster) as client:
481
        _, _, _, _, dX, dy, dw, _ = _create_data("binary-classification", output="array")
482
483

        lightgbm_default_port = 12400
484
        workers_hostname = _get_workers_hostname(cluster)
485
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
486
            s.bind((workers_hostname, lightgbm_default_port))
487
            dask_classifier = lgb.DaskLGBMClassifier(client=client, time_out=5, n_estimators=5, num_leaves=5)
488
489
490
491
492
493
494
            for _ in range(5):
                dask_classifier.fit(
                    X=dX,
                    y=dy,
                    sample_weight=dw,
                )
                assert dask_classifier.booster_
495

496

497
498
499
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("boosting_type", boosting_types)
@pytest.mark.parametrize("tree_learner", distributed_training_algorithms)
500
501
def test_regressor(output, boosting_type, tree_learner, cluster):
    with Client(cluster) as client:
502
        X, y, w, _, dX, dy, dw, _ = _create_data(objective="regression", output=output)
503
504
505
506
507
508
509

        params = {
            "boosting_type": boosting_type,
            "random_state": 42,
            "num_leaves": 31,
            "n_estimators": 20,
        }
510
511
512
513
514
515
516
        if boosting_type == "rf":
            params.update(
                {
                    "bagging_freq": 1,
                    "bagging_fraction": 0.9,
                }
            )
517

518
        dask_regressor = lgb.DaskLGBMRegressor(client=client, time_out=5, tree=tree_learner, **params)
519
520
521
522
523
524
        dask_regressor = dask_regressor.fit(dX, dy, sample_weight=dw)
        p1 = dask_regressor.predict(dX)
        p1_pred_leaf = dask_regressor.predict(dX, pred_leaf=True)

        s1 = _r2_score(dy, p1)
        p1 = p1.compute()
525
526
        p1_raw = dask_regressor.predict(dX, raw_score=True).compute()
        p1_first_iter_raw = dask_regressor.predict(dX, start_iteration=0, num_iteration=1, raw_score=True).compute()
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
        p1_local = dask_regressor.to_local().predict(X)
        s1_local = dask_regressor.to_local().score(X, y)

        local_regressor = lgb.LGBMRegressor(**params)
        local_regressor.fit(X, y, sample_weight=w)
        s2 = local_regressor.score(X, y)
        p2 = local_regressor.predict(X)

        # Scores should be the same
        assert_eq(s1, s2, atol=0.01)
        assert_eq(s1, s1_local)

        # Predictions should be roughly the same.
        assert_eq(p1, p1_local)

        # pref_leaf values should have the right shape
        # and values that look like valid tree nodes
        pred_leaf_vals = p1_pred_leaf.compute()
545
546
        assert pred_leaf_vals.shape == (X.shape[0], dask_regressor.booster_.num_trees())
        assert np.max(pred_leaf_vals) <= params["num_leaves"]
547
        assert np.min(pred_leaf_vals) >= 0
548
        assert len(np.unique(pred_leaf_vals)) <= params["num_leaves"]
549

550
551
        assert_eq(p1, y, rtol=0.5, atol=50.0)
        assert_eq(p2, y, rtol=0.5, atol=50.0)
552

553
        # extra predict() parameters should be passed through correctly
554
        with pytest.raises(AssertionError):  # noqa: PT011
555
556
            assert_eq(p1_raw, p1_first_iter_raw)

557
558
        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
559
560
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
561
            tree_df = dask_regressor.booster_.trees_to_dataframe()
562
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
563
            assert node_uses_cat_col.sum() > 0
564
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
565

566

567
@pytest.mark.parametrize("output", data_output)
568
569
def test_regressor_pred_contrib(output, cluster):
    with Client(cluster) as client:
570
        X, y, w, _, dX, dy, dw, _ = _create_data(objective="regression", output=output)
571

572
        params = {"n_estimators": 10, "num_leaves": 10}
573

574
        dask_regressor = lgb.DaskLGBMRegressor(client=client, time_out=5, tree_learner="data", **params)
575
576
577
578
579
580
581
582
        dask_regressor = dask_regressor.fit(dX, dy, sample_weight=dw)
        preds_with_contrib = dask_regressor.predict(dX, pred_contrib=True).compute()

        local_regressor = lgb.LGBMRegressor(**params)
        local_regressor.fit(X, y, sample_weight=w)
        local_preds_with_contrib = local_regressor.predict(X, pred_contrib=True)

        if output == "scipy_csr_matrix":
583
            preds_with_contrib = preds_with_contrib.toarray()
584
585
586
587
588
589
590
591
592

        # contrib outputs for distributed training are different than from local training, so we can just test
        # that the output has the right shape and base values are in the right position
        num_features = dX.shape[1]
        assert preds_with_contrib.shape[1] == num_features + 1
        assert preds_with_contrib.shape == local_preds_with_contrib.shape

        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
593
594
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
595
            tree_df = dask_regressor.booster_.trees_to_dataframe()
596
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
597
            assert node_uses_cat_col.sum() > 0
598
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
599

600

601
602
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("alpha", [0.1, 0.5, 0.9])
603
604
def test_regressor_quantile(output, alpha, cluster):
    with Client(cluster) as client:
605
        X, y, w, _, dX, dy, dw, _ = _create_data(objective="regression", output=output)
606

607
        params = {"objective": "quantile", "alpha": alpha, "random_state": 42, "n_estimators": 10, "num_leaves": 10}
608

609
        dask_regressor = lgb.DaskLGBMRegressor(client=client, tree_learner_type="data_parallel", **params)
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
        dask_regressor = dask_regressor.fit(dX, dy, sample_weight=dw)
        p1 = dask_regressor.predict(dX).compute()
        q1 = np.count_nonzero(y < p1) / y.shape[0]

        local_regressor = lgb.LGBMRegressor(**params)
        local_regressor.fit(X, y, sample_weight=w)
        p2 = local_regressor.predict(X)
        q2 = np.count_nonzero(y < p2) / y.shape[0]

        # Quantiles should be right
        np.testing.assert_allclose(q1, alpha, atol=0.2)
        np.testing.assert_allclose(q2, alpha, atol=0.2)

        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
625
626
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
627
            tree_df = dask_regressor.booster_.trees_to_dataframe()
628
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
629
            assert node_uses_cat_col.sum() > 0
630
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
631

632

633
@pytest.mark.parametrize("output", data_output)
634
635
def test_regressor_custom_objective(output, cluster):
    with Client(cluster) as client:
636
        X, y, w, _, dX, dy, dw, _ = _create_data(objective="regression", output=output)
637

638
        params = {"n_estimators": 10, "num_leaves": 10, "objective": _objective_least_squares}
639

640
        dask_regressor = lgb.DaskLGBMRegressor(client=client, time_out=5, tree_learner="data", **params)
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
        dask_regressor = dask_regressor.fit(dX, dy, sample_weight=dw)
        dask_regressor_local = dask_regressor.to_local()
        p1 = dask_regressor.predict(dX)
        p1_local = dask_regressor_local.predict(X)
        s1_local = dask_regressor_local.score(X, y)
        s1 = _r2_score(dy, p1)
        p1 = p1.compute()

        local_regressor = lgb.LGBMRegressor(**params)
        local_regressor.fit(X, y, sample_weight=w)
        p2 = local_regressor.predict(X)
        s2 = local_regressor.score(X, y)

        # function should have been preserved
        assert callable(dask_regressor.objective_)
        assert callable(dask_regressor_local.objective_)

        # Scores should be the same
        assert_eq(s1, s2, atol=0.01)
        assert_eq(s1, s1_local)

        # local and Dask predictions should be the same
        assert_eq(p1, p1_local)

        # predictions should be better than random
666
        assert_precision = {"rtol": 0.5, "atol": 50.0}
667
668
669
670
        assert_eq(p1, y, **assert_precision)
        assert_eq(p2, y, **assert_precision)


671
672
673
674
@pytest.mark.parametrize("output", ["array", "dataframe", "dataframe-with-categorical"])
@pytest.mark.parametrize("group", [None, group_sizes])
@pytest.mark.parametrize("boosting_type", boosting_types)
@pytest.mark.parametrize("tree_learner", distributed_training_algorithms)
675
676
def test_ranker(output, group, boosting_type, tree_learner, cluster):
    with Client(cluster) as client:
677
        if output == "dataframe-with-categorical":
678
            X, y, w, g, dX, dy, dw, dg = _create_data(
679
                objective="ranking", output=output, group=group, n_features=1, n_informative=1
680
681
            )
        else:
682
            X, y, w, g, dX, dy, dw, dg = _create_data(objective="ranking", output=output, group=group)
683
684

        # rebalance small dask.Array dataset for better performance.
685
        if output == "array":
686
687
688
689
690
691
692
693
694
695
696
697
698
699
            dX = dX.persist()
            dy = dy.persist()
            dw = dw.persist()
            dg = dg.persist()
            _ = wait([dX, dy, dw, dg])
            client.rebalance()

        # use many trees + leaves to overfit, help ensure that Dask data-parallel strategy matches that of
        # serial learner. See https://github.com/microsoft/LightGBM/issues/3292#issuecomment-671288210.
        params = {
            "boosting_type": boosting_type,
            "random_state": 42,
            "n_estimators": 50,
            "num_leaves": 20,
700
            "min_child_samples": 1,
701
        }
702
703
704
705
706
707
708
709
710
        if boosting_type == "rf":
            params.update(
                {
                    "bagging_freq": 1,
                    "bagging_fraction": 0.9,
                }
            )

        dask_ranker = lgb.DaskLGBMRanker(client=client, time_out=5, tree_learner_type=tree_learner, **params)
711
712
713
714
        dask_ranker = dask_ranker.fit(dX, dy, sample_weight=dw, group=dg)
        rnkvec_dask = dask_ranker.predict(dX)
        rnkvec_dask = rnkvec_dask.compute()
        p1_pred_leaf = dask_ranker.predict(dX, pred_leaf=True)
715
716
        p1_raw = dask_ranker.predict(dX, raw_score=True).compute()
        p1_first_iter_raw = dask_ranker.predict(dX, start_iteration=0, num_iteration=1, raw_score=True).compute()
717
        p1_early_stop_raw = dask_ranker.predict(
718
            dX, pred_early_stop=True, pred_early_stop_margin=1.0, pred_early_stop_freq=2, raw_score=True
719
        ).compute()
720
721
722
723
724
725
726
727
728
729
730
731
732
        rnkvec_dask_local = dask_ranker.to_local().predict(X)

        local_ranker = lgb.LGBMRanker(**params)
        local_ranker.fit(X, y, sample_weight=w, group=g)
        rnkvec_local = local_ranker.predict(X)

        # distributed ranker should be able to rank decently well and should
        # have high rank correlation with scores from serial ranker.
        dcor = spearmanr(rnkvec_dask, y).correlation
        assert dcor > 0.6
        assert spearmanr(rnkvec_dask, rnkvec_local).correlation > 0.8
        assert_eq(rnkvec_dask, rnkvec_dask_local)

733
        # extra predict() parameters should be passed through correctly
734
        with pytest.raises(AssertionError):  # noqa: PT011
735
736
            assert_eq(p1_raw, p1_first_iter_raw)

737
        with pytest.raises(AssertionError):  # noqa: PT011
738
739
            assert_eq(p1_raw, p1_early_stop_raw)

740
741
742
        # pref_leaf values should have the right shape
        # and values that look like valid tree nodes
        pred_leaf_vals = p1_pred_leaf.compute()
743
744
        assert pred_leaf_vals.shape == (X.shape[0], dask_ranker.booster_.num_trees())
        assert np.max(pred_leaf_vals) <= params["num_leaves"]
745
        assert np.min(pred_leaf_vals) >= 0
746
        assert len(np.unique(pred_leaf_vals)) <= params["num_leaves"]
747

748
749
        # be sure LightGBM actually used at least one categorical column,
        # and that it was correctly treated as a categorical feature
750
751
        if output == "dataframe-with-categorical":
            cat_cols = [col for col in dX.columns if dX.dtypes[col].name == "category"]
752
            tree_df = dask_ranker.booster_.trees_to_dataframe()
753
            node_uses_cat_col = tree_df["split_feature"].isin(cat_cols)
754
            assert node_uses_cat_col.sum() > 0
755
            assert tree_df.loc[node_uses_cat_col, "decision_type"].unique()[0] == "=="
756

757

758
@pytest.mark.parametrize("output", ["array", "dataframe", "dataframe-with-categorical"])
759
760
def test_ranker_custom_objective(output, cluster):
    with Client(cluster) as client:
761
        if output == "dataframe-with-categorical":
762
            X, y, w, g, dX, dy, dw, dg = _create_data(
763
                objective="ranking", output=output, group=group_sizes, n_features=1, n_informative=1
764
765
            )
        else:
766
            X, y, w, g, dX, dy, dw, dg = _create_data(objective="ranking", output=output, group=group_sizes)
767
768

        # rebalance small dask.Array dataset for better performance.
769
        if output == "array":
770
771
772
773
774
775
776
777
778
779
780
781
            dX = dX.persist()
            dy = dy.persist()
            dw = dw.persist()
            dg = dg.persist()
            _ = wait([dX, dy, dw, dg])
            client.rebalance()

        params = {
            "random_state": 42,
            "n_estimators": 50,
            "num_leaves": 20,
            "min_child_samples": 1,
782
            "objective": _objective_least_squares,
783
784
        }

785
        dask_ranker = lgb.DaskLGBMRanker(client=client, time_out=5, tree_learner_type="data", **params)
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
        dask_ranker = dask_ranker.fit(dX, dy, sample_weight=dw, group=dg)
        rnkvec_dask = dask_ranker.predict(dX).compute()
        dask_ranker_local = dask_ranker.to_local()
        rnkvec_dask_local = dask_ranker_local.predict(X)

        local_ranker = lgb.LGBMRanker(**params)
        local_ranker.fit(X, y, sample_weight=w, group=g)
        rnkvec_local = local_ranker.predict(X)

        # distributed ranker should be able to rank decently well with the least-squares objective
        # and should have high rank correlation with scores from serial ranker.
        assert spearmanr(rnkvec_dask, y).correlation > 0.6
        assert spearmanr(rnkvec_dask, rnkvec_local).correlation > 0.8
        assert_eq(rnkvec_dask, rnkvec_dask_local)

        # function should have been preserved
        assert callable(dask_ranker.objective_)
        assert callable(dask_ranker_local.objective_)


806
807
808
809
@pytest.mark.parametrize("task", tasks)
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("eval_sizes", [[0.5, 1, 1.5], [0]])
@pytest.mark.parametrize("eval_names_prefix", ["specified", None])
810
def test_eval_set_no_early_stopping(task, output, eval_sizes, eval_names_prefix, cluster):
811
812
    if task == "ranking" and output == "scipy_csr_matrix":
        pytest.skip("LGBMRanker is not currently tested on sparse matrices")
813
814
815
816
817
818
819
820
821
822
823
824
825

    with Client(cluster) as client:
        # Use larger trainset to prevent premature stopping due to zero loss, causing num_trees() < n_estimators.
        # Use small chunk_size to avoid single-worker allocation of eval data partitions.
        n_samples = 1000
        chunk_size = 10
        n_eval_sets = len(eval_sizes)
        eval_set = []
        eval_sample_weight = []
        eval_class_weight = None
        eval_init_score = None

        if eval_names_prefix:
826
            eval_names = [f"{eval_names_prefix}_{i}" for i in range(len(eval_sizes))]
827
828
829
830
        else:
            eval_names = None

        X, y, w, g, dX, dy, dw, dg = _create_data(
831
            objective=task, n_samples=n_samples, output=output, chunk_size=chunk_size
832
833
        )

834
835
        if task == "ranking":
            eval_metrics = ["ndcg"]
836
            eval_at = (5, 6)
837
            eval_metric_names = [f"ndcg@{k}" for k in eval_at]
838
839
840
841
            eval_group = []
        else:
            # test eval_class_weight, eval_init_score on binary-classification task.
            # Note: objective's default `metric` will be evaluated in evals_result_ in addition to all eval_metrics.
842
843
844
            if task == "binary-classification":
                eval_metrics = ["binary_error", "auc"]
                eval_metric_names = ["binary_logloss", "binary_error", "auc"]
845
846
                eval_class_weight = []
                eval_init_score = []
847
848
849
850
851
852
            elif task == "multiclass-classification":
                eval_metrics = ["multi_error"]
                eval_metric_names = ["multi_logloss", "multi_error"]
            elif task == "regression":
                eval_metrics = ["l1"]
                eval_metric_names = ["l2", "l1"]
853
854
855
856
857
858
859
860
861
862
863
864

        # create eval_sets by creating new datasets or copying training data.
        for eval_size in eval_sizes:
            if eval_size == 1:
                y_e = y
                dX_e = dX
                dy_e = dy
                dw_e = dw
                dg_e = dg
            else:
                n_eval_samples = max(chunk_size, int(n_samples * eval_size))
                _, y_e, _, _, dX_e, dy_e, dw_e, dg_e = _create_data(
865
                    objective=task, n_samples=n_eval_samples, output=output, chunk_size=chunk_size
866
867
868
869
                )

            eval_set.append((dX_e, dy_e))
            eval_sample_weight.append(dw_e)
870
            if task == "ranking":
871
872
                eval_group.append(dg_e)

873
            if task == "binary-classification":
874
875
876
877
                n_neg = np.sum(y_e == 0)
                n_pos = np.sum(y_e == 1)
                eval_class_weight.append({0: n_neg / n_pos, 1: n_pos / n_neg})
                init_score_value = np.log(np.mean(y_e) / (1 - np.mean(y_e)))
878
                if "dataframe" in output:
879
                    d_init_score = dy_e.map_partitions(lambda x, val=init_score_value: pd.Series([val] * x.size))
880
                else:
881
                    d_init_score = dy_e.map_blocks(lambda x, val=init_score_value: np.repeat(val, x.size))
882
883
884
885

                eval_init_score.append(d_init_score)

        fit_trees = 50
886
        params = {"random_state": 42, "n_estimators": fit_trees, "num_leaves": 2}
887
888

        model_factory = task_to_dask_factory[task]
889
        dask_model = model_factory(client=client, **params)
890
891

        fit_params = {
892
893
894
895
896
897
898
            "X": dX,
            "y": dy,
            "eval_set": eval_set,
            "eval_names": eval_names,
            "eval_sample_weight": eval_sample_weight,
            "eval_init_score": eval_init_score,
            "eval_metric": eval_metrics,
899
        }
900
901
902
903
        if task == "ranking":
            fit_params.update({"group": dg, "eval_group": eval_group, "eval_at": eval_at})
        elif task == "binary-classification":
            fit_params.update({"eval_class_weight": eval_class_weight})
904
905

        if eval_sizes == [0]:
906
907
908
909
            with pytest.warns(
                UserWarning,
                match="Worker (.*) was not allocated eval_set data. Therefore evals_result_ and best_score_ data may be unreliable.",
            ):
910
911
912
913
914
                dask_model.fit(**fit_params)
        else:
            dask_model = dask_model.fit(**fit_params)

            # total number of trees scales up for ova classifier.
915
            if task == "multiclass-classification":
916
917
918
919
920
921
                model_trees = fit_trees * dask_model.n_classes_
            else:
                model_trees = fit_trees

            # check that early stopping was not applied.
            assert dask_model.booster_.num_trees() == model_trees
922
            assert dask_model.best_iteration_ == 0
923
924
925
926
927
928
929
930
931
932
933
934
935
936

            # checks that evals_result_ and best_score_ contain expected data and eval_set names.
            evals_result = dask_model.evals_result_
            best_scores = dask_model.best_score_
            assert len(evals_result) == n_eval_sets
            assert len(best_scores) == n_eval_sets

            for eval_name in evals_result:
                assert eval_name in dask_model.best_score_
                if eval_names:
                    assert eval_name in eval_names

                # check that each eval_name and metric exists for all eval sets, allowing for the
                # case when a worker receives a fully-padded eval_set component which is not evaluated.
937
                if evals_result[eval_name] != {}:
938
939
940
941
942
943
                    for metric in eval_metric_names:
                        assert metric in evals_result[eval_name]
                        assert metric in best_scores[eval_name]
                        assert len(evals_result[eval_name][metric]) == fit_trees


944
@pytest.mark.parametrize("task", ["binary-classification", "regression", "ranking"])
945
946
947
948
949
def test_eval_set_with_custom_eval_metric(task, cluster):
    with Client(cluster) as client:
        n_samples = 1000
        n_eval_samples = int(n_samples * 0.5)
        chunk_size = 10
950
        output = "array"
951
952

        X, y, w, g, dX, dy, dw, dg = _create_data(
953
            objective=task, n_samples=n_samples, output=output, chunk_size=chunk_size
954
955
        )
        _, _, _, _, dX_e, dy_e, _, dg_e = _create_data(
956
            objective=task, n_samples=n_eval_samples, output=output, chunk_size=chunk_size
957
958
        )

959
        if task == "ranking":
960
            eval_at = (5, 6)
961
962
963
964
965
            eval_metrics = ["ndcg", _constant_metric]
            eval_metric_names = [f"ndcg@{k}" for k in eval_at] + ["constant_metric"]
        elif task == "binary-classification":
            eval_metrics = ["binary_error", "auc", _constant_metric]
            eval_metric_names = ["binary_logloss", "binary_error", "auc", "constant_metric"]
966
        else:
967
968
            eval_metrics = ["l1", _constant_metric]
            eval_metric_names = ["l2", "l1", "constant_metric"]
969
970

        fit_trees = 50
971
        params = {"random_state": 42, "n_estimators": fit_trees, "num_leaves": 2}
972
        model_factory = task_to_dask_factory[task]
973
        dask_model = model_factory(client=client, **params)
974
975

        eval_set = [(dX_e, dy_e)]
976
977
978
        fit_params = {"X": dX, "y": dy, "eval_set": eval_set, "eval_metric": eval_metrics}
        if task == "ranking":
            fit_params.update({"group": dg, "eval_group": [dg_e], "eval_at": eval_at})
979
980
981

        dask_model = dask_model.fit(**fit_params)

982
        eval_name = "valid_0"
983
984
985
986
987
988
989
990
        evals_result = dask_model.evals_result_
        assert len(evals_result) == 1
        assert eval_name in evals_result

        for metric in eval_metric_names:
            assert metric in evals_result[eval_name]
            assert len(evals_result[eval_name][metric]) == fit_trees

991
        np.testing.assert_allclose(evals_result[eval_name]["constant_metric"], 0.708)
992
993


994
@pytest.mark.parametrize("task", tasks)
995
996
def test_training_works_if_client_not_provided_or_set_after_construction(task, cluster):
    with Client(cluster) as client:
997
        _, _, _, _, dX, dy, _, dg = _create_data(objective=task, output="array", group=None)
998
999
        model_factory = task_to_dask_factory[task]

1000
        params = {"time_out": 5, "n_estimators": 1, "num_leaves": 2}
1001
1002
1003
1004

        # should be able to use the class without specifying a client
        dask_model = model_factory(**params)
        assert dask_model.client is None
1005
        with pytest.raises(lgb.compat.LGBMNotFittedError, match="Cannot access property client_ before calling fit"):
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
            dask_model.client_

        dask_model.fit(dX, dy, group=dg)
        assert dask_model.fitted_
        assert dask_model.client is None
        assert dask_model.client_ == client

        preds = dask_model.predict(dX)
        assert isinstance(preds, da.Array)
        assert dask_model.fitted_
        assert dask_model.client is None
        assert dask_model.client_ == client

        local_model = dask_model.to_local()
1020
1021
1022
1023
1024
        no_client_attr_msg = re.compile(
            f"{repr(type(local_model).__name__)} object has no attribute '(client|client_)'"
        )

        with pytest.raises(AttributeError, match=no_client_attr_msg):
1025
            local_model.client
1026
        with pytest.raises(AttributeError, match=no_client_attr_msg):
1027
1028
1029
1030
1031
1032
1033
            local_model.client_

        # should be able to set client after construction
        dask_model = model_factory(**params)
        dask_model.set_params(client=client)
        assert dask_model.client == client

1034
        with pytest.raises(lgb.compat.LGBMNotFittedError, match="Cannot access property client_ before calling fit"):
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
            dask_model.client_

        dask_model.fit(dX, dy, group=dg)
        assert dask_model.fitted_
        assert dask_model.client == client
        assert dask_model.client_ == client

        preds = dask_model.predict(dX)
        assert isinstance(preds, da.Array)
        assert dask_model.fitted_
        assert dask_model.client == client
        assert dask_model.client_ == client

        local_model = dask_model.to_local()
1049
        with pytest.raises(AttributeError, match=no_client_attr_msg):
1050
            local_model.client
1051
        with pytest.raises(AttributeError, match=no_client_attr_msg):
1052
            local_model.client_
1053
1054


1055
1056
1057
1058
1059
1060
@pytest.mark.parametrize("serializer", ["pickle", "joblib", "cloudpickle"])
@pytest.mark.parametrize("task", tasks)
@pytest.mark.parametrize("set_client", [True, False])
def test_model_and_local_version_are_picklable_whether_or_not_client_set_explicitly(
    serializer, task, set_client, tmp_path, cluster, cluster2
):
1061
    with Client(cluster) as client1:
1062
        # data on cluster1
1063
        X_1, _, _, _, dX_1, dy_1, _, dg_1 = _create_data(objective=task, output="array", group=None)
1064

1065
        with Client(cluster2) as client2:
1066
            # create identical data on cluster2
1067
            X_2, _, _, _, dX_2, dy_2, _, dg_2 = _create_data(objective=task, output="array", group=None)
1068

1069
1070
            model_factory = task_to_dask_factory[task]

1071
            params = {"time_out": 5, "n_estimators": 1, "num_leaves": 2}
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084

            # at this point, the result of default_client() is client2 since it was the most recently
            # created. So setting client to client1 here to test that you can select a non-default client
            assert default_client() == client2
            if set_client:
                params.update({"client": client1})

            # unfitted model should survive pickling round trip, and pickling
            # shouldn't have side effects on the model object
            dask_model = model_factory(**params)
            local_model = dask_model.to_local()
            if set_client:
                assert dask_model.client == client1
1085
            else:
1086
1087
                assert dask_model.client is None

1088
1089
1090
            with pytest.raises(
                lgb.compat.LGBMNotFittedError, match="Cannot access property client_ before calling fit"
            ):
1091
1092
1093
1094
1095
                dask_model.client_

            assert "client" not in local_model.get_params()
            assert getattr(local_model, "client", None) is None

1096
            tmp_file = tmp_path / "model-1.pkl"
1097
1098
            pickle_obj(obj=dask_model, filepath=tmp_file, serializer=serializer)
            model_from_disk = unpickle_obj(filepath=tmp_file, serializer=serializer)
1099

1100
            local_tmp_file = tmp_path / "local-model-1.pkl"
1101
1102
            pickle_obj(obj=local_model, filepath=local_tmp_file, serializer=serializer)
            local_model_from_disk = unpickle_obj(filepath=local_tmp_file, serializer=serializer)
1103
1104
1105
1106
1107
1108
1109
1110

            assert model_from_disk.client is None

            if set_client:
                assert dask_model.client == client1
            else:
                assert dask_model.client is None

1111
1112
1113
            with pytest.raises(
                lgb.compat.LGBMNotFittedError, match="Cannot access property client_ before calling fit"
            ):
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
                dask_model.client_

            # client will always be None after unpickling
            if set_client:
                from_disk_params = model_from_disk.get_params()
                from_disk_params.pop("client", None)
                dask_params = dask_model.get_params()
                dask_params.pop("client", None)
                assert from_disk_params == dask_params
            else:
                assert model_from_disk.get_params() == dask_model.get_params()
            assert local_model_from_disk.get_params() == local_model.get_params()

            # fitted model should survive pickling round trip, and pickling
            # shouldn't have side effects on the model object
            if set_client:
                dask_model.fit(dX_1, dy_1, group=dg_1)
            else:
                dask_model.fit(dX_2, dy_2, group=dg_2)
            local_model = dask_model.to_local()

            assert "client" not in local_model.get_params()
1136
1137
1138
1139
            no_client_attr_msg = re.compile(
                f"{repr(type(local_model).__name__)} object has no attribute '(client|client_)'"
            )
            with pytest.raises(AttributeError, match=no_client_attr_msg):
1140
                local_model.client
1141
            with pytest.raises(AttributeError, match=no_client_attr_msg):
1142
1143
                local_model.client_

1144
            tmp_file2 = tmp_path / "model-2.pkl"
1145
1146
            pickle_obj(obj=dask_model, filepath=tmp_file2, serializer=serializer)
            fitted_model_from_disk = unpickle_obj(filepath=tmp_file2, serializer=serializer)
1147

1148
            local_tmp_file2 = tmp_path / "local-model-2.pkl"
1149
1150
            pickle_obj(obj=local_model, filepath=local_tmp_file2, serializer=serializer)
            local_fitted_model_from_disk = unpickle_obj(filepath=local_tmp_file2, serializer=serializer)
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188

            if set_client:
                assert dask_model.client == client1
                assert dask_model.client_ == client1
            else:
                assert dask_model.client is None
                assert dask_model.client_ == default_client()
                assert dask_model.client_ == client2

            assert isinstance(fitted_model_from_disk, model_factory)
            assert fitted_model_from_disk.client is None
            assert fitted_model_from_disk.client_ == default_client()
            assert fitted_model_from_disk.client_ == client2

            # client will always be None after unpickling
            if set_client:
                from_disk_params = fitted_model_from_disk.get_params()
                from_disk_params.pop("client", None)
                dask_params = dask_model.get_params()
                dask_params.pop("client", None)
                assert from_disk_params == dask_params
            else:
                assert fitted_model_from_disk.get_params() == dask_model.get_params()
            assert local_fitted_model_from_disk.get_params() == local_model.get_params()

            if set_client:
                preds_orig = dask_model.predict(dX_1).compute()
                preds_loaded_model = fitted_model_from_disk.predict(dX_1).compute()
                preds_orig_local = local_model.predict(X_1)
                preds_loaded_model_local = local_fitted_model_from_disk.predict(X_1)
            else:
                preds_orig = dask_model.predict(dX_2).compute()
                preds_loaded_model = fitted_model_from_disk.predict(dX_2).compute()
                preds_orig_local = local_model.predict(X_2)
                preds_loaded_model_local = local_fitted_model_from_disk.predict(X_2)

            assert_eq(preds_orig, preds_loaded_model)
            assert_eq(preds_orig_local, preds_loaded_model_local)
1189
1190


1191
1192
1193
1194
1195
def test_warns_and_continues_on_unrecognized_tree_learner(cluster):
    with Client(cluster) as client:
        X = da.random.random((1e3, 10))
        y = da.random.random((1e3, 1))
        dask_regressor = lgb.DaskLGBMRegressor(
1196
            client=client, time_out=5, tree_learner="some-nonsense-value", n_estimators=1, num_leaves=2
1197
        )
1198
        with pytest.warns(UserWarning, match="Parameter tree_learner set to some-nonsense-value"):
1199
            dask_regressor = dask_regressor.fit(X, y)
1200

1201
        assert dask_regressor.fitted_
1202

1203

1204
@pytest.mark.parametrize("tree_learner", ["data_parallel", "voting_parallel"])
1205
1206
def test_training_respects_tree_learner_aliases(tree_learner, cluster):
    with Client(cluster) as client:
1207
1208
        task = "regression"
        _, _, _, _, dX, dy, dw, dg = _create_data(objective=task, output="array")
1209
        dask_factory = task_to_dask_factory[task]
1210
        dask_model = dask_factory(client=client, tree_learner=tree_learner, time_out=5, n_estimators=10, num_leaves=15)
1211
1212
1213
        dask_model.fit(dX, dy, sample_weight=dw, group=dg)

        assert dask_model.fitted_
1214
        assert dask_model.get_params()["tree_learner"] == tree_learner
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224


def test_error_on_feature_parallel_tree_learner(cluster):
    with Client(cluster) as client:
        X = da.random.random((100, 10), chunks=(50, 10))
        y = da.random.random(100, chunks=50)
        X, y = client.persist([X, y])
        _ = wait([X, y])
        client.rebalance()
        dask_regressor = lgb.DaskLGBMRegressor(
1225
            client=client, time_out=5, tree_learner="feature_parallel", n_estimators=1, num_leaves=2
1226
        )
1227
        with pytest.raises(lgb.basic.LightGBMError, match="Do not support feature parallel in c api"):
1228
1229
1230
1231
1232
            dask_regressor = dask_regressor.fit(X, y)


def test_errors(cluster):
    with Client(cluster) as client:
1233

1234
        def f(part):
1235
            raise Exception("foo")
1236
1237
1238

        df = dd.demo.make_timeseries()
        df = df.map_partitions(f, meta=df._meta)
1239
        with pytest.raises(Exception) as info:  # noqa: PT011, PT012 # error message needs to be coerced to a string
1240
1241
            lgb.dask._train(client=client, data=df, label=df.x, params={}, model_factory=lgb.LGBMClassifier)
            assert "foo" in str(info.value)
1242
1243


1244
1245
@pytest.mark.parametrize("task", tasks)
@pytest.mark.parametrize("output", data_output)
1246
def test_training_succeeds_even_if_some_workers_do_not_have_any_data(task, output, cluster_three_workers):
1247
1248
    if task == "ranking" and output == "scipy_csr_matrix":
        pytest.skip("LGBMRanker is not currently tested on sparse matrices")
1249

1250
1251
    with Client(cluster_three_workers) as client:
        _, y, _, _, dX, dy, dw, dg = _create_data(
1252
1253
            objective=task,
            output=output,
1254
1255
1256
            group=None,
            n_samples=1_000,
            chunk_size=200,
1257
1258
1259
1260
        )

        dask_model_factory = task_to_dask_factory[task]

1261
        workers = list(client.scheduler_info()["workers"].keys())
1262
1263
1264
1265
1266
1267
1268
        assert len(workers) == 3
        first_two_workers = workers[:2]

        dX = client.persist(dX, workers=first_two_workers)
        dy = client.persist(dy, workers=first_two_workers)
        dw = client.persist(dw, workers=first_two_workers)
        wait([dX, dy, dw])
1269

1270
1271
1272
1273
1274
1275
        workers_with_data = set()
        for coll in (dX, dy, dw):
            for with_data in client.who_has(coll).values():
                workers_with_data.update(with_data)
                assert workers[2] not in with_data
        assert len(workers_with_data) == 2
1276
1277

        params = {
1278
1279
1280
1281
            "time_out": 5,
            "random_state": 42,
            "num_leaves": 10,
            "n_estimators": 20,
1282
1283
        }

1284
        dask_model = dask_model_factory(tree="data", client=client, **params)
1285
1286
        dask_model.fit(dX, dy, group=dg, sample_weight=dw)
        dask_preds = dask_model.predict(dX).compute()
1287
        if task == "regression":
1288
            score = r2_score(y, dask_preds)
1289
        elif task.endswith("classification"):
1290
            score = accuracy_score(y, dask_preds)
1291
        else:
1292
1293
            score = spearmanr(dask_preds, y).correlation
        assert score > 0.9
1294
1295


1296
@pytest.mark.parametrize("task", tasks)
1297
1298
def test_network_params_not_required_but_respected_if_given(task, listen_port, cluster):
    with Client(cluster) as client:
1299
        _, _, _, _, dX, dy, _, dg = _create_data(objective=task, output="array", chunk_size=10, group=None)
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313

        dask_model_factory = task_to_dask_factory[task]

        # rebalance data to be sure that each worker has a piece of the data
        client.rebalance()

        # model 1 - no network parameters given
        dask_model1 = dask_model_factory(
            n_estimators=5,
            num_leaves=5,
        )
        dask_model1.fit(dX, dy, group=dg)
        assert dask_model1.fitted_
        params = dask_model1.get_params()
1314
1315
        assert "local_listen_port" not in params
        assert "machines" not in params
1316
1317

        # model 2 - machines given
1318
        workers = list(client.scheduler_info()["workers"])
1319
        workers_hostname = _get_workers_hostname(cluster)
1320
1321
1322
        remote_sockets, open_ports = lgb.dask._assign_open_ports_to_workers(client, workers)
        for s in remote_sockets.values():
            s.release()
1323
1324
1325
        dask_model2 = dask_model_factory(
            n_estimators=5,
            num_leaves=5,
1326
            machines=",".join([f"{workers_hostname}:{port}" for port in open_ports.values()]),
1327
1328
1329
1330
1331
        )

        dask_model2.fit(dX, dy, group=dg)
        assert dask_model2.fitted_
        params = dask_model2.get_params()
1332
1333
        assert "local_listen_port" not in params
        assert "machines" in params
1334
1335
1336
1337

        # model 3 - local_listen_port given
        # training should fail because LightGBM will try to use the same
        # port for multiple worker processes on the same machine
1338
        dask_model3 = dask_model_factory(n_estimators=5, num_leaves=5, local_listen_port=listen_port)
1339
1340
1341
        error_msg = "has multiple Dask worker processes running on it"
        with pytest.raises(lgb.basic.LightGBMError, match=error_msg):
            dask_model3.fit(dX, dy, group=dg)
1342
1343


1344
@pytest.mark.parametrize("task", tasks)
1345
def test_machines_should_be_used_if_provided(task, cluster):
1346
    pytest.skip("skipping due to timeout issues discussed in https://github.com/microsoft/LightGBM/issues/5390")
1347
    with Client(cluster) as client:
1348
        _, _, _, _, dX, dy, _, dg = _create_data(objective=task, output="array", chunk_size=10, group=None)
1349
1350

        dask_model_factory = task_to_dask_factory[task]
1351
1352

        # rebalance data to be sure that each worker has a piece of the data
1353
        client.rebalance()
1354

1355
        n_workers = len(client.scheduler_info()["workers"])
1356
        assert n_workers > 1
1357
        workers_hostname = _get_workers_hostname(cluster)
1358
        open_ports = lgb.dask._find_n_open_ports(n_workers)
1359
1360
1361
        dask_model = dask_model_factory(
            n_estimators=5,
            num_leaves=5,
1362
            machines=",".join([f"{workers_hostname}:{port}" for port in open_ports]),
1363
1364
1365
1366
        )

        # test that "machines" is actually respected by creating a socket that uses
        # one of the ports mentioned in "machines"
1367
        error_msg = f"Binding port {open_ports[0]} failed"
1368
        with pytest.raises(lgb.basic.LightGBMError, match=error_msg):  # noqa: PT012
1369
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
1370
                s.bind((workers_hostname, open_ports[0]))
1371
                dask_model.fit(dX, dy, group=dg)
1372

1373
1374
1375
        # The above error leaves a worker waiting
        client.restart()

1376
        # an informative error should be raised if "machines" has duplicates
1377
        one_open_port = lgb.dask._find_n_open_ports(1)
1378
        dask_model.set_params(machines=",".join([f"127.0.0.1:{one_open_port}" for _ in range(n_workers)]))
1379
1380
1381
        with pytest.raises(ValueError, match="Found duplicates in 'machines'"):
            dask_model.fit(dX, dy, group=dg)

1382

1383
@pytest.mark.parametrize(
1384
    ("dask_est", "sklearn_est"),
1385
1386
1387
    [
        (lgb.DaskLGBMClassifier, lgb.LGBMClassifier),
        (lgb.DaskLGBMRegressor, lgb.LGBMRegressor),
1388
1389
        (lgb.DaskLGBMRanker, lgb.LGBMRanker),
    ],
1390
)
1391
1392
1393
1394
1395
def test_dask_classes_and_sklearn_equivalents_have_identical_constructors_except_client_arg(dask_est, sklearn_est):
    dask_spec = inspect.getfullargspec(dask_est)
    sklearn_spec = inspect.getfullargspec(sklearn_est)

    # should not allow for any varargs
1396
    assert dask_spec.varargs == sklearn_spec.varargs
1397
1398
1399
1400
    assert dask_spec.varargs is None

    # the only varkw should be **kwargs,
    # for pass-through to parent classes' __init__()
1401
    assert dask_spec.varkw == sklearn_spec.varkw
1402
    assert dask_spec.varkw == "kwargs"
1403
1404

    # "client" should be the only different, and the final argument
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
    assert dask_spec.kwonlyargs == [*sklearn_spec.kwonlyargs, "client"]

    # default values for all constructor arguments should be identical
    #
    # NOTE: if LGBMClassifier / LGBMRanker / LGBMRegressor ever override
    #       any of LGBMModel's constructor arguments, this will need to be updated
    assert dask_spec.kwonlydefaults == {**sklearn_spec.kwonlydefaults, "client": None}

    # only positional argument should be 'self'
    assert dask_spec.args == sklearn_spec.args
    assert dask_spec.args == ["self"]
    assert dask_spec.defaults is None

    # get_params() should be identical, except for "client"
    assert dask_est().get_params() == {**sklearn_est().get_params(), "client": None}
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430


@pytest.mark.parametrize(
    "methods",
    [
        (lgb.DaskLGBMClassifier.fit, lgb.LGBMClassifier.fit),
        (lgb.DaskLGBMClassifier.predict, lgb.LGBMClassifier.predict),
        (lgb.DaskLGBMClassifier.predict_proba, lgb.LGBMClassifier.predict_proba),
        (lgb.DaskLGBMRegressor.fit, lgb.LGBMRegressor.fit),
        (lgb.DaskLGBMRegressor.predict, lgb.LGBMRegressor.predict),
        (lgb.DaskLGBMRanker.fit, lgb.LGBMRanker.fit),
1431
1432
        (lgb.DaskLGBMRanker.predict, lgb.LGBMRanker.predict),
    ],
1433
1434
1435
1436
1437
1438
)
def test_dask_methods_and_sklearn_equivalents_have_similar_signatures(methods):
    dask_spec = inspect.getfullargspec(methods[0])
    sklearn_spec = inspect.getfullargspec(methods[1])
    dask_params = inspect.signature(methods[0]).parameters
    sklearn_params = inspect.signature(methods[1]).parameters
1439
    assert dask_spec.args == sklearn_spec.args[: len(dask_spec.args)]
1440
1441
    assert dask_spec.varargs == sklearn_spec.varargs
    if sklearn_spec.varkw:
1442
        assert dask_spec.varkw == sklearn_spec.varkw[: len(dask_spec.varkw)]
1443
1444
1445
1446
1447
    assert dask_spec.kwonlyargs == sklearn_spec.kwonlyargs
    assert dask_spec.kwonlydefaults == sklearn_spec.kwonlydefaults
    for param in dask_spec.args:
        error_msg = f"param '{param}' has different default values in the methods"
        assert dask_params[param].default == sklearn_params[param].default, error_msg
1448
1449


1450
@pytest.mark.parametrize("task", tasks)
1451
def test_training_succeeds_when_data_is_dataframe_and_label_is_column_array(task, cluster):
1452
    with Client(cluster):
1453
        _, _, _, _, dX, dy, dw, dg = _create_data(objective=task, output="dataframe", group=None)
1454
1455
1456
1457
1458

        model_factory = task_to_dask_factory[task]

        dy = dy.to_dask_array(lengths=True)
        dy_col_array = dy.reshape(-1, 1)
1459
1460
        assert len(dy_col_array.shape) == 2
        assert dy_col_array.shape[1] == 1
1461

1462
        params = {"n_estimators": 1, "num_leaves": 3, "random_state": 0, "time_out": 5}
1463
1464
1465
        model = model_factory(**params)
        model.fit(dX, dy_col_array, sample_weight=dw, group=dg)
        assert model.fitted_
1466
1467


1468
1469
@pytest.mark.parametrize("task", tasks)
@pytest.mark.parametrize("output", data_output)
1470
def test_init_score(task, output, cluster, rng):
1471
1472
    if task == "ranking" and output == "scipy_csr_matrix":
        pytest.skip("LGBMRanker is not currently tested on sparse matrices")
1473

1474
    with Client(cluster) as client:
1475
        _, _, _, _, dX, dy, dw, dg = _create_data(objective=task, output=output, group=None)
1476

1477
1478
        model_factory = task_to_dask_factory[task]

1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
        params = {
            "n_estimators": 1,
            "num_leaves": 2,
            "time_out": 5,
            "seed": 708,
            "deterministic": True,
            "force_row_wise": True,
            "num_thread": 1,
        }
        num_classes = 1
1489
        if task == "multiclass-classification":
1490
            num_classes = 3
1491

1492
        if output.startswith("dataframe"):
1493
            init_scores = dy.map_partitions(lambda x: pd.DataFrame(rng.uniform(size=(x.size, num_classes))))
1494
        else:
1495
1496
            init_scores = dy.map_blocks(lambda x: rng.uniform(size=(x.size, num_classes)))

1497
        model = model_factory(client=client, **params)
1498
1499
1500
1501
1502
1503
1504
1505
        model.fit(dX, dy, sample_weight=dw, group=dg)
        pred = model.predict(dX, raw_score=True)

        model_init_score = model_factory(client=client, **params)
        model_init_score.fit(dX, dy, sample_weight=dw, init_score=init_scores, group=dg)
        pred_init_score = model_init_score.predict(dX, raw_score=True)

        # check if init score changes predictions
1506
        with pytest.raises(AssertionError):  # noqa: PT011
1507
            assert_eq(pred, pred_init_score)
1508
1509


1510
def sklearn_checks_to_run():
1511
    check_names = ["check_estimator_get_tags_default_keys", "check_get_params_invariance", "check_set_params"]
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
    for check_name in check_names:
        check_func = getattr(sklearn_checks, check_name, None)
        if check_func:
            yield check_func


def _tested_estimators():
    for Estimator in [lgb.DaskLGBMClassifier, lgb.DaskLGBMRegressor]:
        yield Estimator()


@pytest.mark.parametrize("estimator", _tested_estimators())
@pytest.mark.parametrize("check", sklearn_checks_to_run())
1525
def test_sklearn_integration(estimator, check, cluster):
1526
    with Client(cluster):
1527
1528
1529
        estimator.set_params(local_listen_port=18000, time_out=5)
        name = type(estimator).__name__
        check(name, estimator)
1530
1531
1532
1533
1534


# this test is separate because it takes a not-yet-constructed estimator
@pytest.mark.parametrize("estimator", list(_tested_estimators()))
def test_parameters_default_constructible(estimator):
1535
    name = estimator.__class__.__name__
1536
    Estimator = estimator
1537
    sklearn_checks.check_parameters_default_constructible(name, Estimator)
1538
1539


1540
1541
@pytest.mark.parametrize("task", tasks)
@pytest.mark.parametrize("output", data_output)
1542
def test_predict_with_raw_score(task, output, cluster):
1543
1544
    if task == "ranking" and output == "scipy_csr_matrix":
        pytest.skip("LGBMRanker is not currently tested on sparse matrices")
1545

1546
    with Client(cluster) as client:
1547
        _, _, _, _, dX, dy, _, dg = _create_data(objective=task, output=output, group=None)
1548

1549
        model_factory = task_to_dask_factory[task]
1550
        params = {"client": client, "n_estimators": 1, "num_leaves": 2, "time_out": 5, "min_sum_hessian": 0}
1551
1552
1553
1554
1555
1556
        model = model_factory(**params)
        model.fit(dX, dy, group=dg)
        raw_predictions = model.predict(dX, raw_score=True).compute()

        trees_df = model.booster_.trees_to_dataframe()
        leaves_df = trees_df[trees_df.node_depth == 2]
1557
        if task == "multiclass-classification":
1558
1559
            for i in range(model.n_classes_):
                class_df = leaves_df[leaves_df.tree_index == i]
1560
                assert set(raw_predictions[:, i]) == set(class_df["value"])
1561
        else:
1562
            assert set(raw_predictions) == set(leaves_df["value"])
1563

1564
        if task.endswith("classification"):
1565
1566
            pred_proba_raw = model.predict_proba(dX, raw_score=True).compute()
            assert_eq(raw_predictions, pred_proba_raw)
1567
1568


1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
@pytest.mark.parametrize("output", data_output)
@pytest.mark.parametrize("use_init_score", [False, True])
def test_predict_stump(output, use_init_score, cluster, rng):
    with Client(cluster) as client:
        _, _, _, _, dX, dy, _, _ = _create_data(objective="binary-classification", n_samples=1_000, output=output)

        params = {"objective": "binary", "n_estimators": 5, "min_data_in_leaf": 1_000}

        if not use_init_score:
            init_scores = None
        elif output.startswith("dataframe"):
            init_scores = dy.map_partitions(lambda x: pd.DataFrame(rng.uniform(size=x.size)))
        else:
            init_scores = dy.map_blocks(lambda x: rng.uniform(size=x.size))

        model = lgb.DaskLGBMClassifier(client=client, **params)
        model.fit(dX, dy, init_score=init_scores)
        preds_1 = model.predict(dX, raw_score=True, num_iteration=1).compute()
        preds_all = model.predict(dX, raw_score=True).compute()

        if use_init_score:
            # if init_score was provided, a model of stumps should predict all 0s
            all_zeroes = np.full_like(preds_1, fill_value=0.0)
            assert_eq(preds_1, all_zeroes)
            assert_eq(preds_all, all_zeroes)
        else:
            # if init_score was not provided, prediction for a model of stumps should be
            # the "average" of the labels
            y_avg = np.log(dy.mean() / (1.0 - dy.mean()))
            assert_eq(preds_1, np.full_like(preds_1, fill_value=y_avg))
            assert_eq(preds_all, np.full_like(preds_all, fill_value=y_avg))


1602
def test_distributed_quantized_training(tmp_path, cluster):
1603
    with Client(cluster) as client:
1604
        X, y, w, _, dX, dy, dw, _ = _create_data(objective="regression", output="array")
1605

1606
        np.savetxt(tmp_path / "data_dask.csv", np.hstack([np.array([y]).T, X]), fmt="%f,%f,%f,%f,%f")
1607
1608

        params = {
1609
            "boosting_type": "gbdt",
1610
1611
            "n_estimators": 50,
            "num_leaves": 31,
1612
1613
1614
1615
            "use_quantized_grad": True,
            "num_grad_quant_bins": 30,
            "quant_train_renew_leaf": True,
            "verbose": -1,
1616
1617
        }

1618
        quant_dask_classifier = lgb.DaskLGBMRegressor(client=client, time_out=5, **params)
1619
1620
1621
1622
1623
        quant_dask_classifier = quant_dask_classifier.fit(dX, dy, sample_weight=dw)
        quant_p1 = quant_dask_classifier.predict(dX)
        quant_rmse = np.sqrt(np.mean((quant_p1.compute() - y) ** 2))

        params["use_quantized_grad"] = False
1624
        dask_classifier = lgb.DaskLGBMRegressor(client=client, time_out=5, **params)
1625
1626
1627
1628
        dask_classifier = dask_classifier.fit(dX, dy, sample_weight=dw)
        p1 = dask_classifier.predict(dX)
        rmse = np.sqrt(np.mean((p1.compute() - y) ** 2))
        assert quant_rmse < rmse + 7.0