#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2020 The TensorFlow Authors.
End of explanation
"""


!pip install portpicker

#@title
import multiprocessing
import os
import random
import portpicker
import tensorflow as tf

"""
Explanation: Parameter server training with ParameterServerStrategy
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/tutorials/distribute/parameter_server_training"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/parameter_server_training.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs/blob/master/site/en/tutorials/distribute/parameter_server_training.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/distribute/parameter_server_training.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

Overview
Parameter server training is a common data-parallel method to scale up model training on multiple machines.
A parameter server training cluster consists of workers and parameter servers. Variables are created on parameter servers and they are read and updated by workers in each step. By default, workers read and update these variables independently without synchronizing with each other. This is why sometimes parameter server-style training is called asynchronous training.
In TensorFlow 2, parameter server training is powered by the tf.distribute.ParameterServerStrategy class, which distributes the training steps to a cluster that scales up to thousands of workers (accompanied by parameter servers).
Supported training methods
There are two main supported training methods:

The Keras Model.fit API: if you prefer a high-level abstraction and handling of training. This is generally recommended if you are training a tf.keras.Model.
A custom training loop: if you prefer to define the details of your training loop (you can refer to guides on Custom training, Writing a training loop from scratch
 and Custom training loop with Keras and MultiWorkerMirroredStrategy for more details).

A cluster with jobs and tasks
Regardless of the API of choice (Model.fit or a custom training loop), distributed training in TensorFlow 2 involves a 'cluster' with several 'jobs', and each of the jobs may have one or more 'tasks'.
When using parameter server training, it is recommended to have:

One coordinator job (which has the job name chief)
Multiple worker jobs (job name worker)
Multiple parameter server jobs (job name ps)

The coordinator creates resources, dispatches training tasks, writes checkpoints, and deals with task failures. The workers and parameter servers run tf.distribute.Server instances that listen for requests from the coordinator.
Parameter server training with the Model.fit API
Parameter server training with the Model.fit API requires the coordinator to use a tf.distribute.ParameterServerStrategy object. Similar to Model.fit usage with no strategy, or with other strategies, the workflow involves creating and compiling the model, preparing the callbacks, and calling Model.fit.
Parameter server training with a custom training loop
With custom training loops, the tf.distribute.coordinator.ClusterCoordinator class is the key component used for the coordinator.

The ClusterCoordinator class needs to work in conjunction with a tf.distribute.Strategy object.
This tf.distribute.Strategy object is needed to provide the information of the cluster and is used to define a training step, as demonstrated in Custom training with tf.distribute.Strategy.
The ClusterCoordinator object then dispatches the execution of these training steps to remote workers.
For parameter server training, the ClusterCoordinator needs to work with a tf.distribute.ParameterServerStrategy.

The most important API provided by the ClusterCoordinator object is schedule:

The schedule API enqueues a tf.function and returns a future-like RemoteValue immediately.
The queued functions will be dispatched to remote workers in background threads and their RemoteValues will be filled asynchronously.
Since schedule doesn’t require worker assignment, the tf.function passed in can be executed on any available worker.
If the worker it is executed on becomes unavailable before its completion, the function will be retried on another available worker.
Because of this fact and the fact that function execution is not atomic, a single function call may be executed more than once.

In addition to dispatching remote functions, the ClusterCoordinator also helps
to create datasets on all the workers and rebuild these datasets when a worker recovers from failure.
Tutorial setup
The tutorial will branch into Model.fit and custom training loop paths, and you can choose the one that fits your needs. Sections other than "Training with X" are applicable to both paths.
End of explanation
"""


def create_in_process_cluster(num_workers, num_ps):
  """Creates and starts local servers and returns the cluster_resolver."""
  worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)]
  ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)]

  cluster_dict = {}
  cluster_dict["worker"] = ["localhost:%s" % port for port in worker_ports]
  if num_ps > 0:
    cluster_dict["ps"] = ["localhost:%s" % port for port in ps_ports]

  cluster_spec = tf.train.ClusterSpec(cluster_dict)

  # Workers need some inter_ops threads to work properly.
  worker_config = tf.compat.v1.ConfigProto()
  if multiprocessing.cpu_count() < num_workers + 1:
    worker_config.inter_op_parallelism_threads = num_workers + 1

  for i in range(num_workers):
    tf.distribute.Server(
        cluster_spec,
        job_name="worker",
        task_index=i,
        config=worker_config,
        protocol="grpc")

  for i in range(num_ps):
    tf.distribute.Server(
        cluster_spec,
        job_name="ps",
        task_index=i,
        protocol="grpc")

  cluster_resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
      cluster_spec, rpc_layer="grpc")
  return cluster_resolver

# Set the environment variable to allow reporting worker and ps failure to the
# coordinator. This is a workaround and won't be necessary in the future.
os.environ["GRPC_FAIL_FAST"] = "use_caller"

NUM_WORKERS = 3
NUM_PS = 2
cluster_resolver = create_in_process_cluster(NUM_WORKERS, NUM_PS)

"""
Explanation: Cluster setup
As mentioned above, a parameter server training cluster requires a coordinator task that runs your training program, one or several workers and parameter server tasks that run TensorFlow servers—tf.distribute.Server—and possibly an additional evaluation task that runs sidecar evaluation (refer to the sidecar evaluation section below). The requirements to set them up are:

The coordinator task needs to know the addresses and ports of all other TensorFlow servers, except the evaluator.
The workers and parameter servers need to know which port they need to listen to. For the sake of simplicity, you can usually pass in the complete cluster information when creating TensorFlow servers on these tasks.
The evaluator task doesn’t have to know the setup of the training cluster. If it does, it should not attempt to connect to the training cluster.
Workers and parameter servers should have task types as "worker" and "ps", respectively. The coordinator should use "chief" as the task type for legacy reasons.

In this tutorial, you will create an in-process cluster so that the whole parameter server training can be run in Colab. You will learn how to set up real clusters in a later section.
In-process cluster
You will start by creating several TensorFlow servers in advance and you will connect to them later. Note that this is only for the purpose of this tutorial's demonstration, and in real training the servers will be started on "worker" and "ps" machines.
End of explanation
"""


variable_partitioner = (
    tf.distribute.experimental.partitioners.MinSizePartitioner(
        min_shard_bytes=(256 << 10),
        max_shards=NUM_PS))

strategy = tf.distribute.ParameterServerStrategy(
    cluster_resolver,
    variable_partitioner=variable_partitioner)

"""
Explanation: The in-process cluster setup is frequently used in unit testing, such as here.
Another option for local testing is to launch processes on the local machine—check out Multi-worker training with Keras for an example of this approach.
Instantiate a ParameterServerStrategy
Before you dive into the training code, let's instantiate a tf.distribute.ParameterServerStrategy object. Note that this is needed regardless of whether you are proceeding with Model.fit or a custom training loop. The variable_partitioner argument will be explained in the Variable sharding section.
End of explanation
"""


global_batch_size = 64

x = tf.random.uniform((10, 10))
y = tf.random.uniform((10,))

dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(10).repeat()
dataset = dataset.batch(global_batch_size)
dataset = dataset.prefetch(2)

"""
Explanation: In order to use GPUs for training, allocate GPUs visible to each worker. ParameterServerStrategy will use all the available GPUs on each worker, with the restriction that all workers should have the same number of GPUs available.
Variable sharding
Variable sharding refers to splitting a variable into multiple smaller
variables, which are called shards. Variable sharding may be useful to distribute the network load when accessing these shards. It is also useful to distribute computation and storage of a normal variable across multiple parameter servers, for example, when using very large embeddings
that may not fit in a single machine's memory.
To enable variable sharding, you can pass in a variable_partitioner when
constructing a ParameterServerStrategy object. The variable_partitioner will
be invoked every time when a variable is created and it is expected to return
the number of shards along each dimension of the variable. Some out-of-box
variable_partitioners are provided such as
tf.distribute.experimental.partitioners.MinSizePartitioner. It is recommended to use size-based partitioners like
tf.distribute.experimental.partitioners.MinSizePartitioner to avoid
partitioning small variables, which could have a negative impact on model training
speed.
When a variable_partitioner is passed in, and you create a variable directly
under Strategy.scope, the variable will become a container type with a variables
property, which provides access to the list of shards. In most cases, this
container will be automatically converted to a Tensor by concatenating all the
shards. As a result, it can be used as a normal variable. On the other hand,
some TensorFlow methods such as tf.nn.embedding_lookup provide efficient
implementation for this container type and in these methods automatic
concatenation will be avoided.
Refer to the API docs of tf.distribute.ParameterServerStrategy for more details.
Training with Model.fit
<a id="training_with_modelfit"></a>
Keras provides an easy-to-use training API via Model.fit that handles the training loop under the hood, with the flexibility of an overridable train_step, and callbacks which provide functionalities such as checkpoint saving or summary saving for TensorBoard. With Model.fit, the same training code can be used with other strategies with a simple swap of the strategy object.
Input data
Keras Model.fit with tf.distribute.ParameterServerStrategy can take input data in the form of a tf.data.Dataset, tf.distribute.DistributedDataset, or a tf.keras.utils.experimental.DatasetCreator, with Dataset being the recommended option for ease of use. If you encounter memory issues using Dataset, however, you may need to use DatasetCreator with a callable dataset_fn argument (refer to the tf.keras.utils.experimental.DatasetCreator API documentation for details).
If you transform your dataset into a tf.data.Dataset, you should use Dataset.shuffle and Dataset.repeat, as demonstrated in the code example below. 

Keras Model.fit with parameter server training assumes that each worker receives the same dataset, except when it is shuffled differently. Therefore, by calling Dataset.shuffle, you ensure more even iterations over the data.
Because workers do not synchronize, they may finish processing their datasets at different times. Therefore, the easiest way to define epochs with parameter server training is to use Dataset.repeat—which repeats a dataset indefinitely when called without an argument—and specify the steps_per_epoch argument in the Model.fit call. 

Refer to the "Training workflows" section of the tf.data guide for more details on shuffle and repeat.
End of explanation
"""


with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])

  model.compile(tf.keras.optimizers.SGD(), loss="mse", steps_per_execution=10)

"""
Explanation: If you instead create your dataset with tf.keras.utils.experimental.DatasetCreator, the code in dataset_fn will be invoked on the input device, which is usually the CPU, on each of the worker machines.
Model construction and compiling
Now, you will create a tf.keras.Model—a trivial tf.keras.models.Sequential model for demonstration purposes—followed by a Model.compile call to incorporate components, such as an optimizer, metrics, and other parameters such as steps_per_execution:
End of explanation
"""


working_dir = "/tmp/my_working_dir"
log_dir = os.path.join(working_dir, "log")
ckpt_filepath = os.path.join(working_dir, "ckpt")
backup_dir = os.path.join(working_dir, "backup")

callbacks = [
    tf.keras.callbacks.TensorBoard(log_dir=log_dir),
    tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_filepath),
    tf.keras.callbacks.BackupAndRestore(backup_dir=backup_dir),
]

model.fit(dataset, epochs=5, steps_per_epoch=20, callbacks=callbacks)

"""
Explanation: Callbacks and training
<a id="callbacks-and-training"> </a>
Before you call Keras Model.fit for the actual training, prepare any needed callbacks for common tasks, such as:

tf.keras.callbacks.ModelCheckpoint: saves the model at a certain frequency, such as after every epoch.
tf.keras.callbacks.BackupAndRestore: provides fault tolerance by backing up the model and current epoch number, if the cluster experiences unavailability (such as abort or preemption). You can then restore the training state upon a restart from a job failure, and continue training from the beginning of the interrupted epoch.
tf.keras.callbacks.TensorBoard: periodically writes model logs in summary files that can be visualized in the TensorBoard tool.

Note: Due to performance considerations, custom callbacks cannot have batch level callbacks overridden when used with ParameterServerStrategy. Please modify your custom callbacks to make them epoch level calls, and adjust steps_per_epoch to a suitable value. In addition, steps_per_epoch is a required argument for Model.fit when used with ParameterServerStrategy.
End of explanation
"""


feature_vocab = [
    "avenger", "ironman", "batman", "hulk", "spiderman", "kingkong", "wonder_woman"
]
label_vocab = ["yes", "no"]

with strategy.scope():
  feature_lookup_layer = tf.keras.layers.StringLookup(
      vocabulary=feature_vocab,
      mask_token=None)
  label_lookup_layer = tf.keras.layers.StringLookup(
      vocabulary=label_vocab,
      num_oov_indices=0,
      mask_token=None)

  raw_feature_input = tf.keras.layers.Input(
      shape=(3,),
      dtype=tf.string,
      name="feature")
  feature_id_input = feature_lookup_layer(raw_feature_input)
  feature_preprocess_stage = tf.keras.Model(
      {"features": raw_feature_input},
      feature_id_input)

  raw_label_input = tf.keras.layers.Input(
      shape=(1,),
      dtype=tf.string,
      name="label")
  label_id_input = label_lookup_layer(raw_label_input)

  label_preprocess_stage = tf.keras.Model(
      {"label": raw_label_input},
      label_id_input)

"""
Explanation: Direct usage with ClusterCoordinator (optional)
Even if you choose the Model.fit training path, you can optionally instantiate a tf.distribute.coordinator.ClusterCoordinator object to schedule other functions you would like to be executed on the workers. Refer to the Training with a custom training loop section for more details and examples.
Training with a custom training loop
<a id="training_with_custom_training_loop"> </a>
Using custom training loops with tf.distribute.Strategy provides great flexibility to define training loops. With the ParameterServerStrategy defined above (as strategy), you will use a tf.distribute.coordinator.ClusterCoordinator to dispatch the execution of training steps to remote workers.
Then, you will create a model, define a dataset, and define a step function, as you have done in the training loop with other tf.distribute.Strategys. You can find more details in the Custom training with tf.distribute.Strategy tutorial.
To ensure efficient dataset prefetching, use the recommended distributed dataset creation APIs mentioned in the Dispatch training steps to remote workers section below. Also, make sure to call Strategy.run inside worker_fn to take full advantage of GPUs allocated to workers. The rest of the steps are the same for training with or without GPUs.
Let’s create these components in the following steps:
Set up the data
First, write a function that creates a dataset.
If you would like to preprocess the data with Keras preprocessing layers or Tensorflow Transform layers, create these layers outside the dataset_fn and under Strategy.scope, like you would do for any other Keras layers. This is because the dataset_fn will be wrapped into a tf.function and then executed on each worker to generate the data pipeline.
If you don't follow the above procedure, creating the layers might create Tensorflow states which will be lifted out of the tf.function to the coordinator. Thus, accessing them on workers would incur repetitive RPC calls between coordinator and workers, and cause significant slowdown.
Placing the layers under Strategy.scope will instead create them on all workers. Then, you will apply the transformation inside the dataset_fn via tf.data.Dataset.map. Refer to Data preprocessing in the Distributed input tutorial for more information on data preprocessing with distributed input.
End of explanation
"""


def feature_and_label_gen(num_examples=200):
  examples = {"features": [], "label": []}
  for _ in range(num_examples):
    features = random.sample(feature_vocab, 3)
    label = ["yes"] if "avenger" in features else ["no"]
    examples["features"].append(features)
    examples["label"].append(label)
  return examples

examples = feature_and_label_gen()

"""
Explanation: Generate toy examples in a dataset:
End of explanation
"""


def dataset_fn(_):
  raw_dataset = tf.data.Dataset.from_tensor_slices(examples)

  train_dataset = raw_dataset.map(
      lambda x: (
          {"features": feature_preprocess_stage(x["features"])},
          label_preprocess_stage(x["label"])
      )).shuffle(200).batch(32).repeat()
  return train_dataset

"""
Explanation: Then, create the training dataset wrapped in a dataset_fn:
End of explanation
"""


# These variables created under the `Strategy.scope` will be placed on parameter
# servers in a round-robin fashion.
with strategy.scope():
  # Create the model. The input needs to be compatible with Keras processing layers.
  model_input = tf.keras.layers.Input(
      shape=(3,), dtype=tf.int64, name="model_input")

  emb_layer = tf.keras.layers.Embedding(
      input_dim=len(feature_lookup_layer.get_vocabulary()), output_dim=16384)
  emb_output = tf.reduce_mean(emb_layer(model_input), axis=1)
  dense_output = tf.keras.layers.Dense(units=1, activation="sigmoid")(emb_output)
  model = tf.keras.Model({"features": model_input}, dense_output)

  optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.1)
  accuracy = tf.keras.metrics.Accuracy()

"""
Explanation: Build the model
Next, create the model and other objects. Make sure to create all variables under Strategy.scope.
End of explanation
"""


assert len(emb_layer.weights) == 2
assert emb_layer.weights[0].shape == (4, 16384)
assert emb_layer.weights[1].shape == (4, 16384)

print(emb_layer.weights[0].device)
print(emb_layer.weights[1].device)


"""
Explanation: Let's confirm that the use of FixedShardsPartitioner split all variables into two shards and that each shard was assigned to a different parameter server:
End of explanation
"""


@tf.function
def step_fn(iterator):

  def replica_fn(batch_data, labels):
    with tf.GradientTape() as tape:
      pred = model(batch_data, training=True)
      per_example_loss = tf.keras.losses.BinaryCrossentropy(
          reduction=tf.keras.losses.Reduction.NONE)(labels, pred)
      loss = tf.nn.compute_average_loss(per_example_loss)
      gradients = tape.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
    accuracy.update_state(labels, actual_pred)
    return loss

  batch_data, labels = next(iterator)
  losses = strategy.run(replica_fn, args=(batch_data, labels))
  return strategy.reduce(tf.distribute.ReduceOp.SUM, losses, axis=None)

"""
Explanation: Define the training step
Third, create the training step wrapped into a tf.function:
End of explanation
"""


coordinator = tf.distribute.coordinator.ClusterCoordinator(strategy)

"""
Explanation: In the above training step function, calling Strategy.run and Strategy.reduce in the step_fn can support multiple GPUs per worker. If the workers have GPUs allocated, Strategy.run will distribute the datasets on multiple replicas.
Dispatch training steps to remote workers
<a id="dispatch_training_steps_to_remote_workers"> </a>
After all the computations are defined by ParameterServerStrategy, you will use the tf.distribute.coordinator.ClusterCoordinator class to create resources and distribute the training steps to remote workers.
Let’s first create a ClusterCoordinator object and pass in the strategy object:
End of explanation
"""


@tf.function
def per_worker_dataset_fn():
  return strategy.distribute_datasets_from_function(dataset_fn)

per_worker_dataset = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
per_worker_iterator = iter(per_worker_dataset)

"""
Explanation: Then, create a per-worker dataset and an iterator using the ClusterCoordinator.create_per_worker_dataset API, which replicates the dataset to all workers. In the per_worker_dataset_fn below, wrapping the dataset_fn into strategy.distribute_datasets_from_function is recommended to allow efficient prefetching to GPUs seamlessly.
End of explanation
"""


num_epochs = 4
steps_per_epoch = 5
for i in range(num_epochs):
  accuracy.reset_states()
  for _ in range(steps_per_epoch):
    coordinator.schedule(step_fn, args=(per_worker_iterator,))
  # Wait at epoch boundaries.
  coordinator.join()
  print("Finished epoch %d, accuracy is %f." % (i, accuracy.result().numpy()))

"""
Explanation: The final step is to distribute the computation to remote workers using ClusterCoordinator.schedule:

The schedule method enqueues a tf.function and returns a future-like RemoteValue immediately. The queued functions will be dispatched to remote workers in background threads and the RemoteValue will be filled asynchronously.
The join method (ClusterCoordinator.join) can be used to wait until all scheduled functions are executed.
End of explanation
"""


loss = coordinator.schedule(step_fn, args=(per_worker_iterator,))
print("Final loss is %f" % loss.fetch())

"""
Explanation: Here is how you can fetch the result of a RemoteValue:
End of explanation
"""


eval_dataset = tf.data.Dataset.from_tensor_slices(
    feature_and_label_gen(num_examples=16)).map(
          lambda x: (
              {"features": feature_preprocess_stage(x["features"])},
              label_preprocess_stage(x["label"])
          )).batch(8)

eval_accuracy = tf.keras.metrics.Accuracy()

for batch_data, labels in eval_dataset:
  pred = model(batch_data, training=False)
  actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
  eval_accuracy.update_state(labels, actual_pred)

print("Evaluation accuracy: %f" % eval_accuracy.result())

"""
Explanation: Alternatively, you can launch all steps and do something while waiting for
completion:
python
for _ in range(total_steps):
  coordinator.schedule(step_fn, args=(per_worker_iterator,))
while not coordinator.done():
  time.sleep(10)
  # Do something like logging metrics or writing checkpoints.
For the complete training and serving workflow for this particular example, please check out this test.
More about dataset creation
The dataset in the above code is created using the ClusterCoordinator.create_per_worker_dataset API. It creates one dataset per worker and returns a container object. You can call the iter method on it to create a per-worker iterator. The per-worker iterator contains one iterator per worker and the corresponding slice of a worker will be substituted in the input argument of the function passed to the ClusterCoordinator.schedule method before the function is executed on a particular worker.
The ClusterCoordinator.schedule method assumes workers are equivalent and thus assumes the datasets on different workers are the same (except that they may be shuffled differently). Because of this, it is also recommended to repeat datasets, and schedule a finite number of steps instead of relying on receiving an OutOfRangeError from a dataset.
Another important note is that tf.data datasets don’t support implicit serialization and deserialization across task boundaries. So it is important to create the whole dataset inside the function passed to ClusterCoordinator.create_per_worker_dataset. The create_per_worker_dataset API can also directly take a tf.data.Dataset or tf.distribute.DistributedDataset as input.
Evaluation
The two main approaches to performing evaluation with tf.distribute.ParameterServerStrategy training are inline evaluation and sidecar evaluation. Each has its own pros and cons as described below. The inline evaluation method is recommended if you don't have a preference.
Inline evaluation
In this method, the coordinator alternates between training and evaluation, and thus it is called inline evaluation.
There are several benefits of inline evaluation. For example:

It can support large evaluation models and evaluation datasets that a single task cannot hold.
The evaluation results can be used to make decisions for training the next epoch, for example, whether to stop training early.

There are two ways to implement inline evaluation: direct evaluation and distributed evaluation.

Direct evaluation: For small models and evaluation datasets, the coordinator can run evaluation directly on the distributed model with the evaluation dataset on the coordinator:
End of explanation
"""


with strategy.scope():
  # Define the eval metric on parameter servers.
  eval_accuracy = tf.keras.metrics.Accuracy()

@tf.function
def eval_step(iterator):
  def replica_fn(batch_data, labels):
    pred = model(batch_data, training=False)
    actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
    eval_accuracy.update_state(labels, actual_pred)
  batch_data, labels = next(iterator)
  strategy.run(replica_fn, args=(batch_data, labels))

def eval_dataset_fn():
  return tf.data.Dataset.from_tensor_slices(
      feature_and_label_gen(num_examples=16)).map(
          lambda x: (
              {"features": feature_preprocess_stage(x["features"])},
              label_preprocess_stage(x["label"])
          )).shuffle(16).repeat().batch(8)

per_worker_eval_dataset = coordinator.create_per_worker_dataset(eval_dataset_fn)
per_worker_eval_iterator = iter(per_worker_eval_dataset)

eval_steps_per_epoch = 2
for _ in range(eval_steps_per_epoch):
  coordinator.schedule(eval_step, args=(per_worker_eval_iterator,))
coordinator.join()
print("Evaluation accuracy: %f" % eval_accuracy.result())

"""
Explanation: Distributed evaluation: For large models or datasets that are infeasible to run directly on the coordinator, the coordinator task can distribute evaluation tasks to the workers via the ClusterCoordinator.schedule/ClusterCoordinator.join methods:
End of explanation
"""


