使用 Xarray 和 Dask 在 Dataarray 上运行的循环中获得非常慢的迭代答案

【问题标题】：Getting very slow iterations in a loop run over a Datarray using Xarray and Dask使用 Xarray 和 Dask 在 Dataarray 上运行的循环中获得非常慢的迭代
【发布时间】：2021-04-22 19:40:12
【问题描述】：

我正在尝试从 u 和 v 分量以每小时时间步长和 0.1 x 0.1 度分辨率计算 1 年数据的风速，总共 40 年。 1 年的单个 u 和 v netcdf 文件每个约为 5GB。我已经实现了一个基本的for 循环，其中每年的 u 和 v netcdf 文件通过 Xarray open_dataset 打开并重新分块以将它们作为 dask 数组，然后进行计算并将结果导出为新的 netcdf。当循环运行时，第一次迭代几乎是瞬间发生的，但是下一次迭代的循环时间太长（几乎到了似乎停止的地步）。我不明白我的代码的哪一部分在这里遇到瓶颈以及为什么。任何帮助，将不胜感激。此外，我已经正确实现了 dask 调度程序来自适应地请求资源。我附上相关代码 sn-p 以供参考：

    cluster = PBSCluster(cores=1,memory='8GB',queue='standard',project='civil',interface='ib0',walltime='00:20:00')
    cluster.adapt(minimum=1, maximum=8)
    client = Client(cluster) 
    for i in range (1979,2019):
        u_dir = glob.glob('../u_wind/uwind_hourly_'+ str(i)+'*.nc')
        v_dir = glob.glob('../v_wind/vwind_hourly_'+ str(i)+'*.nc')
        w_dir = './wind/wind_hourly_'+str(i)+'-'+str(i)+'.nc'
        u_wind = xr.open_dataset(u_dir[0])
        v_wind = xr.open_dataset(v_dir[0])
        u_wind_rechunk = u_wind.chunk({'time':720})
        v_wind_rechunk = v_wind.chunk({'time':720}) 
        u_var = u_wind_rechunk['UGRD_10m']
        v_var = v_wind_rechunk['VGRD_10m'] 
        wind_speed = xr.Dataset(data_vars=None, coords=None, attrs=None)
        wind_speed=wind_speed.assign(wind_speed=np.sqrt(u_var**2 + v_var**2))
        wind_speed.to_netcdf(w_dir)
        del u_wind
        del v_wind
        del u_wind_rechunk
        del v_wind_rechunk
        del u_var
        del v_var
        del wind_speed
        gc.collect()

【问题讨论】：

通过一些故障排除后，我可以看到使用 cluster.scale() 而不是 cluster.adapt() 时性能有所提高。

标签： python dask netcdf python-xarray

【解决方案1】：

事实上，您的代码仍然看起来是串行的而不是并行的，特别是 wind_speed.to_netcdf(w_dir) 将立即触发计算。下面的代码可能需要一些调整，但重点是并行化您的操作：

def single_run(i):
# nothing is modified in the code below relative
    u_dir = glob.glob('../u_wind/uwind_hourly_'+ str(i)+'*.nc')
    v_dir = glob.glob('../v_wind/vwind_hourly_'+ str(i)+'*.nc')
    w_dir = './wind/wind_hourly_'+str(i)+'-'+str(i)+'.nc'
    u_wind = xr.open_dataset(u_dir[0])
    v_wind = xr.open_dataset(v_dir[0])
    u_wind_rechunk = u_wind.chunk({'time':720})
    v_wind_rechunk = v_wind.chunk({'time':720}) 
    u_var = u_wind_rechunk['UGRD_10m']
    v_var = v_wind_rechunk['VGRD_10m'] 
    wind_speed = xr.Dataset(data_vars=None, coords=None, attrs=None)
    wind_speed=wind_speed.assign(wind_speed=np.sqrt(u_var**2 + v_var**2))
    wind_speed.to_netcdf(w_dir)
    del u_wind
    del v_wind
    del u_wind_rechunk
    del v_wind_rechunk
    del u_var
    del v_var
    del wind_speed
    gc.collect()

# new parts
import dask

run_me = dask.compute([dask.delayed(single_run)(i) for i in range (1979,2019)])

【讨论】：