Ответ 1
Это относительно просто.
В упрощающих предположениях, что вы запрашиваете один процесс на один хост, slurm предоставит вам всю необходимую информацию в переменных среды, в частности, SLURM_PROCID, SLURM_NPROCS и SLURM_NODELIST.
Например, вы можете инициализировать индекс своей задачи, количество задач и нод-лист следующим образом:
from hostlist import expand_hostlist
task_index = int( os.environ['SLURM_PROCID'] )
n_tasks = int( os.environ['SLURM_NPROCS'] )
tf_hostlist = [ ("%s:22222" % host) for host in
expand_hostlist( os.environ['SLURM_NODELIST']) ]
Обратите внимание, что slurm предоставляет вам список хостов в сжатом формате (например, "myhost [11-99]" ), который вам нужно расширить. Я делаю это с помощью хоста-модуля с помощью Kent Engström, доступен здесь https://pypi.python.org/pypi/python-hostlist
В этот момент вы можете пойти прямо вперед и создать свою спецификацию и сервер TensorFlow с информацией, доступной вам, например:
cluster = tf.train.ClusterSpec( {"your_taskname" : tf_hostlist } )
server = tf.train.Server( cluster.as_cluster_def(),
job_name = "your_taskname",
task_index = task_index )
И ты настроен! Теперь вы можете выполнить размещение TensorFlow node на определенном хосте вашего распределения с помощью обычного синтаксиса:
for idx in range(n_tasks):
with tf.device("/job:your_taskname/task:%d" % idx ):
...
Недостаток кода, описанного выше, заключается в том, что все ваши задания будут давать указание Tensorflow устанавливать серверы, прослушивающие фиксированный порт 22222. Если несколько таких заданий будут запланированы на один и тот же node, второй не сможет прослушать 22222.
Лучшее решение - разрешить порты резервирования slurm для каждого задания. Вам нужно привести администратора slurm на борт и попросить его настроить slurm, чтобы он мог запрашивать порты с опцией --resv-ports. На практике это требует, чтобы они добавили следующую строку в свой slurm.conf:
MpiParams=ports=15000-19999
Прежде чем удалять администратор slurm, проверьте, какие параметры уже настроены, например:
scontrol show config | grep MpiParams
Если ваш сайт уже использует старую версию OpenMPI, есть шанс, что такой вариант уже существует.
Затем измените мой первый фрагмент кода следующим образом:
from hostlist import expand_hostlist
task_index = int( os.environ['SLURM_PROCID'] )
n_tasks = int( os.environ['SLURM_NPROCS'] )
port = int( os.environ['SLURM_STEP_RESV_PORTS'].split('-')[0] )
tf_hostlist = [ ("%s:%s" % (host,port)) for host in
expand_hostlist( os.environ['SLURM_NODELIST']) ]
Удачи!