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

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']))