code: def train_func(config: dict):
      per_worker_batch_size = config.get("batch_size", 64)
      epochs = config.get("epochs", 3)
      steps_per_epoch = config.get("steps_per_epoch", 70)
  
      tf_config = json.loads(os.environ["TF_CONFIG"])
      num_workers = len(tf_config["cluster"]["worker"])
  
      strategy = tf.distribute.MultiWorkerMirroredStrategy()
  
      global_batch_size = per_worker_batch_size * num_workers
      multi_worker_dataset = mnist_dataset(global_batch_size)
  
      with strategy.scope():
          # Model building/compiling need to be within `strategy.scope()`.
          multi_worker_model = build_cnn_model()
          learning_rate = config.get("lr", 0.001)
          multi_worker_model.compile(
              loss=tf.keras.losses.SparseCategoricalCrossentropy(
                  from_logits=True
              ),
              optimizer=tf.keras.optimizers.SGD(
                  learning_rate=learning_rate
              ),
              metrics=["accuracy"],
          )
  
      history = multi_worker_model.fit(
          multi_worker_dataset,
          epochs=epochs,
          steps_per_epoch=steps_per_epoch,
          callbacks=[ReportCheckpointCallback()],
      )
      results = history.history
      return results
  
file_path: ray/python/ray/train/examples/tf/tensorflow_mnist_example.py
imports: import argparse
  import json
  import os
  import numpy
  import tensorflow
  from filelock import FileLock
  from ray.air.integrations.keras import ReportCheckpointCallback
  from ray.train import Result, RunConfig, ScalingConfig
  from ray.train.tensorflow import TensorflowTrainer
label_desc: Use image node_id 3 for calling a global function with example usage: train_func(config) and returns: results