resource-backed-dask-array

ResourceBackedDaskArray 是一个实验性的Dask数组子类，在计算时打开/关闭资源（但每个计算调用只打开一次）。

安装

pip install resource-backed-dask-array

此包的动机

考虑以下类，该类模拟一个文件读取器，能够返回Dask数组（使用 dask.array.map_blocks）。文件句柄必须处于打开状态才能读取一个块，否则会崩溃（或产生其他错误）

import dask.array as da
import numpy as np


class FileReader:

    def __init__(self):
        self._closed = False

    def close(self):
        """close the imaginary file"""
        self._closed = True

    @property
    def closed(self):
        return self._closed

    def __enter__(self):
        if self.closed:
            self._closed = False  # open
        return self

    def __exit__(self, *_):
        self.close()

    def to_dask(self) -> da.Array:
        """Method that returns a dask array for this file."""
        return da.map_blocks(
            self._dask_block,
            chunks=((1,) * 4, 4, 4),
            dtype=float,
        )

    def _dask_block(self):
        """simulate getting a single chunk of the file."""
        if self.closed:
            raise RuntimeError("Segfault!")
        return np.random.rand(1, 4, 4)

只要文件保持打开状态，一切正常

>>> fr = FileReader()
>>> dsk_ary = fr.to_dask()
>>> dsk_ary.compute().shape
(4, 4, 4)

然而，如果关闭了文件，从 to_dask 返回的Dask数组现在将失败

>>> fr.close()
>>> dsk_ary.compute()  # RuntimeError: Segfault!

这里的“快速而简单的”解决方案可能是强制 _dask_block 方法在发现文件处于关闭状态时临时重新打开文件，但如果文件打开过程需要任何时间，这将导致显著的开销，因为它会为数组的每个块打开和关闭。

使用方法

ResourceBackedDaskArray.from_array

此库尝试通过一个 ResourceBackedDaskArray 对象解决上述问题。在调用 .compute() 时，它会管理打开/关闭底层资源，并且只在单个计算任务图中对所有块执行一次。

>>> from resource_backed_dask_array import resource_backed_dask_array
>>> safe_dsk_ary = resource_backed_dask_array(dsk_ary, fr)
>>> safe_dsk_ary.compute().shape
(4, 4, 4)

>>> fr.closed  # leave it as we found it
True

传递给 from_array 的第二个参数必须是一个可重用的上下文管理器，并且额外提供 closed 属性（例如，像 io.IOBase）。换句话说，它必须实现以下协议

1. 它必须有一个 __enter__ 方法 来打开底层资源
2. 它必须有一个 __exit__ 方法 来关闭资源，并且可选地处理异常
3. 它必须有一个 closed 属性来报告资源是否已关闭。

在上面的示例中，FileReader 类本身实现了此协议，因此它可以作为 ResourceBackedDaskArray.from_array 上的第二个参数。

重要注意事项

这是为单进程（也许只是单线程？）用例创建的，在这些用例中，Dask 的离核懒加载仍然非常理想。与 dask.distributed 一起使用尚未经过测试，并且可能会失败。在多线程/进程任务中使用有状态的对象（例如这里使用的可重用上下文管理器）是容易出错的。