[train v2] Add Train v2 user-facing callback interface

[justinvyu]

Summary

Adds UserCallback, which is a DeveloperAPI that exposes the reported results from ray.train.report calls as well as exceptions that are raised during training.

For now, it only exposes 2 callback methods:

• after_report(run_context: TrainRunContext, metrics: List[Dict[str, Any]], checkpoint: Optional[Checkpoint])
  • The metrics argument is a list of all reported dicts by ray.train.report, where metrics[i] is the dict reported by rank i.
  • The checkpoint argument is populated if one of the workers reported a checkpoint. There's only one of these since all workers report to the same storage location.
• after_exception(run_context: TrainRunContext, worker_exceptions: Dict[int, Exception])
  • worker_exceptions maps worker rank to an exception. Does not include ranks that do not have errors associated with them.

The TrainRunContext is the run-level context that lives on the Train driver.

Here's a diagram showing the dependencies of the UserCallback.

New CallbackHandler concept

As part of this PR, I consolidated the CheckpointHandler logic to gather checkpoints/metrics reported by all workers into a ReportCallbackHandler, since the UserCallback.on_report method also relies on the same logic of gathering results across multiple poll_status calls.

I also introduced a UserCallbackHandler which implements the logic needed to call the UserCallback interface.

These two new callbacks are examples of "callback handlers" which take in a list of callbacks and are responsible for calling their methods. The naming convention is [SubscriberCallback]Handler.

Follow-up

A follow-up PR here is to implement TuneReportCallback as a provided UserCallback.

Example Usage

def after_report(self, run_context, metrics, checkpoint):
    print(run_context)
    # TrainRunContext(run_config=RunConfig(name='ray_train_run-2025-01-16_18-18-39', storage_path='/private/var/folders/v6/r_z7clls5hl9bhbc2s7yyvjm0000gq/T/pytest-of-justin/pytest-21/test_user_callback0', storage_filesystem=None, failure_config=FailureConfig(), checkpoint_config=CheckpointConfig(), sync_config='DEPRECATED', verbose='DEPRECATED', stop='DEPRECATED', callbacks=[<test_data_parallel_trainer.MyUserCallback object at 0x11d4c1d10>], progress_reporter='DEPRECATED', log_to_file='DEPRECATED', local_dir=None))

    print(metrics)
    # [{'rank': 0}, {'rank': 1}]

def after_exception(self, run_context, worker_exceptions):
    print(worker_exceptions)
    # {0: UserExceptionWithTraceback(ValueError('error'), 'Traceback (most recent call last):\n  File "/Users/justin/Developer/ray/python/ray/train/v2/tests/test_data_parallel_trainer.py", line 201, in _train_fn\n    raise ValueError("error")\nValueError: error\n')}

[justinvyu]

Note: we shouldn't have too many callbacks implementing this method since it might accumulate a lot of overhead every health check period (ex: 2 seconds).

UserCallbackInvoker also implements this now to capture the worker group failures.

[justinvyu]

I refactored the invoke_callbacks_context_managers a bit to fix a cyclical dependency issue.

[justinvyu]

Naming?

[matthewdeng]

UserCallback sounds good & straightforward to me.

[matthewdeng]

nit: Clarify that the checkpoint points to the persisted location, since this is currently a bit ambiguous.

[matthewdeng]

Should we keep this as UserCallback for now to prevent misuse, and expand as needed?

[matthewdeng]

It's a bit confusing if we have a UserCallback class but it's not exposed in the public API.

[justinvyu]

What about making the type UserCallback but then keeping this https://github.com/ray-project/ray/pull/49819/files#diff-730d6551cfa9d3b55d3d413ce0500fad427b6409bc47d7e103f4f586341c3b2eR210-R214

So that workarounds / custom callbacks are still possible for workarounds and in unit tests. It will just break the typechecking on the editor.

[matthewdeng]

Yeah that works for now.

[matthewdeng]

nit: Does this need to be an instance variable or should it just be a local variable? Seems unused outside of the init.

Suggested change

	self._report_handler = ReportHandler(
	report_handler = ReportHandler(

[matthewdeng]

Would prefer if all callbacks have Callback as their suffix, so it's easy to tell if it is a Callback or a related class.

For this one, it can be something like CombinedUserCallback.

[matthewdeng]

Same for ReportHandler -> ReportHandlerCallback.

[justinvyu]

@hongpeng-guo actually thought it was more confusing for callbacks that are not meant to be further subclassed to have a Callback suffix.

Ex: ReportCallback is a callback interface, but UserCallbackInvoker just implements the interface but should not be a new callback interface.

[matthewdeng]

Hmm isn't that what we do for all the other callbacks (e.g. AcceleratorSetupCallback, BackendSetupCallback, ...)?

[justinvyu]

Yeah, the idea was to remove the suffix later but didn't do it yet.

[hongpeng-guo]

Yeah, as the callback structure are getting more complicated, i.e., some callbacks becomes the input of other callbacks, the invoker callbacks are not being subclassed, but the main purpose of them is to invoke other lower level callbacks. I think it's more clear for us to differentiate these two levels of callbacks.

I.e., right now, the UserCallbackInvoker and ReportHandler are the two callbacks that will take other callbacks as input, i.e, these two are the higher level internal callbacks. I think it makes sense to name these two things specially. For other first level callbacks, I think it's fine to have the Callback suffix with them.

Another way to think of it is: Callbacks are relatively stand alone modules that if we modify one callback, it should not influence the effect of another callback. But any changes in the two invokers will change how their subscribers effects. That's why I think these two callbacks should be at least named specially. I may still prefer to have a world of only one level of callbacks if possible.

[hongpeng-guo]

offline discussion result XXXXCallbackHandler takes in a list of XXXXCallback and is responsible for calling its methods.

[matthewdeng]

UserCallback sounds good & straightforward to me.

[matthewdeng]

Add tests?

[hongpeng-guo]

Overall looks good to me. The following comments are not blocking.

It seems the main goal of this PR proposing UserCallback is to enable TuneReportCallback which is a P0 feature to integrate Tune with Train V2. I think we should deliver this PR ASAP to release Train V2 as a whole.

However, as I am thinking about some possible feature on async checkpointing and evaluation, I feel report(metrics: Dict, checkpoint(Optional[CKPT])) may not be an optimal API in the long run. I.e., if the evaluation happens on a separate cluster, it may not be very intuitive for every training worker to report a metric. I think it's possible that we need to revamp the report/ checkpoint related callbacks in some near future. But this should be non-blocking for now.

[justinvyu]

@hongpeng-guo Yeah, I'm also not too confident about how future-proof the current callback API is, so I'm keeping it as a DeveloperAPI and not surfacing it in docs. The only thing in docs will be the TuneReportCallback that is our "out of the box" integration API between Train v2 and Tune.