Run InterGeneratorRouter inside a monarch actor

torchtitan/experiments/rl/actors/generator.py

@@ -644,7 +644,7 @@ class VLLMGenerator(Actor, Configurable):
 
     A weight sync rides the same loop: `pull_model_state_dict` queues a `LoopDecision(LoopAction.PULL_MODEL_STATE_DICT)` applied
     between step bursts. The engine does NOT drain in-flight requests first ("hotswap"). This behavior can be changed
-    on the controller side, by blocking new requests until the engine is drained.
+    in the inter-generator router, by blocking new requests until the engine is drained.
 
     Args:
         config: Generator-specific configuration.

torchtitan/experiments/rl/actors/inter_generator_router.py

@@ -0,0 +1,82 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Actor wrapper for inter-generator routing."""
+
+from collections.abc import Sequence
+from typing import Any
+
+from monarch.actor import Actor, endpoint
+
+from torchtitan.experiments.rl.routing.inter_generator_router import (
+    InterGeneratorRouter,
+)
+from torchtitan.experiments.rl.routing.types import RoutingContext
+from torchtitan.observability import structured_logger as sl
+
+
+class InterGeneratorRouterActor(Actor):
+    """Give mutable inter-generator routing state a single process owner.
+
+    ``InterGeneratorRouter`` tracks mutable generator load and lifecycle state
+    and is not safe to share across threads or processes. Hosting it in a
+    singleton Monarch actor confines that state to one actor process and makes
+    other actors access it through endpoints instead of sharing the router.
+    """
+
+    def __init__(
+        self,
+        config: InterGeneratorRouter.Config,
+        *,
+        generators: Sequence[Any],
+    ) -> None:
+        self._router = InterGeneratorRouter(config, generators=generators)
+
+    @endpoint
+    async def generate(
+        self,
+        prompt_token_ids: list[int],
+        *,
+        request_id: str,
+        routing_session_id: str | None,
+        sampling_config: Any | None,
+        metrics_prefix: str,
+    ) -> Any:
+        """Route one generation call while holding global routing state."""
+        # Dispatches to the chosen generator's rank-0 intake via call_one, so
+        # it returns the Completion directly (no ValueMesh unwrap).
+        return await self._router.route(
+            "generate",
+            prompt_token_ids,
+            request_id=request_id,
+            # VLLMGenerator.generate also requires this field for its
+            # intra-mesh DP routing.
+            routing_session_id=routing_session_id,
+            sampling_config=sampling_config,
+            metrics_prefix=metrics_prefix,
+            # Load is measured as in-flight request count (one unit per call).
+            routing_ctx=RoutingContext(
+                estimated_cost=1,
+                session_id=routing_session_id,
+            ),
+        )
+
+    @endpoint
+    async def start_engine_loop(self) -> None:
+        await self._router.fanout("start_engine_loop")
+
+    @endpoint
+    async def sync_log_step(self, step: int) -> None:
+        sl.set_step(step)
+        await self._router.fanout("sync_log_step", step)
+
+    @endpoint
+    async def pull_model_state_dict(self, policy_version: int) -> None:
+        await self._router.pull_model_state_dict(policy_version=policy_version)
+
+    @endpoint
+    async def close_generators(self) -> list[Any | BaseException]:
+        return await self._router.fanout("close", return_exceptions=True)

torchtitan/experiments/rl/components/weight_sync.py

@@ -6,16 +6,22 @@
 
 """Overlap the trainer->generator weight handoff with the next training step."""
 
+from __future__ import annotations
+
 import asyncio
 import time
+from typing import TYPE_CHECKING
 
 from torchtitan.experiments.rl.components.work_buffer import RolloutGroupWorkBuffer
 from torchtitan.experiments.rl.observability import metrics as m
-from torchtitan.experiments.rl.routing.inter_generator_router import (
-    InterGeneratorRouter,
-)
 from torchtitan.observability import structured_logger as sl
 
+if TYPE_CHECKING:
+    from torchtitan.experiments.rl.actors.inter_generator_router import (
+        InterGeneratorRouterActor,
+    )
+    from torchtitan.experiments.rl.actors.trainer import PolicyTrainer
+
 # dummy no-op for step 0, used in WeightSyncManager
 async def _noop() -> None:
     return None
@@ -47,8 +53,8 @@ class WeightSyncManager:
     def __init__(
         self,
         *,
-        trainer,  # PolicyTrainer actor handle
-        generator_router: InterGeneratorRouter,
+        trainer: PolicyTrainer,
+        generator_router: InterGeneratorRouterActor,
         group_buffer: RolloutGroupWorkBuffer,
         num_prompts_per_train_step: int,
     ) -> None:
@@ -112,7 +118,7 @@ async def _generator_pull_and_release_buffer_slots(
         await push_task
         with sl.log_trace_span("generator_pull_model_state_dict"):
             start = time.perf_counter()
-            await self._generator_router.pull_model_state_dict(policy_version=version)
+            await self._generator_router.pull_model_state_dict.call_one(version)
             self._last_pull_s = time.perf_counter() - start
         # TODO(perf): pull_model_state_dict awaits ALL generators before we release any buffer slots,
         #   so a generator that finishes its pull early idles until the slowest one. Investigate

torchtitan/experiments/rl/controller.py

@@ -102,11 +102,15 @@
 import torch  # noqa: F401
 import torchstore as ts
 import tyro
-from monarch.actor import ProcMesh
+
+from monarch.actor import ProcMesh, this_host
 from monarch.spmd import setup_torch_elastic_env_async
 
 from torchtitan.config import CompileConfig, Configurable
 from torchtitan.experiments.rl.actors.generator import SamplingConfig, VLLMGenerator
+from torchtitan.experiments.rl.actors.inter_generator_router import (
+    InterGeneratorRouterActor,
+)
 from torchtitan.experiments.rl.actors.trainer import PolicyTrainer
 from torchtitan.experiments.rl.components.batcher import Batcher
 from torchtitan.experiments.rl.components.training_sample_builder import (
@@ -134,7 +138,6 @@
 from torchtitan.experiments.rl.routing.inter_generator_router import (
     InterGeneratorRouter,
 )
-from torchtitan.experiments.rl.routing.types import RoutingContext
 from torchtitan.experiments.rl.types import Completion, TrainingBatch
 from torchtitan.observability import structured_logger as sl
 from torchtitan.protocols.model_spec import ModelSpec
@@ -432,7 +435,7 @@ def __post_init__(self):
     def __init__(self, config: Config):
         self.config = config
         self.trainer: PolicyTrainer | None = None
-        self.generator_router: InterGeneratorRouter | None = None
+        self.generator_router: InterGeneratorRouterActor | None = None
         # Resume step (0 = fresh); set in setup_async from the loaded checkpoint.
         self.start_step = 0
         self._proc_meshes = []
@@ -470,19 +473,22 @@ async def close(self):
                 logger.exception("trainer.close failed")
 
         if self.generator_router is not None:
-            close_results = await self.generator_router.fanout(
-                "close", return_exceptions=True
-            )
-            for idx, result in enumerate(close_results):
-                if isinstance(result, BaseException):
-                    actor_name = (
-                        "generator" if len(close_results) == 1 else f"generator[{idx}]"
-                    )
-                    logger.error(
-                        "%s.close failed",
-                        actor_name,
-                        exc_info=(type(result), result, result.__traceback__),
-                    )
+            try:
+                close_results = await self.generator_router.close_generators.call_one()
+                for idx, result in enumerate(close_results):
+                    if isinstance(result, BaseException):
+                        actor_name = (
+                            "generator"
+                            if len(close_results) == 1
+                            else f"generator[{idx}]"
+                        )
+                        logger.error(
+                            "%s.close failed",
+                            actor_name,
+                            exc_info=(type(result), result, result.__traceback__),
+                        )
+            except Exception:
+                logger.exception("generator_router.close_generators failed")
 
         try:
             self.metrics_processor.close()
@@ -519,22 +525,12 @@ async def generate(
             routing_session_id: str | None = None,
             sampling_config: SamplingConfig | None = None,
         ) -> Completion | None:
-            # Dispatches to the chosen generator's rank-0 intake via call_one, so
-            # it returns the Completion directly (no ValueMesh unwrap).
-            return await self.generator_router.route(
-                "generate",
+            return await self.generator_router.generate.call_one(
                 prompt_token_ids,
                 request_id=request_id,
-                # VLLMGenerator.generate also requires this field for its
-                # intra-mesh DP routing.
                 routing_session_id=routing_session_id,
                 sampling_config=sampling_config,
                 metrics_prefix=metrics_prefix,
-                # Load is measured as in-flight request count (one unit per call).
-                routing_ctx=RoutingContext(
-                    estimated_cost=1,
-                    session_id=routing_session_id,
-                ),
             )
 
         return generate
@@ -553,9 +549,9 @@ async def setup_async(
         weight push/pull are all ``await``-based runtime side effects
         that cannot run in a synchronous constructor.
 
-        The trainer and generator meshes are provisioned by the caller
-        (see ``create_proc_mesh``) on disjoint GPUs; this method only
-        spawns the actors on them and synchronizes initial weights from
+        The trainer and generator meshes are provisioned by the caller (see
+        ``spawn_proc_mesh``). The router mesh is created on the controller host.
+        This method spawns the actors and synchronizes initial weights from
         trainer to generator. Must be called before :meth:`run`.
 
         Args:
@@ -605,8 +601,9 @@ async def setup_async(
         # provisioner logic. Pull a PerHostProvisioner.spawn_meshes(...) helper and
         # shrink this span to a single call.
         with sl.log_trace_span("mesh_spawn"):
+            router_mesh = this_host().spawn_procs(per_host={"cpus": 1})
             # Store proc meshes for cleanup
-            self._proc_meshes = [trainer_mesh, *generator_meshes]
+            self._proc_meshes = [router_mesh, trainer_mesh, *generator_meshes]
 
             await setup_torch_elastic_env_async(trainer_mesh)
             for generator_mesh in generator_meshes:
@@ -641,7 +638,12 @@ async def setup_async(
                     output_dir=config.dump_folder,
                 )
                 generators.append(generator)
-            self.generator_router = config.generator_router.build(generators=generators)
+            self.generator_router = router_mesh.spawn(
+                "generator_router",
+                InterGeneratorRouterActor,
+                config.generator_router,
+                generators=generators,
+            )
 
         # Initialize TorchStore for weight sync between trainer and generator.
         # StorageVolumes are spawned on the trainer mesh so they are colocated
@@ -667,15 +669,13 @@ async def setup_async(
         # rank-0-only generate / pull (rank 0 drives the followers through this
         # loop, so every rank must be running it first).
         with sl.log_trace_span("generator_start_engine_loop"):
-            await self.generator_router.fanout("start_engine_loop")
+            await self.generator_router.start_engine_loop.call_one()
 
         # Initial weight sync: only the trainer loads weights; generators pull at start_step.
         with sl.log_trace_span("trainer_push_model_state_dict"):
             await self.trainer.push_model_state_dict.call()
         with sl.log_trace_span("generator_pull_model_state_dict"):
-            await self.generator_router.pull_model_state_dict(
-                policy_version=self.start_step
-            )
+            await self.generator_router.pull_model_state_dict.call_one(self.start_step)
 
     # TODO: fold validation into a Validator(Configurable) the controller attaches, instead of 4 methods.
     @sl.log_trace_span("_collect_validation_rollouts")
@@ -1063,7 +1063,7 @@ async def _trainer_loop(
             sl.set_step(step)  # propagate the step counter to the actors
             with sl.log_trace_span("sync_log_step"):
                 await self.trainer.sync_log_step.call(step)
-                await self.generator_router.fanout("sync_log_step", step)
+                await self.generator_router.sync_log_step.call_one(step)
             step_timer = MetricsTimer()
 
             with sl.log_trace_span("train_step"), step_timer.record(

torchtitan/experiments/rl/routing/inter_generator_router.py

@@ -32,7 +32,7 @@ class _GeneratorState(Enum):
 
 @dataclass(kw_only=True, slots=True)
 class _GeneratorHandle(RoutingCandidate):
-    """Controller-side metadata for one generator mesh."""
+    """Router-side metadata for one generator mesh."""
 
     actor: Any
     """Monarch actor handle for the full generator mesh. Used for fan-out calls

torchtitan/experiments/rl/routing/intra_generator_router.py

@@ -32,7 +32,7 @@ class _DPRankHandle(RoutingCandidate):
 class IntraGeneratorRouter(Configurable):
     """Router that partitions requests across the DP ranks within one generator.
 
-    This is layer 2 of the two-layer routing design: the controller-side
+    This is layer 2 of the two-layer routing design: the router-actor-side
     ``InterGeneratorRouter`` routes a call across generators; this router
     then routes each request across the DP ranks within one generator.
 

torchtitan/experiments/rl/routing/strategies.py

@@ -11,7 +11,7 @@
 ``reserved_load`` field plus object identity -- so the same strategy classes
 serve both routing layers in the RL generator:
 
-- Layer 1: ``InterGeneratorRouter`` (controller side) routes a call across
+- Layer 1: ``InterGeneratorRouter`` (router actor side) routes a call across
   generator *meshes* (replicas). See ``inter_generator_router.py``.
 - Layer 2: ``IntraGeneratorRouter`` (in-mesh, rank-0 side) routes a request
   across the *data-parallel ranks* within one generator mesh. See

torchtitan/experiments/rl/tests/test_shutdown.py

@@ -28,7 +28,12 @@ def __init__(self, config=None):
         self.setup_generator_meshes = None
         self.instances.append(self)
 
-    async def setup_async(self, *, trainer_mesh=None, generator_meshes=None):
+    async def setup_async(
+        self,
+        *,
+        trainer_mesh=None,
+        generator_meshes=None,
+    ):
         self.events.append("setup")
         self.setup_trainer_mesh = trainer_mesh
         self.setup_generator_meshes = generator_meshes
@@ -197,9 +202,10 @@ def stub_mesh_provisioning(monkeypatch):
     monkeypatch.setattr(train, "_compute_generator_world_size", lambda p: 1)
 
     def _spawn_proc_mesh(*args, num_generators=1, **kwargs):
-        return "trainer_mesh", [
-            f"generator_mesh_{idx}" for idx in range(num_generators)
-        ]
+        return (
+            "trainer_mesh",
+            [f"generator_mesh_{idx}" for idx in range(num_generators)],
+        )
 
     monkeypatch.setattr(train, "spawn_proc_mesh", _spawn_proc_mesh)
 
@@ -295,6 +301,19 @@ async def call(self):
             raise RuntimeError(f"{self._name} failed")
 
 
+class _RouterCloseEndpoint:
+    def __init__(self, router):
+        self._router = router
+
+    async def call_one(self):
+        return await self._router.fanout("close", return_exceptions=True)
+
+
+class _StubRouterActor:
+    def __init__(self, router):
+        self.close_generators = _RouterCloseEndpoint(router)
+
+
 class _StubActor:
     def __init__(self, name, events, raises=False):
         self.close = _StubEndpoint(name, events, raises)
@@ -321,9 +340,11 @@ async def stop(self):
 
 
 def _set_generator_router(rl_trainer, generators):
-    rl_trainer.generator_router = InterGeneratorRouter(
-        InterGeneratorRouter.Config(),
-        generators=generators,
+    rl_trainer.generator_router = _StubRouterActor(
+        InterGeneratorRouter(
+            InterGeneratorRouter.Config(),
+            generators=generators,
+        )
     )