torch.distributed.elastic.rendezvous.api 的源代码

# 版权所有 (c) Facebook, Inc. 及其附属公司。
# 保留所有权利。
#
# 本源代码根据在此源树根目录下的LICENSE文件中找到的BSD风格许可证进行许可。

from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, Optional, Tuple

from torch.distributed import Store


class RendezvousError(Exception):
    """表示rendezvous错误的基类型。"""


class RendezvousClosedError(RendezvousError):
    """当rendezvous关闭时引发。"""


class RendezvousTimeoutError(RendezvousError):
    """当rendezvous未按时完成时引发。"""


class RendezvousConnectionError(RendezvousError):
    """当与rendezvous后端的连接失败时引发。"""


class RendezvousStateError(RendezvousError):
    """当rendezvous的状态损坏时引发。"""

class RendezvousGracefulExitError(RendezvousError):
    """当节点未被包含在rendezvous中并优雅退出时引发。

    异常是一种退出堆栈的机制，但这并不意味着失败。
    """

class RendezvousHandler(ABC):
    """主要rendezvous接口。

    注意:
        分布式Torch用户通常**不需要**实现自己的``RendezvousHandler``。
        基于C10d Store的实现已经提供，并推荐给大多数用户。
    """

    @abstractmethod
    def get_backend(self) -> str:
        """返回rendezvous后端的名称。"""

    @abstractmethod
    def next_rendezvous(
        self,
    ) -> Tuple[Store, int, int]:
        """rendezvous屏障的主要入口点。

        阻塞直到rendezvous完成并且当前进程被包含在形成的worker组中，
        或发生超时，或rendezvous被标记为关闭。

        返回:
            一个包含 :py:class:`torch.distributed.Store`、``rank`` 和
            ``world size`` 的元组。

        引发:
            RendezvousClosedError:
                rendezvous已关闭。
            RendezvousConnectionError:
                与rendezvous后端的连接失败。
            RendezvousStateError:
                rendezvous状态损坏。
            RendezvousTimeoutError:
                rendezvous未按时完成。
        """

    @abstractmethod
    def is_closed(self) -> bool:
        """检查rendezvous是否已关闭。

        关闭的rendezvous意味着同一作业中所有未来的重新rendezvous尝试都将失败。

        ``is_closed()`` 和 :py:meth:`set_closed` 具有最终传播的语义，不应
        用于同步。意图是如果至少一个节点决定作业已完成，它将关闭rendezvous，
        其他节点将很快观察到这一点并停止运行。
        """

    @abstractmethod
    def set_closed(self):
        """将rendezvous标记为关闭。"""

    @abstractmethod
    def num_nodes_waiting(self) -> int:
        """返回在rendezvous屏障处迟到的节点数量，因此未被包含在当前worker组中。

        调用者应定期调用此方法以检查是否有新节点等待加入作业，如果有，
        则通过调用 :py:meth:`next_rendezvous()`（重新rendezvous）来接纳它们。
        """

    @abstractmethod
    def get_run_id(self) -> str:
        """返回rendezvous的run id。

        run id是用户定义的id，唯一标识分布式应用程序的实例。它通常映射到作业id，
        并用于允许节点加入正确的分布式应用程序。
        """

    @abstractmethod
    def shutdown(self) -> bool:
        """关闭为rendezvous打开的所有资源。

        示例::

            rdzv_handler = ...
            try:
                store, rank, world_size = rdzv_handler.next_rendezvous()
            finally:
                rdzv_handler.shutdown()
        """


class RendezvousParameters:
    """持有构造 :py:class:`RendezvousHandler` 的参数。

    参数:
        backend:
            用于处理rendezvous的后端的名称。
        endpoint:
            rendezvous的端点，通常为 [:] 形式。
        run_id:
            rendezvous的id。
        min_nodes:
            允许加入rendezvous的最小节点数。
        max_nodes:
            允许加入rendezvous的最大节点数。
        local_addr:
            本地节点的地址。
        **kwargs:
            指定后端的附加参数。
    """

    def __init__(
        self,
        backend: str,
        endpoint: str,
        run_id: str,
        min_nodes: int,
        max_nodes: int,
        local_addr: Optional[str] = None,
        **kwargs,
    ):
        if not backend:
            raise</span