mutiple issues with the new parameter server strategy
> 1. how to shard data by files, how to use the later binding mechanism to shard data by files, which is critical for high performance training using tf.data.Dataset apis, all workers reading the total data is not scalable.

> 3. how to ensure even the workers are with uneven amount of data, the training process with model.fit could end elegantly instead of having to using try catch or data.repeat, as recommendation models generally assume one epoch training.



Due to the nature of the parameter server training environment (unsynchronized workers, higher preemption rate), iterating over the training data exactly once is hard to enforce. For this reason the [default recommendation](https://www.tensorflow.org/tutorials/distribute/parameter_server_training#input_data) is to shuffle data (approximately ensuring workers see different data) and repeat it infinitely, then train for a fixed number of steps. 

There is a [new experimental API](https://www.tensorflow.org/api_docs/python/tf/distribute/coordinator/experimental_get_current_worker_index) to let workers see disjoint datasets, so there would be no data overlap between workers. This should provide the late binding you're looking for. (However if a worker fails and recovers it will still iterate over parts of its data shard more than once). 

> 2. how to control the placement of ops, in one attempt, we build dataset using tf.data.Dataset.from_tensor_slices and.
dataset = dataset.interleave(lambda x: tf.data.TextLineDataset(x).skip(1), cycle_length=15, num_parallel_calls=15)
with parameter server strategy, the dataset related ops runs on workers, however, if I add one more op dataset.repeat(), the dataset ops all runs on chief, which is surprising.

This is surprising, are you able to provide the exact dataset building code that could help reproduce this?  This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. 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/61129">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/61129">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

2.10.1

### Custom Code

Yes

### OS Platform and Distribution

ubuntu 16.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

We meet multiple issues in using the Tensorflow 2.x strategy:
1. how to shard data by files, how to use the later binding mechanism to shard data by files, which is critical for high performance training using tf.data.Dataset apis, all workers reading the total data is not scalable.
2. how to control the placement of ops, in one attempt, we build dataset using tf.data.Dataset.from_tensor_slices and.     
     dataset = dataset.interleave(lambda x: tf.data.TextLineDataset(x).skip(1),        cycle_length=15,       num_parallel_calls=15)
   with parameter server strategy, the dataset related ops runs on workers, however, if I add one more op dataset.repeat(), the dataset ops all runs on chief, which is surprising.
4. how to ensure even the workers are with uneven amount of data, the training process with model.fit could end elegantly instead of having to using try catch or data.repeat, as recommendation models generally assume one epoch training.

[yuefengz@google.com](mailto:yuefengz@google.com)
[rchao@google.com](mailto:rchao@google.com)

### Standalone code to reproduce the issue

```shell
# for chief: 
file_list = tf.io.gfile.glob(input_pattern)
dataset = tf.data.Dataset.from_tensor_slices(file_list)
dataset = dataset.shard(worker_num, worker_id)  # how to later bind worker_id?
dataset = dataset.interleave(lambda x: tf.data.TextLineDataset(x).skip(1),        cycle_length=15,       num_parallel_calls=15)

def _parse_csv(line):
  record_defaults = []
  for x in all_columns:
    record_defaults.append(0)
  with tf.control_dependencies([tf.print(tf.shape(line), line[0], output_stream=sys.stderr)]):
    fields = tf.io.decode_csv(
        line,
        field_delim=',',
        record_defaults=record_defaults,
        name='decode_csv')
    return fields

dataset = dataset.map(_parse_csv, num_parallel_calls=8)
dataset = dataset.repeat()

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

with strategy.scope():
  model = build_model(...)

model.fit(dataset, epochs=num_epoch,
     callbacks=[
         tf.keras.callbacks.ProgbarLogger(count_mode='steps'),
         tf.keras.callbacks.TensorBoard(log_dir='/train/tensorboard/', histogram_freq=0)
     ],
     steps_per_epoch=steps_per_epoch)

# for workers and ps:
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()
```


### Relevant log output

```shell
# tf.print outputs are all on chief stdout, and no improvements if we place the dataset built process under strategy.scope() or using ops.device() operations.
```
</details>