Add bound method tests for async_execution with RRef helper (pytorch#44716)

Summary: Pull Request resolved: pytorch#44716

Test Plan: Imported from OSS

Reviewed By: rohan-varma

Differential Revision: D23707326

Pulled By: mrshenli

fbshipit-source-id: a2f8db17447e9f82c9f6ed941ff1f8cb9090ad74

Author: mrshenli

torch/distributed/rpc/functions.py

@@ -115,6 +115,12 @@ def async_execution(fn):
         >>>         )
         >>>         return ret_fut
         >>>
+        >>>     @rpc.functions.async_execution
+        >>>     def bound_async_add(self, to, x, y, z):
+        >>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
+        >>>             lambda fut: fut.wait() + z
+        >>>         )
+        >>>
         >>> # On worker0
         >>> ret = rpc.rpc_sync(
         >>>     "worker1",
@@ -129,6 +135,26 @@ def async_execution(fn):
         >>>     args=("worker2", torch.ones(2), 1, 2)
         >>> )
         >>> print(ret)  # prints tensor([4., 4.])
+
+        This decorator also works with RRef helpers, i.e., .
+        :meth:`torch.distributed.rpc.RRef.rpc_sync`,
+        :meth:`torch.distributed.rpc.RRef.rpc_async`, and
+        :meth:`torch.distributed.rpc.RRef.remote`.
+
+        >>> from torch.distributed import rpc
+        >>>
+        >>> # reuse the AsyncExecutionClass class above
+        >>> rref = rpc.remote("worker1", AsyncExecutionClass)
+        >>> ret = rref.rpc_sync().static_async_add("worker2", torch.ones(2), 1, 2)
+        >>> print(ret)  # prints tensor([4., 4.])
+        >>>
+        >>> rref = rpc.remote("worker1", AsyncExecutionClass)
+        >>> ret = rref.rpc_async().static_async_add("worker2", torch.ones(2), 1, 2).wait()
+        >>> print(ret)  # prints tensor([4., 4.])
+        >>>
+        >>> rref = rpc.remote("worker1", AsyncExecutionClass)
+        >>> ret = rref.remote().static_async_add("worker2", torch.ones(2), 1, 2).to_here()
+        >>> print(ret)  # prints tensor([4., 4.])
     """
     @functools.wraps(fn)
     def wrapper(*args, **kwargs):

torch/testing/_internal/distributed/rpc/rpc_test.py

@@ -448,6 +448,12 @@ def class_async_add(cls, to, x, y, z):
         )
         return ret_fut
 
+    @rpc.functions.async_execution
+    def bound_async_add(self, to, x, y, z):
+        return rpc.rpc_async(to, torch.add, args=(x, y)).then(
+            lambda fut: fut.wait() + z
+        )
+
 
 def return_future():
     return torch.futures.Future()
@@ -3051,14 +3057,17 @@ def _test_test_async_class_rref_proxy(self, mode=RPCExecMode.SYNC):
         if mode == RPCExecMode.SYNC:
             ret = rref.rpc_sync().static_async_add(dst2, x, x, y)
             ret += rref.rpc_sync().class_async_add(dst2, x, x, y)
+            ret += rref.rpc_sync().bound_async_add(dst2, x, x, y)
         elif mode == RPCExecMode.ASYNC:
             ret = rref.rpc_async().static_async_add(dst2, x, x, y).wait()
             ret += rref.rpc_async().class_async_add(dst2, x, x, y).wait()
+            ret += rref.rpc_async().bound_async_add(dst2, x, x, y).wait()
         elif mode == RPCExecMode.REMOTE:
             ret = rref.remote().static_async_add(dst2, x, x, y).to_here()
             ret += rref.remote().class_async_add(dst2, x, x, y).to_here()
+            ret += rref.remote().bound_async_add(dst2, x, x, y).to_here()
 
-        self.assertEqual(ret, 2 * 4 * x)
+        self.assertEqual(ret, 3 * 4 * x)
 
     @dist_init
     def test_async_class_rref_proxy(self):