python - Dask add_done_callback 与其他参数？

Question

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

根据文档:

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.

这并没有为我提供我需要的东西，因为回调 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)

事情像这样按预期工作。 但我更喜欢从外面这样做 我错过或没有看到什么使这项工作有效？

谢谢

Answer 1

您也可以考虑在客户端使用 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