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

Code before:

                self._worker_group_state = worker_group_state_builder.build()

                after_wg_start_cb_start = time_monotonic()
                for callback in self._callbacks:
                    callback.after_worker_group_start(self)
                logger.debug(
                    "[Train Worker Initialization] after_worker_group_start "
                    f"callbacks completed in {time_monotonic() - after_wg_start_cb_start:.2f}s."
                )

            except RayActorError as actor_error:
                error_msg = "At least one of the worker actors failed to initialize."
                raise WorkerGroupStartupFailedError(error_msg) from actor_error

        # Launch the training function on each worker.
        # This task should start a worker thread and return immediately.
        launch_start = time_monotonic()
        ray_get_safe(
            [
                worker.actor.run_train_fn.remote(worker_group_context.train_fn_ref)
                for worker in workers
            ]
        )
        logger.debug(
            "[Train Worker Initialization] Train function launched on "
            f"workers in {time_monotonic() - launch_start:.2f}s."
        )

        workers_info = "\n".join(
            [
                f"- (ip={w.metadata.node_ip}, pid={w.metadata.pid}) "
                f"world_rank={w.distributed_context.world_rank}, "
                f"local_rank={w.distributed_context.local_rank}, "
                f"node_rank={w.distributed_context.node_rank}"
                for w in workers
            ]
        )
        logger.info(
            f"Started training worker group of size {len(workers)}: \n{workers_info}"
        )

        after_training_start_cb_start = time_monotonic()
        for callback in self._callbacks:
            callback.after_worker_group_training_start(self)
        logger.debug(
            "[Train Worker Initialization] after_worker_group_training_start "
            f"callbacks completed in {time_monotonic() - after_training_start_cb_start:.2f}s."
        )

        logger.debug(

Review comment (suggestion):
can you also show the per-callback breakdown for this one? Also happy for this to be a followup PR.

something like:

```python
[Train Worker Initialization] after_worker_group_start callbacks completed in 0.88s.
   Individual callback time breakdown:
     ReportCallbackHandler: 0.01s
     AcceleratorSetupCallback: xxx
     BackendSetupCallback: xxx
```

The main one I want to showcase is the BackendSetupCallback time, since that is what runs the torch init_process_group.



Code after:
                    workers, sync_actor, train_context_args
                )
                logger.debug(
                    "[Worker Group Initialization] Train context initialized "
                    f"on workers in {time_monotonic() - init_ctx_start:.2f}s."
                )

                self._worker_group_state = worker_group_state_builder.build()

                after_wg_start_cb_start = time_monotonic()
                after_wg_start_cb_times = {}
                for callback in self._callbacks:
                    cb_start = time_monotonic()
                    callback.after_worker_group_start(self)
                    after_wg_start_cb_times[type(callback).__name__] = (
                        time_monotonic() - cb_start
                    )
                after_wg_start_cb_breakdown = "\n".join(
                    f"    {name}: {elapsed:.2f}s"
                    for name, elapsed in after_wg_start_cb_times.items()
                )
                logger.debug(
                    "[Worker Group Initialization] after_worker_group_start "
                    f"callbacks completed in {time_monotonic() - after_wg_start_cb_start:.2f}s.\n"
                    f"  Individual callback time breakdown:\n{after_wg_start_cb_breakdown}"
                )

            except RayActorError as actor_error:
                error_msg = "At least one of the worker actors failed to initialize."
                raise WorkerGroupStartupFailedError(error_msg) from actor_error

        # Launch the training function on each worker.
        # This task should start a worker thread and return immediately.
        launch_start = time_monotonic()
        ray.get(
            [
                worker.actor.run_train_fn.remote(worker_group_context.train_fn_ref)
                for worker in workers
            ]
        )
        logger.debug(
            "[Worker Group Initialization] Train function launched on "
            f"workers in {time_monotonic() - launch_start:.2f}s."
        )

        workers_info = "\n".join(
            [
                f"- (ip={w.metadata.node_ip}, pid={w.metadata.pid}) "
                f"world_rank={w.distributed_context.world_rank}, "
                f"local_rank={w.distributed_context.local_rank}, "
                f"node_rank={w.distributed_context.node_rank}"