与其他参数Dask add_done_callback?

时间:2019-03-08 21:29:48

标签: python distributed-computing dask

我希望在完成后向其添加回调。

根据文档:

  

当回调完成时,将来调用回调。

     

回调fn应该将Future作为唯一参数。无论将来是否成功完成,出错或被取消,都将调用它。

     

回调是在单独的线程中执行的。

这不能满足我的需要,因为要求fn回调必须以future作为唯一参数。

这是我要执行的操作的示例代码:

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

这里的想法是让客户端从队列中获取可用的GPU,然后在完成使用后将其放回队列中,以便另一个客户端可以使用它。

一种解决方法是将gpu_queue传递给客户端:

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)

事情如预期般运作。 但是我更希望能够从外部进行此操作,我缺少或看不到要完成这项工作的什么?

谢谢

1 个答案:

答案 0 :(得分:0)

您还可以考虑使用as_completed迭代器在客户端上进行处理

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