自定义初始化
目录
自定义初始化¶
通常，我们希望在启动或关闭调度器或 worker 时运行自定义代码。我们可以手动使用 Client.run 或 Client.run_on_scheduler 等函数来实现，但这容易出错且难以自动化。
为解决此问题，Dask 提供了几种机制，用于在 Scheduler、Worker、Nanny 或 Client 的生命周期中运行任意代码。
预加载脚本¶
dask-scheduler 和 dask-worker 都支持 --preload 选项，该选项允许对各自的 scheduler/worker 进行自定义初始化。作为 --preload 值传递的模块或 Python 文件，保证在建立任何连接之前被导入。如果找到一个 dask_setup(service) 函数，则会被调用，并以 Scheduler, Worker, Nanny 或 Client 的实例作为参数。当服务停止时，如果存在 dask_teardown(service)，则会调用该函数。
为了支持额外的配置，单个 --preload 模块可以通过将 dask_setup 公开为 Click 命令来注册额外的命令行参数。此命令将用于解析提供给 dask-worker 或 dask-scheduler 的额外参数，并在服务初始化之前调用。
示例¶
作为一个示例，请考虑以下文件，它创建了一个调度器插件并将其注册到调度器
# scheduler-setup.py
import click
from distributed.diagnostics.plugin import SchedulerPlugin
class MyPlugin(SchedulerPlugin):
def __init__(self, print_count):
self.print_count = print_count
super().__init__()
def add_worker(self, scheduler=None, worker=None, **kwargs):
print("Added a new worker at:", worker)
if self.print_count and scheduler is not None:
print("Total workers:", len(scheduler.workers))
@click.command()
@click.option("--print-count/--no-print-count", default=False)
def dask_setup(scheduler, print_count):
plugin = MyPlugin(print_count)
scheduler.add_plugin(plugin)
然后，当启动调度器时，可以通过引用其文件名（如果它在路径中，也可以是模块名）来运行此预加载脚本
dask-scheduler --preload scheduler-setup.py --print-count
类型¶
预加载可以指定为以下任何形式
脚本的路径，例如
/path/to/myfile.py在路径中的模块名，例如
my_module.initializePython 脚本的文本，例如
import os; os.environ["A"] = "value"
配置¶
预加载也可以在配置中以下列值进行注册
distributed:
scheduler:
preload:
- "import os; os.environ['A'] = 'b'" # use Python text
- /path/to/myfile.py # or a filename
- my_module # or a module name
preload-argv:
- [] # Pass optional keywords
- ["--option", "value"]
- []
worker:
preload: []
preload-argv: []
nanny:
preload: []
preload-argv: []
client:
preload: []
preload-argv: []
注意
由于 dask-worker 命令需要同时接受 Worker 和 Nanny（如果使用了 nanny）的关键字，因此它同时具有 --preload 和 --preload-nanny 关键字。所有额外的关键字（例如上面的 --print-count）将发送给 worker 而不是 nanny。在命令行上无法为 nanny 预加载脚本指定额外的关键字。如果需要，我们建议使用更灵活的配置。
Worker 生命周期插件¶
您还可以创建一个包含 setup、teardown 和 transition 方法的类，并使用 Client.register_worker_plugin 方法将其注册到调度器，以便提供给每个 worker。
|
为所有当前和未来的 worker 注册生命周期 worker 插件。 |
- Client.register_worker_plugin(plugin: distributed.diagnostics.plugin.NannyPlugin | distributed.diagnostics.plugin.WorkerPlugin, name: str | None = None, nanny: bool | None = None)[source]
为所有当前和未来的 worker 注册生命周期 worker 插件。
自 2023.9.2 版本起已弃用: 请改用
Client.register_plugin()。这会注册一个新对象，用于处理此集群中 worker 的设置、任务状态转换和关闭。该插件将在所有当前连接的 worker 上实例化自身。它也将在将来连接的任何 worker 上运行。
该插件可能包含
setup、teardown、transition和release_key方法。有关接口和文档字符串，请参阅dask.distributed.WorkerPlugin类或下面的示例。它必须可以使用 pickle 或 cloudpickle 模块进行序列化。如果插件具有
name属性，或者使用了name=关键字，则会控制幂等性。如果具有该名称的插件已注册，则它将被移除并替换为新的插件。对于插件的替代方案，您也可以考虑预加载脚本。
- 参数
- pluginWorkerPlugin 或 NannyPlugin
要注册的 WorkerPlugin 或 NannyPlugin 实例。
- namestr，可选
插件的名称。注册同名插件将无效。如果插件没有 name 属性，则使用随机名称。
- nannybool，可选
是否将插件注册到 worker 或 nanny。
另请参阅
distributed.WorkerPluginunregister_worker_plugin
示例
>>> class MyPlugin(WorkerPlugin): ... def __init__(self, *args, **kwargs): ... pass # the constructor is up to you ... def setup(self, worker: dask.distributed.Worker): ... pass ... def teardown(self, worker: dask.distributed.Worker): ... pass ... def transition(self, key: str, start: str, finish: str, ... **kwargs): ... pass ... def release_key(self, key: str, state: str, cause: str | None, reason: None, report: bool): ... pass
>>> plugin = MyPlugin(1, 2, 3) >>> client.register_plugin(plugin)
您可以使用
get_worker函数访问插件>>> client.register_plugin(other_plugin, name='my-plugin') >>> def f(): ... worker = get_worker() ... plugin = worker.plugins['my-plugin'] ... return plugin.my_state
>>> future = client.run(f)