dask.py 17.7 KB
Newer Older
1
2
3
# coding: utf-8
"""Distributed training with LightGBM and Dask.distributed.

4
5
6
7
This module enables you to perform distributed training with LightGBM on
Dask.Array and Dask.DataFrame collections.

It is based on dask-lightgbm, which was based on dask-xgboost.
8
"""
9
import socket
10
from collections import defaultdict
11
from copy import deepcopy
12
from typing import Dict, Iterable
13
14
15
16
from urllib.parse import urlparse

import numpy as np
import pandas as pd
17
18
import scipy.sparse as ss

19
20
21
from dask import array as da
from dask import dataframe as dd
from dask import delayed
22
from dask.distributed import Client, default_client, get_worker, wait
23

24
from .basic import _ConfigAliases, _LIB, _log_warning, _safe_call
25
from .sklearn import LGBMClassifier, LGBMRegressor, LGBMRanker
26
27


28
29
def _find_open_port(worker_ip: str, local_listen_port: int, ports_to_skip: Iterable[int]) -> int:
    """Find an open port.
30

31
32
    This function tries to find a free port on the machine it's run on. It is intended to
    be run once on each Dask worker, sequentially.
33

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
    Parameters
    ----------
    worker_ip : str
        IP address for the Dask worker.
    local_listen_port : int
        First port to try when searching for open ports.
    ports_to_skip: Iterable[int]
        An iterable of integers referring to ports that should be skipped. Since multiple Dask
        workers can run on the same physical machine, this method may be called multiple times
        on the same machine. ``ports_to_skip`` is used to ensure that LightGBM doesn't try to use
        the same port for two worker processes running on the same machine.

    Returns
    -------
    result : int
        A free port on the machine referenced by ``worker_ip``.
    """
    max_tries = 1000
    out_port = None
    found_port = False
    for i in range(max_tries):
        out_port = local_listen_port + i
        if out_port in ports_to_skip:
            continue
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind((worker_ip, out_port))
            found_port = True
            break
        # if unavailable, you'll get OSError: Address already in use
        except OSError:
            continue
    if not found_port:
        msg = "LightGBM tried %s:%d-%d and could not create a connection. Try setting local_listen_port to a different value."
        raise RuntimeError(msg % (worker_ip, local_listen_port, out_port))
    return out_port


def _find_ports_for_workers(client: Client, worker_addresses: Iterable[str], local_listen_port: int) -> Dict[str, int]:
    """Find an open port on each worker.

    LightGBM distributed training uses TCP sockets by default, and this method is used to
    identify open ports on each worker so LightGBM can reliable create those sockets.
77
78
79

    Parameters
    ----------
80
81
82
83
    client : dask.distributed.Client
        Dask client.
    worker_addresses : Iterable[str]
        An iterable of addresses for workers in the cluster. These are strings of the form ``<protocol>://<host>:port``
84
    local_listen_port : int
85
        First port to try when searching for open ports.
86
87
88

    Returns
    -------
89
90
    result : Dict[str, int]
        Dictionary where keys are worker addresses and values are an open port for LightGBM to use.
91
    """
92
93
94
95
96
97
98
99
100
101
102
103
104
105
    lightgbm_ports = set()
    worker_ip_to_port = {}
    for worker_address in worker_addresses:
        port = client.submit(
            func=_find_open_port,
            workers=[worker_address],
            worker_ip=urlparse(worker_address).hostname,
            local_listen_port=local_listen_port,
            ports_to_skip=lightgbm_ports
        ).result()
        lightgbm_ports.add(port)
        worker_ip_to_port[worker_address] = port

    return worker_ip_to_port
106
107
108
109
110
111
112
113
114
115
116
117
118


def _concat(seq):
    if isinstance(seq[0], np.ndarray):
        return np.concatenate(seq, axis=0)
    elif isinstance(seq[0], (pd.DataFrame, pd.Series)):
        return pd.concat(seq, axis=0)
    elif isinstance(seq[0], ss.spmatrix):
        return ss.vstack(seq, format='csr')
    else:
        raise TypeError('Data must be one of: numpy arrays, pandas dataframes, sparse matrices (from scipy). Got %s.' % str(type(seq[0])))


119
def _train_part(params, model_factory, list_of_parts, worker_address_to_port, return_model,
120
                time_out=120, **kwargs):
121
122
123
124
125
126
127
128
129
130
131
132
    local_worker_address = get_worker().address
    machine_list = ','.join([
        '%s:%d' % (urlparse(worker_address).hostname, port)
        for worker_address, port
        in worker_address_to_port.items()
    ])
    network_params = {
        'machines': machine_list,
        'local_listen_port': worker_address_to_port[local_worker_address],
        'time_out': time_out,
        'num_machines': len(worker_address_to_port)
    }
133
134
    params.update(network_params)

135
136
    is_ranker = issubclass(model_factory, LGBMRanker)

137
138
139
140
141
142
143
    # Concatenate many parts into one
    parts = tuple(zip(*list_of_parts))
    data = _concat(parts[0])
    label = _concat(parts[1])

    try:
        model = model_factory(**params)
144
145
146

        if is_ranker:
            group = _concat(parts[-1])
147
148
149
150
            if len(parts) == 4:
                weight = _concat(parts[2])
            else:
                weight = None
151
152
            model.fit(data, y=label, sample_weight=weight, group=group, **kwargs)
        else:
153
154
155
156
            if len(parts) == 3:
                weight = _concat(parts[2])
            else:
                weight = None
157
158
            model.fit(data, y=label, sample_weight=weight, **kwargs)

159
160
161
162
163
164
165
166
167
    finally:
        _safe_call(_LIB.LGBM_NetworkFree())

    return model if return_model else None


def _split_to_parts(data, is_matrix):
    parts = data.to_delayed()
    if isinstance(parts, np.ndarray):
168
169
170
171
        if is_matrix:
            assert parts.shape[1] == 1
        else:
            assert parts.ndim == 1 or parts.shape[1] == 1
172
173
174
175
        parts = parts.flatten().tolist()
    return parts


176
def _train(client, data, label, params, model_factory, sample_weight=None, group=None, **kwargs):
177
178
179
180
181
182
183
184
185
186
    """Inner train routine.

    Parameters
    ----------
    client: dask.Client - client
    X : dask array of shape = [n_samples, n_features]
        Input feature matrix.
    y : dask array of shape = [n_samples]
        The target values (class labels in classification, real numbers in regression).
    params : dict
187
    model_factory : lightgbm.LGBMClassifier, lightgbm.LGBMRegressor, or lightgbm.LGBMRanker class
188
    sample_weight : array-like of shape = [n_samples] or None, optional (default=None)
189
190
191
192
193
194
195
        Weights of training data.
    group : array-like or None, optional (default=None)
        Group/query data.
        Only used in the learning-to-rank task.
        sum(group) = n_samples.
        For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
        where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
196
    """
197
198
    params = deepcopy(params)

199
    # Split arrays/dataframes into parts. Arrange parts into tuples to enforce co-locality
200
201
202
203
204
205
206
207
208
209
210
211
    data_parts = _split_to_parts(data=data, is_matrix=True)
    label_parts = _split_to_parts(data=label, is_matrix=False)

    if sample_weight is not None:
        weight_parts = _split_to_parts(data=sample_weight, is_matrix=False)
    else:
        weight_parts = None

    if group is not None:
        group_parts = _split_to_parts(data=group, is_matrix=False)
    else:
        group_parts = None
212
213
214
215
216
217
218
219

    # choose between four options of (sample_weight, group) being (un)specified
    if weight_parts is None and group_parts is None:
        parts = zip(data_parts, label_parts)
    elif weight_parts is not None and group_parts is None:
        parts = zip(data_parts, label_parts, weight_parts)
    elif weight_parts is None and group_parts is not None:
        parts = zip(data_parts, label_parts, group_parts)
220
    else:
221
        parts = zip(data_parts, label_parts, weight_parts, group_parts)
222
223

    # Start computation in the background
224
    parts = list(map(delayed, parts))
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    parts = client.compute(parts)
    wait(parts)

    for part in parts:
        if part.status == 'error':
            return part  # trigger error locally

    # Find locations of all parts and map them to particular Dask workers
    key_to_part_dict = dict([(part.key, part) for part in parts])
    who_has = client.who_has(parts)
    worker_map = defaultdict(list)
    for key, workers in who_has.items():
        worker_map[next(iter(workers))].append(key_to_part_dict[key])

    master_worker = next(iter(worker_map))
    worker_ncores = client.ncores()

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    tree_learner = None
    for tree_learner_param in _ConfigAliases.get('tree_learner'):
        tree_learner = params.get(tree_learner_param)
        if tree_learner is not None:
            break

    allowed_tree_learners = {
        'data',
        'data_parallel',
        'feature',
        'feature_parallel',
        'voting',
        'voting_parallel'
    }
    if tree_learner is None:
257
        _log_warning('Parameter tree_learner not set. Using "data" as default')
258
        params['tree_learner'] = 'data'
259
    elif tree_learner.lower() not in allowed_tree_learners:
260
        _log_warning('Parameter tree_learner set to %s, which is not allowed. Using "data" as default' % tree_learner)
261
262
263
264
265
266
267
268
        params['tree_learner'] = 'data'

    local_listen_port = 12400
    for port_param in _ConfigAliases.get('local_listen_port'):
        val = params.get(port_param)
        if val is not None:
            local_listen_port = val
            break
269

270
271
272
273
274
275
276
277
278
    # find an open port on each worker. note that multiple workers can run
    # on the same machine, so this needs to ensure that each one gets its
    # own port
    worker_address_to_port = _find_ports_for_workers(
        client=client,
        worker_addresses=worker_map.keys(),
        local_listen_port=local_listen_port
    )

279
280
281
282
    # num_threads is set below, so remove it and all aliases of it from params
    for num_thread_alias in _ConfigAliases.get('num_threads'):
        params.pop(num_thread_alias, None)

283
    # Tell each worker to train on the parts that it has locally
284
285
286
287
288
289
290
291
292
293
294
295
296
    futures_classifiers = [
        client.submit(
            _train_part,
            model_factory=model_factory,
            params={**params, 'num_threads': worker_ncores[worker]},
            list_of_parts=list_of_parts,
            worker_address_to_port=worker_address_to_port,
            time_out=params.get('time_out', 120),
            return_model=(worker == master_worker),
            **kwargs
        )
        for worker, list_of_parts in worker_map.items()
    ]
297
298
299
300
301
302

    results = client.gather(futures_classifiers)
    results = [v for v in results if v]
    return results[0]


303
def _predict_part(part, model, raw_score, pred_proba, pred_leaf, pred_contrib, **kwargs):
304
305
306
307
    data = part.values if isinstance(part, pd.DataFrame) else part

    if data.shape[0] == 0:
        result = np.array([])
308
309
310
311
312
313
314
315
    elif pred_proba:
        result = model.predict_proba(
            data,
            raw_score=raw_score,
            pred_leaf=pred_leaf,
            pred_contrib=pred_contrib,
            **kwargs
        )
316
    else:
317
318
319
320
321
322
323
        result = model.predict(
            data,
            raw_score=raw_score,
            pred_leaf=pred_leaf,
            pred_contrib=pred_contrib,
            **kwargs
        )
324
325

    if isinstance(part, pd.DataFrame):
326
        if pred_proba or pred_contrib:
327
328
329
330
331
332
333
            result = pd.DataFrame(result, index=part.index)
        else:
            result = pd.Series(result, index=part.index, name='predictions')

    return result


334
335
def _predict(model, data, raw_score=False, pred_proba=False, pred_leaf=False, pred_contrib=False,
             dtype=np.float32, **kwargs):
336
337
338
339
    """Inner predict routine.

    Parameters
    ----------
340
    model : lightgbm.LGBMClassifier, lightgbm.LGBMRegressor, or lightgbm.LGBMRanker class
341
342
    data : dask array of shape = [n_samples, n_features]
        Input feature matrix.
343
344
345
346
347
348
    pred_proba : bool, optional (default=False)
        Should method return results of ``predict_proba`` (``pred_proba=True``) or ``predict`` (``pred_proba=False``).
    pred_leaf : bool, optional (default=False)
        Whether to predict leaf index.
    pred_contrib : bool, optional (default=False)
        Whether to predict feature contributions.
349
    dtype : np.dtype
350
        Dtype of the output.
351
352
    kwargs : dict
        Other parameters passed to ``predict`` or ``predict_proba`` method.
353
354
    """
    if isinstance(data, dd._Frame):
355
356
357
358
359
360
361
362
363
        return data.map_partitions(
            _predict_part,
            model=model,
            raw_score=raw_score,
            pred_proba=pred_proba,
            pred_leaf=pred_leaf,
            pred_contrib=pred_contrib,
            **kwargs
        ).values
364
    elif isinstance(data, da.Array):
365
        if pred_proba:
366
367
368
            kwargs['chunks'] = (data.chunks[0], (model.n_classes_,))
        else:
            kwargs['drop_axis'] = 1
369
370
371
372
373
374
375
376
377
378
        return data.map_blocks(
            _predict_part,
            model=model,
            raw_score=raw_score,
            pred_proba=pred_proba,
            pred_leaf=pred_leaf,
            pred_contrib=pred_contrib,
            dtype=dtype,
            **kwargs
        )
379
380
381
382
383
384
    else:
        raise TypeError('Data must be either Dask array or dataframe. Got %s.' % str(type(data)))


class _LGBMModel:

385
    def _fit(self, model_factory, X, y=None, sample_weight=None, group=None, client=None, **kwargs):
386
387
388
389
390
        """Docstring is inherited from the LGBMModel."""
        if client is None:
            client = default_client()

        params = self.get_params(True)
391
392
393
394
395
396
397
398
399
400
401

        model = _train(
            client=client,
            data=X,
            label=y,
            params=params,
            model_factory=model_factory,
            sample_weight=sample_weight,
            group=group,
            **kwargs
        )
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425

        self.set_params(**model.get_params())
        self._copy_extra_params(model, self)

        return self

    def _to_local(self, model_factory):
        model = model_factory(**self.get_params())
        self._copy_extra_params(self, model)
        return model

    @staticmethod
    def _copy_extra_params(source, dest):
        params = source.get_params()
        attributes = source.__dict__
        extra_param_names = set(attributes.keys()).difference(params.keys())
        for name in extra_param_names:
            setattr(dest, name, attributes[name])


class DaskLGBMClassifier(_LGBMModel, LGBMClassifier):
    """Distributed version of lightgbm.LGBMClassifier."""

    def fit(self, X, y=None, sample_weight=None, client=None, **kwargs):
426
        """Docstring is inherited from the lightgbm.LGBMClassifier.fit."""
427
428
429
430
431
432
433
434
435
        return self._fit(
            model_factory=LGBMClassifier,
            X=X,
            y=y,
            sample_weight=sample_weight,
            client=client,
            **kwargs
        )

436
437
438
439
    fit.__doc__ = LGBMClassifier.fit.__doc__

    def predict(self, X, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMClassifier.predict."""
440
441
442
443
444
445
446
        return _predict(
            model=self.to_local(),
            data=X,
            dtype=self.classes_.dtype,
            **kwargs
        )

447
448
449
450
    predict.__doc__ = LGBMClassifier.predict.__doc__

    def predict_proba(self, X, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMClassifier.predict_proba."""
451
452
453
454
455
456
457
        return _predict(
            model=self.to_local(),
            data=X,
            pred_proba=True,
            **kwargs
        )

458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
    predict_proba.__doc__ = LGBMClassifier.predict_proba.__doc__

    def to_local(self):
        """Create regular version of lightgbm.LGBMClassifier from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMClassifier
        """
        return self._to_local(LGBMClassifier)


class DaskLGBMRegressor(_LGBMModel, LGBMRegressor):
    """Docstring is inherited from the lightgbm.LGBMRegressor."""

    def fit(self, X, y=None, sample_weight=None, client=None, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMRegressor.fit."""
475
476
477
478
479
480
481
482
483
        return self._fit(
            model_factory=LGBMRegressor,
            X=X,
            y=y,
            sample_weight=sample_weight,
            client=client,
            **kwargs
        )

484
485
486
487
    fit.__doc__ = LGBMRegressor.fit.__doc__

    def predict(self, X, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMRegressor.predict."""
488
489
490
491
492
493
        return _predict(
            model=self.to_local(),
            data=X,
            **kwargs
        )

494
495
496
497
498
499
500
501
502
503
    predict.__doc__ = LGBMRegressor.predict.__doc__

    def to_local(self):
        """Create regular version of lightgbm.LGBMRegressor from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMRegressor
        """
        return self._to_local(LGBMRegressor)
504
505
506
507
508
509
510
511
512
513


class DaskLGBMRanker(_LGBMModel, LGBMRanker):
    """Docstring is inherited from the lightgbm.LGBMRanker."""

    def fit(self, X, y=None, sample_weight=None, init_score=None, group=None, client=None, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMRanker.fit."""
        if init_score is not None:
            raise RuntimeError('init_score is not currently supported in lightgbm.dask')

514
515
516
517
518
519
520
521
522
523
        return self._fit(
            model_factory=LGBMRanker,
            X=X,
            y=y,
            sample_weight=sample_weight,
            group=group,
            client=client,
            **kwargs
        )

524
525
526
527
528
    fit.__doc__ = LGBMRanker.fit.__doc__

    def predict(self, X, **kwargs):
        """Docstring is inherited from the lightgbm.LGBMRanker.predict."""
        return _predict(self.to_local(), X, **kwargs)
529

530
531
532
533
534
535
536
537
538
539
    predict.__doc__ = LGBMRanker.predict.__doc__

    def to_local(self):
        """Create regular version of lightgbm.LGBMRanker from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMRanker
        """
        return self._to_local(LGBMRanker)