从 Future 创建惰性 xarray 对象

Question

我有一个dask.delayed 函数，它接受xarray.Dataarray 作为参数并返回一个。

我正在创建一些延迟任务，并使用dask.distributed 将它们传递给client.compute。每次调用计算都会返回一个distributed.client.Future，表示将返回的数据数组。

我的问题是：

有没有办法从未来再次构建一个“惰性”数据数组而不从工作人员加载实际数据？我的意图是根据第一次计算的输出构建第二个任务图。

我见过client.gather，但这似乎将所有数据拉回客户端，这不是我想要的。

这是一个小例子：

import dask
from distributed import Client
import xarray as xr

# load example data
x = xr.tutorial.open_dataset("air_temperature")

# use first timestep
x_t0 = x.isel(time=0)

# delayed 'processing' function
@dask.delayed
def fun(x):
    return x*2

# init client
client = Client()

# compute on worker
future = client.compute(fun(x_t0))

# when done
print(future)
# <Future: finished, type: xarray.Dataset, key: fun-96cd56f4-4ed3-4eac-ade9-fe3f17e4b8c6>

## now how to get back to lazy xarray from future?

Answer 1

我不知道你到底想达到什么目的。可能有比从未来创建一个新数组更好的方法来做到这一点。话虽如此，这将从您的数据创建一个新的数据数组： 你不必调用 compute 来保持它的懒惰。

（如果您想要一个 dask 数组而不是 xarray 数组，请删除 xr.DataArray）

import dask
from distributed import Client
import xarray as xr

# load example data
x = xr.tutorial.open_dataset("air_temperature")

# use first timestep
x_t0 = x.isel(time=0)

# delayed 'processing' function
@dask.delayed
def fun(x):
    return x*2

# init client
client = Client()

# Create lazy xarray object from future:
import dask.array as da

new_ds = xr.DataArray(da.from_delayed(client.persist(fun(x_t0)), shape=x_t0.air.shape, meta='f8'), coords=x.coords)

编辑：添加了 client.persist 以将数据留在客户端

输出：