使用--preload在dask worker中全局初始化任务模块?

时间:2019-06-13 00:58:03

标签: python dask dask-distributed

我正在尝试实现与以下问题类似的东西(Initializing state on dask-distributed workersSetting up Dask worker with variable),在这里我有一个(相对)较大的模型,我希望对该子集的工人进行预初始化接受需要模型的任务。理想情况下,我甚至不希望客户端计算机具有该模型。

在发现这些问题之前,我的最初尝试是在共享模块delayed中定义一个worker_task.model任务,并在worker中分配一个模块全局变量(例如worker_tasks.model.model) '--preload脚本供该任务使用;但是,由于某些原因,此操作不起作用-在预加载脚本中设置了变量,但是在调用任务时仍为None

init_model_worker.py:

import logging
from uuid import uuid4

from worker_tasks import model


def dask_setup(worker):
    model.model = f'<mock model {uuid4()}>'

    logger = logging.getLogger('distributed')
    logger.warning(f'model = {model.model}')

worker_tasks / model.py:

import logging
import random
from time import sleep
from uuid import uuid4

import dask

model = None


@dask.delayed
def compute_clinical(inp):        
    if model is None:
        raise RuntimeError('Model not initialized.')

    sleep(random.uniform(3, 17))

    return {
        'result': random.choice((True, False)),
        'confidence': random.uniform(0, 1)
        }

这是我启动它并将其提交给调度程序时的工作日志:

> dask-worker --preload init_model_worker.py tcp://scheduler:8786 --name model-worker
distributed.utils - INFO - Reload module init_model_worker from .py file                                  
distributed.nanny - INFO -         Start Nanny at: 'tcp://172.28.0.4:41743'                         
distributed.diskutils - INFO - Found stale lock file and directory '/worker-epptq9sh', purging      
distributed.utils - INFO - Reload module init_model_worker from .py file                                  
distributed - WARNING - model = <mock model faa41af0-d925-46ef-91c9-086093d37c71>                   
distributed.worker - INFO -       Start worker at:     tcp://172.28.0.4:37973                       
distributed.worker - INFO -          Listening to:     tcp://172.28.0.4:37973                       
distributed.worker - INFO -              nanny at:           172.28.0.4:41743                       
distributed.worker - INFO -              bokeh at:           172.28.0.4:37766                       
distributed.worker - INFO - Waiting to connect to:       tcp://scheduler:8786                       
distributed.worker - INFO - -------------------------------------------------                       
distributed.worker - INFO -               Threads:                          4                       
distributed.worker - INFO -                Memory:                    1.93 GB                       
distributed.worker - INFO -       Local Directory:           /worker-mhozo9ru                       
distributed.worker - INFO - -------------------------------------------------                       
distributed.worker - INFO -         Registered to:       tcp://scheduler:8786                       
distributed.worker - INFO - -------------------------------------------------                       
distributed.core - INFO - Starting established connection                                           
distributed.worker - WARNING -  Compute Failed                                                      
Function:  compute_clinical                                                                         
args:      ('mock')                                                                                 
kwargs:    {}                                                                                       
Exception: RuntimeError('Model not initialized.')                                                   

您可以看到,重新加载预加载脚本后,model<mock model faa41af0-d925-46ef-91c9-086093d37c71>;但是当我尝试从任务中调用它时,我得到了None

我将尝试根据其他问题的答案实施解决方案,但我有几个与工人预载有关的问题:

  1. 在预加载脚本中分配任务后,为什么调用None模型时会出现问题?
  2. 是否通常建议避免在工作程序--preload脚本中执行类似的操作?从客户端调用工作程序状态的初始化更好吗? 如果是,为什么

1 个答案:

答案 0 :(得分:1)

我怀疑通过Python序列化函数,模型变量会立即绑定到函数中。您可以尝试以下方法:

@dask.delayed
def compute_clinical(inp):       
    from worker_tasks.model import model

    if model is None:
        raise RuntimeError('Model not initialized.')

或者,而不是将变量分配给全局模块范围(在Python中可能很难理解),或者尝试将其分配给工作程序本身。

from dask.distributed import get_worker

def dask_setup(worker):
    worker.model = f'<mock model {uuid4()}>'

@dask.delayed
def compute_clinical(inp):       
    if get_worker().model is None:
        raise RuntimeError('Model not initialized.')