frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add API usage to torch.RPC (#67515)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67515

Adding API usage to torch.rpc to better understand usage of this API.
ghstack-source-id: 141877028

Reviewed By: rohan-varma

Differential Revision: D32011465

fbshipit-source-id: 34d006ece307ae4a90fbcc6cb44fc0b7edca611e
This commit is contained in:
Pritam Damania
2021-10-29 10:35:58 -07:00
committed by Facebook GitHub Bot
parent ddc9bd335b
commit 285d5a55b9
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torch/distributed/rpc/api.py
View File

@ -583,6 +583,7 @@ def remote(to, func, args=None, kwargs=None, timeout=UNSET_RPC_TIMEOUT):
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()
"""
torch._C._log_api_usage_once("torch.distributed.rpc_remote")
qualified_name = torch.jit._builtins._find_builtin(func)
dst_worker_info = _to_worker_info(to)
should_profile = torch.autograd._profiler_enabled()
@ -761,6 +762,7 @@ def rpc_sync(to, func, args=None, kwargs=None, timeout=UNSET_RPC_TIMEOUT):
>>> rpc.shutdown()
"""
torch._C._log_api_usage_once("torch.distributed.rpc_sync")
fut = _invoke_rpc(to, func, RPCExecMode.SYNC, args, kwargs, timeout)
return fut.wait()
@ -853,6 +855,7 @@ def rpc_async(to, func, args=None, kwargs=None, timeout=UNSET_RPC_TIMEOUT):
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()
"""
torch._C._log_api_usage_once("torch.distributed.rpc_async")
fut = _invoke_rpc(to, func, RPCExecMode.ASYNC, args, kwargs, timeout)
if hasattr(_thread_local_var, "future_list"):
_thread_local_var.future_list.append(fut)