From d3561fd8e4c1bcf809fa1c6aec7fa2c2f6e89fa3 Mon Sep 17 00:00:00 2001 From: Justin Yu Date: Fri, 17 Jan 2025 11:59:41 -0800 Subject: [PATCH] add comment Signed-off-by: Justin Yu --- python/ray/train/v2/_internal/callbacks/user_callback.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/ray/train/v2/_internal/callbacks/user_callback.py b/python/ray/train/v2/_internal/callbacks/user_callback.py index 82ef232a04e6..517c58b2b93b 100644 --- a/python/ray/train/v2/_internal/callbacks/user_callback.py +++ b/python/ray/train/v2/_internal/callbacks/user_callback.py @@ -6,10 +6,14 @@ WorkerGroupCallback, ) from ray.train.v2._internal.execution.context import TrainRunContext +from ray.train.v2._internal.execution.worker_group import WorkerGroupStatus from ray.train.v2.api.callback import UserCallback class UserCallbackHandler(WorkerGroupCallback, ReportCallback): + """Responsible for calling methods of subscribers implementing + the `UserCallback` interface. + """ def __init__( self, user_callbacks: List[UserCallback], train_run_context: TrainRunContext ): @@ -34,7 +38,7 @@ def after_report( # WorkerGroupCallback # -------------------------- - def after_worker_group_poll_status(self, worker_group_status): + def after_worker_group_poll_status(self, worker_group_status: WorkerGroupStatus): if not worker_group_status.errors: return