python - 确定分配给作业的处理器数量

Question

我正在尝试修正一个 python 程序，该程序将使用工作池在 Linux hpc 计算环境中执行任务。我想将池中的工作线程数设置为等于我为我的作业请求的线程数。

multiprocessing.cup_count() 和 num_procs 等命令返回 64，这是节点上的线程数，但不是分配给我的作业的线程数.

我目前的解决方案是grep批处理文件中指定线程数的命令，然后检查输出:

os.system('grep -i "batch -n" * > output')

file_ = open('output','r')
file_lines = file_.readlines()
file_.close()

for line in file_lines:
  elems = line.strip().split()
  if '-n' in elems:
    position = elems.index('-n')
    num_procs = elems[position+1]

有没有更简洁直接的方法来做到这一点？

Answer 1

根据您使用的调度程序，您可以查询您的作业从调度程序请求的核心数。对于 SGE - $NSLOTS；对于 SLURM - $SLURM_CPUS_PER_TASK；对于 PBS - $PBS_NP。

例如 SGE:

import os
#for SGE:
nslots = int(os.environ['NSLOTS']))

# for SLURM:
nslots = int(os.environ['SLURM_CPUS_PER_TASK']))