Python PyTorch shutdown用法及代码示例

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

用法:

torch.distributed.rpc.shutdown(graceful=True)

参数：

graceful(bool) -是否进行正常关机。如果为 True，这将 1) 等到 UserRRefs 没有待处理的系统消息并删除它们； 2) 阻塞，直到所有本地和远程 RPC 进程都到达此方法并等待所有未完成的工作完成。

关闭 RPC 代理，然后销毁 RPC 代理。这会阻止本地代理接受未完成的请求，并通过终止所有 RPC 线程来关闭 RPC 框架。如果 graceful=True ，这将阻塞，直到所有本地和远程 RPC 进程到达此方法并等待所有未完成的工作完成。否则，如果 graceful=False ，这是本地关闭，它不会等待其他 RPC 进程到达此方法。

警告

对于 rpc_async() 返回的 Future 对象，不应在 shutdown() 之后调用 future.wait() 。

例子：

确保 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)
>>> # do some work
>>> result = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(1), 1))
>>> # ready to shutdown
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> # wait for worker 0 to finish work, and then shutdown.
>>> rpc.shutdown()