Repository: ray-project/ray
File: python/ray/train/v2/_internal/execution/callback.py (Python)

Code before:

    def before_controller_execute_resize_decision(
        self,
        resize_decision: "ResizeDecision",
    ):
        """Called before the controller executes a resize decision."""
        pass

    def after_controller_finish(self, result: "Result"):
        """Called after the training run completes, providing access to the final result.

        Args:
            result: The final training result containing metrics and checkpoint.
        """
        pass

    def on_callback_hook_exception(
        self, hook_name: str, error: Exception, **context
    ) -> CallbackHookExceptionHandlerResult:
        """Called when a callback hook raises an exception.

        This is invoked after the exception is caught by the callback manager.

        Return:
            A tuple of (action, error):
            - (CallbackErrorAction.SUPPRESS, None) to continue.
            - (CallbackErrorAction.SUPPRESS, TrainingFailedError) to log the exception and continue train.
            - (CallbackErrorAction.RAISE, TrainingFailedError) to propagate the exception to the training control flow.
        """
        return (CallbackErrorAction.SUPPRESS, ControllerError(error))


# TODO: consider consolidating all metrics into one dict, possibly with UDF
@DeveloperAPI
class ReportCallback(RayTrainCallback):
    def after_report(
        self,
        training_report: _TrainingReport,
        metrics: List[Dict[str, Any]],
    ):
        """Called after all workers have reported a training result.

        Note that this differs from `after_worker_group_poll_status`,
        which may only contain a subset of workers that have reported.
        For example, if only rank 0 is performing checkpointing, then
        rank 0 would report a training result the slowest.
        """
        pass


@DeveloperAPI

Review comment (suggestion):
Are these the only three allowed combinations? If so, we should clarify it in the docstrings- right now they may be confused for examples- and enforce it in your validation function `_is_callback_hook_exception_handler_result`

Code after:
        which may only contain a subset of workers that have reported.
        For example, if only rank 0 is performing checkpointing, then
        rank 0 would report a training result the slowest.
        """
        pass


@DeveloperAPI
class WorkerCallback(RayTrainCallback):
    """
    Callbacks that are hooked to the worker event.

    These callbacks are created on the train driver process and then
    copied and passed to all the workers.
    The execution of these callbacks happens on each of the workers,
    not on the train driver process.
    """

    def after_init_train_context(self):
        pass

    def before_worker_shutdown(self):
        pass


@DeveloperAPI
class TrainContextCallback(RayTrainCallback):
    """
    Callbacks that are hooked to the train context event.

    These callbacks are created on the train driver process and then
    copied and passed to all the workers.
    The execution of these callbacks happens on the train context of the workers.
    """

    @contextmanager
    def on_report(self):
        yield