Python PyTorch rpc_async用法及代码示例

本文简要介绍python语言中 torch.distributed.rpc.rpc_async 的用法。

用法:

torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None, timeout=- 1.0)

参数：

• to(str或者WorkerInfo或者int) -目标工作人员的名称/等级/WorkerInfo。

• func(可调用的) -可调用函数，例如 Python 可调用函数、内置运算符(例如 add() )和带注释的 TorchScript 函数。

• args(tuple) -func 调用的参数元组。

• kwargs(dict) -是func 调用的关键字参数字典。

• timeout(float,可选的) -用于此 RPC 的超时秒数。如果 RPC 在这段时间内没有完成，则会引发一个表明它已超时的异常。值 0 表示无限超时，即永远不会引发超时错误。如果未提供，则使用在初始化期间或使用_set_rpc_timeout 设置的默认值。

返回：

返回可以等待的 Future 对象。完成后，可以从 Future 对象中检索args 和kwargs 上func 的返回值。

进行非阻塞 RPC 调用以在 worker to 上运行函数 func。 RPC 消息与 Python 代码的执行并行发送和接收。此方法是线程安全的。此方法将立即返回一个可以等待的 Future 。

警告

不支持使用 GPU 张量作为参数或 func 的返回值，因为我们不支持通过网络发送 GPU 张量。在将 GPU 张量用作参数或返回 func 值之前，您需要将它们显式复制到 CPU。

警告

rpc_async API 不会复制参数张量的存储，直到通过网络发送它们，这可以由不同的线程完成，具体取决于 RPC 后端类型。调用者应确保这些张量的内容保持不变，直到返回的 Future 完成。

例子：

确保 MASTER_ADDR 和 MASTER_PORT 在两个 worker 上都正确设置。有关详细信息，请参阅 init_process_group() API。例如，

>>> export MASTER_ADDR=localhost
>>> export MASTER_PORT=5678

然后在两个不同的进程中运行以下代码：

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut1 = rpc.rpc_async("worker1", torch.add, args=(torch.ones(2), 3))
>>> fut2 = rpc.rpc_async("worker1", min, args=(1, 2))
>>> result = fut1.wait() + fut2.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

下面是使用 RPC 运行 TorchScript 函数的示例。

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(t1, t2):
>>>    return torch.add(t1, t2)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut = rpc.rpc_async("worker1", my_script_add, args=(torch.ones(2), 3))
>>> ret = fut.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()