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()