Skip to content

Commit

Permalink
[Core] Shell escape worker process command in RuntimeEnvContext.exec_…
Browse files Browse the repository at this point in the history
…worker (#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>
  • Loading branch information
jjyao authored Jan 12, 2024
1 parent 943490a commit 3a306ef
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 24 deletions.
47 changes: 24 additions & 23 deletions python/ray/_private/runtime_env/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import os
import subprocess
import shlex
import sys
from typing import Any, Dict, List, Optional

Expand Down Expand Up @@ -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 = []
Expand All @@ -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
Expand All @@ -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)])
1 change: 0 additions & 1 deletion python/ray/_private/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}",
Expand Down
19 changes: 19 additions & 0 deletions python/ray/tests/test_storage.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import sys
import subprocess
import urllib
from pathlib import Path
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 3a306ef

Please sign in to comment.