为数据框操作定义 dask worker 资源

Question

我正在对一个 dask 数据帧应用多个操作。我可以为特定操作定义分布式工作人员资源要求吗？

例如我这样称呼：

df.fillna(value="").map_partitions(...).map(...)

我想为 map_partitions() 指定资源需求（可能与 map() 的需求不同） , 但似乎该方法不接受 resources 参数。

PS。或者，我发现我可以在 map_partitions() 之后调用 client.persist() 并在此调用中指定资源，但这会立即触发计算。

Answer 1

您可以在调用 compute 或 persist 时通过提供中间集合来指定计算的特定部分的资源限制。

x = dd.read_csv(...)
y = x.map_partitions(func)
z = y.map(func2)

z.compute(resources={tuple(y._keys()): {'GPU': 1}})

谢谢你的问题，我去包含一个 link 到关于这个功能的文档，发现它没有记录。我会尽快修复。

今天看来有一个错误，在某些情况下中间键可能会被优化掉（尽管这不太可能用于数据帧操作），因此您可能还想传递 optimize_graph=False 关键字。

z.compute(resources={tuple(y._keys()): {'GPU': 1}}, optimize_graph=False)

Defining dask worker resources for a dataframe operation