Dask add_done_callback с другими аргументами? - PullRequest
Новая
       32

Dask add_done_callback с другими аргументами?

0 голосов
/ 09 марта 2019

Я хочу добавить обратный вызов в будущее после его завершения.

Согласно документации:

Обратный вызов в будущем, когда обратный вызов завершится.

Обратный вызов fn должен принимать будущее как единственный аргумент. Это будет вызвано независимо от того, завершится ли будущее успешно, ошибочно или отменено.

Обратный вызов выполняется в отдельном потоке.

Это не дает мне того, что мне нужно, из-за требования обратного вызова fn, чтобы принимать в качестве единственного аргумента будущее.

Вот пример кода части того, что я пытаюсь сделать: 

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_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 голосов
/ 10 марта 2019

Вы можете также рассмотреть возможность обработки этого на стороне клиента с помощью 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
...