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

问题描述

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

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()

解决方法

实际上，您的代码仍然是串行而不是并行的，特别是 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)])