[runtime_env] Allow specifying runtime env in @ray.remote decorator with Ray Client (#19626)

doc/source/advanced.rst

@@ -468,8 +468,6 @@ Or specify per-actor or per-task in the ``@ray.remote()`` decorator or by using
    :start-after: __per_task_per_actor_start__
    :end-before: __per_task_per_actor_end__
 
-Note: specifying within the ``@ray.remote()`` decorator is currently unsupported while using Ray Client; please use ``.options()`` instead in this case.
-
 The ``runtime_env`` is a Python dictionary including one or more of the following arguments:
 
 - ``working_dir`` (Path): Specifies the working directory for your job. This must be an existing local directory.

python/ray/_private/client_mode_hook.py

@@ -158,7 +158,7 @@ def client_mode_convert_function(func_cls, in_args, in_kwargs, **kwargs):
     The common case for this is to transparently convert that RemoteFunction
     to a ClientRemoteFunction. This happens in circumstances where the
     RemoteFunction is declared early, in a library and only then is Ray used in
-    client mode -- nescessitating a conversion.
+    client mode -- necessitating a conversion.
     """
     from ray.util.client import ray
 

python/ray/actor.py

@@ -493,7 +493,13 @@ class ActorClass:
         # Parse local pip/conda config files here. If we instead did it in
         # .remote(), it would get run in the Ray Client server, which runs on
         # a remote node where the files aren't available.
-        new_runtime_env = ParsedRuntimeEnv(runtime_env or {})
+        if runtime_env is not None:
+            new_runtime_env = ParsedRuntimeEnv(runtime_env)
+        else:
+            # Keep the runtime_env as None.  In .remote(), we need to know if
+            # runtime_env is None to know whether or not to fall back to the
+            # runtime_env specified in the @ray.remote decorator.
+            new_runtime_env = runtime_env
 
         class ActorOptionWrapper:
             def remote(self, *args, **kwargs):

python/ray/remote_function.py

@@ -171,7 +171,13 @@ class RemoteFunction:
         # Parse local pip/conda config files here. If we instead did it in
         # .remote(), it would get run in the Ray Client server, which runs on
         # a remote node where the files aren't available.
-        new_runtime_env = ParsedRuntimeEnv(runtime_env or {})
+        if runtime_env is not None:
+            new_runtime_env = ParsedRuntimeEnv(runtime_env)
+        else:
+            # Keep the runtime_env as None.  In .remote(), we need to know if
+            # runtime_env is None to know whether or not to fall back to the
+            # runtime_env specified in the @ray.remote decorator.
+            new_runtime_env = runtime_env
 
         class FuncWrapper:
             def remote(self, *args, **kwargs):

python/ray/tests/test_runtime_env.py

@@ -43,9 +43,29 @@ def test_get_release_wheel_url():
                 assert requests.head(url).status_code == 200, url
 
 
+@pytest.fixture(scope="function", params=["ray_client", "no_ray_client"])
+def start_cluster(ray_start_cluster, request):
+    assert request.param in {"ray_client", "no_ray_client"}
+    use_ray_client: bool = request.param == "ray_client"
+
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=4)
+    if use_ray_client:
+        cluster.head_node._ray_params.ray_client_server_port = "10003"
+        cluster.head_node.start_ray_client_server()
+        address = "ray://localhost:10003"
+    else:
+        address = cluster.address
+
+    yield cluster, address
+
+
 @pytest.mark.skipif(
     sys.platform == "win32", reason="runtime_env unsupported on Windows.")
-def test_decorator_task(ray_start_cluster_head):
+def test_decorator_task(start_cluster):
+    cluster, address = start_cluster
+    ray.init(address)
+
     @ray.remote(runtime_env={"env_vars": {"foo": "bar"}})
     def f():
         return os.environ.get("foo")
@@ -55,7 +75,10 @@ def test_decorator_task(ray_start_cluster_head):
 
 @pytest.mark.skipif(
     sys.platform == "win32", reason="runtime_env unsupported on Windows.")
-def test_decorator_actor(ray_start_cluster_head):
+def test_decorator_actor(start_cluster):
+    cluster, address = start_cluster
+    ray.init(address)
+
     @ray.remote(runtime_env={"env_vars": {"foo": "bar"}})
     class A:
         def g(self):
@@ -67,12 +90,9 @@ def test_decorator_actor(ray_start_cluster_head):
 
 @pytest.mark.skipif(
     sys.platform == "win32", reason="runtime_env unsupported on Windows.")
-def test_decorator_complex(shutdown_only):
-    ray.init(
-        job_config=ray.job_config.JobConfig(
-            runtime_env={"env_vars": {
-                "foo": "job"
-            }}))
+def test_decorator_complex(start_cluster):
+    cluster, address = start_cluster
+    ray.init(address, runtime_env={"env_vars": {"foo": "job"}})
 
     @ray.remote
     def env_from_job():

python/ray/util/client/worker.py

@@ -724,6 +724,7 @@ class Worker:
                 "object_store_memory": md.object_store_memory,
                 "resources": md.resources,
                 "accelerator_type": md.accelerator_type,
+                "runtime_env": md.runtime_env
             })
         return key
 
@@ -741,7 +742,8 @@ class Worker:
                 "resources": func._resources,
                 "accelerator_type": func._accelerator_type,
                 "num_returns": func._num_returns,
-                "memory": func._memory
+                "memory": func._memory,
+                "runtime_env": func._runtime_env,
             })
         return key