扩展 sizeof

扩展 sizeof

当Dask需要计算一个对象的字节大小,例如确定哪些对象需要溢出到磁盘时,它会使用 dask.sizeof.sizeof 注册机制。需要为自己的对象定义 sizeof 实现的用户可以使用 sizeof.register

>>> import numpy as np
>>> from dask.sizeof import sizeof
>>> @sizeof.register(np.ndarray)
>>> def sizeof_numpy_like(array):
...     return array.nbytes

此代码可以通过将其放置在库的模块中(例如 __init__.py)来执行,以便向 Dask 注册实现。然而,这给这些库的开发者带来了维护负担,并且如果这些库不接受补丁,则必须在所有工作节点上手动导入。

因此,Dask 也在 dask.sizeof 组下公开了一个 入口点,以允许第三方库开发和维护这些 sizeof 实现。

对于一个虚构的库 numpy_sizeof_dask.py,必要的 setup.cfg 配置如下:

[options.entry_points]
dask.sizeof =
   numpy = numpy_sizeof_dask:sizeof_plugin

numpy_sizeof_dask.py 将包含

>>> import numpy as np
>>> def sizeof_plugin(sizeof):
...    @sizeof.register(np.ndarray)
...    def sizeof_numpy_like(array):
...        return array.nbytes

在首次导入 dask.sizeof 时,Dask 会调用入口点(sizeof_plugin),并传入 dask.sizeof.sizeof 对象,该对象随后可用于注册 sizeof 实现。