async_execution and RRef helpers do not work together

[mrshenli]

The following code does not work

class Model:
    ...
    @staticmethod
    @rpc.functions.async_execution
    def func(x):
        ...

model_rref = rpc.remote("worker1", Model)
model_rref.rpc_sync().func(x)

This is because the _local_invoke function ignored the attribute installed by the decorator. To fix this, we will need to have another version of _local_invoke which is decorated by @rpc.functions.async_execution and returns the Future object from the wrapped function.

pytorch/torch/distributed/rpc/rref_proxy.py

Lines 8 to 13 in 03eec07

	def _invoke_rpc(rref, rpc_api, func_name, *args, **kwargs):
	return rpc_api(
	rref.owner(),
	_local_invoke,
	args=(rref, func_name, args, kwargs)
	)

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse

[mrshenli]

closed by #44666 and #44716