tensorflow2 on slurm cluster not work correctly - tensorflow

Question

I have seen similar questions but there were for tensorflow1.x or didn't work for me.
I want to train models on multiple threads on multiple CPUs across the cluster.
For now I have example code for simple MNIST classification :
import os
import tensorflow as tf
import numpy as np
import json
def mnist_dataset(batch_size):
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train / np.float32(255)
y_train = y_train.astype(np.int64)
train_dataset = tf.data.Dataset.from_tensor_slices(
(x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
return train_dataset
def build_and_compile_cnn_model():
model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10)
])
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
metrics=['accuracy'])
return model
from hostlist import expand_hostlist
task_index = int( os.environ['SLURM_PROCID'] )
n_tasks = int( os.environ['SLURM_NPROCS'] )
tf_hostlist = [ ("%s:33333" % host) for host in
expand_hostlist( os.environ['SLURM_NODELIST']) ]
node = os.environ['SLURMD_NODENAME']
print('id',task_index)
index = tf_hostlist.index(f'{node}:33333')
print('index', index)
tf_config = {
'cluster': {
'worker': [tf_hostlist[1], tf_hostlist[2]],
'chief': [tf_hostlist[0]]
},
'task': {'type': 'worker', 'index':index}
}
if index == 0:
tf_config['task']['type'] = 'chief'
tf_config['task']['index'] = index
else:
tf_config['task']['index'] = index-1
per_worker_batch_size = 64
os.environ['TF_CONFIG'] = json.dumps(tf_config)
num_workers = len(tf_config['cluster']['worker'])
print('num_workers', num_workers)
strategy = tf.distribute.MultiWorkerMirroredStrategy()
global_batch_size = per_worker_batch_size * num_workers
multi_worker_dataset = mnist_dataset(global_batch_size)
with strategy.scope():
multi_worker_model = build_and_compile_cnn_model()
multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
It's written to run on three dfferent CPUs
I have also script for run on slurm cluster:
#!/bin/bash
#SBATCH --job-name=mnist_tf_distributed
#SBATCH --nodes=3
#SBATCH --cpus-per-task=10
#SBATCH --time=00:05:00
#SBATCH --exclusive
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
srun python keras_multi.py
The problem is that training start three times, not once disturbed on three CPUs.
Can anyone know how to change code to run above training once across different machines.