Question

TL; DR
我想在数据集启动时将数据集预加载到Dask Distributed调度程序中。

背景
我使用Dask以实时查询方式使用较小的内存数据集。因为它是实时的，所以工作人员可以相信调度程序始终具有某些可用的数据集 - 即使在启动后立即也是如此。工作人员始终将整个数据集保存在内存中。

传统上，我通过连接客户端，散布df并发布数据集来完成此任务：

df = dd.read_parquet('df.parq')
df = client.persist(df)
client.publish_dataset(flights=dfa)

但是这使得调度程序重新启动并且数据集没有被加载的可能性。

我知道您可以使用--preload在启动时执行脚本，如下所示：

dask-scheduler --preload=scheduler-startup.py

锅炉板代码如下所示：

from distributed.diagnostics.plugin import SchedulerPlugin

class MyPlugin(SchedulerPlugin):
    def add_worker(self, scheduler=None, worker=None, **kwargs):
        print("Added a new worker at", worker)

def dask_setup(scheduler):
    plugin = MyPlugin()
    scheduler.add_plugin(plugin)

但是如何在不使用外部客户端的情况下说服调度程序加载我的数据集？

理论上我可能会删除启动预填充客户端的子流程，但感觉不太理想：）

调度程序启动时的普通客户端
尝试在调度程序启动时作为客户端进行连接：

from distributed.diagnostics.plugin import SchedulerPlugin
from dask.distributed import Client

class MyPlugin(SchedulerPlugin):
    def add_worker(self, scheduler=None, worker=None, **kwargs):
        print("Added a new worker at", worker)

def dask_setup(scheduler):
    c = Client(scheduler.address)
    df = dd.read_parquet('df.parq')
    df = c.persist(df)
    c.publish_dataset(flights=dfa)

挂起c = Client(scheduler.address)并且必须被强行杀死（kill -9）

Answer 1

您可以考虑在事件循环上运行的异步函数中添加客户端代码。这将允许预加载脚本完成，让调度程序启动，然后运行您的客户端代码。您可能需要以下内容：

async def f(scheduler):
    client =  await Client(scheduler.address)
    df = dd.read_parquet(...)
    await client.publish_dataset(flights=df)

def dask_setup(scheduler):
    scheduler.loop.add_callback(f, scheduler)

Answer 2

@ MRocklin的回答让我走上了正确的道路，但我确实需要放到另一个线程：

from concurrent.futures import ThreadPoolExecutor

def load_dataset():
    client = Client('127.0.0.1:8786')
    df = dd.read_parquet(...)
    df = client.persist(df)
    client.publish_dataset(flights=df)

async def f(scheduler):
    executor = ThreadPoolExecutor(max_workers=1)
    executor.submit(load_dataset)

def dask_setup(scheduler):
    scheduler.loop.add_callback(f, scheduler)

缺点是它在加载数据时不会阻止工人连接，但我认为必须在工人端进行管理（如果数据集不可用则重试）

在启动时自动将数据集添加到Dask调度程序

2 个答案: