diff --git a/python/ray/_private/client_mode_hook.py b/python/ray/_private/client_mode_hook.py index 4fbc568c8..9029b05c0 100644 --- a/python/ray/_private/client_mode_hook.py +++ b/python/ray/_private/client_mode_hook.py @@ -1,5 +1,6 @@ import os from contextlib import contextmanager +from functools import wraps client_mode_enabled = os.environ.get("RAY_CLIENT_MODE", "0") == "1" @@ -38,6 +39,7 @@ def client_mode_hook(func): """ from ray.experimental.client import ray + @wraps(func) def wrapper(*args, **kwargs): global _client_hook_enabled if client_mode_enabled and _client_hook_enabled: