dask-lgbm.py

from urllib.parse import urlparse
import lightgbm as lgb
import threading


def _parse_machines(workers, listen_port):
    """From dask worker info to LightGBM mlist"""
    # TODO: Assert that we're using TCP?
    mlist = ['127.0.0.1:{}'.format(listen_port)] + [urlparse(worker).netloc
                                                    for worker in workers]
    return mlist


def train(params, train):
    bst = lgb.train(params, train)
    return bst


def main():
    from sklearn.datasets import make_classification
    from distributed import Client

    # setup
    port = 12400
    c = Client()
    print('scheduler', c.scheduler_info()['address'])
    machines = _parse_machines(c.scheduler_info()['workers'], port)
    params = {"machines": machines}

    # Data
    X, y = make_classification()
    dset = lgb.Dataset(X, label=y)

    # Train
    train(params, dset)


if __name__ == '__main__':
    main()