# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ctypes
import functools
import inspect
import logging
import os
import signal
import sys
import threading
import time
import traceback
import warnings
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar

import ray
import ray.util.state
import torch

from ..cluster import (
    Cluster,
    ClusterEnvVar,
    load_user_extension_module,
    without_http_proxies,
)
from ..hardware import AcceleratorType, AcceleratorUtil, HardwareInfo
from ..manager import WorkerAddress
from .routing import split_channel_message

if TYPE_CHECKING:
    from ..collective import CollectiveGroupOptions
    from ..manager import WorkerInfo
    from .worker_group import WorkerGroup

WorkerClsType = TypeVar("WorkerClsType")


class WorkerMeta(type):
    """Metaclass to capture failures in worker classes."""

    def __new__(cls, name: str, bases: tuple[type], attrs: dict[str, Any]):
        """Wrap the function to catch SystemExit exceptions."""
        for attr_name, attr_value in attrs.items():
            if callable(attr_value):
                attrs[attr_name] = cls._catch_failure_for_cls_func(
                    name, attr_name, attr_value
                )
        return super().__new__(cls, name, bases, attrs)

    @classmethod
    def _catch_failure_for_cls_func(cls, cls_name, func_name: str, func: Callable):
        """Wrap a try...except SystemExit block around the class function calls."""
        # Get all callable methods of the WorkerGroup class and the Worker class
        if func_name.startswith("_") and func_name != "__init__":
            return func

        def func_wrapper(func: Callable):
            @functools.wraps(func)
            def sync_func(*args, **kwargs):
                try:
                    return func(*args, **kwargs)
                except SystemExit:
                    # Catch SystemExit and log the error
                    raise RuntimeError(
                        f"SystemExit caught in {cls_name}'s function {func.__name__}, traceback is below: {traceback.format_exc()}"
                    )

            @functools.wraps(func)
            async def async_func(*args, **kwargs):
                try:
                    return await func(*args, **kwargs)
                except SystemExit:
                    # Catch SystemExit and log the error
                    raise RuntimeError(
                        f"SystemExit caught in {cls_name}'s function {func.__name__}, traceback is below: {traceback.format_exc()}"
                    )

            if inspect.iscoroutinefunction(func):
                return async_func
            elif inspect.isasyncgenfunction(func):
                raise NotImplementedError(
                    f"Async generator function {func.__name__} is not supported when CATCH_FAILURE is enabled."
                )
            else:
                return sync_func

        return func_wrapper(func)


class Worker(metaclass=WorkerMeta):
    """Class representing a remote process or worker.

    Inheriting `Worker` will grant your worker or processor class the ability to run remotely and communicate with other workers in the cluster.
    Also, essential environment variables like MASTER_ADDR, MASTER_PORT, RANK, LOCAL_RANK, WORLD_SIZE will be set automatically.
    This allows easy creation of torch process groups and distributed training.

    The following example shows how to use the Worker class to create a simple distributed worker that can run on multiple GPUs and nodes.

    Example::

        >>> import torch
        >>> from rlinf.scheduler import Cluster, Worker
        >>>
        >>> class MyWorker(Worker):
        ...     def __init__(self):
        ...         super().__init__()
        ...
        ...     def initialize(self):
        ...         torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
        ...         if not torch.distributed.is_initialized():
        ...             torch.distributed.init_process_group(backend="nccl")
        ...
        ...         test_tensor = torch.ones(
        ...             size=(1, 1), dtype=torch.float32, device=torch.cuda.current_device()
        ...         )
        ...         torch.distributed.all_reduce(test_tensor)
        ...         return test_tensor
        ...
        ...     def hello(self):
        ...         return self._rank
        >>>
        >>> cluster = Cluster(num_nodes=1)
        >>> my_worker_group = MyWorker.create_group().launch(cluster=cluster, name="my_worker_group")
        >>> my_worker_group.initialize().wait()[0]
        tensor([[4.]], device='cuda:0')
        >>> # This will execute the hello method only on ranks 0 and 1.
        >>> my_worker_group.execute_on(0, 3).hello().wait()
        [0, 3]

    The following example shows the communication capabilities of the Worker class.

    Example::

        >>> import asyncio
        >>> import torch
        >>> from rlinf.scheduler import Cluster, Worker
        >>> SEND_GROUP_NAME = "send_worker_group"
        >>> RECV_GROUP_NAME = "recv_worker_group"
        >>>
        >>> class SendWorker(Worker):
        ...     def __init__(self):
        ...         super().__init__()
        ...
        ...     def hello_recv(self):
        ...         # 1. Send a message (string or any serializable object) to the RecvWorker group with the same rank as this SendWorker worker.
        ...         msg = f"Hello from SendWorker Rank {self._rank}!"
        ...         self.send(msg, dst_group_name=RECV_GROUP_NAME, dst_rank=self._rank)
        ...
        ...         # 2. Receive a reply from the RecvWorker group with the same rank.
        ...         reply = self.recv(
        ...             src_group_name=RECV_GROUP_NAME, src_rank=self._rank
        ...         )
        ...
        ...         # 3. The send/recv APIs can also handle tensor, list of tensors and dict of tensors.
        ...         torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
        ...         dst_rank = (
        ...             self._rank + 1
        ...         ) % self._world_size  # Send to the next rank in the group
        ...         tensor = torch.ones(
        ...             size=(1, 1),
        ...             dtype=torch.float32,
        ...             device=torch.cuda.current_device(),
        ...         )
        ...         self.send(tensor, dst_group_name=RECV_GROUP_NAME, dst_rank=dst_rank)
        ...
        ...         tensor_list = [
        ...             torch.tensor(
        ...                 1.0, dtype=torch.float32, device=torch.cuda.current_device()
        ...             )
        ...             for _ in range(4)
        ...         ]
        ...         self.send(
        ...             tensor_list, dst_group_name=RECV_GROUP_NAME, dst_rank=dst_rank
        ...         )
        ...
        ...         tensor_dict = {
        ...             "tensor1": torch.tensor(
        ...                 2.0, dtype=torch.float32, device=torch.cuda.current_device()
        ...             ),
        ...             "tensor2": torch.tensor(
        ...                 3.0, dtype=torch.float32, device=torch.cuda.current_device()
        ...             ),
        ...         }
        ...         self.send(
        ...             tensor_dict, dst_group_name=RECV_GROUP_NAME, dst_rank=dst_rank
        ...         )
        ...
        ...         # 4. Send tensor directly without metadata overhead if you already know the tensor shape and dtype at the recv side
        ...         tensor = torch.ones(
        ...             size=(2, 1),
        ...             dtype=torch.float32,
        ...             device=torch.cuda.current_device(),
        ...         )
        ...         self.send_tensor(
        ...             tensor, dst_group_name=RECV_GROUP_NAME, dst_rank=dst_rank
        ...         )
        ...
        ...     def hello_recv_async(self):
        ...         # 1. Send a tensor asynchronously to the RecvWorker group with the next rank.
        ...         dst_rank = (self._rank + 1) % self._world_size
        ...         tensor = torch.ones(
        ...             size=(3, 1),
        ...             dtype=torch.float32,
        ...             device=torch.cuda.current_device(),
        ...         )
        ...         async_send_work = self.send(
        ...             tensor,
        ...             dst_group_name=RECV_GROUP_NAME,
        ...             dst_rank=dst_rank,
        ...             async_op=True,
        ...         )
        ...         async_send_work.wait()  # Wait for the async send to complete
        ...
        ...         # 2. Send a tensor asynchronously and use asyncio to wait for the operation to complete.
        ...         async def send_tensor_async():
        ...             dst_rank = (self._rank + 1) % self._world_size
        ...             tensor = torch.ones(
        ...                 size=(4, 1),
        ...                 dtype=torch.float32,
        ...                 device=torch.cuda.current_device(),
        ...             )
        ...             async_send_work = self.send(
        ...                 tensor,
        ...                 dst_group_name=RECV_GROUP_NAME,
        ...                 dst_rank=dst_rank,
        ...                 async_op=True,
        ...             )
        ...             await async_send_work.async_wait()
        ...
        ...         asyncio.run(send_tensor_async())
        >>>
        >>> class RecvWorker(Worker):
        ...     def __init__(self):
        ...         super().__init__()
        ...
        ...     def hello_recv(self):
        ...         # 1. Receive a message from the SendWorker worker group with the same rank.
        ...         msg = self.recv(src_group_name=SEND_GROUP_NAME, src_rank=self._rank)
        ...
        ...         # 2. Send a reply back to the SendWorker worker group with the same rank.
        ...         reply = f"Hello from RecvWorker Rank {self._rank}!"
        ...         self.send(
        ...             reply, dst_group_name=SEND_GROUP_NAME, dst_rank=self._rank
        ...         )
        ...
        ...         # 3. Receive a tensor, tensor list and tensor dict from the SendWorker worker group with the same rank.
        ...         torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
        ...         src_rank = (
        ...             self._rank - 1
        ...         ) % self._world_size  # Receive from the previous rank in the group
        ...         tensor = self.recv(
        ...             src_group_name=SEND_GROUP_NAME, src_rank=src_rank
        ...         )
        ...         tensor_list = self.recv(
        ...             src_group_name=SEND_GROUP_NAME, src_rank=src_rank
        ...         )
        ...         tensor_dict = self.recv(
        ...             src_group_name=SEND_GROUP_NAME, src_rank=src_rank
        ...         )
        ...
        ...         # 4. In-place receive tensor directly without metadata overhead
        ...         tensor = torch.empty(
        ...             size=(2, 1),
        ...             dtype=torch.float32,
        ...             device=torch.cuda.current_device(),
        ...         )
        ...         self.recv_tensor(
        ...             tensor, src_group_name=SEND_GROUP_NAME, src_rank=src_rank
        ...         )
        ...
        ...     def hello_recv_async(self):
        ...         # 1. Receive a tensor asynchronously from the SendWorker group with the next rank.
        ...         src_rank = (self._rank - 1) % self._world_size
        ...         async_recv_work = self.recv(
        ...             src_group_name=SEND_GROUP_NAME, src_rank=src_rank, async_op=True
        ...         )
        ...         tensor = async_recv_work.wait()
        ...
        ...         # 2. Receive a tensor asynchronously and use asyncio to wait for the operation to complete.
        ...         async def recv_tensor_async():
        ...             src_rank = (self._rank - 1) % self._world_size
        ...             async_recv_work = self.recv(
        ...                 src_group_name=SEND_GROUP_NAME,
        ...                 src_rank=src_rank,
        ...                 async_op=True,
        ...             )
        ...             tensor = await async_recv_work.async_wait()
        ...
        ...         asyncio.run(recv_tensor_async())
        >>>
        >>> cluster = Cluster(num_nodes=1)
        >>> send_group = SendWorker.create_group().launch(cluster=cluster, name=SEND_GROUP_NAME)
        >>> recv_group = RecvWorker.create_group().launch(cluster=cluster, name=RECV_GROUP_NAME)
        >>> res = send_group.hello_recv()
        >>> res = recv_group.hello_recv().wait()
        >>> res = send_group.hello_recv_async()
        >>> res = recv_group.hello_recv_async().wait()

    """

    PID = None
    current_worker = None
    logging.basicConfig()
    logger = logging.getLogger(Cluster.SYS_NAME)
    logger.setLevel(Cluster.LOGGING_LEVEL)
    accelerator_type = AcceleratorUtil.get_accelerator_type()
    torch_platform = AcceleratorUtil.get_torch_platform(accelerator_type)
    torch_device_type = AcceleratorUtil.get_device_type(accelerator_type)

    def __new__(cls, *args, **kwargs):
        """Create a new instance of the Worker class."""
        instance = super().__new__(cls)

        cluster_node_rank = os.environ.get("CLUSTER_NODE_RANK", None)

        # ray.remote initializes the class with the ActorClass wrapper locally first (not in a remote process),
        # which doesn't have the environment variables set.
        if cluster_node_rank is not None and "ActorClass(" not in cls.__name__:
            instance._env_setup_before_init()
            # Handle OS signals for better debuggability
            # Ray new the class in main thread but call __init__ in worker thread if it's an Actor with async functions
            # Since signal handlers must be registered in main thread, we call the registration in __new__
            instance._register_signal_handlers()
            instance._enable_ptrace()

        return instance

    def _env_setup_before_init(self):
        """Set up distributed Worker environments."""
        # These are required env_vars necessary for both Ray Worker and non-Ray Worker
        # For non-ray workers, these are reset in the __init__ method
        self._rank = int(os.environ.get("RANK", "-1"))
        self._worker_name = os.environ.get("WORKER_NAME", None)
        self._world_size = int(os.environ.get("WORLD_SIZE", "-1"))
        if self._worker_name is not None:
            self._worker_address = WorkerAddress.from_name(self._worker_name)

        # These are not required env_vars, but are set by Ray Worker for convenience
        self._cluster_node_rank = int(os.environ.get("CLUSTER_NODE_RANK", -1))
        self._accelerator_type = AcceleratorType(
            os.environ.get("ACCELERATOR_TYPE", str(AcceleratorType.NO_ACCEL.value))
        )
        self._accelerator_model = os.environ.get("ACCELERATOR_MODEL", "")
        self._local_accelerator_rank = int(os.environ.get("LOCAL_ACCELERATOR_RANK", -1))
        self._node_local_rank = int(os.environ.get("NODE_LOCAL_RANK", -1))
        self._node_local_world_size = int(os.environ.get("NODE_LOCAL_WORLD_SIZE", -1))
        Worker.accelerator_type = self._accelerator_type
        Worker.torch_device_type = AcceleratorUtil.get_device_type(
            self._accelerator_type
        )
        Worker.torch_platform = AcceleratorUtil.get_torch_platform(
            self._accelerator_type
        )
        self.accelerator_type = Worker.accelerator_type
        self.torch_device_type = Worker.torch_device_type
        self.torch_platform = Worker.torch_platform

        self._actor = None
        self._has_initialized = False
        self._timer_metrics: dict[str, float] = {}

        # Load user-provided extension modules (e.g., for registering custom envs/models)
        self._load_user_extensions()

    def _load_user_extensions(self):
        """Load extension modules specified via EXT_MODULE environment variable.

        This allows users to register custom environments, models, or other extensions
        without patching.
        The extension module should have a `register()` function that performs the necessary registrations.

        The module's register() function will be called once per Worker process.
        """
        load_user_extension_module(logger=Worker.logger)

    def __init__(
        self,
        parent_address: Optional[WorkerAddress] = None,
        world_size: Optional[int] = None,
        rank: Optional[int] = None,
    ):
        """Initialize the Worker with the given parent address and world size.

        Only non-Ray workers should provide parent_address, world_size and rank. For example, when a Worker is created via multiprocessing by another Worker, the parent address, world size and rank should be provided.

        Args:
            parent_address (Optional[WorkerAddress]): The address of the parent worker. This is used to set up the WorkerAddress for this worker.
            world_size (Optional[int]): The total number of workers in the group. If not provided, it will be set to the environment variable WORLD_SIZE.
            rank (Optional[int]): The rank of this worker in the group. If not provided, it will be set to the environment variable RANK.

        """
        if rank is not None and parent_address is not None and world_size is not None:
            # The Worker is not a Ray actor
            self._rank = rank
            self._worker_address = parent_address.get_child_address(rank)
            self._world_size = world_size
            self._worker_name = self._worker_address.get_name()
            # Forked process might inherit the environment variable RAY_ACTOR, but it is not a Ray actor.
            self._is_ray_actor = False
        else:
            self._is_ray_actor = True

        if self._is_ray_actor and not hasattr(self, "_local_accelerator_rank"):
            raise RuntimeError(
                "You may have mistakenly initialized the Worker class directly without `create_group` and `launch`. Please ensure a worker class is not instantiated on the main process directly like `Worker()`, but `Worker.create_group().launch()`."
            )

        Worker.PID = os.getpid()
        self._thread = threading.current_thread()
        self._stacklevel = 4 if self._is_ray_actor else 3

        # Reset Cluster.NAMESPACE for this Worker process according to the environment variable
        namespace = os.environ.get("CLUSTER_NAMESPACE", None)
        assert namespace is not None, (
            "CLUSTER_NAMESPACE environment variable must be set before initializing Worker."
        )
        Cluster.NAMESPACE = namespace

        # Initialize Ray if not already initialized
        if not ray.is_initialized():
            ray.init(
                address="auto",
                namespace=Cluster.NAMESPACE,
                logging_level=Cluster.LOGGING_LEVEL,
            )

        if self._is_ray_actor and parent_address is not None:
            # The Worker is a Ray actor launched inside a Worker
            self._worker_address = parent_address.get_child_address(self._rank)
            self._worker_name = self._worker_address.get_name()
            os.environ["WORKER_NAME"] = self._worker_name
        self._group_name = self._worker_address.get_parent_address().get_name()

        # Initialize global locks
        from .lock import DeviceLock, PortLock

        self._device_lock = DeviceLock(self)
        self._port_lock = PortLock(self)

        # Setup local rank and world size
        self._setup_local_rank_world_size()

        # Setup accelerator ID
        self._setup_accelerator_info()

        # Configure logging
        self._setup_logging()

        # Setup node group and hardware ranks
        self._setup_hardware()

        # Setup worker info
        self._setup_worker_info()

        # Init ray and managers
        self._manager_proxy = None
        self._collective = None
        self._setup_managers()

        # Setup MASTER_ADDR and MASTER_PORT
        self._setup_master_address_and_port()

        # Setup communication envs
        self._setup_comm_envs()

        self._lock = threading.Lock()
        Worker.current_worker = self
        self._has_initialized = True

    @property
    def has_accelerator(self) -> bool:
        """Whether the worker has been allocated with accelerators."""
        return self._accelerator_type != AcceleratorType.NO_ACCEL

    @property
    def worker_address(self) -> WorkerAddress:
        """Get the WorkerAddress of the worker.

        This is used to identify the worker in the WorkerGroup.
        """
        return self._worker_address

    @property
    def worker_info(self) -> "WorkerInfo":
        """Get the WorkerInfo of the worker."""
        return self._worker_info

    @property
    def manager_proxy(self):
        """Get the SchedulerProxy instance for this worker.

        This is used to interact with the scheduler and register the worker.
        """
        return self._manager_proxy

    @property
    def device_lock(self):
        """Get the DeviceLock instance for this worker."""
        return self._device_lock

    @property
    def hardware_type(self) -> str:
        """Get the hardware type of the current worker.

        Returns:
            str: The hardware type of the current worker.
        """
        return self._node_group.hardware_type

    @property
    def hardware_infos(self) -> list[HardwareInfo]:
        """Get the hardware information of the current worker.

        Returns:
            list[HardwareInfo]: The list hardware information assigned to the current worker.
        """
        infos = []
        for local_hw_rank in self._local_hardware_ranks:
            infos.append(
                self._node_group.get_hardware_infos(self._cluster_node_rank)[
                    local_hw_rank
                ]
            )
        return infos

    @classmethod
    def create_group(
        cls: type[WorkerClsType], *args, **kwargs
    ) -> "WorkerGroup[WorkerClsType] | WorkerClsType":
        """Create a worker group with the class arguments.

        Args:
            args: The positional arguments of the class.
            kwargs: The keyword arguments of the class.
        """
        from .worker_group import WorkerGroup

        return WorkerGroup(cls, args, kwargs)

    def send(
        self,
        object: torch.Tensor | list[torch.Tensor] | dict[str, torch.Tensor] | Any,
        dst_group_name: str,
        dst_rank: int | list[int],
        async_op: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
        piggyback_payload: Optional[Any] = None,
    ):
        """Send an object to a specific worker address in the collective group.

        The function is specially optimized for torch.Tensor, List of torch.Tensor, Dict of torch.Tensor, and dataclass containing torch.Tensor, which go through NCCL when the contained tensors are on GPU. Otherwise, all communications go through GLOO.

        .. note::
            Do not mix send with recv_tensor

        .. note::
            We only use NCCL primitives when the list or dict values only contain GPU tensors. We also see complex dicts with deep hierarchy as common Python objects, which will be serialized into a CPU tensor and sent through GLOO.

        .. note::
            When transferring GPU objects, the first send needs to be paired with a recv at the other end. Calling async send or recv first at both ends will result in communication hang, because NCCL communicators are established in a lazy manner when the first pair of send/recv is called.

        .. note::
            Do not mix CPU and GPU tensors in a list or dict.

        .. note::
            This method is not thread safe.

        Args:
            object (torch.Tensor | List[torch.Tensor] | Dict[str, torch.Tensor] | Any): The object to send.
            dst_group_name (str): The name of the destination worker group.
            dst_rank (int | List[int]): The rank or list of ranks in the destination worker group to send the object to. For SPMD-like workers, this should be a single rank. For SPSD-like workers forked by parent workers, this can be a list of ranks that forms a path from the root worker to the target worker.
            async_op (bool): Whether to perform the operation asynchronously.
            options (Optional[CollectiveGroupOptions]): The options for the collective group. The options will only take effect when two workers first communicate with each other, and will be ignored for subsequent communications. This option must match the options of the recv side.
            piggyback_payload (Optional[Any]): The payload to piggyback on the send operation. This payload will be sent to the recv side and can be used to pass additional information to the recv side without disrupting the object's data structure, e.g., list/dict of tensors that are optimized for sending.

        Returns:
            Optional[AsyncWork]: An AsyncWork object if async_op is True, otherwise None.

        """
        dst_addr = WorkerAddress(dst_group_name, ranks=dst_rank)
        group = self._get_collective_group(dst_addr)
        return group.send(
            object=object,
            async_op=async_op,
            options=options,
            piggyback_payload=piggyback_payload,
        )

    def recv(
        self,
        src_group_name: str,
        src_rank: int | list[int],
        async_op: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
    ):
        """Out-of-place receive of an object from a specific worker address in the collective group.

        .. note::
            Do not mix recv with send_tensor

        .. note::
            When transferring GPU objects, the first send needs to be paired with a recv at the other end. Calling async send or recv first at both ends will result in communication hang, because NCCL communicators are established in a lazy manner when the first pair of send/recv is called.

        .. note::
            This method is not thread safe.

        Args:
            async_op (bool): Whether to perform the operation asynchronously.
            src_group_name (str): The name of the source worker group.
            src_rank (int | List[int]): The rank or list of ranks in the source worker group to receive the object from. For SPMD-like workers, this should be a single rank. For SPSD-like workers forked by parent workers, this can be a list of ranks that forms a path from the root worker to the target worker.
            options (Optional[CollectiveGroupOptions]): The options for the collective group. The options will only take effect when two workers first communicate with each other, and will be ignored for subsequent communications. This option must match the options of the send side.

        Returns:
            AsyncWork | torch.Tensor | List[torch.Tensor] | Dict[str, torch.Tensor] | Any: An AsyncWork object if async_op is True, otherwise the received object. If the send side sends a piggyback payload, the received object will be a tuple of the received object and the piggyback payload.
        """
        src_addr = WorkerAddress(src_group_name, ranks=src_rank)
        group = self._get_collective_group(src_addr)
        return group.recv(async_op=async_op, options=options)

    def send_tensor(
        self,
        tensor: torch.Tensor,
        dst_group_name: str,
        dst_rank: int | list[int],
        async_op: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
    ):
        """Send a tensor to a specific worker address in the collective group. This function is optimized for sending a single tensor and does not introduce metadata communication overhead like send. But it needs to be paired with the in-place recv_tensor function which requires apriori knowledge of the tensor shape and dtype.

        .. note::
            Do not mix send_tensor with recv

        .. note::
            When transferring GPU objects, the first send_tensor needs to be paired with a recv_tensor at the other end. Calling async send_tensor or recv_tensor first at both ends will result in communication hang, because NCCL communicators are established in a lazy manner when the first pair of send/recv is called.

        .. note::
            This method is not thread safe.

        Args:
            tensor (torch.Tensor): The tensor to send.
            dst_group_name (str): The name of the destination worker group.
            dst_rank (int | List[int]): The rank or list of ranks in the destination worker group to send the tensor to. For SPMD-like workers, this should be a single rank. For SPSD-like workers forked by parent workers, this can be a list of ranks that forms a path from the root worker to the target worker.
            async_op (bool): Whether to perform the operation asynchronously.
            options (Optional[CollectiveGroupOptions]): The options for the collective group. The options will only take effect when two workers first communicate with each other, and will be ignored for subsequent communications. This option must match the options of the send side.

        Returns:
            Optional[AsyncWork]: An AsyncWork object if async_op is True, otherwise None.

        """
        dst_addr = WorkerAddress(dst_group_name, ranks=dst_rank)
        group = self._get_collective_group(dst_addr)
        return group.send_tensor(tensor=tensor, async_op=async_op, options=options)

    def recv_tensor(
        self,
        tensor: torch.Tensor,
        src_group_name: str,
        src_rank: int | list[int],
        async_op: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
    ):
        """In-place receive of a tensor from a specific worker address in the collective group. This function is optimized for receiving a single tensor and does not introduce metadata communication overhead like recv. But it requires preallocation of the tensor with the correct shape and dtype.

        .. note::
            Do not mix recv_tensor with send

        .. note::
            When transferring GPU objects, the first send_tensor needs to be paired with a recv_tensor at the other end. Calling async send_tensor or recv_tensor first at both ends will result in communication hang, because NCCL communicators are established in a lazy manner when the first pair of send/recv is called.

        .. note::
            This method is not thread safe.

        Args:
            tensor (torch.Tensor): The tensor to receive. It must be preallocated with the correct shape and dtype.
            src_group_name (str): The name of the source worker group.
            src_rank (int | List[int]): The rank or list of ranks in the source worker group to receive the tensor from. For SPMD-like workers, this should be a single rank. For SPSD-like workers forked by parent workers, this can be a list of ranks that forms a path from the root worker to the target worker.
            async_op (bool): Whether to perform the operation asynchronously.
            options (Optional[CollectiveGroupOptions]): The options for the collective group. The options will only take effect when two workers first communicate with each other, and will be ignored for subsequent communications. This option must match the options of the send side.

        Returns:
            Optional[AsyncWork]: An AsyncWork object if async_op is True, otherwise None.

        """
        src_addr = WorkerAddress(src_group_name, ranks=src_rank)
        group = self._get_collective_group(src_addr)
        return group.recv_tensor(tensor=tensor, async_op=async_op, options=options)

    def broadcast(
        self,
        object: Optional[Any] = None,
        groups: Optional[
            list[tuple[str, list[int] | list[tuple[int]] | tuple[int] | int]]
        ] = None,
        src: Optional[tuple[str, tuple[int] | int]] = None,
        async_op: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
    ):
        """Broadcast an object across workers in one or more groups.

        The source is the first worker address in the expanded group list.
        The index in the expanded list is the rank in the communication group.
        All participating workers must call this method with identical arguments.

        Args:
            object (Any): The object to broadcast on the source worker. For non-src ranks, this is typically None.
            groups: The participating groups with ranks. Each element must be a (group_name, ranks) tuple where ranks is either a single int (one worker of the rank), a list of ints (multiple workers of the same group), a tuple of ints (one worker of the rank path), or a list of tuples of ints (multiple workers of the rank paths of the same group).
            src: The source group and rank. If not provided, the source will be the first worker address in the expanded group list.
            async_op (bool): Whether to perform the operation asynchronously.
            options (Optional[CollectiveGroupOptions]): The options for the collective group.

        Returns:
            AsyncWork | Any: An AsyncWork object if async_op is True, otherwise the
            broadcast object.
        """
        if groups is None:
            raise ValueError("groups must be provided with explicit ranks.")
        if not isinstance(groups, list):
            raise TypeError("groups must be a list of (group_name, rank) tuples.")
        if len(groups) == 0:
            raise ValueError("groups must contain at least one entry.")

        worker_addresses: list[WorkerAddress] = []
        for entry in groups:
            if not isinstance(entry, tuple) or len(entry) != 2:
                raise TypeError(
                    "Each groups entry must be a (group_name, ranks) tuple."
                )
            group_name, ranks = entry
            if not isinstance(group_name, str):
                raise TypeError(
                    f"group_name must be a string. But got {type(group_name)}."
                )
            if isinstance(ranks, list):
                if len(ranks) == 0:
                    raise ValueError("ranks list must not be empty.")
                if not all(
                    isinstance(rank, int) or isinstance(rank, tuple) for rank in ranks
                ):
                    raise TypeError(
                        f"All ranks must be integers or tuples. But got {type(ranks)}."
                    )
                for rank in ranks:
                    worker_addresses.append(WorkerAddress(group_name, ranks=rank))
            elif isinstance(ranks, int) or isinstance(ranks, tuple):
                worker_addresses.append(WorkerAddress(group_name, ranks=ranks))
            else:
                raise TypeError(
                    f"ranks must be an int, tuple, list[int], list[tuple[int]]. But got {type(ranks)}."
                )

        if not worker_addresses:
            return object

        if self._worker_address not in worker_addresses:
            raise ValueError(
                f"Worker {self._worker_address.get_name()} is not part of the broadcast group."
            )

        # Get the src addr before sorting
        if src is not None:
            src_group_name, src_ranks = src
            if not isinstance(src_group_name, str):
                raise TypeError(
                    f"src_group_name must be a string. But got {type(src_group_name)}."
                )
            if not isinstance(src_ranks, int) and not isinstance(src_ranks, tuple):
                raise TypeError(
                    f"src_ranks must be an int or tuple. But got {type(src_ranks)}."
                )
            src_addr = WorkerAddress(src_group_name, ranks=src_ranks)
        else:
            src_addr = worker_addresses[0]
        with self._lock:
            worker_addresses.sort()
            group = self._collective.create_collective_group(worker_addresses)

        return group.broadcast(
            object=object,
            src_addr=src_addr,
            async_op=async_op,
            options=options,
        )

    def create_channel(
        self,
        channel_name: str,
        maxsize: int = 0,
        distributed: bool = False,
        node_rank: int = 0,
        local: bool = False,
    ):
        """Create a new channel with the specified placement rank and maximum size.

        Args:
            channel_name (str): The name of the channel.
            maxsize (int): The maximum size of the channel queue. Defaults to 0 (unbounded).
            distributed (bool): Whether the channel should be distributed. A distributed channel creates a distributed worker on each node, and routes communications to the channel worker on the same node as the current worker, benefitting from the locality of the data. The routing is based on the key of the put/get APIs. So if you expect the key to be randomly distributed, you should set this to False to avoid unnecessary routing overhead.
            node_rank (int): The node rank of the current worker. Only valid when distributed is False.
            local (bool): Create the channel for intra-process communication. A local channel cannot be connected by other workers, and its data cannot be shared among different processes.

        Returns:
            Channel: A new instance of the Channel class.

        """
        from ..channel.channel import Channel

        return Channel.create(
            name=channel_name,
            maxsize=maxsize,
            distributed=distributed,
            node_rank=node_rank,
            local=local,
        )

    def connect_channel(self, channel_name: str):
        """Connect to an existing channel.

        Args:
            channel_name (str): The name of the channel to connect to.

        Returns:
            Channel: An instance of the Channel class connected to the specified channel.

        """
        from ..channel.channel import Channel

        return Channel.connect(name=channel_name, current_worker=self)

    def send_to(
        self,
        group_name: str,
        channel: Any | None,
        data: Any,
        *,
        route_key: Any = None,
        mode: str | None = None,
        tag: str | None = None,
        async_op: bool = False,
        batch_size: int | None = None,
        split_fn: Optional[Callable[[Any, list[int]], list[Any]]] = None,
        enable_p2p: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
        decoupled_mode: bool = False,
        send_queue_size: int = 0,
    ):
        """Send a payload to another worker group according to a routing plan.

        The payload is split into shards based on the generated send plan. Each shard
        is then sent either through ``channel.put`` or, when ``enable_p2p`` is True,
        through the collective ``send`` API.

        In normal mode, the routing plan is built from the current worker rank and the
        destination group size. In decoupled mode, each shard is wrapped with route
        metadata:

            {
                "batch_index": <route metadata>,
                "batch": <payload shard>,
            }

        When ``tag`` exists in ``self.batch_router``, the stored batch indices are used
        to route the response back to the workers that originally sent the batches.
        Those stored routes are cleared after the send plan is built. Otherwise, the
        current worker rank is encoded in the generated ``batch_index`` so a later
        receiver can return results to this worker.

        Args:
            group_name: Destination worker group name.
            channel: Channel used for routed transfer. Required unless ``enable_p2p``
                is True.
            data: Payload to split and send.
            route_key: Optional key used to separate independent routed streams.
            mode: Optional phase name used only in decoupled routing. It is encoded in
            ``batch_index`` so the response path can distinguish streams that share the
            same base ``tag``. For example, ``mode="train"`` and
            ``tag="rollout_results"`` make the return message use
            ``"train_rollout_results"``.
            tag: Optional routing tag used to build channel keys and decoupled
                batch indices.
            async_op: If True, return an ``AsyncRouteWork`` wrapping async send/put
                operations.
            batch_size: Total logical batch size to route. If omitted, it is inferred
                from ``data`` and multiplied by the current worker group size.
            split_fn: Optional custom function for splitting ``data`` by planned shard
                sizes. If omitted, ``split_batch`` is used.
            enable_p2p: If True, send shards with collective ``send`` instead of
                ``channel.put``. Not supported with ``decoupled_mode``.
            options: Optional collective options forwarded to ``send``.
            decoupled_mode: If True, use decoupled routing and wrap each shard
                with ``batch_index`` metadata.
            send_queue_size: Number of send queue entries used when building the
                decoupled send plan.

        Returns:
            ``AsyncRouteWork`` if ``async_op`` is True; otherwise ``None``.
        """
        from ..collective import AsyncRouteWork
        from .routing import (
            build_send_plan,
            decoupled_build_send_plan,
            get_group_world_size,
            infer_batch_size,
            split_batch,
        )

        if decoupled_mode:
            assert not enable_p2p, (
                "Now, enable_p2p dpn't support when decoupled_mode is True."
            )
            assert getattr(self, "batch_router", None) is not None, (
                "batch_router must be provided when decoupled_mode is True."
            )

        if not enable_p2p and channel is None:
            raise ValueError("send_to requires ``channel`` when enable_p2p is False.")

        world_size = get_group_world_size(self._manager_proxy, group_name)
        if batch_size is None:
            batch_size = infer_batch_size(data) * self._world_size

        if not decoupled_mode:
            plan = build_send_plan(
                src_group_name=self.worker_address.root_group_name,
                dst_group_name=group_name,
                src_rank=self._rank,
                src_world_size=self._world_size,
                dst_world_size=world_size,
                tag=tag if mode is None else f"{mode}_{tag}",
                route_key=route_key,
                batch_size=batch_size,
            )
        else:
            if tag in self.batch_router:
                # if the batch_router has this tag
                # The sending and receiving logic of batch_router is as follows:
                # recv_from the worker -> Save the data's batch_index to batch_router[tag] ->
                # Use the batch_router[tag] to get the batch_index to create the send plan ->
                # Send data to the worker that originally sent it.
                # the src_rank get from the batch_router[tag]
                tag_batch_router = self.batch_router[tag]
                src_rank = None
            else:
                # if the batch_router does not have this tag
                # The sending and receiving logic is as follows:
                # Save the send_rank in batch_index -> send the data and batch_index to channel ->
                # Any workers can recv the data and save the batch_index -> handle the data in the worker ->
                # Send the data to the worker that originally sent it by the batch_index.
                # the src_rank is the current worker's rank
                tag_batch_router = None
                src_rank = self._rank

            plan = decoupled_build_send_plan(
                src_group_name=self.worker_address.root_group_name,
                dst_group_name=group_name,
                src_rank=src_rank,
                src_world_size=self._world_size,
                dst_world_size=world_size,
                tag=tag,
                route_key=route_key,
                batch_size=batch_size,
                tag_batch_router=tag_batch_router,
                send_queue_size=send_queue_size,
                mode=mode,
            )
            if tag_batch_router is not None:
                # delete the used batch_router item to avoid duplicate sending
                self.batch_router[tag] = []

        split_sizes = [entry.batch_size for entry in plan.entries]
        payloads = (
            split_fn(data, split_sizes)
            if split_fn is not None
            else split_batch(data, split_sizes)
        )

        works = []
        for entry, payload in zip(plan.entries, payloads):
            if decoupled_mode:
                # After enabling decoupled_mode, the data sending format is as follows:
                # {
                #     "batch_index": batch_index,
                #     "batch": batch,
                # }
                # The batch_index is the index of the batch in the data.
                # The batch is the data to send.
                # batch_index: {send_rank}_{batch_idx}_{mode}_{tag}
                # The send_rank is the rank of the worker that originally sent the data.
                # The batch_idx is the index of the batch in the data.
                # The tag is the tag of the data.
                senditem = {
                    "batch_index": entry.batch_index,
                    "batch": payload,
                }
                work = channel.put(
                    item=senditem,
                    key=entry.key,
                    async_op=async_op,
                )
            elif enable_p2p:
                work = self.send(
                    payload,
                    dst_group_name=group_name,
                    dst_rank=entry.peer_rank,
                    async_op=async_op,
                    options=options,
                )
            else:
                work = channel.put(
                    item=payload,
                    key=entry.key,
                    async_op=async_op,
                )
            if async_op and work is not None:
                works.append(work)

        if async_op:
            return AsyncRouteWork(works, lambda _: None)
        return None

    def recv_from(
        self,
        group_name: str,
        channel: Any | None,
        *,
        route_key: Any = None,
        tag: str | None = None,
        async_op: bool = False,
        batch_size: int | None = None,
        merge_fn: Optional[Callable[[list[Any]], Any]] = None,
        infer_batch_size_fn: Optional[Callable[[Any], int]] = None,
        enable_p2p: bool = False,
        options: Optional["CollectiveGroupOptions"] = None,
        decoupled_mode: bool = False,
        recv_queue_size: int = 0,
    ):
        """Receive routed payload shards from another worker group.

        This method builds a receive plan, receives one or more shards, validates their
        batch sizes, and returns either a single payload or a merged payload. Shards can
        be received through ``channel.get`` or, when ``enable_p2p`` is True, through the
        collective ``recv`` API.

        In decoupled mode, each channel item is expected to have the form:

            {
                "batch_index": <route metadata>,
                "batch": <payload shard>,
            }

        When ``tag`` exists in ``self.batch_router``, the received ``batch_index`` values
        are recorded in ``self.batch_router[tag]`` so a later ``send_to`` call can route
        responses back to the original source workers. Otherwise, received shards are
        sorted by the batch index encoded in ``batch_index`` before being merged.

        Args:
            group_name: Source worker group name.
            channel: Channel used for routed transfer. Required unless ``enable_p2p``
                is True.
            route_key: Optional key used to separate independent routed streams.
            tag: Optional routing tag used to build channel keys and decoupled
                batch indices.
            async_op: If True, return an ``AsyncRouteWork`` that finalizes the received
                shards after all async receive operations complete.
            batch_size: Expected local batch size after routing.
            merge_fn: Optional custom function for merging received shards. If omitted,
                ``merge_batches`` is used when more than one shard is received.
            infer_batch_size_fn: Optional function used to infer shard batch size during
                validation.
            enable_p2p: If True, receive shards with collective ``recv`` instead of
                ``channel.get``. Not supported with ``decoupled_mode``.
            options: Optional collective options forwarded to ``recv``.
            decoupled_mode: If True, use decoupled routing and unwrap
                ``batch_index`` metadata from channel items.
            recv_queue_size: Number of receive queue entries used when building the
                decoupled receive plan.

        Returns:
            If ``async_op`` is True, an ``AsyncRouteWork``. Otherwise, returns ``None``
            when no items are received, a single payload when one shard is received, or
            the merged payload when multiple shards are received.
        """
        if decoupled_mode:
            assert not enable_p2p, (
                "Now, enable_p2p dpn't support when decoupled_mode is True."
            )
            assert getattr(self, "batch_router", None) is not None, (
                "batch_router must be provided when decoupled_mode is True."
            )

        from ..collective import AsyncRouteWork
        from .routing import (
            build_recv_plan,
            decoupled_build_recv_plan,
            get_group_world_size,
            merge_batches,
            validate_batch_size,
        )

        if not enable_p2p and channel is None:
            raise ValueError("recv_from requires ``channel`` when enable_p2p is False.")

        world_size = get_group_world_size(self._manager_proxy, group_name)

        if not decoupled_mode:
            plan = build_recv_plan(
                src_group_name=group_name,
                dst_group_name=self.worker_address.root_group_name,
                dst_rank=self._rank,
                src_world_size=world_size,
                dst_world_size=self._world_size,
                tag=tag,
                route_key=route_key,
                batch_size=batch_size,
            )
        else:
            if tag in self.batch_router:
                # if the batch_router has this tag
                # The sending and receiving logic of batch_router is as follows:
                # recv_from the worker -> Save the data's batch_index to batch_router[tag] ->
                # Use the batch_router[tag] to get the batch_index to create the send plan ->
                # Send data to the worker that originally sent it.
                # The recv_rank don't need to be provided
                # The current worker will directly and arbitrarily fetch any data from the channel.
                recv_rank = None
            else:
                # if the batch_router does not have this tag
                # The sending and receiving logic is as follows:
                # Save the send_rank in batch_index -> send the data and batch_index to channel ->
                # Any workers can recv the data and save the batch_index -> handle the data in the worker ->
                # Send the data to the worker that originally sent it by the batch_index.
                # In this recv_from, the processed data will be received, so recv_rank needs to be set to the current worker rank
                recv_rank = self._rank
            plan = decoupled_build_recv_plan(
                src_group_name=group_name,
                dst_group_name=self.worker_address.root_group_name,
                recv_rank=recv_rank,
                src_world_size=self._world_size,
                dst_world_size=world_size,
                tag=tag,
                route_key=route_key,
                batch_size=batch_size,
                recv_queue_size=recv_queue_size,
            )

        def _finalize(received_items: list[Any]):
            if not received_items:
                return None
            if decoupled_mode:
                # get the tag from the received_items
                _, _, _, tag = split_channel_message(received_items[0]["batch_index"])
                if tag in self.batch_router:
                    # If the batch_router is provided,
                    # Save the batch_index to the batch_router.
                    list_received_items = []
                    for item in received_items:
                        batch_index = item["batch_index"]
                        received_item = item["batch"]
                        list_received_items.append(received_item)
                        # Save the batch_index to the batch_router.
                        self.batch_router[tag].append(batch_index)
                    received_items = list_received_items
                else:
                    # Otherwise, the worker get the batch_index from the channel.
                    # Sort the works by the batch_index.
                    sorted_received_items = []
                    for item in received_items:
                        batch_index = item["batch_index"]
                        received_item = item["batch"]
                        # divide the batch_index into send_rank, batch_idx and tag
                        # batch_index: {send_rank}_{batch_idx}_{mode}_{tag}
                        _, batch_idx, _, _ = split_channel_message(batch_index)
                        sorted_received_items.append((batch_idx, received_item))
                    sorted_received_items.sort(key=lambda x: x[0])
                    received_items = [x[1] for x in sorted_received_items]

            for item, entry in zip(received_items, plan.entries):
                validate_batch_size(
                    data=item,
                    expected_batch_size=entry.batch_size,
                    infer_batch_size_fn=infer_batch_size_fn,
                )
            if merge_fn is not None:
                return merge_fn(received_items)
            if len(received_items) == 1:
                return received_items[0]
            return merge_batches(received_items)

        if async_op:
            if enable_p2p:
                works = [
                    self.recv(
                        src_group_name=group_name,
                        src_rank=entry.peer_rank,
                        async_op=True,
                        options=options,
                    )
                    for entry in plan.entries
                ]
            else:
                works = [
                    channel.get(key=entry.key, async_op=True) for entry in plan.entries
                ]
            return AsyncRouteWork(works, _finalize)

        if enable_p2p:
            received_items = [
                self.recv(
                    src_group_name=group_name,
                    src_rank=entry.peer_rank,
                    async_op=False,
                    options=options,
                )
                for entry in plan.entries
            ]
        else:
            received_items = [channel.get(key=entry.key) for entry in plan.entries]
        return _finalize(received_items)

    def get_name(self) -> str:
        """Convert the WorkerAddress to a string representation.

        Returns:
            str: The string representation of the worker name.

        """
        return self._worker_address.get_name()

    def get_parent_rank(self) -> int:
        """Get the rank of the parent worker in the WorkerAddress.

        Returns:
            int: The rank of the parent worker, or 0 if this is the root worker.

        """
        return self._worker_address.get_parent_rank()

    def acquire_free_port(self):
        """Safely acquire a free port on the current node without causing conflicts within the node."""
        max_tries = 10000  # Retry up to 10000 times to find a free port
        for _ in range(max_tries):
            port = Cluster.find_free_port()
            success = self._port_lock.acquire(port)
            if success:
                return port
        raise RuntimeError(f"Failed to acquire a free port after {max_tries} attempts.")

    def log_on_first_rank(self, msg):
        """Log a message only on the first rank of the worker group."""
        if self._rank == 0:
            self._logger.info(msg, stacklevel=self._stacklevel)

    def log_on_last_rank(self, msg):
        """Log a message only on the last rank of the worker group."""
        if self._rank == self._world_size - 1:
            self._logger.info(msg, stacklevel=self._stacklevel)

    def log_debug(self, msg):
        """Log at the debug level."""
        self._logger.debug(msg, stacklevel=self._stacklevel)

    def log_info(self, msg):
        """Log at the info level."""
        self._logger.info(msg, stacklevel=self._stacklevel)

    def log_warning(self, msg):
        """Log at the warning level."""
        self._logger.warning(msg, stacklevel=self._stacklevel)

    def log_error(self, msg):
        """Log at the error level."""
        self._logger.error(msg, stacklevel=self._stacklevel)

    def pop_execution_time(self, tag: str):
        """Retrieve the execution time of a function.

        Args:
            tag (str): The name of the timer to retrieve the execution time for.
        """
        if tag not in self._timer_metrics:
            raise ValueError(f"Timer '{tag}' has not been recorded.")
        return self._timer_metrics.pop(tag)

    def pop_execution_times(self) -> dict[str, float]:
        """Retrieve and clear all execution times."""
        metrics = dict(self._timer_metrics)
        self._timer_metrics.clear()
        return metrics

    @contextmanager
    def worker_timer(self, tag: Optional[str] = None):
        """Context manager to time the execution of a worker function.

        Args:
            tag (str): The name of the timer to record the execution time for. Default is the current function name.
        """
        if tag is None:
            frame_num = 2
            frame = inspect.stack()[frame_num]
            tag = frame.function
        assert tag is not None, "Timer tag must be provided."
        try:
            start_time = time.perf_counter()
            yield
        finally:
            duration = time.perf_counter() - start_time
            self._timer_metrics[tag] = self._timer_metrics.get(tag, 0.0) + duration

    @staticmethod
    def timer(tag: Optional[str] = None):
        """Decorator to time a worker function and emit a profiling annotation."""

        def decorator(func):
            label = tag or func.__name__
            if inspect.iscoroutinefunction(func):

                @functools.wraps(func)
                async def wrapper(self, *args, **kwargs):
                    with self.worker_timer(label):
                        with AcceleratorUtil.profiling_range(
                            self._accelerator_type, label
                        ):
                            return await func(self, *args, **kwargs)

                return wrapper

            @functools.wraps(func)
            def wrapper(self, *args, **kwargs):
                with self.worker_timer(label):
                    with AcceleratorUtil.profiling_range(self._accelerator_type, label):
                        return func(self, *args, **kwargs)

            return wrapper

        return decorator

    def start_profile(self, step_idx: int) -> None:
        """Open a profiling capture window for the given step on this worker."""
        AcceleratorUtil.start_profiling(self._accelerator_type, step_idx)

    def stop_profile(self) -> None:
        """Close the current profiling capture window on this worker."""
        AcceleratorUtil.stop_profiling(self._accelerator_type)

    @staticmethod
    def check_worker_alive(worker_name: str) -> bool:
        """Check if a worker is alive.

        Args:
            worker_name (str): The name of the worker to check.

        Returns:
            bool: True if the worker is alive, False otherwise.
        """
        try:
            with without_http_proxies():
                actors = ray.util.state.list_actors(
                    filters=[("NAME", "=", worker_name), ("STATE", "!=", "DEAD")]
                )

            return len(actors) > 0
        except Exception:
            # Simply treat the worker as alive if any unexpected error occurs during state query
            return True

    def _check_initialized(self):
        """Check if the Worker has been initialized.

        This is used to ensure that the Worker is ready to be used.
        """
        if not self._has_initialized:
            raise RuntimeError(
                "Worker has not been initialized. Please call Worker.__init__(self) in your class's __init__ method."
            )

    def _setup_managers(self):
        """When the Worker is not a Ray actor, we need to initialize Ray if it is not already initialized."""
        from ..collective import Collective
        from ..manager import WorkerManager

        if (
            self._manager_proxy is None
            or self._collective is None
            or Worker.PID != os.getpid()
        ):
            self._manager_proxy = WorkerManager.get_proxy()
            self._manager_proxy.register_worker(self._worker_address, self._worker_info)
            self._collective = Collective(self)

            Worker.PID = os.getpid()

    def _setup_local_rank_world_size(self):
        if self._is_ray_actor:
            if os.environ.get("ISOLATE_ACCELERATOR", "0") == "1":
                # Ray limits the number of accelerators per worker to 1, so when calling torch.cuda.set_device(), we must ensure that 0 is passed as the local rank.
                os.environ["LOCAL_RANK"] = "0"
                os.environ["LOCAL_WORLD_SIZE"] = "1"
                self._isolate_gpu = True
            else:
                os.environ["LOCAL_RANK"] = str(
                    self._local_accelerator_rank
                )  # Must use the actual device ID
                os.environ["LOCAL_WORLD_SIZE"] = str(self._node_local_world_size)
                self._isolate_gpu = False

            self._local_rank = int(os.environ["LOCAL_RANK"])
            self._local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
        else:
            # These are not set for non-Ray workers
            self._local_rank = -1
            self._local_world_size = -1

    def _setup_master_address_and_port(self):
        # Executed after _init_ray_and_proxies
        from ..manager import WorkerInfo

        if self._is_ray_actor:
            master_worker_address = (
                self._worker_address.get_parent_address().get_child_address(0)
            )
            worker_info: WorkerInfo = None
            count = 0
            while worker_info is None:
                worker_info = self._manager_proxy.get_worker_info(master_worker_address)
                time.sleep(0.001)
                count += 1
                if count % Cluster.TIMEOUT_WARN_TIME == 0:
                    self._logger.warning(
                        f"Waiting for rank 0 of group {self._worker_address.root_group_name} to be up for {count // 1000} seconds"
                    )
            self._master_address = worker_info.node_ip
            self._master_port = worker_info.node_port
            os.environ["MASTER_ADDR"] = self._master_address
            os.environ["MASTER_PORT"] = str(self._master_port)

    def _setup_accelerator_info(self) -> int:
        cluster = Cluster()
        visible_devices = AcceleratorUtil.get_visible_devices(self._accelerator_type)
        node_accelerator_ranks = cluster.accelerator_ranks[self._cluster_node_rank]
        self.global_accelerator_ids = [
            node_accelerator_ranks[local_id] for local_id in visible_devices
        ]

        if not self._is_ray_actor:
            if len(visible_devices) > 0:
                self._local_accelerator_rank = visible_devices[0]
            else:
                self._local_accelerator_rank = -1

    def _setup_hardware(self):
        cluster = Cluster()
        hardware_ranks_str = os.environ.get("LOCAL_HARDWARE_RANKS", "")
        if hardware_ranks_str == "":
            self._local_hardware_ranks = []
        else:
            self._local_hardware_ranks = list(
                map(int, hardware_ranks_str.strip().split(","))
            )
        node_group_label = os.environ.get("NODE_GROUP_LABEL", None)
        self._node_group = cluster.get_node_group(node_group_label)
        assert self._node_group is not None, (
            f"Node group {node_group_label} not found in cluster. Available node groups: {[node_group.label for node_group in cluster._node_groups]}"
        )

    def _setup_comm_envs(self):
        # Communication devices
        self._comm_devices = Cluster.get_sys_env_var(
            ClusterEnvVar.COMM_NET_DEVICES, None
        )
        if self._comm_devices is not None:
            self.log_info(
                f"Using communication devices for worker {self._worker_name}: {self._comm_devices}"
            )
            # Validate the format of comm devices
            if os.getenv("GLOO_SOCKET_IFNAME") is None:
                os.environ["GLOO_SOCKET_IFNAME"] = self._comm_devices
            elif self._comm_devices != os.environ["GLOO_SOCKET_IFNAME"]:
                self.log_warning(
                    f"GLOO_SOCKET_IFNAME is already set to {os.environ['GLOO_SOCKET_IFNAME']}, ignoring {Cluster.get_full_env_var_name(ClusterEnvVar.COMM_NET_DEVICES)}={self._comm_devices}"
                )

            ccl_socket_env_var = AcceleratorUtil.get_ccl_socket_ifname_env_var(
                self._accelerator_type
            )
            if os.environ.get(ccl_socket_env_var) is None:
                os.environ[ccl_socket_env_var] = self._comm_devices
            elif self._comm_devices != os.environ[ccl_socket_env_var]:
                self.log_warning(
                    f"{ccl_socket_env_var} is already set to {os.environ[ccl_socket_env_var]}, ignoring {Cluster.get_full_env_var_name(ClusterEnvVar.COMM_NET_DEVICES)}={self._comm_devices}"
                )

    def _setup_logging(self):
        self._logger = logging.getLogger(self._worker_name)
        logging_level = Cluster.get_sys_env_var(ClusterEnvVar.LOG_LEVEL, "INFO").upper()
        if logging_level == "DEBUG":
            self._logging_level = logging.DEBUG
        elif logging_level == "INFO":
            self._logging_level = logging.INFO
        elif logging_level == "WARNING":
            self._logging_level = logging.WARNING
        elif logging_level == "ERROR":
            self._logging_level = logging.ERROR
        self._logger.setLevel(self._logging_level)

        self._logger.propagate = False
        for handler in self._logger.handlers:
            self._logger.removeHandler(handler)
        handler = logging.StreamHandler(sys.stdout)
        formatter = logging.Formatter(
            fmt=f"[%(levelname)s %(asctime)s {self._worker_address.get_parent_address().get_name()}-Rank-{self._rank}][%(filename)s:%(lineno)d] %(message)s",
            datefmt="%H:%M:%S",
        )
        handler.setFormatter(formatter)
        self._logger.addHandler(handler)
        Worker.logger = self._logger

    def _register_signal_handlers(self):
        """Register signal handlers for this worker process for more informative debugging."""

        def signal_handler(signum, frame):
            self._logger.error(
                f"Received signal {signum} ({signal.strsignal(signum)}) in worker {self._worker_address}, traceback is below:"
            )
            traceback.print_stack(frame)

            if self._thread is not threading.main_thread():
                # If the Worker is running in a worker thread (e.g., Worker with async functions), print the stack trace of the worker thread
                self._logger.error("Worker thread traceback is below:")
                traceback.print_stack(sys._current_frames()[self._thread.ident])
            os.kill(os.getpid(), signal.SIGKILL)

        should_catch_system_failure = os.environ.get("CATCH_SYSTEM_FAILURE", "0") == "1"
        if not should_catch_system_failure:
            # If the environment variable CATCH_SYSTEM_FAILURE is "0", do not register signal handlers
            return

        try:
            # Register signal handlers for common signals
            signal.signal(signal.SIGINT, signal_handler)  # Handle Ctrl+C
            signal.signal(signal.SIGTERM, signal_handler)  # Handle termination signal
            signal.signal(signal.SIGSEGV, signal_handler)  # Handle segmentation fault
            signal.signal(signal.SIGABRT, signal_handler)  # Handle abort signal
            signal.signal(signal.SIGQUIT, signal_handler)  # Handle quit signal
            signal.signal(
                signal.SIGUSR1, signal_handler
            )  # Handle user-defined signal 1
            signal.signal(
                signal.SIGUSR2, signal_handler
            )  # Handle user-defined signal 2
        except ValueError:
            self._logger.warning(
                "Failed to register signal handlers. This may happen if the Worker is not running in the main thread."
            )

    def _enable_ptrace(self):
        """Enable ptrace from any same-UID process. Used for enable CUDAIPC when PTRACE is disabled.

        https://gist.github.com/youkaichao/8f87555bdeaaf68f4492b0dc96fbd206
        """
        # https://github.com/torvalds/linux/blob/24d479d26b25bce5faea3ddd9fa8f3a6c3129ea7/include/uapi/linux/prctl.h#L155
        PR_SET_PTRACER = 0x59616D61
        PR_SET_PTRACER_ANY = -1

        try:
            libc = ctypes.CDLL("libc.so.6")

            result = libc.prctl(PR_SET_PTRACER, PR_SET_PTRACER_ANY, 0, 0, 0)
            if result != 0:
                warnings.warn("prctl(PR_SET_PTRACER, ANY) failed!")
        except Exception as e:
            warnings.warn(f"Failed to enable ptrace from any same-UID process: {e}")

    def _get_collective_group(self, peer_addr: WorkerAddress):
        """Get a collective group for communication with a peer worker."""
        workers = [self._worker_address, peer_addr]
        # Ensure the order is the same with the same two ranks
        workers = sorted(workers, key=lambda x: x.get_name())
        self._setup_managers()
        with self._lock:
            return self._collective.create_collective_group(workers)

    def _setup_worker_info(self):
        """Get the worker information for local access.

        This method is used to retrieve the worker properties without calling remote functions.
        """
        if self._actor is None and self._is_ray_actor:
            self._actor = ray.get_actor(self._worker_name, namespace=Cluster.NAMESPACE)

        node_ip = ray.util.get_node_ip_address()
        node_port = self.acquire_free_port()

        from ..manager import WorkerInfo

        self._worker_info = WorkerInfo(
            address=self._worker_address,
            rank=self._rank,
            group_world_size=self._world_size,
            cluster_node_rank=self._cluster_node_rank,
            accelerator_type=self._accelerator_type,
            accelerator_model=self._accelerator_model,
            accelerator_rank=self._local_accelerator_rank,
            node_ip=node_ip,
            node_port=node_port,
            available_accelerators=self.global_accelerator_ids,
            hardware_infos=self.hardware_infos,
        )

    def __repr__(self):
        """Return a string representation of the Worker."""
        return f"{self._group_name}(rank={self._rank})"
