xarray/dask - 限制线程/CPU数量

Question

我对xarray比较陌生，目前正在尝试使用它来对一些NetCDF文件进行子集处理。我在共享服务器上运行，希望知道如何最好地限制xarray使用的处理能力，以便与其他应用程序协调。我已经阅读了一些dask和xarray文档，但似乎不清楚如何设置CPU/线程数量上限。以下是空间子集的示例：

import glob
import os
import xarray as xr

from multiprocessing.pool import ThreadPool
import dask

wd = os.getcwd()

test_data = os.path.join(wd, 'test_data')
lat_bnds = (43, 50)
lon_bnds = (-67, -80)
output = 'test_data_subset'

def subset_nc(ncfile, lat_bnds, lon_bnds, output):
    if not glob.os.path.exists(output):
        glob.os.makedirs(output)
    outfile = os.path.join(output, os.path.basename(ncfile).replace('.nc', '_subset.nc'))

    with dask.config.set(scheduler='threads', pool=ThreadPool(5)):
        ds = xr.open_dataset(ncfile, decode_times=False)

        ds_sub = ds.where(
            (ds.lon >= min(lon_bnds)) & (ds.lon <= max(lon_bnds)) & (ds.lat >= min(lat_bnds)) & (ds.lat <= max(lat_bnds)),
            drop=True)
        comp = dict(zlib=True, complevel=5)
        encoding = {var: comp for var in ds.data_vars}
        ds_sub.to_netcdf(outfile, format='NETCDF4', encoding=encoding)

list_files = glob.glob(os.path.join(test_data, '*'))
print(list_files)

for i in list_files:
    subset_nc(i, lat_bnds, lon_bnds, output)

我尝试了几种变化，移动了ThreadPool的配置，但服务器的top仍然显示有太多活动（>3000%的CPU活动）。 我不确定问题出在哪里。

- Trevor J. Smith

---

你尝试过将dask的工作进程数和每个工作进程的线程数设置为dask设置吗？也许你可以尝试类似这样的代码：with dd.LocalCluster(n_workers=1, threads_per_worker=5, memory_limit='15GiB') as cluster, dd.Client(cluster, set_as_default=True) as client: 并根据你的机器调整这些值？好了，我刚看到，任何评论现在都太晚了:D - Helmut

">3000% CPU 使用率" - 你有这么多个CPU吗，还是报告有误？ - mdurant

Answer 1

这个问题实际上在这里得到了解决 https://github.com/pydata/xarray/issues/2417#issuecomment-460298993（GitHub的问题似乎是由提问者发起的）。

{{link2:@jhamman（再次非常匹配的个人资料）提出的解决方案是将环境变量OMP_NUM_THREADS设置为所需的线程数（建议大约是期望核心数的两倍，可能是为了充分利用现代Intel/AMD的线程技术）。

- ti7