[Core] Add initial foundational configurations and flags for fork server

python/ray/_private/ray_constants.py

@@ -813,3 +813,7 @@ def get_num_threads_env_vars() -> list[str]:
 
 # Name of the environment variable for the Redis password.
 RAY_REDIS_PASSWORD_ENV = "RAY_REDIS_PASSWORD"
+
+# Modules that initialize an internal thread pool at import time and require an
+# explicit .set_num_threads() call if imported before fork().
+THREAD_POOL_MODULES = ["torch"]

python/ray/_private/runtime_env/context.py

@@ -64,13 +64,15 @@ def __init__(
         override_worker_entrypoint: Optional[str] = None,
         java_jars: List[str] = None,
         container: Optional[Dict[str, Any]] = None,
+        fork_server_socket_path: Optional[str] = None,
     ):
         self.command_prefix = command_prefix or []
         self.env_vars = env_vars or {}
         self.py_executable = py_executable or sys.executable
         self.override_worker_entrypoint: Optional[str] = override_worker_entrypoint
         self.java_jars = java_jars or []
         self.container = container or {}
+        self.fork_server_socket_path: Optional[str] = fork_server_socket_path
 
     def serialize(self) -> str:
         return json.dumps(self.__dict__)

python/ray/_private/runtime_env/validation.py

@@ -1,3 +1,4 @@
+import keyword
 import logging
 import sys
 from collections import OrderedDict
@@ -452,6 +453,38 @@ def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str,
     return env_vars
 
 
+def parse_and_validate_preload_modules(
+    preload_modules: List[str],
+) -> Optional[List[str]]:
+    """Parses and validates a user-provided 'preload_modules' option."""
+    assert preload_modules is not None
+    if not isinstance(preload_modules, list):
+        raise TypeError(
+            "runtime_env['preload_modules'] must be of type List[str], "
+            f"got {type(preload_modules)}"
+        )
+    if len(preload_modules) == 0:
+        return None
+    validated: List[str] = []
+    for mod in preload_modules:
+        if not isinstance(mod, str):
+            raise TypeError(
+                "Elements of runtime_env['preload_modules'] must be of type str, "
+                f"got {type(mod)}"
+            )
+        stripped = mod.strip()
+        if not stripped or not all(
+            part.isidentifier() and not keyword.iskeyword(part)
+            for part in stripped.split(".")
+        ):
+            raise ValueError(
+                f"Invalid module name '{mod}' in runtime_env['preload_modules']."
+            )
+        if stripped not in validated:
+            validated.append(stripped)
+    return validated
+
+
 # Dictionary mapping runtime_env options with the function to parse and
 # validate them.
 OPTION_TO_VALIDATION_FN = {
@@ -463,6 +496,7 @@ def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str,
     "uv": parse_and_validate_uv,
     "env_vars": parse_and_validate_env_vars,
     "container": parse_and_validate_container,
+    "preload_modules": parse_and_validate_preload_modules,
 }
 
 # RuntimeEnv can be created with local paths

python/ray/_private/utils.py

@@ -277,6 +277,16 @@ def set_num_threads_env_vars_if_unset() -> List[str]:
     num_threads = str(max(math.floor(num_assigned_cpus), 1))
     for name in unset_env_vars:
         os.environ[name] = num_threads
+    if "OMP_NUM_THREADS" in unset_env_vars:
+        # If modules like torch were pre-imported before fork(), their internal thread
+        # pool is already initialized and ignores subsequent os.environ["OMP_NUM_THREADS"]
+        # updates.
+        for module_name in ray_constants.THREAD_POOL_MODULES:
+            if module_name in sys.modules:
+                try:
+                    sys.modules[module_name].set_num_threads(int(num_threads))
+                except Exception:
+                    pass
     return unset_env_vars
 
 

python/ray/_private/workers/default_worker.py

@@ -200,12 +200,9 @@
     ),
 )
 
-if __name__ == "__main__":
-    # NOTE(sang): For some reason, if we move the code below
-    # to a separate function, tensorflow will capture that method
-    # as a step function. For more details, check out
-    # https://github.com/ray-project/ray/pull/12225#issue-525059663.
-    args = parser.parse_args()
+
+def initialize_default_worker(argv=None):
+    args = parser.parse_args(argv)
     ray._private.ray_logging.setup_logger(args.logging_level, args.logging_format)
     worker_launched_time_ms = time.time_ns() // 1e6
     if args.worker_type == "WORKER":
@@ -319,6 +316,15 @@
         if error is not None:
             worker.core_worker.drain_and_exit_worker("system", error)
 
+    return mode, worker
+
+
+if __name__ == "__main__":
+    # NOTE(sang): For some reason, if we move the main_loop call below
+    # to a separate function, tensorflow will capture that method
+    # as a step function. For more details, check out
+    # https://github.com/ray-project/ray/pull/12225#issue-525059663.
+    mode, worker = initialize_default_worker()
     if mode == ray.WORKER_MODE:
         worker.main_loop()
     elif mode in [ray.RESTORE_WORKER_MODE, ray.SPILL_WORKER_MODE]:

python/ray/_raylet.pyx

@@ -2591,6 +2591,15 @@ cdef execute_task_with_cancellation_handler(
             # Reset the num threads env vars if they were set.
             for env_var in num_threads_env_vars_set:
                 os.environ.pop(env_var, None)
+            if "OMP_NUM_THREADS" in num_threads_env_vars_set:
+                # Restore default single-thread fallback (set_num_threads_env_vars_if_unset()
+                # sets OMP_NUM_THREADS when unset) if modules were imported pre-fork.
+                for module_name in ray_constants.THREAD_POOL_MODULES:
+                    if module_name in sys.modules:
+                        try:
+                            sys.modules[module_name].set_num_threads(1)
+                        except Exception:
+                            pass
 
 
     if execution_info.max_calls != 0:

python/ray/autoscaler/_private/constants.py

@@ -137,6 +137,7 @@ def env_integer(key, default):
     ["ray.util.client.server", False],
     ["default_worker.py", False],  # Python worker.
     ["setup_worker.py", False],  # Python environment setup worker.
+    ["fork_server.py", False],  # Python worker fork server.
     # For mac osx, setproctitle doesn't change the process name returned
     # by psutil but only cmdline.
     [

python/ray/runtime_env/runtime_env.py

@@ -321,6 +321,8 @@ class MyClass:
             or with the ``config`` or ``env_vars`` parameters.
         uv: Either a list of pip packages, or a Python dictionary that has one field:
             1) ``packages`` (required, List[str]).
+        preload_modules: (Experimental) List of Python module names to preload in the
+            worker fork server before forking worker processes.
         **kwargs: Additional runtime environment fields that are forwarded
             verbatim to the underlying runtime environment dictionary.
     """
@@ -344,6 +346,7 @@ class MyClass:
         "_nsight",
         "_rocprof_sys",
         "image_uri",
+        "preload_modules",
     }
 
     extensions_fields: Set[str] = {
@@ -369,6 +372,7 @@ def __init__(
         _validate: bool = True,
         image_uri: Optional[str] = None,
         uv: Optional[List[str]] = None,
+        preload_modules: Optional[List[str]] = None,
         **kwargs,
     ):
         super().__init__()
@@ -400,6 +404,8 @@ def __init__(
             runtime_env["worker_process_setup_hook"] = worker_process_setup_hook
         if image_uri is not None:
             runtime_env["image_uri"] = image_uri
+        if preload_modules is not None:
+            runtime_env["preload_modules"] = preload_modules
 
         self.update(runtime_env)
 
@@ -847,6 +853,16 @@ def image_uri(self) -> Optional[str]:
         """
         return self.get("image_uri")
 
+    def preload_modules(self) -> List[str]:
+        """Return the entries of the ``preload_modules`` field.
+
+        Returns:
+            A list of the ``preload_modules`` entries, or an empty list if it isn't set.
+        """
+        if "preload_modules" in self:
+            return list(self["preload_modules"])
+        return []
+
     def plugins(self) -> List[Tuple[str, Any]]:
         """Return the fields that aren't built-in runtime environment fields.
 

python/ray/tests/unit/test_runtime_env_validation.py

@@ -857,6 +857,33 @@ def test_type_validation(self):
             validation.parse_and_validate_env_vars({1.23: "hi"})
 
 
+class TestValidatePreloadModules:
+    @pytest.mark.parametrize(
+        "preload_modules,expected_result,expected_error",
+        [
+            (["os", "sys", "urllib.parse"], ["os", "sys", "urllib.parse"], None),
+            ([" os ", "os", "sys"], ["os", "sys"], None),
+            ([], [], None),
+            ("os", None, TypeError),
+            ([123], None, TypeError),
+            ([""], None, ValueError),
+            (["123invalid"], None, ValueError),
+            (["bad..module"], None, ValueError),
+            (["import"], None, ValueError),
+            (["foo.class"], None, ValueError),
+        ],
+    )
+    def test_preload_modules_validation(
+        self, preload_modules, expected_result, expected_error
+    ):
+        if expected_error is None:
+            env = RuntimeEnv(preload_modules=preload_modules)
+            assert env.preload_modules() == expected_result
+        else:
+            with pytest.raises(expected_error):
+                RuntimeEnv(preload_modules=preload_modules)
+
+
 def test_validate_no_local_paths_raises_exceptions_on_type_mismatch():
     with pytest.raises(TypeError):
         _validate_no_local_paths(1)

python/ray/tests/unit/test_utils.py

@@ -0,0 +1,55 @@
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import ray
+from ray._private.utils import set_num_threads_env_vars_if_unset
+
+
+@pytest.mark.parametrize(
+    "num_threads_env_vars,preset_omp,expected_ret,expected_torch_threads",
+    [
+        (None, None, ["OMP_NUM_THREADS"], 4),
+        ("OMP_NUM_THREADS,RAYON_NUM_THREADS", "2", ["RAYON_NUM_THREADS"], None),
+    ],
+)
+def test_set_num_threads_env_vars_if_unset_updates_thread_pool_modules(
+    monkeypatch,
+    num_threads_env_vars,
+    preset_omp,
+    expected_ret,
+    expected_torch_threads,
+):
+    """Test that set_num_threads_env_vars_if_unset calls .set_num_threads() on loaded
+    THREAD_POOL_MODULES only when OMP_NUM_THREADS is set by Ray.
+    """
+    if num_threads_env_vars is None:
+        monkeypatch.delenv("RAY_NUM_THREADS_ENV_VARS", raising=False)
+    else:
+        monkeypatch.setenv("RAY_NUM_THREADS_ENV_VARS", num_threads_env_vars)
+    if preset_omp is None:
+        monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
+    else:
+        monkeypatch.setenv("OMP_NUM_THREADS", preset_omp)
+    monkeypatch.delenv("RAYON_NUM_THREADS", raising=False)
+
+    mock_ctx = MagicMock()
+    mock_ctx.worker.mode = ray._private.worker.WORKER_MODE
+    mock_ctx.get_assigned_resources.return_value = {"CPU": 4}
+    mock_torch = MagicMock()
+    with patch("ray.get_runtime_context", return_value=mock_ctx), patch.dict(
+        sys.modules, {"torch": mock_torch}
+    ):
+        assert set_num_threads_env_vars_if_unset() == expected_ret
+        if expected_torch_threads is not None:
+            assert os.environ["OMP_NUM_THREADS"] == str(expected_torch_threads)
+            mock_torch.set_num_threads.assert_called_once_with(expected_torch_threads)
+        else:
+            assert os.environ["OMP_NUM_THREADS"] == preset_omp
+            mock_torch.set_num_threads.assert_not_called()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

src/ray/common/ray_config.cc

@@ -33,6 +33,11 @@ void RayConfig::initialize(const std::string &config_list) {
   name##_ = ReadEnv<type>("RAY_" #name, #type, default_value);
 
 #include "ray/common/ray_config_def.h"
+// Re-read internal flags (e.g., RAY_JOB_ID) when initialize() is called
+// post-fork in a worker process after RayConfig was already constructed.
+#define RAY_INTERNAL_FLAG RAY_CONFIG
+#include "ray/common/ray_internal_flag_def.h"
+#undef RAY_INTERNAL_FLAG
 #undef RAY_CONFIG
 
   if (config_list.empty()) {

src/ray/common/ray_config_def.h

@@ -1300,3 +1300,11 @@ RAY_CONFIG(uint64_t, gcs_resource_broadcast_max_batch_delay_ms, 0)
 // Whether to enable/disable multiple gRPC connections to improve object transfer
 // throughput.
 RAY_CONFIG(bool, experimental_object_manager_enable_multiple_connections, true)
+
+/// If true, Python workers are forked from a warm fork server process
+/// instead of starting a new Python interpreter from scratch.
+RAY_CONFIG(bool, enable_worker_fork_server, false)
+
+/// Timeout in milliseconds for a fork server socket request before falling back to
+/// starting a fresh worker process.
+RAY_CONFIG(int64_t, worker_fork_server_timeout_ms, 500)

src/ray/common/ray_internal_flag_def.h

@@ -17,10 +17,11 @@
 // RAY_{VARIABLE NAME}.
 //
 // The format is used to avoid code duplication.
-// It can be included multiple times in ray_config.h, and each inclusion
-// could use a different definition of the RAY_INTERNAL_FLAG macro.
+// It can be included multiple times in ray_config.h and ray_config.cc, and each
+// inclusion could use a different definition of the RAY_INTERNAL_FLAG macro.
 // Macro definition format: RAY_INTERNAL_FLAG(type, name, default_value).
-// NOTE: This file should NOT be included in any file other than ray_config.h.
+// NOTE: This file should NOT be included in any file other than ray_config.h
+// and ray_config.cc.
 
 // WARNING: if you update this file, please also update RAY_INTERNAL_FLAGS in
 // ray_constants.py.

src/ray/common/tests/BUILD.bazel

@@ -18,6 +18,7 @@ ray_cc_test(
     deps = [
         "//src/ray/common:grpc_util",
         "//src/ray/common:ray_config",
+        "//src/ray/util:env",
     ],
 )
 
@@ -40,6 +41,7 @@ ray_cc_test(
     deps = [
         "//src/ray/common:id",
         "//src/ray/protobuf:common_cc_proto",
+        "@com_google_absl//absl/random",
     ],
 )