SLURM: запуск заданий графического процессора на узлах из Python ⇐ Python

[Anonymous]

Я пытаюсь запустить скрипт Python, который оценивает другую функцию для разных значений параметров, и делает это N раз. Вычисления функций независимы, и я хочу запустить каждую из них на одном графическом процессоре. Я запускаю это в кластере, где сценарий запуска для моего основного файла Python выглядит примерно так

Код: Выделить всё

#!/bin/bash

#SBATCH -N 12
#SBATCH -C gpu
#SBATCH -G 48
#SBATCH -q regular

module load pytorch
srun --exact -N 1 -G 0 -n 1 --cpu_bind=cores python run_main.py

Затем в run_main.py у меня есть цикл

Код: Выделить всё

for i in range(N):
processes = []
for k in range(48): #i.e. one task per gpu
cmd = "srun --exact --ntasks=1 -G 1 -N 1 --cpu-bind=none python run_task.py (with a bunch of arguments depending on k, etc)"
print(cmd)
process = Popen(cmd, shell=True)
processes.append(process)
for process in processes:
process.wait()
#do something else depending on run_task outcomes, and then launch over GPUs again

Кажется, это слишком медленно для того, насколько быстро должны выполняться отдельные задачи, поэтому я предполагаю, что либо это большие накладные расходы, либо я не использую все свои графические процессоры одновременно время. Как я могу это проверить? Как я могу сделать это лучше?

Подробнее здесь: https://stackoverflow.com/questions/792 ... hin-python