Как раздавать кастомный код через SLURM manager?

Question

У меня есть доступ к компьютерному кластеру с менеджером SLURM. Я хочу добиться, чтобы разные узлы выполняли разные части моего кода. Если я правильно понял, этого можно достичь с помощью SLURM с помощью команды srun, если ваш код написан правильно. Здесь должно быть что-то вроде этого примера MPI https://support.ceci-hpc.be/doc/_contents/QuickStart/SubmittingJobs/SlurmTutorial.html.

Но я не понимаю, как создать этот код в TF. Для версии 1 TF гораздо больше информации. Если я попробую что-то вроде этого

jobs={'worker': 4}
cluster=tf.distribute.cluster_resolver.SlurmClusterResolver(jobs=jobs)
server0 = tf.distribute.Server(cluster.cluster_spec(), job_name='worker', task_index=0)
server1 = tf.distribute.Server(cluster.cluster_spec(), job_name='worker', task_index=1)
server2 = tf.distribute.Server(cluster.cluster_spec(), job_name='worker', task_index=2)
server2 = tf.distribute.Server(cluster.cluster_spec(), job_name='worker', task_index=3)

и запустил это с помощью SLURM, я получаю сообщение об ошибке, в котором я вижу, что только первый сервер был запущен, но второй пытался использовать тот же адрес, который является локальный: 8888. По сути, я не знаю, как создавать серверы на разных узлах, которые затем могут обмениваться данными. Должен ли я запускать одновременно разные скрипты? Должен ли я использовать командную строку с флагами или чем-то в этом роде?

Впоследствии моя идея состоит в том, чтобы использовать

with tf.device("/job:worker/task:0"):
#some code
with tf.device("/job:worker/task:1"):
#some other code

для распределения работы. Любая помощь? Я не думаю, что смогу использовать какие-либо стратегии распространения, которые предлагает TF.

Answer 1

Вроде нашел решение, поэтому выкладываю, может кому поможет. Похоже, что

cluster = tf.compat.v1.train.ClusterSpec({'worker': ['n03:2222', 'n04:2223'] })

вместо cluster_resolver решает проблему одинаковых адресов. Позже мне нужно было вызвать сеанс, и это должен был быть сеанс с целью сервера, относящегося к задаче 1 (не знаю, почему, вероятно, это как-то связано с главным узлом), например:

with tf.compat.v1.Session(server1.target) as sess:
    x=tf.Variable(...)
    for k in range(n):
        y1=f1(x)
        y2=f2(x)
        y1=y1.eval()
        y2=y2.eval()

, где f1(x) - это tf.функция, в которой происходит раздача работникам, например:

@tf.function
def f1(x):
    with tf.device("/job:worker/task:0"):
         y=...
         x.assign(x+1)
    return y

и f2(x) аналогичны, только с задачей 1. Это было все в одном скрипт, который я вызываю в файле. sh.