chore: import upstream snapshot with attribution

python/paddle/distributed/rpc/__init__.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+from paddle.distributed.rpc.rpc import (
+    get_all_worker_infos,
+    get_current_worker_info,
+    get_worker_info,
+    init_rpc,
+    rpc_async,
+    rpc_sync,
+    shutdown,
+)
+
+__all__ = [
+    "init_rpc",
+    "shutdown",
+    "rpc_async",
+    "rpc_sync",
+    "get_worker_info",
+    "get_all_worker_infos",
+    "get_current_worker_info",
+]

python/paddle/distributed/rpc/internal.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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 pickle
+from collections import namedtuple
+
+PythonFunc = namedtuple("PythonFunc", ["func", "args", "kwargs"])
+"""Some Python code interfaces called in C++"""
+
+
+def _serialize(obj):
+    return pickle.dumps(obj)
+
+
+def _deserialize(obj):
+    return pickle.loads(obj)
+
+
+def _run_py_func(python_func):
+    result = python_func.func(*python_func.args, **python_func.kwargs)
+    return result

python/paddle/distributed/rpc/rpc.py

@@ -0,0 +1,456 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+from __future__ import annotations
+
+import datetime
+import os
+import pickle
+import time
+from collections import namedtuple
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
+
+from paddle.base import core
+from paddle.distributed.launch.context import Node
+from paddle.distributed.rpc.internal import PythonFunc, _serialize
+from paddle.distributed.utils.launch_utils import logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    _RetT = TypeVar("_RetT", covariant=True)
+
+    class _FutureWrapper(Protocol[_RetT]):
+        def wait(self) -> _RetT: ...
+
+
+WorkerInfo = namedtuple("WorkerInfo", ["name", "rank", "ip", "port"])
+
+_DEFAULT_RPC_TIMEOUT = -1
+_MAX_RPC_TIMEOUT_MS = 0x7FFFFFFF
+_BARRIER_TIMEOUT_MAX_DAYS = 99999999
+# tcp store for `_barrier_never_timeout`
+_barrier_store = None
+# count the number of `_barrier_never_timeout` is called and
+# ensure that the barrier key is unique
+_barrier_count = 0
+
+
+def _set_barrier_store(store):
+    global _barrier_store
+    _barrier_store = store
+
+
+def _del_barrier_store():
+    global _barrier_store
+    del _barrier_store
+
+
+def _set_self_info(name, rank, ip, port):
+    self_info = pickle.dumps(WorkerInfo(name, rank, ip, port))
+    _barrier_store.set(str(rank), self_info)
+
+
+def _exchange_all_service_infos(world_size):
+    all_infos = []
+    s = set()
+    for rank in range(world_size):
+        info = pickle.loads(_barrier_store.get(str(rank)))
+        assert info.name not in s, (
+            "The Worker name must be unique, but name `{}` is repeated."
+        )
+        s.add(info.name)
+        all_infos.append(info)
+    return all_infos
+
+
+def _gen_endpoint():
+    node = Node()
+    ip = node.get_host_ip()
+    free_port = node.get_free_port()
+    return f"{ip}:{free_port}"
+
+
+def init_rpc(
+    name: str,
+    rank: int | None = None,
+    world_size: int | None = None,
+    master_endpoint: str | None = None,
+) -> None:
+    """
+    init rpc. Warning: All RPC API should only be used internally within a secure network environment and
+    must not be accessible via the public internet.
+
+    Args:
+        name (str): worker name.
+        rank (int, optional): worker id, default is None.
+        world_size (int, optional): number of workers, default is None.
+        master_endpoint (str, optional): id address of master, other nodes communicate with the master to
+            get the information of all worker nodes, default is None.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8001",
+            ... )
+
+            >>> rpc.shutdown()
+
+    """
+    rank = int(os.environ["PADDLE_TRAINER_ID"]) if rank is None else rank
+    world_size = (
+        int(os.environ["PADDLE_TRAINERS_NUM"])
+        if world_size is None
+        else world_size
+    )
+    worker_endpoint = os.getenv("PADDLE_WORKER_ENDPOINT", None)
+    if worker_endpoint is None:
+        worker_endpoint = _gen_endpoint()
+    logger.info(f"Trainer {rank}: worker endpoint: {worker_endpoint}")
+    master_endpoint = (
+        master_endpoint
+        if master_endpoint is not None
+        else os.environ["PADDLE_MASTER_ENDPOINT"]
+    )
+    master_addr, master_port = master_endpoint.split(":")
+    master_port = int(master_port)
+    stop_check_timeout = int(os.getenv("FLAGS_stop_check_timeout", "900"))
+    store = core.TCPStore(
+        master_addr,
+        master_port,
+        rank == 0,
+        world_size,
+        timeout=stop_check_timeout,
+    )
+    _set_barrier_store(store)
+    ip, port = worker_endpoint.split(":")
+    port = int(port)
+    _set_self_info(name, rank, ip, port)
+    all_infos = _exchange_all_service_infos(world_size)
+    c_infos = []
+    for node_info in all_infos:
+        info = core.WorkerInfo(
+            node_info.name, node_info.rank, node_info.ip, node_info.port
+        )
+        c_infos.append(info)
+    core.init_and_set_agent_instance(name, c_infos)
+    core.rpc_start_worker()
+    # ensure that all the workers are started
+    _barrier_never_timeout(rank, world_size)
+    core.rpc_start_client()
+    logger.info(f"Trainer {rank}: Init RPC done!")
+
+
+def rpc_sync(
+    to: str,
+    fn: Callable[..., _RetT],
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    timeout: int = _DEFAULT_RPC_TIMEOUT,
+) -> _RetT:
+    """
+    Make a blocking RPC call to run function ``fn`` on worker ``to``. Warning: All RPC API should
+    only be used internally within a secure network environment and must not be accessible via
+    the public internet.
+
+    Args:
+        to (str): name of the destination worker.
+        fn (fn): a callable function, such as Python callables.
+        args (tuple, optional): the argument tuple for the ``fn`` invocation, default is None.
+        kwargs (dict, optional): is a dictionary of keyword arguments for the ``fn``
+                       invocation, default is None.
+        timeout (int, optional): timeout in seconds to use for this RPC. If
+                                   the RPC does not complete in this amount of
+                                   time, an exception indicating it has
+                                   timed out will be raised. A value less than or equal to 0
+                                   indicates an infinite timeout, i.e. a timeout
+                                   error will never be raised. The default value is -1.
+
+    Returns:
+        Returns the result of running ``fn`` with ``args`` and ``kwargs``.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+
+            >>> def add(a, b):
+            ...     return a + b
+
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8002",
+            ... )
+
+            >>> ret = rpc.rpc_sync("worker0", add, args=(2, 3))
+            >>> rpc.shutdown()
+
+    """
+    fut = _invoke_rpc(to, fn, args, kwargs, timeout)
+    return fut.wait()
+
+
+def rpc_async(
+    to: str,
+    fn: Callable[..., _RetT],
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    timeout: int = _DEFAULT_RPC_TIMEOUT,
+) -> _FutureWrapper[_RetT]:
+    """
+    Make a non-blocking RPC call to run function ``fn`` on worker ``to``. Warning: All RPC API should
+    only be used internally within a secure network environment and must not be accessible via the public internet.
+
+    Args:
+        to (str): name of the destination worker.
+        fn (fn): a callable function, such as Python callables.
+        args (tuple, optional): the argument tuple for the ``fn`` invocation, default is None.
+        kwargs (dict, optional): is a dictionary of keyword arguments for the ``fn``
+                       invocation, default is None.
+        timeout (int, optional): timeout in seconds to use for this RPC. If
+                                   the RPC does not complete in this amount of
+                                   time, an exception indicating it has
+                                   timed out will be raised. A value less than or equal to 0
+                                   indicates an infinite timeout, i.e. a timeout
+                                   error will never be raised. The default value is -1.
+
+    Returns:
+        Returns a :class:`FutureWrapper` object that can be waited
+        on. When completed, the return value of ``fn`` on ``args`` and
+        ``kwargs`` can be got by `fut.wait()`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+
+            >>> def add(a, b):
+            ...     return a + b
+
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8003",
+            ... )
+
+            >>> fut = rpc.rpc_async("worker0", add, args=(2, 3))
+            >>> print(fut.wait())
+            5
+
+            >>> rpc.shutdown()
+
+    """
+    return _invoke_rpc(to, fn, args, kwargs, timeout)
+
+
+def _invoke_rpc(to, fn, args, kwargs, timeout):
+    args = args if args else ()
+    kwargs = kwargs if kwargs else {}
+    serial_obj = _serialize(PythonFunc(fn, args, kwargs))
+    timeout_ms = timeout * 1000
+    timeout_ms = _MAX_RPC_TIMEOUT_MS if timeout_ms <= 0 else timeout_ms
+    future = core.invoke_rpc(to, serial_obj, timeout_ms)
+    return future
+
+
+def _barrier_never_timeout(global_rank, global_world_size):
+    # max timeout
+    timeout = datetime.timedelta(days=_BARRIER_TIMEOUT_MAX_DAYS)
+
+    if global_world_size < 2:
+        return
+
+    global _barrier_count
+    barrier_prefix = "Barrier/" + str(_barrier_count) + "/"
+    _barrier_count += 1
+    is_master = global_rank == 0
+
+    def _check_keys_ready(wait_keys):
+        start_time = time.time()
+        while len(wait_keys) > 0:
+            time.sleep(0.1)
+            elapse_time = time.time() - start_time
+            if datetime.timedelta(seconds=elapse_time) > timeout:
+                raise RuntimeError(
+                    f"Keys {wait_keys} are not ready since rank {global_rank} is waiting them."
+                )
+            wait_keys = list(
+                filter(lambda key: int(_barrier_store.get(key)) != 1, wait_keys)
+            )
+
+    if is_master:
+        # the master will add key, wait for all workers'exiting key and exit in the end.
+        # Note: the master must exit in the end to ensure that the TcpServer is destroyed in the end.
+        wait_keys = [
+            barrier_prefix + str(rank) for rank in range(1, global_world_size)
+        ]
+        _barrier_store.add(barrier_prefix + str(0), 1)
+        _check_keys_ready(wait_keys)
+    else:
+        wait_keys = [barrier_prefix + str(0)]
+        _check_keys_ready(wait_keys)
+        _barrier_store.add(barrier_prefix + str(global_rank), 1)
+
+
+def shutdown() -> None:
+    """
+    Perform a shutdown of the RPC agent, stop the worker and destroy the agent.
+    This will block until all local and remote RPC processes reach this method
+    and wait for all outstanding work to complete. Warning: All RPC API should
+    only be used internally within a secure network environment and must not be
+    accessible via the public internet.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8004",
+            ... )
+
+            >>> rpc.shutdown()
+
+    """
+    info = get_current_worker_info()
+    rank = info.rank
+    world_size = len(get_all_worker_infos())
+    # master will exit in the end
+    _barrier_never_timeout(rank, world_size)
+    core.rpc_stop_worker()
+    _del_barrier_store()
+    logger.info(f"Trainer {rank}: rpc shutdown!")
+
+
+def get_worker_info(name: str) -> WorkerInfo:
+    """
+    Get worker information by worker name. Warning: All RPC API should
+    only be used internally within a secure network environment and must
+    not be accessible via the public internet.
+
+    Args:
+        name (str): name of the worker.
+
+    Returns:
+        class `WorkerInfo` with attribute `name`, `rank`, `ip` and `port`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+            >>> import os
+
+            >>> os.environ["PADDLE_WORKER_ENDPOINT"] = "127.0.0.1:9002"
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8005",
+            ... )
+
+            >>> print(rpc.get_worker_info("worker0"))
+            {name: worker0, rank: 0, ip: 127.0.0.1, port: 9002}
+
+            >>> rpc.shutdown()
+
+    """
+    return core.rpc_get_worker_info(name)
+
+
+def get_all_worker_infos() -> list[WorkerInfo]:
+    """
+    Get all worker information. Warning: All RPC API should only be used
+    internally within a secure network environment and must not be
+    accessible via the public internet.
+
+    Returns:
+        List[WorkerInfo].
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+            >>> import os
+
+            >>> os.environ["PADDLE_WORKER_ENDPOINT"] = "127.0.0.1:9003"
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8006",
+            ... )
+
+            >>> print(rpc.get_all_worker_infos())
+            [{name: worker0, rank: 0, ip: 127.0.0.1, port: 9003}]
+
+            >>> rpc.shutdown()
+
+    """
+    return core.rpc_get_all_worker_infos()
+
+
+def get_current_worker_info() -> WorkerInfo:
+    """
+    Get current worker information. Warning: All RPC API should only be used internally
+    within a secure network environment and must not be accessible via the public internet.
+
+    Returns:
+        class `WorkerInfo` with attribute `name`, `rank`, `ip` and `port`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle.distributed.rpc as rpc
+            >>> import os
+
+            >>> os.environ["PADDLE_WORKER_ENDPOINT"] = "127.0.0.1:9004"
+            >>> rpc.init_rpc(
+            ...     "worker0",
+            ...     rank=0,
+            ...     world_size=1,
+            ...     master_endpoint="127.0.0.1:8007",
+            ... )
+
+            >>> print(rpc.get_current_worker_info())
+            {name: worker0, rank: 0, ip: 127.0.0.1, port: 9004}
+
+            >>> rpc.shutdown()
+
+    """
+    return core.rpc_get_current_worker_info()