Distributed training with parameter servers example using a single binary
@ali-raza-tariq,
**ClusterResolvers** are a way for TensorFlow to communicate with various cluster management systems (e.g. GCE, AWS, etc...) and gives TensorFlow necessary information to set up distributed training.

By letting TensorFlow communicate with these systems, we will be able to automatically discover and resolve IP addresses for various TensorFlow workers. This will eventually allow us to automatically recover from underlying machine failures and scale TensorFlow worker clusters up and down.

Could you please take a look at the official document for the information on cluster_resolver.
https://www.tensorflow.org/api_docs/python/tf/distribute/cluster_resolver/ClusterResolver
https://www.tensorflow.org/api_docs/python/tf/distribute/cluster_resolver/TFConfigClusterResolver

Also provide the complete error log to debug the issue. Thank you! Hi @tilakrayal thank you for your response. First of all, does this mean the `Chief` no longer participates in the training? Previously, `Chief` (also referred to as `worker0`) used to be just another worker with some extra coordination responsibilities. Secondly, kindly let me know if there are any end to end examples using using `Keras` new style of code (specifically for the `parameterServerStrategy`). This code is just an example using snippets of code I took from the tensorflow documentation from the link above. Which is why I shared the code to make sure I am not doing something fundamentally wrong in the training job. To be fair I am not even sure if its an error or not but this is what I get on my Chief, kindly look below:

```
2023-08-05 23:40:02.230291: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-08-05 23:40:03.951682: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:449] Started server with target: grpc://localhost:12345
/usr/local/lib/python3.8/dist-packages/tensorflow/python/data/ops/dataset_ops.py:468: UserWarning: To make it possible to preserve tf.data options across serialization boundaries, their implementation has moved to be part of the TensorFlow graph. As a consequence, the options value is in general no longer known at graph construction time. Invoking this method in graph mode retains the legacy behavior of the original implementation, but note that the returned value might not reflect the actual value of the options.
  warnings.warn("To make it possible to preserve tf.data options across "
Epoch 1/12
2023-08-05 23:40:06.968495: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 3s - loss: 0.7186 - accuracy: 0.7920 - 3s/epoch - 28ms/step
Epoch 2/12
2023-08-05 23:40:08.341929: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.2847 - accuracy: 0.9103 - 1s/epoch - 14ms/step
Epoch 3/12
2023-08-05 23:40:09.810759: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.2160 - accuracy: 0.9369 - 1s/epoch - 15ms/step
Epoch 4/12
2023-08-05 23:40:11.204296: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.1960 - accuracy: 0.9413 - 1s/epoch - 14ms/step
Epoch 5/12
WARNING:tensorflow:5 out of the last 5 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7fdd00686d30> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:5 out of the last 5 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7fdd00686d30> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
2023-08-05 23:40:12.640156: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.1419 - accuracy: 0.9598 - 1s/epoch - 14ms/step
Epoch 6/12
WARNING:tensorflow:6 out of the last 6 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7fdd668989d0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:6 out of the last 6 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7fdd668989d0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
2023-08-05 23:40:13.990166: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.1409 - accuracy: 0.9558 - 1s/epoch - 13ms/step
Epoch 7/12
2023-08-05 23:40:15.301438: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.1303 - accuracy: 0.9634 - 1s/epoch - 13ms/step
Epoch 8/12
2023-08-05 23:40:16.636748: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.1158 - accuracy: 0.9672 - 1s/epoch - 13ms/step
Epoch 9/12
2023-08-05 23:40:17.985054: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.0969 - accuracy: 0.9745 - 1s/epoch - 13ms/step
Epoch 10/12
2023-08-05 23:40:19.361131: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.0949 - accuracy: 0.9725 - 1s/epoch - 14ms/step
Epoch 11/12
2023-08-05 23:40:20.846003: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.0965 - accuracy: 0.9727 - 1s/epoch - 15ms/step
Epoch 12/12
2023-08-05 23:40:22.234380: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
100/100 - 1s - loss: 0.0703 - accuracy: 0.9784 - 1s/epoch - 14ms/step
2023-08-05 23:40:22.723441: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_65'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.723257314","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_65'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.724180: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_65'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.724041713","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_65'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.724904: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_54'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.724815013","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_54'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.725483: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_54'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.725363814","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_54'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.726046: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_25'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.725959729","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_25'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.726620: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_25'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.726483537","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_25'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.727303: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_22'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.727133327","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_22'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.727894: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_22'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.727780581","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_22'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.762456: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_133'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.762253070","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_133'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.763108: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_133'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.762942628","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_lookup_nest_133'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.763966: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_122'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.763808342","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_122'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.764617: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_122'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.764451669","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_map_parse_and_decode_122'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.765414: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_93'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.765246613","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_93'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.766108: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_93'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.765934285","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_interleave_classfunctools.partial_93'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.766876: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_90'.
Additional GRPC error information from remote target /job:ps/replica:0/task:0:
:{"created":"@1691278822.766725231","description":"Error received from peer ipv4:127.0.0.1:34567","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_90'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
2023-08-05 23:40:22.767543: E tensorflow/core/common_runtime/eager/context.cc:878] Failed to remove function remotely due to Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_90'.
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1691278822.767393731","description":"Error received from peer ipv4:127.0.0.1:23456","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Tried to remove non-existent function '__inference_Dataset_flat_map_read_one_file_90'.","grpc_status":3}
This could happen if the remote target has been disconnected from the client.
```

For some reason, the worker (separate shell) does not show any thing related to training progress. Worker log:
```
2023-08-05 23:40:03.944832: I tensorflow/core/distributed_runtime/eager/eager_service_impl.cc:311] Creating sync eager service context with rendezvous_id on host tensorflow4 /job:worker/replica:0/task:0
2023-08-05 23:40:03.944929: I tensorflow/core/distributed_runtime/eager/eager_service_impl.cc:317] SessionOptions: device_count {
  key: "CPU"
  value: 1
}
device_count {
  key: "GPU"
  value: 0
}
gpu_options {
  experimental {
  }
}
allow_soft_placement: true
experimental {
  collective_group_leader: "/job:ps/replica:0/task:0"
}

2023-08-05 23:40:04.238759: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:551] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed.
2023-08-05 23:40:04.389447: W tensorflow/core/framework/dataset.cc:956] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
2023-08-05 23:40:22.752175: I tensorflow/core/common_runtime/eager/kernel_and_device.cc:94] Ignoring error status when releasing multi-device function handle UNIMPLEMENTED: Releasing a multi-device component handle on a remote device is not yet implemented.
``` > To my understanding, all workers and paramter-servers will start and wait for chief to assign the tasks. Chief or coordinator (documentation uses them interchangeably but is there any difference between the two?) will automatically divide the work based on the information it gets from cluster_resolver (let me know if that's wrong interpretation).

Yes this is accurate. The chief (AKA the coordinator; they are different names for the same thing) is the only task running `Model.fit`, which under the hood schedules the function executions and dispatches them to workers and parameter servers. 

> First of all, does this mean the Chief no longer participates in the training? Previously, Chief (also referred to as worker0) used to be just another worker with some extra coordination responsibilities.

This is also accurate. Unlike MultiWorkerMirroredStrategy, the chief task is not a worker and doesn't run any training steps itself. It's only called the "chief" so that the `TF_CONFIG` can work the same way. "Coordinator" is the most meaningful term for that task -- we can look into consolidating on that term.

> Secondly, kindly let me know if there are any end to end examples using using Keras new style of code (specifically for the parameterServerStrategy). This code is just an example using snippets of code I took from the tensorflow documentation from the link above. 

The tutorial is the most complete example we have. I don't see anything fundamentally wrong with your code. It looks like the training completes 12/12 epochs. Those errors look like something that happens at the end of training when TF is internally calling destructors on some resources. Sorry they are spammy but I'm pretty confident they can be ignored. You could confirm by adding a final `model.evaluate` or `model.predict` calls after `model.fit` in your code.

> For some reason, the worker (separate shell) does not show any thing related to training progress. Worker log:

Yes, the workers and parameter servers won't have many useful logs during training. If they encounter errors those would be propagated to the coordinator which will log them. 

Overall your code looks to be working appropriately and hopefully these answers clear things up. Will close this now -- feel free to reopen or open a new issue if you encounter other errors. Thanks! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61457">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61457">No</a>
 Hi @jamesmullenbach, thank you so much for the detailed response! yes it has cleared things up for me.
Hello everyone! I am sorry if this is a duplicate issue but from my considerable search - I could not find a single end-to-end distributed parameter-server example to run using tensorflow (using the keras api with `.fit()` method). Also, for some reason - the documentation for parameter-server strategy seems a lot more confusing and difficult to get started with, compared to multi-worker strategy. 

I have been running training jobs using the estimator api before and now trying to update it to TF2.x style distributed training job with parameter-server training strategy using a single binary file for all workers and parameter-servers. I started with the example in documentation here (https://www.tensorflow.org/tutorials/distribute/parameter_server_training) and modified the code to be used as a single binary. 

Code:
```
import tensorflow_datasets as tfds
import tensorflow as tf

import os

cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()
if cluster_resolver.task_type in ("worker", "ps"):
  # Start a TensorFlow server and wait.
  server = tf.distribute.Server(cluster_resolver.cluster_spec(),
                                      job_name=cluster_resolver.task_type,
                                      task_index=cluster_resolver.task_id,
                                      protocol=cluster_resolver.rpc_layer or "grpc",
                                      start=True)
  server.join()
else:
  ## parameter-server
  strategy = tf.distribute.ParameterServerStrategy(cluster_resolver=cluster_resolver)
  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)
  with strategy.scope():
    model = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
    model.compile(tf.keras.optimizers.legacy.SGD(), loss="mse", steps_per_execution=10)
  working_dir = "./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)
```

To my understanding, all workers and paramter-servers will start and wait for chief to assign the tasks. Chief or coordinator (documentation uses them interchangeably but is there any difference between the two?) will automatically divide the work based on the information it gets from `cluster_resolver` (let me know if that's wrong interpretation). In any case, I would highly appreciate if someone can point out what I am doing wrong in this example because I have not been able to get it to work!