diff --git a/python/ray/_private/ray_constants.py b/python/ray/_private/ray_constants.py index ec370a63c4f6..3454d78af045 100644 --- a/python/ray/_private/ray_constants.py +++ b/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"] diff --git a/python/ray/_private/runtime_env/context.py b/python/ray/_private/runtime_env/context.py index be77d5b8d098..e51780b3e016 100644 --- a/python/ray/_private/runtime_env/context.py +++ b/python/ray/_private/runtime_env/context.py @@ -64,6 +64,7 @@ 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 {} @@ -71,6 +72,7 @@ def __init__( 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__) diff --git a/python/ray/_private/runtime_env/validation.py b/python/ray/_private/runtime_env/validation.py index ba1efd0db376..a48a1d62c7e7 100644 --- a/python/ray/_private/runtime_env/validation.py +++ b/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 diff --git a/python/ray/_private/utils.py b/python/ray/_private/utils.py index 7e7755a6224a..f6f9b14c53e7 100644 --- a/python/ray/_private/utils.py +++ b/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 diff --git a/python/ray/_private/workers/default_worker.py b/python/ray/_private/workers/default_worker.py index 8a082d502302..9a9d95e42922 100644 --- a/python/ray/_private/workers/default_worker.py +++ b/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]: diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 270a7e7850d8..f1a09a7c3d27 100644 --- a/python/ray/_raylet.pyx +++ b/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: diff --git a/python/ray/autoscaler/_private/constants.py b/python/ray/autoscaler/_private/constants.py index a997a48b06ec..da0751caf2f9 100644 --- a/python/ray/autoscaler/_private/constants.py +++ b/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. [ diff --git a/python/ray/runtime_env/runtime_env.py b/python/ray/runtime_env/runtime_env.py index a69acb5fd7d6..373d04b27dd4 100644 --- a/python/ray/runtime_env/runtime_env.py +++ b/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. diff --git a/python/ray/tests/unit/test_runtime_env_validation.py b/python/ray/tests/unit/test_runtime_env_validation.py index b95182b9972b..8ada737f6215 100644 --- a/python/ray/tests/unit/test_runtime_env_validation.py +++ b/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) diff --git a/python/ray/tests/unit/test_utils.py b/python/ray/tests/unit/test_utils.py new file mode 100644 index 000000000000..c88a565e20a4 --- /dev/null +++ b/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__])) diff --git a/src/ray/common/ray_config.cc b/src/ray/common/ray_config.cc index 9a78bc62ee22..80a90ed44eb3 100644 --- a/src/ray/common/ray_config.cc +++ b/src/ray/common/ray_config.cc @@ -33,6 +33,11 @@ void RayConfig::initialize(const std::string &config_list) { name##_ = ReadEnv("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()) { diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h index 2beeb0450eef..5d3cb8812d16 100644 --- a/src/ray/common/ray_config_def.h +++ b/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) diff --git a/src/ray/common/ray_internal_flag_def.h b/src/ray/common/ray_internal_flag_def.h index ed478ef76216..7107fb42b66a 100644 --- a/src/ray/common/ray_internal_flag_def.h +++ b/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. diff --git a/src/ray/common/tests/BUILD.bazel b/src/ray/common/tests/BUILD.bazel index efbd779bf5e9..387afbc5cf6b 100644 --- a/src/ray/common/tests/BUILD.bazel +++ b/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", ], ) diff --git a/src/ray/common/tests/id_test.cc b/src/ray/common/tests/id_test.cc index 3cbf782eb5a3..9efd5e4fc7bd 100644 --- a/src/ray/common/tests/id_test.cc +++ b/src/ray/common/tests/id_test.cc @@ -15,10 +15,17 @@ #include "ray/common/id.h" #include +#ifndef _WIN32 +#include +#include +#endif +#include +#include #include #include "absl/container/flat_hash_set.h" +#include "absl/random/random.h" namespace ray { @@ -124,6 +131,74 @@ TEST(ObjectIDTest, TestObjectID) { } } +#ifndef _WIN32 +TEST(ObjectIDTest, TestFromRandomForkSafety) { + // Initialize both Abseil's process-global RandenPool and FillRandom's thread-local + // PRNG before fork(), so any reliance on cloned RandenPool state across sibling + // forked processes or newly spawned child threads will collide. + (void)absl::BitGen{}; + const ObjectID pre_fork_id = ObjectID::FromRandom(); + ASSERT_FALSE(pre_fork_id.IsNil()); + + constexpr int kNumChildren = 4; + int pipefds[kNumChildren][2]; + pid_t child_pids[kNumChildren]; + + for (int i = 0; i < kNumChildren; ++i) { + ASSERT_EQ(pipe(pipefds[i]), 0); + pid_t pid = fork(); + ASSERT_GE(pid, 0); + if (pid == 0) { + for (int j = 0; j <= i; ++j) { + close(pipefds[j][0]); + } + const std::string main_binary = ObjectID::FromRandom().Binary(); + std::string thread_binary; + std::thread child_thread( + [&thread_binary]() { thread_binary = ObjectID::FromRandom().Binary(); }); + child_thread.join(); + + const std::string payload = main_binary + thread_binary; + ssize_t written = write(pipefds[i][1], payload.data(), payload.size()); + close(pipefds[i][1]); + _exit(written == static_cast(payload.size()) ? 0 : 1); + } + close(pipefds[i][1]); + child_pids[i] = pid; + } + + const ObjectID parent_id = ObjectID::FromRandom(); + absl::flat_hash_set seen_ids = {pre_fork_id, parent_id}; + ASSERT_EQ(seen_ids.size(), 2u); + + for (int i = 0; i < kNumChildren; ++i) { + std::string buf(2 * ObjectID::Size(), '\0'); + size_t total_read = 0; + while (total_read < buf.size()) { + ssize_t n = read(pipefds[i][0], buf.data() + total_read, buf.size() - total_read); + if (n <= 0) { + break; + } + total_read += static_cast(n); + } + close(pipefds[i][0]); + + int status = 0; + ASSERT_EQ(waitpid(child_pids[i], &status, 0), child_pids[i]); + ASSERT_TRUE(WIFEXITED(status)); + ASSERT_EQ(WEXITSTATUS(status), 0); + ASSERT_EQ(total_read, buf.size()); + + const ObjectID child_main_id = ObjectID::FromBinary(buf.substr(0, ObjectID::Size())); + const ObjectID child_thread_id = + ObjectID::FromBinary(buf.substr(ObjectID::Size(), ObjectID::Size())); + ASSERT_TRUE(seen_ids.insert(child_main_id).second); + ASSERT_TRUE(seen_ids.insert(child_thread_id).second); + } + ASSERT_EQ(seen_ids.size(), static_cast(2 + 2 * kNumChildren)); +} +#endif + TEST(NilTest, TestIsNil) { ASSERT_TRUE(TaskID().IsNil()); ASSERT_TRUE(TaskID::Nil().IsNil()); diff --git a/src/ray/common/tests/ray_config_test.cc b/src/ray/common/tests/ray_config_test.cc index 5584ab647eb8..7a0ae3c99b79 100644 --- a/src/ray/common/tests/ray_config_test.cc +++ b/src/ray/common/tests/ray_config_test.cc @@ -19,6 +19,7 @@ #include "gtest/gtest.h" #include "ray/common/grpc_util.h" +#include "ray/util/env.h" namespace ray { class RayConfigTest : public ::testing::Test {}; @@ -31,4 +32,18 @@ TEST_F(RayConfigTest, ConvertValueTrimsVectorElements) { ASSERT_EQ(output, expected_output); } +TEST_F(RayConfigTest, InitializeRereadsJobIdFromEnv) { + SetEnv("RAY_JOB_ID", "01000000"); + RayConfig::instance().initialize(""); + EXPECT_EQ(RayConfig::instance().JOB_ID(), "01000000"); + + SetEnv("RAY_JOB_ID", "02000000"); + RayConfig::instance().initialize(""); + EXPECT_EQ(RayConfig::instance().JOB_ID(), "02000000"); + + UnsetEnv("RAY_JOB_ID"); + RayConfig::instance().initialize(""); + EXPECT_EQ(RayConfig::instance().JOB_ID(), ""); +} + } // namespace ray diff --git a/src/ray/raylet/agent_manager.cc b/src/ray/raylet/agent_manager.cc index f24897d59ede..a8507b91600b 100644 --- a/src/ray/raylet/agent_manager.cc +++ b/src/ray/raylet/agent_manager.cc @@ -20,6 +20,7 @@ #include #include +#include "absl/strings/str_join.h" #include "ray/common/ray_config.h" #include "ray/util/logging.h" #include "ray/util/process.h" @@ -56,6 +57,12 @@ void AgentManager::StartAgent(AddProcessToCgroupHook add_to_cgroup) { RayConfig::instance().enable_pipe_based_agent_to_parent_health_check() ? "1" : "0"}); + env.insert({"RAY_enable_worker_fork_server", + RayConfig::instance().enable_worker_fork_server() ? "1" : "0"}); + if (!RayConfig::instance().preload_python_modules().empty()) { + env.insert({"RAY_preload_python_modules", + absl::StrJoin(RayConfig::instance().preload_python_modules(), ",")}); + } // Launch the process to create the agent. std::error_code ec; diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 93323ba7b2c1..d887841cd251 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -878,17 +878,25 @@ Status WorkerPool::RegisterWorker(const std::shared_ptr &worker std::unique_ptr process = std::make_unique(pid); worker->SetProcess(std::move(process)); #if !defined(_WIN32) - // Save the worker's actual PGID at registration for safe cleanup later. - // If setpgrp() succeeded in the child, pgid will equal pid; otherwise it will be the - // parent's PGID. We save whatever the OS reports and will validate again at cleanup. - pid_t pgid = -1; - errno = 0; - pgid = getpgid(pid); - if (pgid != -1) { - worker->SetSavedProcessGroupId(pgid); - } else { - RAY_LOG(WARNING) << "getpgid(" << pid - << ") failed at registration: " << strerror(errno); + if (RayConfig::instance().process_group_cleanup_enabled()) { + // Save the worker's actual PGID at registration for safe cleanup later. + // Only save if the worker is its own process group leader (pgid == pid), so cleanup + // never signals the raylet's, fork server's, or runtime_env_agent's process group. + pid_t pgid = -1; + errno = 0; + pgid = getpgid(pid); + if (pgid != -1) { + if (pgid == pid) { + worker->SetSavedProcessGroupId(pgid); + } else { + RAY_LOG(WARNING) << "Worker " << pid << " pgid (" << pgid + << ") != pid; skipping SetSavedProcessGroupId to avoid " + << "signaling parent process group."; + } + } else { + RAY_LOG(WARNING) << "getpgid(" << pid + << ") failed at registration: " << strerror(errno); + } } #endif diff --git a/src/ray/util/random.h b/src/ray/util/random.h index c21b19f08531..650048f550b5 100644 --- a/src/ray/util/random.h +++ b/src/ray/util/random.h @@ -14,17 +14,63 @@ #pragma once +#include +#include #include +#include +#ifndef _WIN32 +#include +#include +#endif #include "absl/random/random.h" -/// A helper function to fill random bytes into the `data`. -/// Warning: this is not fork-safe, we need to re-seed after that. +namespace ray::internal { + +/** + * @brief Create an `absl::BitGen` seeded from OS entropy instead of Abseil's global + * `RandenPool`, which `fork()` copies into child processes. + */ +inline absl::BitGen CreateForkSafeBitGen() { +#ifdef _WIN32 + std::random_device rd; + std::seed_seq seq{rd(), rd(), rd(), rd(), rd(), rd(), rd(), rd()}; +#else + uint32_t seed[8] = {}; + const int fd = open("/dev/urandom", O_RDONLY | O_CLOEXEC); + RAY_CHECK(fd >= 0) << "Failed to open /dev/urandom"; + const ssize_t bytes_read = read(fd, seed, sizeof(seed)); + close(fd); + RAY_CHECK(bytes_read == static_cast(sizeof(seed))) + << "Failed to read seed from /dev/urandom"; + std::seed_seq seq(std::begin(seed), std::end(seed)); +#endif + return absl::BitGen(seq); +} + +} // namespace ray::internal + +/** + * @brief A helper function to fill random bytes into `data`. + * + * Seeds from OS entropy rather than Abseil's `RandenPool` and automatically reseeds + * the thread-local PRNG if the PID changes across `fork()`. + * + * @param[out] data Pointer to the byte container to fill with random bytes. + */ template void FillRandom(T *data) { RAY_CHECK(data != nullptr); - thread_local absl::BitGen generator; + thread_local absl::BitGen generator = ::ray::internal::CreateForkSafeBitGen(); +#ifndef _WIN32 + thread_local pid_t generator_pid = getpid(); + const pid_t current_pid = getpid(); + if (generator_pid != current_pid) { + generator = ::ray::internal::CreateForkSafeBitGen(); + generator_pid = current_pid; + } +#endif for (size_t i = 0; i < data->size(); i++) { (*data)[i] = static_cast( absl::Uniform(generator, 0, std::numeric_limits::max()));