Data wrangling using CPU workers and training xgboost using GPU workers with dask

Viewed 101

I am trying to read 200 parquet files from hdfs and then try to train a model using 4 GPUs. I have 48 vcores available on the machine as well. If I start the cluster with just the GPU workers then reading part is going to be very slow (since it just uses 4 cpu workers assigned to the gpu workers and you can't really run more workers than the number of gpus you have unless you run them on separate shells and then it gets nasty because you are on your own for memory management issues.) I would like to read the files using CPU workers, play with the data with the cpu workers and then train an xgboost model using GPU workers. I read the documentation here about how to start and assign workers with different resources to different tasks. Also I have seen this question, but I am confused a bit.

Here is the the code I am trying to run to read the .parquet files:

import dask.dataframe as dd 
df = dd \
    .read_parquet(
        "hdfs://address/to/the/*.parquet",
        storage_options = {
            "user":user,
            "kerb_ticket":kerb_ticket},
        engine='pyarrow') \
    .persist()

This will automatically use all the cpu and gpu workers which is fine. After this I need to create my training data and label. Let's say I have X_train, y_train, and params. Here I convert them to dask_cudf:

X_train = dask_cudf.from_dask_dataframe(X_train)
y_train = dask_cudf.from_dask_dataframe(y_train)

Here is the part that I need to use just GPU workers:

Xy = dxgb.DaskDMatrix(client, X_train, y_train)

in order to follow the document I should convert it to this:

Xy = client.submit(dxgb.DaskDMatrix, client, X_train, y_train, resources={'GPU': 1})

But then I'll get this error:

distributed.protocol.pickle - INFO - Failed to serialize (<Client: 'tcp://169.68.236.35:8786' processes=52 threads=52, memory=1.97 TiB>, <dask_cudf.DataFrame | 19200 tasks | 200 npartitions>, <dask_cudf.Series | 600 tasks | 200 npartitions>). Exception: cannot pickle 'socket' object
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/envs/dask/lib/python3.8/site-packages/distributed/protocol/pickle.py in dumps(x, buffer_callback, protocol)
     48         buffers.clear()
---> 49         result = pickle.dumps(x, **dump_kwargs)
     50         if len(result) < 1000:

/envs/dask/lib/python3.8/socket.py in __getstate__(self)
    271     def __getstate__(self):
--> 272         raise TypeError(f"cannot pickle {self.__class__.__name__!r} object")
    273 

TypeError: cannot pickle 'socket' object

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-12-0d6a943365a9> in <module>
      1 # Xy = dxgb.DaskDMatrix(client, X_train, y_train)
      2 # Xy = dxgb.DaskDeviceQuantileDMatrix(client, X_train, y_train)
----> 3 Xy = client.submit(dxgb.DaskDMatrix, client, X_train, y_train, resources={'GPU': 1})
      4 # Xy_valid = dxgb.DaskDMatrix(client, X_valid, y_valid)

/envs/dask/lib/python3.8/site-packages/distributed/client.py in submit(self, func, key, workers, resources, retries, priority, fifo_timeout, allow_other_workers, actor, actors, pure, *args, **kwargs)
   1629             dsk = {skey: (func,) + tuple(args)}
   1630 
-> 1631         futures = self._graph_to_futures(
   1632             dsk,
   1633             [skey],

/envs/dask/lib/python3.8/site-packages/distributed/client.py in _graph_to_futures(self, dsk, keys, workers, allow_other_workers, priority, user_priority, resources, retries, fifo_timeout, actors)
   2646             # Pack the high level graph before sending it to the scheduler
   2647             keyset = set(keys)
-> 2648             dsk = dsk.__dask_distributed_pack__(self, keyset, annotations)
   2649 
   2650             # Create futures before sending graph (helps avoid contention)

/envs/dask/lib/python3.8/site-packages/dask/highlevelgraph.py in __dask_distributed_pack__(self, client, client_keys, annotations)
   1045                     "__module__": layer.__module__,
   1046                     "__name__": type(layer).__name__,
-> 1047                     "state": layer.__dask_distributed_pack__(
   1048                         self.get_all_external_keys(),
   1049                         self.key_dependencies,

/envs/dask/lib/python3.8/site-packages/dask/highlevelgraph.py in __dask_distributed_pack__(self, all_hlg_keys, known_key_dependencies, client, client_keys)
    424             for k, v in dsk.items()
    425         }
--> 426         dsk = toolz.valmap(dumps_task, dsk)
    427         return {"dsk": dsk, "dependencies": dependencies}
    428 

/envs/dask/lib/python3.8/site-packages/cytoolz/dicttoolz.pyx in cytoolz.dicttoolz.valmap()

/envs/dask/lib/python3.8/site-packages/cytoolz/dicttoolz.pyx in cytoolz.dicttoolz.valmap()

/envs/dask/lib/python3.8/site-packages/distributed/worker.py in dumps_task(task)
   3784             return d
   3785         elif not any(map(_maybe_complex, task[1:])):
-> 3786             return {"function": dumps_function(task[0]), "args": warn_dumps(task[1:])}
   3787     return to_serialize(task)
   3788 

/envs/dask/lib/python3.8/site-packages/distributed/worker.py in warn_dumps(obj, dumps, limit)
   3793 def warn_dumps(obj, dumps=pickle.dumps, limit=1e6):
   3794     """Dump an object to bytes, warn if those bytes are large"""
-> 3795     b = dumps(obj, protocol=4)
   3796     if not _warn_dumps_warned[0] and len(b) > limit:
   3797         _warn_dumps_warned[0] = True

/envs/dask/lib/python3.8/site-packages/distributed/protocol/pickle.py in dumps(x, buffer_callback, protocol)
     58         try:
     59             buffers.clear()
---> 60             result = cloudpickle.dumps(x, **dump_kwargs)
     61         except Exception as e:
     62             logger.info("Failed to serialize %s. Exception: %s", x, e)

/envs/dask/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py in dumps(obj, protocol, buffer_callback)
     71                 file, protocol=protocol, buffer_callback=buffer_callback
     72             )
---> 73             cp.dump(obj)
     74             return file.getvalue()
     75 

/envs/dask/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py in dump(self, obj)
    561     def dump(self, obj):
    562         try:
--> 563             return Pickler.dump(self, obj)
    564         except RuntimeError as e:
    565             if "recursion" in e.args[0]:

/envs/dask/lib/python3.8/socket.py in __getstate__(self)
    270 
    271     def __getstate__(self):
--> 272         raise TypeError(f"cannot pickle {self.__class__.__name__!r} object")
    273 
    274     def dup(self):

TypeError: cannot pickle 'socket' object

Anyone knows how to fix this issue?

1 Answers

The problem is that dask.Client is not serializable, so you can't submit it.

You can work around this problem accessing dask.Client within a task by using dask.distributed.get_client:

from dask.distributed import get_client

def create_dmatrix(X_train, y_train):
    client = get_client()

    return dxgb.DaskDMatrix(client, X_train, y_train)

Xy = client.submit(create_dmatrix, X_train, y_train, resources={'GPU': 1})
Related