diff --git a/python/ray/remote_function.py b/python/ray/remote_function.py index 44d2777a2900..9ff6994b8d42 100644 --- a/python/ray/remote_function.py +++ b/python/ray/remote_function.py @@ -4,6 +4,7 @@ import copy import logging +from functools import wraps from ray.function_manager import FunctionDescriptor import ray.signature @@ -74,15 +75,18 @@ def __init__(self, function, num_cpus, num_gpus, resources, self._last_driver_id_exported_for = None + # Override task.remote's signature and docstring + @wraps(function) + def _remote_proxy(*args, **kwargs): + return self._remote(args=args, kwargs=kwargs) + + self.remote = _remote_proxy + def __call__(self, *args, **kwargs): raise Exception("Remote functions cannot be called directly. Instead " "of running '{}()', try '{}.remote()'.".format( self._function_name, self._function_name)) - def remote(self, *args, **kwargs): - """This runs immediately when a remote function is called.""" - return self._remote(args=args, kwargs=kwargs) - def _submit(self, args=None, kwargs=None,