使用 xarray dask apply_ufunc 时避免内存溢出答案

【问题标题】：Avoiding memory overflow while using xarray dask apply_ufunc使用 xarray dask apply_ufunc 时避免内存溢出
【发布时间】：2020-10-30 21:27:28
【问题描述】：

我需要沿此形状的 xarray dask 数组的时间维度应用一个函数：

<xarray.DataArray 'tasmax' (time: 14395, lat: 1801, lon: 3600)>
dask.array<rechunk-merge, shape=(14395, 1801, 3600), dtype=float32, chunksize=(14395, 1801, 600), chunktype=numpy.ndarray>
Coordinates:
  * lat      (lat) float64 90.0 89.9 89.8 89.7 89.6 ... -89.7 -89.8 -89.9 -90.0
  * time     (time) datetime64[ns] 1981-01-01 1981-01-02 ... 2020-05-31
  * lon      (lon) float64 -180.0 -179.9 -179.8 -179.7 ... 179.7 179.8 179.9

进程的输出将是一个小得多的数组，大小为(time=365, lat=1801, lon=3600)，但如上所示，输入数组内存大小约为 360 GB。我有一台具有 16 个 CPU 内核和 126 GB RAM 的机器。我正在尝试通过使用apply_ufunc 和dask='parallelized' 参数来优化该过程，但是由于使用了所有126GB RAM，它会导致内存错误。我可以避免并行化，但完成这个过程需要很长时间。有什么方法可以控制 apply_ufunc 的内存使用以将进程包含在我的可用 RAM 中？

【问题讨论】：

你试过改变chunksize吗？ (14395, 1801, 600) 似乎有点太大了。每个块的大小约为 60GB，但这并不意味着在计算期间您不会超过 126GB 的分配内存（取决于您应用的函数）。例如，您可以使用 DataArray.chunk({'time':14395, 'lat':180, 'lon':360)。在使用apply_ufunc之前。 xarray.pydata.org/en/stable/generated/…
是的，我尝试了不同的块大小，无论是在纬度还是经度上。它仍然超出了我的记忆。 :(
chunksize=(14395, 1801, 600) 绝对太大了，因为 dask 工作人员会一次使用多个块。根据经验，块的大小应约为 100mb。另外，看看仪表板发生了什么。

标签： python-3.6 dask python-xarray dask-distributed numpy-ufunc

【解决方案1】：

您使用的chunksize 太大。正如 Val 在 cmets 中提到的，每个块的推荐大小约为 100mb。

您可以将 1801 更改为 1-4 之间的值以获得更好的性能：

chunksize=(14395, 4, 600)

【讨论】：