Please start by defining the specialized terms you will use, e.g. coordinator. Make a list of the different terms and their definition (this will also enable us to check whether we have a unified terminology)

Sounds good. Added a "Glossary" section.

Add return statement

Added.

We should generally ask users to return finite datasets and avoid steps_per_epoch, no? Or it is different for PS?

Since the ClusterCoordinator.schedule can schedule a function to any worker. So it would be less error-prone to just ask users to create the same repeated, but shuffled differently, datasets on different workers. In this case, steps_per_epoch is needed.

It is right now difficult to achieve visitation guarantee with ClusterCoordinator.schedule but may be easier with tf.data service in the future.

Does this mean we should require the user to return an infinite dataset in the PS case? It sounds like the Dataset might hit OutOfRange error otherwise. In that case, we should throw an error upfront if the user passes a finite dataset

Yes, we have explicitly documented the requirement: https://www.tensorflow.org/tutorials/distribute/parameter_server_training#more_about_dataset_creation

There is not error checking in ParameterServerStrategy now, but would be nice to add it.

With the current design of parameter server training APIs, error handling allows OutOfRangeError to be reported to user, but the caveat is that it can mean only one worker has exhausted the dataset, while others still have some remaining to be processed. The behavior may likely be different from what the users would expect, especially when dataset is sharded, so the recommendation is to provide an infinite dataset that is shuffled with a different seed.

SGTM, IMO then in the DataAdapter class for the DatasetFactory, we should add an error if the dataset isn't infinite that explains this

Yes - added this in DataAdapter class part.

It would also be useful to cover a few crucial details in this snippet: TensorBoard remote monitoring, fault tolerance, and saving artifacts

@fchollet It's important that we discuss what is preferred solution here, i.e., dataset, or dataset_fn

"...instance" (both examples use a dataset)

In this case, the dataset instance is serialized and deserialized on each worker?

Yeah, from what I understand though, there are currently issues with trying to serialize and deserialize the Dataset onto each worker, which is why IMO we'll probably have to start with only supporting a "Dataset factory" input type for PS

Done, and I'm thinking the same @omalleyt12

The dataset instance path would require the library to serialize and deserialize the dataset on workers.

We should not have callbacks include PS-specific logic. The task of synchronizing (and deciding when it's appropriate) should be done in the Model.

Note that it is necessary for us to provide a way to synchronize and run callbacks at higher frequency than once per epoch. This is critical. We could use steps_per_execution to configure this (with a high default value for the PS use case).

It is desirable that some callbacks like checkpointing callbacks will be invoked in the middle of an epoch. But increasing the steps_per_execution to a large value has disadvantages: these steps inside a single function call cannot be distributed across workers for load-balancing purposes; worker failure (leading to the interruption of a function execution) will retry all these steps_per_execution steps.

Also checkpointing doesn't necessary need synchronization in asynchronous training. IMO it would be useful for Keras to support time-based callback for checkpointing that doesn't need synchronization. The checkpointing, especially in the middle of an epoch, is mainly for failure recovery purposes and therefore time-based callbacks would easily allow users to know that at most X minutes of training would be lost if there is PS failure.

We should not have callbacks include PS-specific logic

Agree, really the only change needed is in tf_utils.to_numpy_or_python_type, which is used by CallbackList when it detects that it needs to convert Tensor logs into NumPy for a batch. A user would write Callbacks without having to know anything about this, they'll still see NumPy data.

The change would roughly be:

def to_numpy_or_python_type(data):
  # Sync workers and get NumPy values
  if isinstance(data, RemoteValue):
    get_strategy().join()
    return data.fetch()
  else:
    # Current logic stays the same for other strategies.
    ...

Re: The task of synchronizing (and deciding when it's appropriate) should be done in the Model, one option is having logs resolved in model.fit, after the steps are done, and before epoch_end callbacks are called. Concretely, in the current code, we do:

      for epoch, iterator in data_handler.enumerate_epochs():
        ...
        with data_handler.catch_stop_iteration():  # This will sync at end of steps
          for step in data_handler.steps():
              tmp_logs = self.train_function(iterator)
              logs = tmp_logs
              ...
        logs = data_handler.resolve_logs(logs)  # convert remote to concrete
        if logs is None:
          raise ValueError('Expect x to be a non-empty array or dataset.')
        epoch_logs = copy.copy(logs)

        # Run validation.
        ...

        callbacks.on_epoch_end(epoch, epoch_logs)  # epoch_logs are concrete values

With this, we would just need to provide an implementation for resolve_logs. In the case with ClusterCoordinatorDataHandler, it would be logs.fetch(), whereas the default is an no-op. I believe this could avoid a change in tf_utils.to_numpy_or_python_type.

I think it's OK to have the logic in tf_utils.to_numpy_or_python_type (this utility is basically already just "the thing that takes values returned by Model methods and converts them to what the user should see"

IIUC, Francois's concern was that (please correct me if I'm wrong Francois) users who write Callbacks shouldn't have to be aware of any of this, shouldn't have to know what a RemoteValue is, etc.

As long as we achieve that, either data_handler.resolve_logs or tf_utils.to_numpy_or_python_type seems fine to me, with a slight preference for tf_utils because it consolidates all of our logs munging into one place

Sounds good Tom - and thanks for updating the callback section which reflects the current consensus.

This should be done in the Model. The model should initiate the sync every steps_per_execution steps, and call the callbacks at that time with a dict of plain values.

Callbacks must be agnostic to PS.

In addition to the steps_per_execution logic in Model, the CallbackList object has this mechanism so that it only blocks when absolutely necessary. This saves a lot of time for TPUs. This is complimentary to steps_per_execution

But we can re-use this logic while still fully supporting batch-level Callbacks. This logic is just so that, for the users who don't pass batch-level Callbacks, we don't end up unnecessarily slowing down training

For users who pass batch-level Callbacks, the CallbackList object can detect this and pass the logs every batch, with the modifications to tf_utils.to_numpy_or_python_type noted above.

I disagree, batch-level callbacks are useful and important. We just need to make sure that they're called at an appropriate frequency (especially by default)

No to this and the remaining 2 options below.

Sounds good.

It is desirable that some callbacks like checkpointing callbacks will be invoked in the middle of an epoch. But increasing the steps_per_execution to a large value has disadvantages: these steps inside a single function call cannot be distributed across workers for load-balancing purposes; worker failure (leading to the interruption of a function execution) will retry all these steps_per_execution steps.

Also checkpointing doesn't necessary need synchronization in asynchronous training. IMO it would be useful for Keras to support time-based callback for checkpointing that doesn't need synchronization. The checkpointing, especially in the middle of an epoch, is mainly for failure recovery purposes and therefore time-based callbacks would easily allow users to know that at most X minutes of training would be lost if there is PS failure.

+1

Since the ClusterCoordinator.schedule can schedule a function to any worker. So it would be less error-prone to just ask users to create the same repeated, but shuffled differently, datasets on different workers. In this case, steps_per_epoch is needed.

It is right now difficult to achieve visitation guarantee with ClusterCoordinator.schedule but may be easier with tf.data service in the future.

Is the side-car eval API part of the proposal? If so, is there any alternative API you have considered?

I think it's here: "Initially, we aim to have evaluate and predict to only be carried out on the coordinator"

@tanzhenyu you mean this is not part of the proposal? Then we should make it clear here.

SidecarEvaluator API will be covered in a separate RFC (and is just briefly mentioned in this RFC).

Is this API part of the proposal?

Updated this part. We expect users to use try-except.

Question @omalleyt12 - does TPU have ProgressBar by default? It looks like it blocks async when verbose=1, which seems to be the default.

This option would make the total steps dependent on the number of workers, which is not desirable from users' point of view. It is at odds with the programming model provided by the CTL + PSStrategy. The proposed mental model would also break down if we support elasticity in the future.

In this option, the name of the steps_per_epoch argument should be changed to steps_per_epoch_per_worker.

Alternatively, you can sync every W steps only run steps_per_epoch/W steps, effectively turning it into synchronous parameter server training.

@tanzhenyu you mean this is not part of the proposal? Then we should make it clear here.

I think some of the scheduled functions in the same epoch can still pick up the updated variables because the variable read happens at function execution, but not at scheduling time (unless scheduled functions execute super quickly and callback's variable updates are super slow).

Agreeing with this.

Is it still a hard requirement if the data comes from tf.data service processing in distributed_epoch mode?

I think having tf.data service support would be ideal; workers then would just fetch the next example without having to worry that OutOfRangeError occurs. There would not be sharding needed either.