[Core] Shell escape worker process command in RuntimeEnvContext.exec_worker (ray-project#42332)

RuntimeEnvContext.exec_worker used " ".join(cmds) to construct the worker process command but it didn't do any shell escape. This will cause the worker process fail to start if there is any special character (e.g. ?) in the command. Instead, we should use shlex.join.

Signed-off-by: Jiajun Yao <jeromeyjj@gmail.com>

Author: jjyao

python/ray/_private/runtime_env/context.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import subprocess
+import shlex
 import sys
 from typing import Any, Dict, List, Optional
 
@@ -48,11 +49,11 @@ def exec_worker(self, passthrough_args: List[str], language: Language):
         update_envs(self.env_vars)
 
         if language == Language.PYTHON and sys.platform == "win32":
-            executable = self.py_executable
+            executable = [self.py_executable]
         elif language == Language.PYTHON:
-            executable = f"exec {self.py_executable}"
+            executable = ["exec", self.py_executable]
         elif language == Language.JAVA:
-            executable = "java"
+            executable = ["java"]
             ray_jars = os.path.join(get_ray_jars_dir(), "*")
 
             local_java_jars = []
@@ -63,9 +64,9 @@ def exec_worker(self, passthrough_args: List[str], language: Language):
             class_path_args = ["-cp", ray_jars + ":" + str(":".join(local_java_jars))]
             passthrough_args = class_path_args + passthrough_args
         elif sys.platform == "win32":
-            executable = ""
+            executable = []
         else:
-            executable = "exec "
+            executable = ["exec"]
 
         # By default, raylet uses the path to default_worker.py on host.
         # However, the path to default_worker.py inside the container
@@ -79,29 +80,29 @@ def exec_worker(self, passthrough_args: List[str], language: Language):
             )
             passthrough_args[0] = default_worker_path
 
-        passthrough_args = [s.replace(" ", r"\ ") for s in passthrough_args]
-        exec_command = " ".join([f"{executable}"] + passthrough_args)
-        command_str = " ".join(self.command_prefix + [exec_command])
-        # TODO(SongGuyang): We add this env to command for macOS because it doesn't
-        # work for the C++ process of `os.execvp`. We should find a better way to
-        # fix it.
-        MACOS_LIBRARY_PATH_ENV_NAME = "DYLD_LIBRARY_PATH"
-        if MACOS_LIBRARY_PATH_ENV_NAME in os.environ:
-            command_str = (
-                MACOS_LIBRARY_PATH_ENV_NAME
-                + "="
-                + os.environ.get(MACOS_LIBRARY_PATH_ENV_NAME)
-                + " "
-                + command_str
-            )
-        logger.debug(f"Exec'ing worker with command: {command_str}")
         if sys.platform == "win32":
-            cmd = [*self.command_prefix, executable, *passthrough_args]
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            logger.debug(f"Exec'ing worker with command: {cmd}")
             subprocess.Popen(cmd, shell=True).wait()
         else:
+            # We use shlex to do the necessary shell escape
+            # of special characters in passthrough_args.
+            passthrough_args = [shlex.quote(s) for s in passthrough_args]
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            # TODO(SongGuyang): We add this env to command for macOS because it doesn't
+            # work for the C++ process of `os.execvp`. We should find a better way to
+            # fix it.
+            MACOS_LIBRARY_PATH_ENV_NAME = "DYLD_LIBRARY_PATH"
+            if MACOS_LIBRARY_PATH_ENV_NAME in os.environ:
+                cmd.insert(
+                    0,
+                    f"{MACOS_LIBRARY_PATH_ENV_NAME}="
+                    f"{os.environ[MACOS_LIBRARY_PATH_ENV_NAME]}",
+                )
+            logger.debug(f"Exec'ing worker with command: {cmd}")
             # PyCharm will monkey patch the os.execvp at
             # .pycharm_helpers/pydev/_pydev_bundle/pydev_monkey.py
             # The monkey patched os.execvp function has a different
             # signature. So, we use os.execvp("executable", args=[])
             # instead of os.execvp(file="executable", args=[])
-            os.execvp("bash", args=["bash", "-c", command_str])
+            os.execvp("bash", args=["bash", "-c", " ".join(cmd)])

python/ray/_private/services.py

@@ -1661,7 +1661,6 @@ def start_raylet(
             f"--object-store-name={plasma_store_name}",
             f"--raylet-name={raylet_name}",
             f"--redis-address={redis_address}",
-            f"--temp-dir={temp_dir}",
             f"--metrics-agent-port={metrics_agent_port}",
             f"--runtime-env-agent-port={runtime_env_agent_port}",
             f"--logging-rotate-bytes={max_bytes}",

python/ray/tests/test_storage.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import subprocess
 import urllib
 from pathlib import Path
@@ -64,6 +65,24 @@ def test_get_filesystem_s3(shutdown_only):
         assert isinstance(fs, pyarrow.fs.S3FileSystem), fs
 
 
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="The issue is not fixed for windows yet"
+)
+def test_escape_storage_uri_with_runtime_env(shutdown_only):
+    # https://github.com/ray-project/ray/issues/41568
+    # Test to make sure we can successfully start worker process
+    # when storage uri contains ? and we use runtime env.
+    with simulate_storage("s3") as s3_uri:
+        assert "?" in s3_uri
+        ray.init(storage=s3_uri, runtime_env={"env_vars": {"TEST_ENV": "1"}})
+
+        @ray.remote
+        def f():
+            return 1
+
+        assert ray.get(f.remote()) == 1
+
+
 def test_get_filesystem_invalid(shutdown_only, tmp_path):
     with pytest.raises(pyarrow.lib.ArrowInvalid):
         ray.init(storage="blahblah://bad")