code: def train_func(use_ray: bool, config: dict):
      local_start_time = time.monotonic()
  
      per_worker_batch_size = config.get("batch_size", 64)
      epochs = config.get("epochs", 3)
      steps_per_epoch = config.get("steps_per_epoch", None)
      learning_rate = config.get("lr", 0.001)
  
      tf_config = json.loads(os.environ["TF_CONFIG"])
      num_workers = len(tf_config["cluster"]["worker"])
      local_rank = tf_config["task"]["index"]
  
      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()
          multi_worker_model.compile(
              loss=tf.keras.losses.SparseCategoricalCrossentropy(
                  from_logits=True
              ),
              optimizer=tf.keras.optimizers.SGD(
                  learning_rate=learning_rate
              ),
              metrics=["accuracy"],
          )
  
      if use_ray:
          from ray.air.integrations.keras import (
              ReportCheckpointCallback,
          )
  
          class CustomReportCallback(ReportCheckpointCallback):
              def _handle(self, logs: dict, when: str = None):
                  logs["local_time_taken"] = (
                      time.monotonic() - local_start_time
                  )
                  super()._handle(logs, when)
  
          # NOTE: We shouldn't checkpoint to be identical to the vanilla TF run.
          callbacks = [CustomReportCallback(checkpoint_on=[])]
      else:
          callbacks = []
  
      history = multi_worker_model.fit(
          multi_worker_dataset,
          epochs=epochs,
          steps_per_epoch=steps_per_epoch,
          callbacks=callbacks,
          verbose=2,  # Disables progress bar in remote actors.
      )
      results = history.history
      loss = results["loss"][-1]
  
      if not use_ray:
          local_time_taken = time.monotonic() - local_start_time
          print(f"Reporting loss: {loss:.4f}")
          if local_rank == 0:
              with open(VANILLA_RESULT_JSON, "w") as f:
                  json.dump(
                      {
                          "loss": loss,
                          "local_time_taken": local_time_taken,
                      },
                      f,
                  )
  
      return results
  
file_path: ray/release/air_tests/air_benchmarks/workloads/tensorflow_benchmark.py
imports: import json
  import os
  import time
  from pathlib import Path
  import click
  import numpy
  import tensorflow
  from typing import List, Tuple
label_desc: Use image node_id 3 for calling a global function with example usage: train_func(use_ray, config) and returns: results