Dask add_done_callback with other args?

Question

I'm looking to add a callback to a future once it is finished.

Per the documentation:

Call callback on future when callback has finished.

The callback fn should take the future as its only argument. This will be called regardless of if the future completes successfully, errs, or is cancelled.

The callback is executed in a separate thread.

This doesn't provide me with what I need because of the requirement for callback fn to take future as its only argument.

Here's a sample portion code of what I am tying to do:

def method(cu_device_id):
    print("Hello world, I'm going to use GPU %i" % cu_device_id)

def callback_fn(cu_device_id)
    gpu_queue.put(cu_device_id)

cu_device_id = gpu_queue.get()
future = client.submit(method, cu_device_id)
#gpu_queue.put(cu_device_id) # Does not work, clients will shortly end up piled onto the slowest GPU
result.add_done_callback(callback_fn) # Crash / no way to pass in cu_device_id

The idea here is to have a client take an available GPU from the queue, and then once it is finished using it, put it back into the queue so that another client can use it.

One way to get around this is to pass gpu_queue into client:

def method(gpu_queue):
    cu_device_id = gpu_queue.get()
    print("Hello world, I'm going to use GPU %i" % cu_device_id)
    gpu_queue.put(cu_device_id)

future = client.submit(method, gpu_queue)

Things work as expected like this. But I prefer to be able to do this from outside What I missing or not seeing to make this work?

Thanks

score 2 · Answer 1 · answered Jun 02 '22 at 19:08

2

Use functools.partial. It helps to send parameters in such scenarios. Find more about it here. Besides, Now the official documentation also includes information about how to send parameters to add_done_callback

answered Jun 02 '22 at 19:08

Abhilash Reddy Yammanuru

83
7

score 1 · Answer 2 · answered Sep 22 '21 at 01:02

This should be handled on the server-side. When the task is done, a Future class instance (future) is passed to the callback_fn. You can pass an argument by wrapping this function definition by another function that returns the callback in the right format:

def method(cu_device_id):
    print("Hello world, I'm going to use GPU %i" % cu_device_id)

def get_callback_fn(cu_device_id):
    def callback_fn(future):
        gpu_queue.put(cu_device_id)
    return callback_fn

cu_device_id = gpu_queue.get()
future = client.submit(method, cu_device_id)
future.add_done_callback(get_callback_fn(cu_device_id))

score 0 · Answer 3 · answered Mar 09 '19 at 23:48

You might also consider handling this on the client side with an as_completed iterator

data = iter(data)
futures = []
using_gpu = {}

for i in range(n_gpus):
    future = client.submit(process, next(data), use_gpu=i)
    using_gpu[future] = i

seq = as_completed(futures)
for future in seq:
    gpu = using_gpu.pop(future)
    new = client.submit(process, next(data), use_gpu=gpu)  # TODO: handle end of data sequence gracefully
    using_gpu[new] = gpu
    seq.add(new)  # add this into the sequence

Dask add_done_callback with other args?

3 Answers3