如何正确使用dask的upload_file()将本地代码传递给工作人员

时间:2019-07-19 18:54:45

标签: python dask dask-distributed

我在local_code.py文件中有一些函数,希望通过dask传递给工作人员。我在这里已经看到问题的答案,说可以使用upload_file()函数来完成此操作,但是由于我仍然在使用ModuleNotFoundError,因此似乎无法正常工作。

代码的相关部分如下。

from dask.distributed import Client
from dask_jobqueue import SLURMCluster

from local_code import *
helper_file = '/absolute/path/to/local_code.py'

def main():
    with SLURMCluster(**slurm_params) as cluster:

        cluster.scale(n_workers)

        with Client(cluster) as client:
            client.upload_file(helper_file)
            mapping = client.map(myfunc, data)
            client.gather(mapping)

if __name__ == '__main__':
    main()

请注意,myfunc是从local_code导入的,将其导入地图没有错误。函数myfunc还取决于local_code中定义的其他函数。

使用此代码,我仍然遇到此错误

distributed.protocol.pickle - INFO - Failed to deserialize b'\x80\x04\x95+\x00\x00\x00\x00\x00\x00\x00\x8c\x11local_code\x94\x8c\x$
Traceback (most recent call last):
  File "/home/gallagher.r/.local/lib/python3.7/site-packages/distributed/protocol/pickle.py", line 61, in loads
    return pickle.loads(x)
ModuleNotFoundError: No module named 'local_code'

使用upload_file()似乎很简单,以至于我不确定自己在做什么错。我必须把它放在错误的位置,否则就无法正确理解传递给它的内容。

对此我将不胜感激。如果您需要任何其他信息,或者我可以从错误文件中提供其他信息,请告诉我。

2 个答案:

答案 0 :(得分:2)

当有工作人员进/出时,另一种选择是在注册/添加新工作人员时使用Client.register_worker_callbacks进行挂接。一个警告是您将需要在回调部分中序列化文件:

fname = ...
with open(fname, 'rb') as f:
  data = f.read()

client.register_worker_callbacks(
  setup=functools.partial(
    _worker_upload, data=data, fname=fname,
  )
)

def _worker_upload(dask_worker, *, data, fname):
  dask_worker.loop.add_callback(
    callback=dask_worker.upload_file,
    comm=None,  # not used
    filename=fname,
    data=data,
    load=True)

这还将在首次注册回调时上传文件,因此您可以避免完全调用client.upload_file

答案 1 :(得分:1)

upload_file方法仅将文件上传到当前可用的工作程序。如果工作人员在您致电upload_file之后到达,则该工作人员将没有提供的文件。

如果您的情况最简单,可能是等到所有工作人员到达后再调用上传文件

cluster.scale(n)
with Client(cluster) as client:
    client.wait_for_workers(n)
    client.upload_file(...)