[Core] Add worker fork server daemon and post-fork child setup

python/ray/_private/utils.py

@@ -281,15 +281,32 @@ def set_num_threads_env_vars_if_unset() -> List[str]:
         # 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
+        set_thread_pool_module_threads(int(num_threads))
     return unset_env_vars
 
 
+def set_thread_pool_module_threads(num_threads: int) -> None:
+    """Set the thread count of already-imported THREAD_POOL_MODULES (e.g. torch).
+
+    These modules size their thread pool at import time and ignore later changes to
+    OMP_NUM_THREADS, so modules imported before fork() must be resized explicitly.
+
+    Args:
+        num_threads: The number of threads to set on each imported module.
+    """
+    for module_name in ray_constants.THREAD_POOL_MODULES:
+        module = sys.modules.get(module_name)
+        if module is not None:
+            try:
+                module.set_num_threads(num_threads)
+            except Exception:
+                pass
+
+
+_cached_torch_backend_eps = None
+_torch_backends_autoload_pending: bool = False
+
+
 def set_visible_accelerator_ids() -> Mapping[str, Optional[str]]:
     """Set (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, HIP_VISIBLE_DEVICES,
     NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS , HABANA_VISIBLE_MODULES ,...)
@@ -328,6 +345,19 @@ def set_visible_accelerator_ids() -> Mapping[str, Optional[str]]:
         for key in env_before.keys() | os.environ.keys():
             if env_before.get(key) != os.environ.get(key):
                 original_visible_accelerator_env_vars[key] = env_before.get(key)
+    global _torch_backends_autoload_pending
+    if _torch_backends_autoload_pending:
+        _torch_backends_autoload_pending = False
+        torch_mod = sys.modules.get("torch")
+        if (
+            torch_mod is not None
+            and os.environ.get("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") != "0"
+        ):
+            if _cached_torch_backend_eps is not None:
+                for ep in _cached_torch_backend_eps:
+                    ep.load()()
+            elif hasattr(torch_mod, "_import_device_backends"):
+                torch_mod._import_device_backends()
     return original_visible_accelerator_env_vars
 
 

python/ray/_private/worker.py

@@ -464,12 +464,11 @@ def __init__(self):
         # tensor_transport, to avoid circular import and because it imports
         # third-party dependencies like PyTorch.
         self._rdt_manager = None
-        # When the worker is constructed. Record the original value of the
+        # Lazily record the initial value of the
         # (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, HIP_VISIBLE_DEVICES,
-        # NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS, ..) environment variables.
-        self.original_visible_accelerator_ids = (
-            ray._private.utils.get_visible_accelerator_ids()
-        )
+        # NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS, ..) environment variables
+        # so importing ray._private.worker does not probe accelerator SDKs pre-fork.
+        self._original_visible_accelerator_ids = None
         # A dictionary that maps from driver id to SerializationContext
         # TODO: clean up the SerializationContext once the job finished.
         self.serialization_context_map = {}
@@ -524,6 +523,18 @@ def rdt_manager(self) -> "ray.experimental.RDTManager":
             self._rdt_manager = RDTManager()
         return self._rdt_manager
 
+    @property
+    def original_visible_accelerator_ids(self):
+        if self._original_visible_accelerator_ids is None:
+            self._original_visible_accelerator_ids = (
+                ray._private.utils.get_visible_accelerator_ids()
+            )
+        return self._original_visible_accelerator_ids
+
+    @original_visible_accelerator_ids.setter
+    def original_visible_accelerator_ids(self, value):
+        self._original_visible_accelerator_ids = value
+
     @property
     def connected(self):
         """bool: True if Ray has been started and False otherwise."""

python/ray/_private/workers/_cgroup_bootstrap.py

@@ -0,0 +1,22 @@
+"""Move this process into a cgroup, then run the given command.
+
+This runs before Ray is imported, so the fork server's memory counts toward that cgroup.
+"""
+
+import os
+import sys
+
+
+def main() -> None:
+    if len(sys.argv) < 3:
+        raise SystemExit(f"Usage: {sys.argv[0]} <cgroup_path> <executable> [args...]")
+    fd = os.open(os.path.join(sys.argv[1], "cgroup.procs"), os.O_WRONLY)
+    try:
+        os.write(fd, str(os.getpid()).encode())
+    finally:
+        os.close(fd)
+    os.execvp(sys.argv[2], sys.argv[2:])
+
+
+if __name__ == "__main__":
+    main()