fix(client): monitor proxier death for specific servers

Signed-off-by: JerryLee <223425819+Jerry2003826@users.noreply.github.com>

Author: Jerry2003826

python/ray/_private/services.py

@@ -750,7 +750,7 @@ def extract_ip_port(bootstrap_address: str):
     ip_port = parse_address(bootstrap_address)
     if ip_port is None:
         raise ValueError(
-            f"Malformed address {bootstrap_address}. " f"Expected '<host>:<port>'."
+            f"Malformed address {bootstrap_address}. Expected '<host>:<port>'."
         )
     ip, port = ip_port
     try:
@@ -759,8 +759,7 @@ def extract_ip_port(bootstrap_address: str):
         raise ValueError(f"Malformed address port {port}. Must be an integer.")
     if port < 1024 or port > 65535:
         raise ValueError(
-            f"Invalid address port {port}. Must be between 1024 "
-            "and 65535 (inclusive)."
+            f"Invalid address port {port}. Must be between 1024 and 65535 (inclusive)."
         )
     return ip, port
 
@@ -1053,21 +1052,35 @@ def preexec_fn():
         total_chrs = sum([len(x) for x in command])
         if total_chrs > 31766:
             raise ValueError(
-                f"command is limited to a total of 31767 characters, "
-                f"got {total_chrs}"
+                f"command is limited to a total of 31767 characters, got {total_chrs}"
             )
 
-    process = ConsolePopen(
-        command,
-        env=modified_env,
-        cwd=cwd,
-        stdout=stdout_file,
-        stderr=stderr_file,
-        stdin=subprocess.PIPE if pipe_stdin else None,
-        preexec_fn=(None if sys.platform == "win32" or use_posix_spawn else preexec_fn),
-        close_fds=close_fds,
-        creationflags=CREATE_SUSPENDED if win32_fate_sharing else 0,
+    previous_sigmask = None
+    should_block_sigint_for_spawn = (
+        use_posix_spawn
+        and hasattr(signal, "pthread_sigmask")
+        and hasattr(signal, "SIG_BLOCK")
+        and hasattr(signal, "SIG_SETMASK")
     )
+    if should_block_sigint_for_spawn:
+        previous_sigmask = signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+    try:
+        process = ConsolePopen(
+            command,
+            env=modified_env,
+            cwd=cwd,
+            stdout=stdout_file,
+            stderr=stderr_file,
+            stdin=subprocess.PIPE if pipe_stdin else None,
+            preexec_fn=(
+                None if sys.platform == "win32" or use_posix_spawn else preexec_fn
+            ),
+            close_fds=close_fds,
+            creationflags=CREATE_SUSPENDED if win32_fate_sharing else 0,
+        )
+    finally:
+        if previous_sigmask is not None:
+            signal.pthread_sigmask(signal.SIG_SETMASK, previous_sigmask)
 
     if win32_fate_sharing:
         try:
@@ -2495,9 +2508,12 @@ def start_ray_client_server(
     # multi-threaded gRPC server. Avoid a fork+preexec path there: gRPC may have
     # active poller threads and can skip fork handlers, leaving the child to
     # crash before it opens its channel. Specific servers self-terminate after
-    # being idle and are also cleaned up by the proxier, so they can trade
-    # kernel fate sharing and the preexec SIGINT mask for a fork-safe spawn path.
+    # being idle, monitor stdin EOF from setup_worker for abnormal parent death,
+    # and inherit a temporarily-blocked SIGINT mask from the spawning thread, so
+    # they can trade kernel fate sharing for a fork-safe spawn path.
     process_fate_share = False if use_posix_spawn else fate_share
+    if use_posix_spawn:
+        command.append("--monitor-parent-pipe")
 
     process_info = start_ray_process(
         command,
@@ -2506,6 +2522,7 @@ def start_ray_client_server(
         stderr_file=stderr_file,
         fate_share=process_fate_share,
         env_updates=env_updates,
+        pipe_stdin=use_posix_spawn,
         use_posix_spawn=use_posix_spawn,
     )
     return process_info

python/ray/_private/workers/setup_worker.py

@@ -1,5 +1,8 @@
 import argparse
 import logging
+import os
+import sys
+import threading
 
 from ray._private.ray_constants import LOGGER_FORMAT, LOGGER_LEVEL
 from ray._private.ray_logging import setup_logger
@@ -20,10 +23,43 @@
 
 parser.add_argument("--language", type=str, help="the language type of the worker")
 
+parser.add_argument(
+    "--monitor-parent-pipe",
+    required=False,
+    action="store_true",
+    help="Internal: exit when inherited stdin pipe reaches EOF.",
+)
+
+
+def _start_parent_pipe_monitor(enabled: bool, exit_func=os._exit):
+    if not enabled:
+        return None
+    stdin = getattr(sys.stdin, "buffer", sys.stdin)
+
+    def monitor_parent_pipe():
+        try:
+            while stdin.read(1):
+                pass
+        except Exception:
+            logger.exception("Failed while monitoring parent pipe.")
+        logger.info("Parent pipe closed; exiting worker setup process.")
+        exit_func(0)
+
+    thread = threading.Thread(
+        target=monitor_parent_pipe,
+        name="ray-worker-parent-pipe-monitor",
+        daemon=True,
+    )
+    thread.start()
+    return thread
+
 
 if __name__ == "__main__":
     setup_logger(LOGGER_LEVEL, LOGGER_FORMAT)
     args, remaining_args = parser.parse_known_args()
+    _start_parent_pipe_monitor(args.monitor_parent_pipe)
+    if args.monitor_parent_pipe:
+        remaining_args.append("--monitor-parent-pipe")
     # NOTE(edoakes): args.serialized_runtime_env_context is only None when
     # we're starting the main Ray client proxy server. That case should
     # probably not even go through this codepath.

python/ray/tests/test_debug_tools.py

@@ -1,4 +1,5 @@
 import os
+import signal
 import subprocess
 import sys
 from pathlib import Path
@@ -8,6 +9,7 @@
 import ray
 import ray._private.ray_constants as ray_constants
 import ray._private.services as services
+import ray.util.client.server.server as ray_client_server
 from ray._common.test_utils import wait_for_condition
 
 
@@ -230,10 +232,91 @@ def fake_start_ray_process(command, process_type, **kwargs):
         assert process_info is expected_process_info
         assert captured["process_type"] == ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER
         assert "--mode=specific-server" in captured["command"]
+        assert "--monitor-parent-pipe" in captured["command"]
         assert captured["kwargs"]["fate_share"] is False
+        assert captured["kwargs"]["pipe_stdin"] is True
         assert captured["kwargs"]["use_posix_spawn"] is True
 
 
+def test_setup_worker_parent_pipe_monitor_exits_on_eof(monkeypatch):
+    from ray._private.workers import setup_worker
+
+    class FakeStdin:
+        def read(self, size):
+            return ""
+
+    exit_codes = []
+    with monkeypatch.context() as m:
+        m.setattr(setup_worker.sys, "stdin", FakeStdin())
+        thread = setup_worker._start_parent_pipe_monitor(True, exit_codes.append)
+        assert thread is not None
+        thread.join(timeout=1)
+
+    assert exit_codes == [0]
+    assert setup_worker._start_parent_pipe_monitor(False, exit_codes.append) is None
+
+
+def test_ray_client_specific_server_parent_pipe_monitor_exits_on_eof(monkeypatch):
+    class FakeStdin:
+        def read(self, size):
+            return ""
+
+    exit_codes = []
+    with monkeypatch.context() as m:
+        m.setattr(ray_client_server.sys, "stdin", FakeStdin())
+        thread = ray_client_server._start_parent_pipe_monitor(True, exit_codes.append)
+        assert thread is not None
+        thread.join(timeout=1)
+
+    assert exit_codes == [0]
+    assert (
+        ray_client_server._start_parent_pipe_monitor(False, exit_codes.append) is None
+    )
+
+
+def test_ray_client_specific_server_blocks_sigint(monkeypatch):
+    calls = []
+    sig_block = object()
+
+    def fake_pthread_sigmask(how, signals):
+        calls.append((how, signals))
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            ray_client_server.signal,
+            "pthread_sigmask",
+            fake_pthread_sigmask,
+            raising=False,
+        )
+        m.setattr(ray_client_server.signal, "SIG_BLOCK", sig_block, raising=False)
+
+        ray_client_server._block_sigint_for_specific_server()
+
+    assert calls == [(sig_block, {signal.SIGINT})]
+
+
+def test_ray_client_specific_server_sigint_block_noops_without_posix_signal_support(
+    monkeypatch,
+):
+    calls = []
+
+    def fake_pthread_sigmask(how, signals):
+        calls.append((how, signals))
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            ray_client_server.signal,
+            "pthread_sigmask",
+            fake_pthread_sigmask,
+            raising=False,
+        )
+        m.delattr(ray_client_server.signal, "SIG_BLOCK", raising=False)
+
+        ray_client_server._block_sigint_for_specific_server()
+
+    assert calls == []
+
+
 def test_start_ray_process_posix_spawn_close_fds_when_supported(monkeypatch):
     captured = {}
     expected_process = object()
@@ -260,6 +343,46 @@ def fake_console_popen(command, **kwargs):
         assert captured["kwargs"]["close_fds"] is True
 
 
+def test_start_ray_process_posix_spawn_blocks_sigint_for_child(monkeypatch):
+    captured = {}
+    expected_process = object()
+    calls = []
+    previous_mask = {signal.SIGTERM}
+    sig_block = object()
+    sig_setmask = object()
+
+    def fake_console_popen(command, **kwargs):
+        captured["command"] = command
+        captured["kwargs"] = kwargs
+        return expected_process
+
+    def fake_pthread_sigmask(how, signals):
+        calls.append((how, signals))
+        return previous_mask
+
+    with monkeypatch.context() as m:
+        m.setattr(services.sys, "platform", "linux")
+        m.setattr(services.os, "POSIX_SPAWN_CLOSEFROM", object(), raising=False)
+        m.setattr(services.signal, "pthread_sigmask", fake_pthread_sigmask)
+        m.setattr(services.signal, "SIG_BLOCK", sig_block, raising=False)
+        m.setattr(services.signal, "SIG_SETMASK", sig_setmask, raising=False)
+        m.setattr(services, "ConsolePopen", fake_console_popen)
+
+        process_info = services.start_ray_process(
+            [sys.executable],
+            ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER,
+            fate_share=False,
+            use_posix_spawn=True,
+        )
+
+    assert process_info.process is expected_process
+    assert captured["kwargs"]["preexec_fn"] is None
+    assert calls == [
+        (sig_block, {signal.SIGINT}),
+        (sig_setmask, previous_mask),
+    ]
+
+
 def test_start_ray_process_posix_spawn_leaves_fds_open_for_older_runtime(
     monkeypatch,
 ):

python/ray/util/client/server/server.py

@@ -8,6 +8,8 @@
 import os
 import pickle
 import queue
+import signal
+import sys
 import threading
 import time
 from collections import defaultdict
@@ -851,6 +853,34 @@ def ray_connect_handler(job_config: JobConfig = None, **ray_init_kwargs):
     return ray_connect_handler
 
 
+def _block_sigint_for_specific_server() -> None:
+    if hasattr(signal, "pthread_sigmask") and hasattr(signal, "SIG_BLOCK"):
+        signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+
+
+def _start_parent_pipe_monitor(enabled: bool, exit_func=os._exit):
+    if not enabled:
+        return None
+    stdin = getattr(sys.stdin, "buffer", sys.stdin)
+
+    def monitor_parent_pipe():
+        try:
+            while stdin.read(1):
+                pass
+        except Exception:
+            logger.exception("Failed while monitoring specific-server parent pipe.")
+        logger.info("Specific Ray Client server parent pipe closed; shutting down.")
+        exit_func(0)
+
+    thread = threading.Thread(
+        target=monitor_parent_pipe,
+        name="ray-client-specific-server-parent-monitor",
+        daemon=True,
+    )
+    thread.start()
+    return thread
+
+
 def try_create_gcs_client(address: Optional[str]) -> Optional[GcsClient]:
     """
     Try to create a gcs client based on the command line args or by
@@ -900,9 +930,20 @@ def main():
         default=None,
         help="The hex ID of this node.",
     )
+    parser.add_argument(
+        "--monitor-parent-pipe",
+        required=False,
+        action="store_true",
+        help="Internal: exit when inherited stdin pipe reaches EOF.",
+    )
     args, _ = parser.parse_known_args()
     redis_password = os.environ.get(ray_constants.RAY_REDIS_PASSWORD_ENV)
     setup_logger(ray_constants.LOGGER_LEVEL, ray_constants.LOGGER_FORMAT)
+    _start_parent_pipe_monitor(
+        args.mode == "specific-server" and args.monitor_parent_pipe
+    )
+    if args.mode == "specific-server":
+        _block_sigint_for_specific_server()
 
     ray_connect_handler = create_ray_handler(
         args.address, redis_password, args.redis_username