xarray / dask-限制线程数/ cpus

时间:2018-09-17 20:03:42

标签: python multiprocessing dask python-xarray netcdf4

我对xarray还是很陌生,我目前正在尝试利用它来对某些NetCDF进行子集化。我正在共享服务器上运行此脚本,并且想知道如何最好地限制xarray使用的处理能力,以便它与他人很好地配合使用。我已经阅读了一下dask和xarray文档,但是对我来说如何设置cpus / threads上限似乎还不太清楚。这是一个空间子集的示例:

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活动)。我不确定问题出在哪里。

0 个答案:

没有答案