code: def train_loop_for_worker(config):
      ray.data.DataContext.get_current().execution_options.verbose_progress = (
          True
      )
  
      epoch_times = []
      throughputs = []
      if config["train_sleep_time_ms"] >= 0:
          model = None
      else:
          strategy = (
              tf.distribute.experimental.MultiWorkerMirroredStrategy()
          )
          with strategy.scope():
              model = build_model()
              # model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")
              model.compile(
                  optimizer="Adam",
                  loss="mean_squared_error",
                  metrics=["mse"],
              )
  
      dataset_shard = train.get_dataset_shard("train")
      _tf_dataset = None
      synthetic_dataset = None
      if config["data_loader"] == TF_DATA:
          assert dataset_shard is None
          logger.info("Building tf.dataset...")
          filenames = get_tfrecords_filenames(
              config["data_root"],
              config["num_images_per_epoch"],
              config["num_images_per_input_file"],
          )
          _tf_dataset = build_tf_dataset(
              filenames,
              config["batch_size"],
              config["num_images_per_epoch"],
              config["num_epochs"],
              shuffle_buffer=config["shuffle_buffer_size"],
          )
      elif config["data_loader"] == SYNTHETIC:
          # Build an empty batch and repeat it.
          synthetic_dataset = build_synthetic_dataset(
              config["batch_size"]
          )
      # TODO(swang): We should use the TorchTrainer and iter_torch_batches to
      # compare properly against TORCH_DATALOADER.
      # elif config["data_loader"] == TORCH_DATALOADER:
      #     assert dataset_shard is None
      #     logger.info("Building torch.DataLoader...")
      #     # TODO(swang): pass in shuffle buffer size.
      #     # NOTE(swang): There is no way to .limit() the number of images read
      #     # for torch.
      #     torch_dataset = build_torch_dataset(
      #         config["data_root"],
      #         config["batch_size"],
      #     )
  
      def build_synthetic_tf_dataset(
          dataset, batch_size, num_steps_per_epoch
      ):
          batch = list(
              dataset.iter_tf_batches(
                  batch_size=batch_size, dtypes=tf.float32
              )
          )[0]
          batch = (batch["image"], batch["label"])
  
          # TODO(swang): Might generate a few more records than expected if
          # batches don't divide evenly into num_images_per_epoch.
          def to_tensor_iterator():
              for _ in range(num_steps_per_epoch):
                  yield batch
  
          output_signature = (
              tf.TensorSpec(shape=IMAGE_DIMS, dtype=tf.uint8),
              tf.TensorSpec(shape=(None,), dtype=tf.int32),
          )
          tf_dataset = tf.data.Dataset.from_generator(
              to_tensor_iterator, output_signature=output_signature
          )
          return prepare_dataset_shard(tf_dataset)
  
      num_steps_per_epoch = (
          config["num_images_per_epoch"] // config["batch_size"]
      )
      if config["num_images_per_epoch"] % config["batch_size"]:
          # Assuming batches will respect epoch boundaries.
          num_steps_per_epoch += 1
  
      for epoch in range(config["num_epochs"]):
          tf_dataset = None
          if config["data_loader"] == TF_DATA:
              assert _tf_dataset is not None
              tf_dataset = _tf_dataset
          elif config["data_loader"] == RAY_DATA:
              assert dataset_shard is not None
              tf_dataset = dataset_shard.to_tf(
                  feature_columns="image",
                  label_columns="label",
                  batch_size=config["batch_size"],
              )
          elif config["data_loader"] == SYNTHETIC:
              tf_dataset = build_synthetic_tf_dataset(
                  synthetic_dataset,
                  batch_size=config["batch_size"],
                  num_steps_per_epoch=num_steps_per_epoch,
              )
  
          epoch_start_time_s = time.perf_counter()
  
          if model:
              model.fit(tf_dataset, steps_per_epoch=num_steps_per_epoch)
          else:
              num_rows_read = 0
              for i, batch in enumerate(tf_dataset):
                  num_rows_read += len(batch[0])
                  if i >= num_steps_per_epoch:
                      break
                  time.sleep(config["train_sleep_time_ms"] / 1000)
                  if i % 10 == 0:
                      print("Step", i)
  
              assert num_rows_read >= config["num_images_per_epoch"], (
                  num_rows_read,
                  config["num_images_per_epoch"],
              )
  
          epoch_time_s = time.perf_counter() - epoch_start_time_s
          epoch_times.append(epoch_time_s)
          throughputs.append(
              config["num_images_per_epoch"] / epoch_time_s
          )
  
          total_tput = config["num_images_per_epoch"] / epoch_time_s
          # Drop the first epoch to remove warmup time.
          if len(epoch_times) > 1:
              total_tput = (
                  (epoch)
                  * config["num_images_per_epoch"]
                  / sum(epoch_times[1:])
              )
          logger.info(
              "Epoch time: {epoch_time_s}s, images/s: {throughput}".format(
                  epoch_time_s=epoch_time_s,
                  throughput=config["num_images_per_epoch"]
                  / epoch_time_s,
              )
          )
  
          train.report(
              {
                  "all_epoch_times_s": epoch_times,
                  "all_throughputs_imgs_s": throughputs,
                  "tput_images_per_s": total_tput,
              }
          )
  
          if config["data_loader"] == RAY_DATA:
              print_dataset_stats(dataset_shard)
              print("epoch time", epoch, epoch_time_s)
  
imports: from collections import defaultdict
  import tensorflow
  import numpy
  import os
  import pandas
  import time
  import logging
  import csv
  import json
  import torchvision
  import torch
  import ray
  from ray.train.tensorflow import prepare_dataset_shard, TensorflowTrainer
  from ray.train import DataConfig, ScalingConfig
  from ray import train, tune
  from ray.tune import Tuner
  from ray.data.datasource.partitioning import Partitioning
  from tf_utils import DEFAULT_IMAGE_SIZE, NUM_CHANNELS, preprocess_image, build_tf_dataset
  from metric_utils import determine_if_memory_monitor_is_enabled_in_latest_session, get_ray_spilled_and_restored_mb, MaxMemoryUtilizationTracker
label_desc: Use image node_id 3 for calling a global function with example usage: train_loop_for_worker(config) and returns: prepare_dataset_shard