add comment

Signed-off-by: Justin Yu <[email protected]>

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