code: def train_fn(
      compute_config: TfDataServiceConfig,
      reuse_dataset: bool = False,
      round_robin: bool = False,
  ):
      # Horovod: initialize Horovod.
      hvd.init()
  
      # Horovod: pin GPU to be used to process local rank (one GPU per process)
      gpus = tf.config.experimental.list_physical_devices("GPU")
      for gpu in gpus:
          tf.config.experimental.set_memory_growth(gpu, True)
      if gpus:
          tf.config.experimental.set_visible_devices(
              gpus[hvd.local_rank()], "GPU"
          )
  
      # this lock guarantees only one training task downloads the dataset
      with FileLock(os.path.expanduser("~/.horovod_lock")):
          (
              mnist_images,
              mnist_labels,
          ), _ = tf.keras.datasets.mnist.load_data()
  
      dataset = tf.data.Dataset.from_tensor_slices(
          (
              tf.cast(
                  mnist_images[..., tf.newaxis] / 255.0, tf.float32
              ),
              tf.cast(mnist_labels, tf.int64),
          )
      )
  
      # Allow tf.data service to pre-process the pipeline
      dataset = (
          dataset.repeat()
          .shuffle(10000)
          .batch(128)
          .send_to_data_service(
              compute_config,
              hvd.rank(),
              hvd.size(),
              reuse_dataset=reuse_dataset,
              round_robin=round_robin,
          )
          .prefetch(tf.data.experimental.AUTOTUNE)
      )
  
      mnist_model = tf.keras.Sequential(
          [
              tf.keras.layers.Conv2D(32, [3, 3], activation="relu"),
              tf.keras.layers.Conv2D(64, [3, 3], activation="relu"),
              tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
              tf.keras.layers.Dropout(0.25),
              tf.keras.layers.Flatten(),
              tf.keras.layers.Dense(128, activation="relu"),
              tf.keras.layers.Dropout(0.5),
              tf.keras.layers.Dense(10, activation="softmax"),
          ]
      )
  
      # Horovod: adjust learning rate based on number of GPUs.
      scaled_lr = 0.001 * hvd.size()
      opt = tf.optimizers.Adam(scaled_lr)
  
      # Horovod: add Horovod DistributedOptimizer.
      opt = hvd.DistributedOptimizer(
          opt,
          backward_passes_per_step=1,
          average_aggregated_gradients=True,
      )
  
      # Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow
      # uses hvd.DistributedOptimizer() to compute gradients.
      mnist_model.compile(
          loss=tf.losses.SparseCategoricalCrossentropy(),
          optimizer=opt,
          metrics=["accuracy"],
          experimental_run_tf_function=False,
      )
  
      callbacks = [
          hvd.callbacks.BroadcastGlobalVariablesCallback(0),
          hvd.callbacks.LearningRateWarmupCallback(
              initial_lr=scaled_lr, warmup_epochs=3, verbose=1
          ),
      ]
  
      # Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them.
      if hvd.rank() == 0:
          callbacks.append(
              tf.keras.callbacks.ModelCheckpoint(
                  "./checkpoint-{epoch}.h5"
              )
          )
  
      # Horovod: write logs on worker 0.
      verbose = 1 if hvd.rank() == 0 else 0
  
      # Train the model.
      # Horovod: adjust number of steps based on number of GPUs.
      mnist_model.fit(
          dataset,
          steps_per_epoch=32 // hvd.size(),
          callbacks=callbacks,
          epochs=24,
          verbose=verbose,
      )
  
imports: import os
  import tensorflow
  from filelock import FileLock
  import horovod.tensorflow.keras
  from horovod.tensorflow.data.compute_service import TfDataServiceConfig
label_desc: Use image node_id 1 for calling a global function with example usage: train_fn(compute_config, reuse_dataset, round_robin) without return types