Table of Contents
Shortcuts

# Source code for torch.distributed.rpc

    from datetime import timedelta
    import logging
    import os
    import threading
    import warnings
    from typing import Generator, Tuple
    from urllib.parse import urlparse

    import torch
    import torch.distributed as dist

    logger = logging.getLogger(__name__)

    _init_counter = 0
    _init_counter_lock = threading.Lock()

    __all__ = ["is_available"]

    def is_available() -> bool:
        return hasattr(torch._C, "_rpc_init")

    if is_available() and not torch._C._rpc_init():
        raise RuntimeError("Failed to initialize torch.distributed.rpc")

    if is_available():
        from torch._C._distributed_c10d import Store
        from torch._C._distributed_rpc import (
            _disable_jit_rref_pickle,
            _enable_jit_rref_pickle,
            _disable_server_process_global_profiler,
            _enable_server_process_global_profiler,
            _set_and_start_rpc_agent,
            _reset_current_rpc_agent,
            _delete_all_user_and_unforked_owner_rrefs,
            _destroy_rref_context,
            _set_profiler_node_id,
            _is_current_rpc_agent_set,
            _rref_context_get_debug_info,
            _cleanup_python_rpc_handler,
            _invoke_rpc_builtin,
            _invoke_rpc_python_udf,
            _invoke_rpc_torchscript,
            _invoke_remote_builtin,
            _invoke_remote_python_udf,
            _invoke_remote_torchscript,
            _set_rpc_timeout,
            _get_current_rpc_agent,
            get_rpc_timeout,
            enable_gil_profiling,
            RpcBackendOptions,
            _TensorPipeRpcBackendOptionsBase,
            RpcAgent,
            PyRRef,
            TensorPipeAgent,
            RemoteProfilerManager,
            WorkerInfo,
            _DEFAULT_INIT_METHOD,
            _DEFAULT_NUM_WORKER_THREADS,
            _UNSET_RPC_TIMEOUT,
            _DEFAULT_RPC_TIMEOUT_SEC,
        )  # noqa: F401

        from . import api, backend_registry, functions
        from .api import *  # noqa: F401,F403
        import numbers

        import torch.distributed.autograd as dist_autograd

        from .backend_registry import BackendType
        from .options import TensorPipeRpcBackendOptions  # noqa: F401
        from .server_process_global_profiler import (
            _server_process_global_profile,
        )

        rendezvous_iterator: Generator[Tuple[Store, int, int], None, None]

        __all__ += ["init_rpc", "BackendType", "TensorPipeRpcBackendOptions"]
        __all__ = __all__ + api.__all__ + backend_registry.__all__  # noqa: PLE0605

        def init_rpc(
            name,
            backend=None,
            rank=-1,
            world_size=None,
            rpc_backend_options=None,
        ):
            r"""
            Initializes RPC primitives such as the local RPC agent
            and distributed autograd, which immediately makes the current
            process ready to send and receive RPCs.

            Args:
                name (str): a globally unique name of this node. (e.g.,
                    ``Trainer3``, ``ParameterServer2``, ``Master``, ``Worker1``)
                    Name can only contain number, alphabet, underscore, colon,
                    and/or dash, and must be shorter than 128 characters.
                backend (BackendType, optional): The type of RPC backend
                    implementation. Supported values is
                    ``BackendType.TENSORPIPE`` (the default).
                    See :ref:`rpc-backends` for more information.
                rank (int): a globally unique id/rank of this node.
                world_size (int): The number of workers in the group.
                rpc_backend_options (RpcBackendOptions, optional): The options
                    passed to the RpcAgent constructor. It must be an agent-specific
                    subclass of :class:`~torch.distributed.rpc.RpcBackendOptions`
                    and contains agent-specific initialization configurations. By
                    default, for all agents, it sets the default timeout to 60
                    seconds and performs the rendezvous with an underlying process
                    group initialized using ``init_method = "env://"``,
                    meaning that environment variables ``MASTER_ADDR`` and
                    ``MASTER_PORT`` need to be set properly. See
                    :ref:`rpc-backends` for more information and find which options
                    are available.
            """
            torch._C._log_api_usage_once("torch.distributed.init_rpc")
            if backend is not None and not isinstance(
                backend, backend_registry.BackendType
            ):
                raise TypeError("Argument backend must be a member of BackendType")

            if rpc_backend_options is not None and not isinstance(
                rpc_backend_options, RpcBackendOptions
            ):
                raise TypeError(
                    "Argument rpc_backend_options must be an instance of RpcBackendOptions"
                )

            # Try to detect the backend from the options
            if backend is None and rpc_backend_options is not None:
                for candidate_backend in BackendType:
                    if isinstance(
                        rpc_backend_options,
                        type(
                            backend_registry.construct_rpc_backend_options(
                                candidate_backend
                            )
                        ),
                    ):
                        backend = candidate_backend
                        break
                else:
                    raise TypeError(
                        f"Could not infer backend for options {rpc_backend_options}"
                    )
                # Ignore type error because mypy doesn't handle dynamically generated type objects (#4865)
                if backend != BackendType.TENSORPIPE:  # type: ignore[attr-defined]
                    logger.warning(
                        "RPC was initialized with no explicit backend but with options "  # type: ignore[attr-defined]
                        "corresponding to %(backend)s, hence that backend will be used "
                        "instead of the default BackendType.TENSORPIPE. To silence this "
                        "warning pass `backend=%(backend)s` explicitly.",
                        {'backend': backend}
                    )

            if backend is None:
                backend = BackendType.TENSORPIPE  # type: ignore[attr-defined]

            if rpc_backend_options is None:
                # default construct a set of RPC backend options.
                rpc_backend_options = backend_registry.construct_rpc_backend_options(
                    backend
                )

            # Create store, performs rendezvous for static RPC group.
            if not world_size:
                # If world_size is not set in construction and also not set in environment variables
                # The store will be created for the dynamic group setting
                store = dist._create_store_from_options(rpc_backend_options, rank)
            else:
                # This rendezvous state sometimes is destroyed before all processes
                # finishing handshaking. To avoid that issue, we make it global to
                # keep it alive.
                global rendezvous_iterator
                rendezvous_iterator = dist.rendezvous(
                    rpc_backend_options.init_method, rank=rank, world_size=world_size
                )
                store, _, _ = next(rendezvous_iterator)
            # Use same timeout as RPC.
            store.set_timeout(timedelta(seconds=rpc_backend_options.rpc_timeout))

            # Use a PrefixStore to distinguish multiple invocations.
            with _init_counter_lock:
                global _init_counter
                store = dist.PrefixStore(str(f"rpc_prefix_{_init_counter}"), store)
                _init_counter += 1

            # Initialize autograd before RPC since _init_rpc_backend guarantees all
            # processes sync via the store. If we initialize autograd after RPC,
            # there could be a race where some nodes might have initialized autograd
            # and others might not have. As a result, a node calling
            # torch.distributed.autograd.backward() would run into errors since
            # other nodes might not have been initialized.
            dist_autograd._init(rank)

            _set_profiler_node_id(rank)
            # Initialize RPC.
            _init_rpc_backend(backend, store, name, rank, world_size, rpc_backend_options)

        def _validate_rpc_args(backend, store, name, rank, world_size, rpc_backend_options):
            type_mapping = {
                backend: backend_registry.BackendType,
                store: dist.Store,
                name: str,
                rank: numbers.Integral,
                # world_size can be None for a dynamic group
                world_size: (numbers.Integral, type(None)),
                rpc_backend_options: RpcBackendOptions,
            }
            for arg, arg_type in type_mapping.items():
                if not isinstance(arg, arg_type):  # type: ignore[arg-type]
                    raise RuntimeError(
                        f"Argument {arg} must be of type {arg_type} but got type {type(arg)}"
                    )

        def _init_rpc_backend(
            backend=BackendType.TENSORPIPE,  # type: ignore[attr-defined]
            store=None,
            name=None,
            rank=-1,
            world_size=None,
            rpc_backend_options=None,
        ):

            _validate_rpc_args(backend, store, name, rank, world_size, rpc_backend_options)

            if _is_current_rpc_agent_set():
                raise RuntimeError("RPC is already initialized")

            # Initialize RPC.
            rpc_agent = backend_registry.init_backend(
                backend,
                store=store,
                name=name,
                rank=rank,
                world_size=world_size,
                rpc_backend_options=rpc_backend_options,
            )

            api._init_rpc_states(rpc_agent)

        @api._require_initialized
        def _get_debug_info():
            info = _rref_context_get_debug_info()
            info.update(api._get_current_rpc_agent().get_debug_info())
            info.update(dist_autograd._get_debug_info())
            return info

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.