use subprocess to fork colocated workers and centralize the usage of threadpoolexecutor

Author: garrett4wade

areal/api/cli_args.py

@@ -630,6 +630,14 @@ class SchedulingStrategy:
     target: str | None = field(
         default=None, metadata={"help": "The target role to be colocated with"}
     )
+    fork: bool = field(
+        default=True,
+        metadata={
+            "help": "When True with colocation, the target worker spawns a new "
+            "process on the same node/GPUs instead of sharing its process. "
+            "Provides process isolation while sharing GPU resources."
+        },
+    )
 
 
 @dataclass

areal/api/scheduler_api.py

@@ -36,7 +36,6 @@ class Job:
     replicas: int = 0
     tasks: list[SchedulingSpec] = field(default_factory=list)
     scheduling_strategy: SchedulingStrategy = field(default_factory=SchedulingStrategy)
-    shared_placement_group: bool = False
 
 
 class Scheduler(abc.ABC):

areal/controller/rollout_callback.py

@@ -1,6 +1,4 @@
-import atexit
-import threading
-from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent.futures import Future
 from dataclasses import dataclass
 from typing import Any
 
@@ -9,39 +7,10 @@
 from areal.api.io_struct import ParamSpec, WeightUpdateMeta
 from areal.scheduler.rpc.serialization import serialize_value
 from areal.utils import logging
+from areal.utils.concurrent import get_executor
 
 logger = logging.getLogger(__name__)
 
-# Lazy-initialized thread pool for async HTTP requests
-_executor: ThreadPoolExecutor | None = None
-_executor_lock = threading.Lock()
-
-
-def _get_executor() -> ThreadPoolExecutor:
-    """Get or create the shared thread pool executor."""
-    global _executor
-    if _executor is None:
-        with _executor_lock:
-            if _executor is None:
-                _executor = ThreadPoolExecutor(
-                    max_workers=4, thread_name_prefix="rollout_callback"
-                )
-                # Register cleanup on process exit
-                atexit.register(_shutdown_executor)
-    return _executor
-
-
-def _shutdown_executor() -> None:
-    """Shutdown the shared thread pool executor if it exists.
-
-    Called via atexit at process exit, when no other threads should be
-    accessing the executor.
-    """
-    global _executor
-    if _executor is not None:
-        _executor.shutdown(wait=False)
-        _executor = None
-
 
 @dataclass
 class RolloutCallback:
@@ -110,7 +79,7 @@ def _post_nowait(
         Future[dict]
             Future that completes when the HTTP response is received
         """
-        return _get_executor().submit(self._post, endpoint, payload)
+        return get_executor().submit(self._post, endpoint, payload)
 
     def _post_nowait_void(
         self, endpoint: str, payload: dict[str, Any] | None = None

areal/controller/rollout_controller.py

@@ -29,6 +29,7 @@
 from areal.core.workflow_executor import BatchTaskDispatcher, TaskIdGenerator
 from areal.scheduler.rpc.serialization import deserialize_value
 from areal.utils import logging, perf_tracer
+from areal.utils.concurrent import run_async_task
 from areal.utils.data import concat_padded_tensors, cycle_dataloader
 from areal.utils.dynamic_import import import_from_string
 from areal.utils.network import find_free_ports, gethostip
@@ -139,18 +140,11 @@ def initialize(
             tasks=[sch_spec for _ in range(alloc_mode.gen.dp_size)],
             scheduling_strategy=self.config.scheduling_strategy,
             role=self._worker_role,
-            shared_placement_group=False,
         )
 
-        # Use asyncio.run to call async scheduler methods synchronously
-        asyncio.run(
-            self._async_initialize(
-                job,
-                server_args,
-                server_infos,
-                *args,
-                **kwargs,
-            )
+        # Call async scheduler methods synchronously
+        run_async_task(
+            self._async_initialize, job, server_args, server_infos, *args, **kwargs
         )
 
         # Initialize staleness manager for global capacity control
@@ -385,7 +379,7 @@ def _resolve_task_future(self, task_id: int):
             future.get_loop().call_soon_threadsafe(future.set_result, None)
 
     def _collective_rpc(self, method: str, *args, **kwargs) -> list[Any]:
-        return asyncio.run(self._collective_rpc_async(method, *args, **kwargs))
+        return run_async_task(self._collective_rpc_async, method, *args, **kwargs)
 
     async def _collective_rpc_async(self, method: str, *args, **kwargs) -> list[Any]:
         tasks = [
@@ -764,7 +758,7 @@ async def _call():
             ]
             return await asyncio.gather(*tasks)
 
-        asyncio.run(_call())
+        run_async_task(_call)
 
     def save_perf_tracer(self, step: int | None = None, force: bool = False) -> None:
         self._collective_rpc("save_perf_tracer", step=step, force=force)

areal/controller/train_controller.py

@@ -19,6 +19,7 @@
 from areal.controller.rollout_controller import RolloutController
 from areal.scheduler.rpc.rtensor import RTensor
 from areal.utils import logging, stats_tracker
+from areal.utils.concurrent import run_async_task
 from areal.utils.network import find_free_ports
 
 logger = logging.getLogger("TrainController")
@@ -132,7 +133,6 @@ def initialize(
             tasks=list(self.config.scheduling_spec),
             scheduling_strategy=self.config.scheduling_strategy,
             role=self._worker_role,
-            shared_placement_group=True,
         )
 
         # Create workers via scheduler
@@ -164,21 +164,16 @@ def initialize(
         engine_class = self.train_engine
 
         # Create and initialize engines on workers
-        self._run_async_task(
-            self._async_create_engines(
-                f"{engine_class.__module__}.{engine_class.__name__}"
-            )
+        run_async_task(
+            self._async_create_engines,
+            f"{engine_class.__module__}.{engine_class.__name__}",
         )
-        self._run_async_task(self._async_initialize_engines(ft_spec, **kwargs))
+        run_async_task(self._async_initialize_engines, ft_spec, **kwargs)
 
         # Identify DP head workers
         self._identify_dp_heads()
         logger.info("TrainController initialization complete")
 
-    def _run_async_task(self, task):
-        """Run an async task synchronously."""
-        return asyncio.run(task)
-
     def _engine_name(self, rank: int) -> str:
         """Generate engine name for a worker rank.
 
@@ -255,7 +250,7 @@ async def _get_dp_head():
             ]
             return await asyncio.gather(*tasks)
 
-        self.workers_is_dp_head = self._run_async_task(_get_dp_head())
+        self.workers_is_dp_head = run_async_task(_get_dp_head)
 
     def destroy(self):
         """Destroy the controller and release GPU memory of models.
@@ -280,7 +275,7 @@ async def _destroy_all_engines():
                     ]
                     await asyncio.gather(*tasks, return_exceptions=True)
 
-                self._run_async_task(_destroy_all_engines())
+                run_async_task(_destroy_all_engines)
                 logger.info("Engines destroyed")
             except Exception as e:
                 logger.error(f"Error destroying engines: {e}")
@@ -306,8 +301,8 @@ def _custom_function_call(self, method: str, *args, **kwargs):
         dp_split_args, dp_split_kwargs, group_indices = self._dispatch_inputs(
             *args, **kwargs
         )
-        results = self._run_async_task(
-            self._call_with_dispatched_inputs(method, dp_split_args, dp_split_kwargs)
+        results = run_async_task(
+            self._call_with_dispatched_inputs, method, dp_split_args, dp_split_kwargs
         )
         # Filter to only keep results from DP head workers
         results = [r for idx, r in enumerate(results) if self.workers_is_dp_head[idx]]
@@ -527,7 +522,7 @@ async def _call():
             ]
             return await asyncio.gather(*tasks)
 
-        self._run_async_task(_call())
+        run_async_task(_call)
 
     def save_perf_tracer(self, step: int | None = None, force: bool = False) -> None:
         self._custom_function_call("save_perf_tracer", step=step, force=force)
@@ -588,4 +583,4 @@ async def _async_clear_batches(self, *targets: dict[str, RTensor]):
 
     def clear_batches(self, *targets: dict[str, RTensor]):
         """Clear distributed batch shards from workers to free memory."""
-        self._run_async_task(self._async_clear_batches(*targets))
+        run_async_task(self._async_clear_batches, *targets)

areal/core/workflow_executor.py

@@ -7,7 +7,6 @@
 import threading
 import time
 from collections.abc import Awaitable, Callable
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, TypeVar, Generic, Protocol
 from collections.abc import Generator
@@ -33,6 +32,7 @@
 from areal.core.workflow_context import WorkflowContext
 from areal.experimental.openai.types import InteractionWithTokenLogpReward
 from areal.utils import logging, perf_tracer, stats_tracker
+from areal.utils.concurrent import get_executor
 from areal.utils.data import concat_padded_tensors, cycle_dataloader
 from areal.utils.dynamic_import import import_from_string
 from areal.utils.perf_tracer import trace_perf, trace_session_event
@@ -361,8 +361,6 @@ def __init__(
 
         # Callback support: task_id -> callback_addr
         self._task_callbacks: dict[int, str] = {}
-        # Thread pool for sending callbacks (avoids creating threads per callback)
-        self._callback_executor: ThreadPoolExecutor | None = None
 
     def _set_thread_exception(self, exc: Exception):
         """Store exception from background thread for fail-fast behavior."""
@@ -407,7 +405,7 @@ def post():
             except requests.RequestException as e:
                 self.logger.error(f"Callback to {addr} failed: {e}")
 
-        self._callback_executor.submit(post)
+        get_executor().submit(post)
 
     def _commit_loop(self) -> None:
         """Producer thread - continuously submits tasks based on capacity."""
@@ -504,9 +502,6 @@ def initialize(self, logger: Logger):
         self.runner.initialize(logger=logger)
 
         self._shutdown_event.clear()
-        self._callback_executor = ThreadPoolExecutor(
-            max_workers=4, thread_name_prefix="callback"
-        )
 
         self._commit_thread = threading.Thread(target=self._commit_loop, daemon=True)
         self._commit_thread.start()
@@ -538,11 +533,6 @@ def destroy(self):
         # Clear pending callbacks to prevent memory leak
         self._task_callbacks.clear()
 
-        # Shutdown callback thread pool
-        if self._callback_executor is not None:
-            self._callback_executor.shutdown(wait=False)
-            self._callback_executor = None
-
         # Shutdown the async task runner
         self.runner.destroy()