如何确定在python脚本中的哪一点SLURM中超出了step内存

Question

我有一个python脚本，正在SLURM集群上运行，用于多个输入文件：

#!/bin/bash

#SBATCH -p standard
#SBATCH -A overall 
#SBATCH --time=12:00:00
#SBATCH --output=normalize_%A.out
#SBATCH --error=normalize_%A.err
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=20
#SBATCH --mem=240000

HDF5_DIR=...
OUTPUT_DIR=...
NORM_SCRIPT=...

norm_func () {
  local file=$1
  echo "$file"
  python $NORM_SCRIPT -data $file -path $OUTPUT_DIR
}

# Doing normalization in parallel
for file in $HDF5_DIR/*; do norm_func "$file" & done
wait

python脚本只是加载一个数据集（scRNAseq），进行规范化并保存为.csv文件。其中的一些主要代码行包括：

        f = h5py.File(path_to_file, 'r')
        rawcounts = np.array(rawcounts)

        unique_code = np.unique(split_code)
        for code in unique_code:
            mask = np.equal(split_code, code)
            curr_counts = rawcounts[:,mask]

            # Actual TMM normalization
            mtx_norm = gmn.tmm_normalization(curr_counts)

            # Writing the results into .csv file
            csv_path = path_to_save + "/" + file_name + "_" + str(code) + ".csv"
            with open(csv_path,'w', encoding='utf8') as csvfile:
                writer = csv.writer(csvfile, delimiter=',')
                writer.writerow(["", cell_ids])
                for idx, row in enumerate(mtx_norm):
                    writer.writerow([gene_symbols[idx], row])

对于高于10Gb的数据集，我不断得到step memory exceeded错误，我不知道为什么。如何更改.slurm脚本或python代码以减少其内存使用？我如何才能真正确定是什么导致了memory问题，在这种情况下是否有一种特殊的内存调试方法？如有任何建议，将不胜感激。你知道吗

Answer 1

通过使用srun启动python脚本，可以获得精确的信息：

srun python $NORM_SCRIPT -data $file -path $OUTPUT_DIR

然后，Slurm将为python脚本的每个实例创建一个“步骤”，并在记帐中独立地报告每个步骤的信息（错误、返回代码、使用的内存等），您可以使用sacct命令查询这些信息。你知道吗

如果由管理员配置，请使用 profile选项获取每个步骤的内存使用时间线。你知道吗

在python脚本中，可以使用memory_profile模块获得脚本内存使用情况的反馈。你知道吗