Distributed computing with Dask

[mrocklin]

Hello, I am an author of Dask, a library for parallel and distributed computing in Python. I am curious if there is interest within this community to collaborate on distributing XGBoost on Dask either for parallel training or for ETL.

There are probably two components of Dask that are relevant for this project:

1. A generic system for parallel and distributed computing, built on arbitrary dynamic task scheduling. The relevant APIs here are probably dask.delayed and concurrent.futures
2. A parallel and distributed subset of the Pandas API, dask.dataframe useful for feature engineering and data pre-processing. This doesn't implement the entire Pandas API, but comes decently close.

Is there interest in collaborating here?

[terrytangyuan]

@mrocklin I thought Dask has integrations with sklearn. Did you take a look at our sklearn wrapper to see if it will work with that?

[mrocklin]

Meaningfully integrating with a distributed system typically has to be done at the per-algorithm level rather than at the library level. There are some ways in which SKLearn and Dask can help each other, yes, but they're not particularly deep.

[terrytangyuan]

Dask dataframe would be a good start. In our code base, we have a check for pandas dataframe. That might be where dask dataframe would fit as a start.

[mrocklin]

So what happens if someone arrives with a multi-terabyte dask dataframe? Do you just convert it to Pandas and proceed? Or is there a way to parallelize XGBoost intelligently across a cluster, pointing to the various pandas dataframes that make up a dask dataframe?

[terrytangyuan]

Users can specify batch size? I'd imagine users can be benefited through partial_fit.

cc @tqchen who's more familiar with the distributed part of the code.

[tqchen]

The distributed version of xgboost can be hooked into a distributed job launcher, ideally get data partition feed into xgboost then continues.

@mrocklin I think the most relevant part is the xgboost-spark and xgboost-flink module, which embeds xgboost into mapPartition function of spark/flink. I guess there would be something similar in Dask

The requirement from xgboost side is that XGBoost handles inter process connection by rabit, and will need to start a tracker (that connects each job) from the client side.

[tqchen]

see relevant code in https://github.com/dmlc/xgboost/blob/master/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala#L112

Rabit is designed to be embedded into other distributed system, so I think it might not be too hard to do the adjustment in python side.

[mrocklin]

Launching other distributed systems from Dask is usually pretty doable. How do you move data from the hosting distributed system (spark/flink/dask) to xg-boost? Or is this for distributed training on small-data?

[mrocklin]

More concretely, I expect to build a system as follows:

• On every dask worker I start a Rabit server. Dask gives these Rabit servers enough information to find each other.
• I create some local XGBoost state on every worker that represents the currently training model
• I repeatedly feed this per-worker object pandas dataframes or numpy arrays
• I listen for some signal from XGBoost that tells me to stop

Does this match your expectation? Is it easy for you to point me to relevant Python API?

[tqchen]

Yes, see relevant information here https://github.com/dmlc/xgboost/blob/master/tests/distributed/ for python API.

What you will need to do additionally is to start a rabit tracker at driver side(likely to be the place which drives dask), this is done in the dmlc-submit script here https://github.com/dmlc/dmlc-core/tree/master/tracker/dmlc_tracker

[mrocklin]

OK, filling out my outline from before:

Before running any XGBoost code we set up a Rabit network

On the driver/scheduler node we start a rabit tracker

envs = {'DMLC_NUM_WORKER' : nworker,
        'DMLC_NUM_SERVER' : nserver}

rabit = RabitTracker(hostIP=ip_address, nslave=num_workers)
envs.update(rabit.slave_envs())
rabit.start(args.num_workers)  # manages connections in background thread

I may also go through a similar process to start a PSTracker. Should this be on the same centralized machine or should it be elsewhere within the network? Should there be a few of these? Should this be user configurable?

Eventually I have my tracker (and pstrackers?) join the rabit network and block.

rabit.join()  # join network

On worker nodes I need to dump these environment variables (which I'll move through normal dask channels) into the local environment. Then just calling xgboost.rabit.init() should suffice

import os
os.environ.update(envs)
xgboost.rabit.init()

Looking at the Rabit code it looks like environment variables are the only way to provide this information. Can you verify this? Is there a way to supply tracker host/port information as direct inputs?

Training

Then I convert my numpy arrays / pandas dataframes / scipy sparse arrays to DMatrix objects, this seems relatively straightforward. However I'm likely to have several batches of data per worker. Is there a clean way to call train several times with more data as it becomes available? I'm concerned about the comments on these lines:

# Run training, all the features in training API is available.
# Currently, this script only support calling train once for fault recovery purpose.
bst = xgb.train(param, dtrain, num_round, watchlist, early_stopping_rounds=2)

Do we need to wait for all data to arrive before starting training?

Example dataset / problem

Assuming that I have everything above correct then is there a standard distributed training example that people use for demonstration?

[tqchen]

There is no need to start pstracker.

• Tracker only need to be started in one place, likely on scheduler(driver), it has no data heavy job, and only serves to connect the works.
• The env args can be passed as kwargs in rabit.init
• Since tree boosting is a batched algorithm, we do need to wait for all data to be ingested before start training.
  • Note however each worker only need to take a shard(subset of rows) of data.
  • Ideally, we should use the data iter interface to pass the data into DMatrix as mini-batch fashion, so the entire dataset does not have to sit in memory
  • This is done via https://github.com/dmlc/xgboost/blob/master/include/xgboost/c_api.h#L117, which do not yet have a python wrapper.
  • For the first solution, I would recommend to directly pass by array

[mrocklin]

I had some time to play with this this morning. Results here: https://github.com/mrocklin/dask-xgboost

So far it only handles the distributed learning of a single in-memory dataset. Some questions arose:

1. What is the best way to serialize and pass around DMatrix objects?
2. What is the best way to serialize and return a Booster result?
3. How do the environment variables listed above map to arguments in rabit.init? What precisely is the expected form of inputs to rabit.init? Passing the result of slave_envs() to rabit.init obviously won't work because it expects a list. Should we convert each keyname to --key, perhaps dropping the DMLC prefix and converting to lowercase?
4. Is there a good way to test correctness? How do we compare two Booster objects? Should we expect distributed training to produce precisely the same result and sequential training?

[tqchen]

• You don't normally serialize DMatrix, it is more like a training time data holder, I assume data is passed around and shared by dask (array/dataframe), then passed to xgboost
  • We can explore better ways to pass data other than directly through in-memory array, possibly by exposing an data iterator to xgboost
• You can pickle Booster, as long as xgboost is installed in both side.
• Sorry about not elaborating how things are passed, it should be

rabit.init(['DMLC_KEY1=VALUE1', 'DMLC_KEY2=VALUE2']

• Normally the booster trained from distributed and single machine are not the same, but here are a few things to check
  • The booster returned from all worker should be identical
  • Looking for the predictive validation error, it should be nroughly as low as single machine case

[mrocklin]

Two more questions generally about how this gets used (I have no experience with XGBoost and only a little experience with machine learning, please forgive my ignorance).

1. Is it reasonable to use multiple workers on the same input data? (XGBoost is computationally bound?)
2. If we operate on larger datasets do I have to do anything special to tell each XGBoost worker that its data differs from its peers?

Which use case is more common?