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

Code before:
            )

            # TODO: Figure out ordering between these different calls/callbacks.
            worker_group_state_builder.with_placement_group_handle(pg_handle)

            # Initialize the synchronization actor on the driver node
            sync_actor = SynchronizationActor.options(
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    node_id=ray.get_runtime_context().get_node_id(),
                    soft=False,
                )
            ).remote(
                timeout_s=self._collective_timeout_s,
                warn_interval_s=self._collective_warn_interval_s,
            )
            worker_group_state_builder.with_sync_actor(sync_actor)

            create_workers_start = time_monotonic()
            workers = self._create_workers(
                worker_group_context.num_workers,
                pg_handle.placement_group,
                worker_group_context.resources_per_worker,
            )
            worker_group_state_builder.with_workers(workers)
            logger.debug(
                "[Train Worker Initialization] Workers created in "
                f"{time_monotonic() - create_workers_start:.2f}s."
            )

            # All the ray.get calls in this try block can possibly error if the
            # worker actors die during initialization.
            # To prevent the driver from crashing, catch all `RayActorError`s and
            # raise a specially handled error to the controller.
            try:
                before_init_train_context_cb_start = time_monotonic()
                train_context_args = {}
                for callable in self._callbacks:
                    args = callable.before_init_train_context(workers)
                    for arg, arg_values in args.items():
                        assert len(arg_values) == worker_group_context.num_workers, (
                            f"Callback {callable} returned {arg} with "
                            f"{len(arg_values)} values, expected {worker_group_context.num_workers}."
                        )
                        assert (
                            arg not in train_context_args
                        ), f"Callback {callable} returned {arg} which is already set."
                        train_context_args[arg] = arg_values
                logger.debug(
                    "[Train Worker Initialization] before_init_train_context "
                    f"callbacks completed in {time_monotonic() - before_init_train_context_cb_start:.2f}s."
                )

Review comment (suggestion):
```suggestion
                "[Train Worker Initialization] {num_workers} actors created in "
```

Code after:
                f"{time_monotonic() - pg_wait_start:.2f}s."
            )

            # TODO: Figure out ordering between these different calls/callbacks.
            worker_group_state_builder.with_placement_group_handle(pg_handle)

            # Initialize the synchronization actor on the driver node
            sync_actor = SynchronizationActor.options(
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    node_id=ray.get_runtime_context().get_node_id(),
                    soft=False,
                )
            ).remote(
                timeout_s=self._collective_timeout_s,
                warn_interval_s=self._collective_warn_interval_s,
            )
            worker_group_state_builder.with_sync_actor(sync_actor)

            create_workers_start = time_monotonic()
            workers = self._create_workers(
                worker_group_context.num_workers,
                pg_handle.placement_group,
                worker_group_context.resources_per_worker,
            )
            worker_group_state_builder.with_workers(workers)
            logger.debug(
                f"[Worker Group Initialization] {worker_group_context.num_workers} worker actors created in "
                f"{time_monotonic() - create_workers_start:.2f}s."
            )

            self._replica_groups = [
                ReplicaGroup([worker], worker_group_context.resources_per_worker)
                for worker in workers
            ]

            # All the ray.get calls in this try block can possibly error if the
            # worker actors die during initialization.
            # To prevent the driver from crashing, catch all `RayActorError`s and
            # raise a specially handled error to the controller.
            try:
                before_init_train_context_cb_start = time_monotonic()
                train_context_args = {}
                for callable in self._callbacks:
                    args = callable.before_init_train_context(workers)
                    for arg, arg_values in args.items():
                        assert len(arg_values) == worker_group_context.num_workers, (
                            f"Callback {callable} returned {arg} with "
                            f"{len(arg_values)} values, expected {worker_group_context.num_workers}."
                        )
                        assert (
                            arg not in train_context_args