//  [python]
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)