Rename the trainer group and cell to TrainerController and TrainerCell

miles/ray/placement_group.py

@@ -7,7 +7,7 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
 from miles.ray.specs.train import compute_critic_args
-from miles.ray.train.group import RayTrainGroup
+from miles.ray.train.group import TrainerController
 from ..utils.ray_utils import compute_ray_pin_head_options
 from .rollout.inference_controller import InferenceController
 from .rollout.rollout_executor import RolloutExecutor
@@ -125,7 +125,7 @@ def create_placement_groups(args) -> dict[str, PlacementGroupInfo]:
 
 
 async def create_training_models(args, inference_controller, rollout_executor):
-    actor_model = RayTrainGroup(
+    actor_model = TrainerController(
         args=args,
         role="actor",
         with_ref=args.kl_coef != 0 or args.use_kl_loss,
@@ -136,7 +136,7 @@ async def create_training_models(args, inference_controller, rollout_executor):
     actor_start_rollout_ids = await actor_model.init()
 
     if args.use_critic:
-        critic_model = RayTrainGroup(
+        critic_model = TrainerController(
             args=compute_critic_args(args),
             role="critic",
             with_ref=False,

miles/ray/train/cell.py

@@ -26,7 +26,7 @@
 CONFIRM_DEAD_TIMEOUT_S = 120.0
 
 
-class RayTrainCell:
+class TrainerCell:
     def __init__(
         self,
         *,

miles/ray/train/cell_monitor.py

@@ -11,12 +11,12 @@
 from miles.utils.ft_utils.health_checker import ActiveAndEpoch, SimpleHealthChecker, SimpleHealthCheckerConfig
 
 if TYPE_CHECKING:
-    from miles.ray.train.cell import RayTrainCell
+    from miles.ray.train.cell import TrainerCell
 
 
 def create_trainer_cell_health_checker(
     *,
-    cell: "RayTrainCell",
+    cell: "TrainerCell",
     config: SimpleHealthCheckerConfig,
     get_activeness: Callable[[], ActiveAndEpoch],
 ) -> SimpleHealthChecker:

miles/ray/train/group.py

@@ -6,7 +6,7 @@
 
 from miles.backends.megatron_utils.ft.types import TrainStepOutcome
 from miles.ray.specs.train import compute_trainer_num_cells, compute_trainer_pool_id
-from miles.ray.train.cell import RayTrainCell
+from miles.ray.train.cell import TrainerCell
 from miles.ray.train.cell_monitor import create_trainer_cell_health_checker
 from miles.utils.async_utils import AsyncioGatherUtils
 from miles.utils.audit_utils.checksum_utils import flatten_inference_engine_checksums
@@ -36,7 +36,7 @@
 _CELLS_READY_TIMEOUT_SECONDS = 3600.0
 
 
-class RayTrainGroup:
+class TrainerController:
     def __init__(
         self,
         args,
@@ -66,15 +66,15 @@ def __init__(
 
         self._health_checker_activeness = ActivenessTracker(active=True)
 
-        self._cells_by_id: dict[str, RayTrainCell] = {}
+        self._cells_by_id: dict[str, TrainerCell] = {}
 
         self._witness_allocator: WitnessIdAllocator | None = (
             WitnessIdAllocator(buffer_size=args.witness_buffer_size) if args.enable_witness else None
         )
         if self._witness_allocator is not None and args.save_debug_event_data is not None:
             self._witness_allocator.resume(read_persisted_witness_counter(Path(args.save_debug_event_data)))
 
-        self._test_action_executor = FTTestActionControllerExecutor.from_args(args, group=self)
+        self._test_action_executor = FTTestActionControllerExecutor.from_args(args, controller=self)
 
     @property
     def pool_id(self) -> str:
@@ -89,7 +89,7 @@ def _expected_num_cells(self) -> int:
         return compute_trainer_num_cells(self.args, role=self._role)
 
     @property
-    def _cells(self) -> list[RayTrainCell]:
+    def _cells(self) -> list[TrainerCell]:
         return sorted(self._cells_by_id.values(), key=lambda cell: cell.cell_index)
 
     @property
@@ -130,8 +130,8 @@ async def _remove_cell(self, cell_id: str) -> None:
         cell = self._cells_by_id.pop(cell_id)
         cell.health_checker.stop()
 
-    def _create_cell(self, cell_id: str, *, cell_index: int, workers_hash: str) -> RayTrainCell:
-        cell = RayTrainCell(
+    def _create_cell(self, cell_id: str, *, cell_index: int, workers_hash: str) -> TrainerCell:
+        cell = TrainerCell(
             args=self.args,
             role=self._role,
             with_ref=self._with_ref,
@@ -242,7 +242,7 @@ def _log_step_end_event(self, *, rollout_id: int, snapshot_alive_cells: list, re
             )
 
     def _check_train_one_attempt(self, snapshot_alive_cells, results):
-        outcomes = RayTrainGroup._compute_attempt_outcomes(snapshot_alive_cells, results)
+        outcomes = TrainerController._compute_attempt_outcomes(snapshot_alive_cells, results)
         if not outcomes["normal"] and not outcomes["discarded"]:
             log_structured(
                 logger.error, tag="ft", op="check", **outcomes, decision="retry", reason="all alive cells failed"

miles/utils/audit_utils/event_analyzer/rules/witness.py

@@ -39,7 +39,7 @@ def check(events: list[Event]) -> list[WitnessIssue]:
     * WitnessAllocateIdEvent: when allocating `witness_id` to `sample_index`
     * WitnessSnapshotParamEvent: near the end of each train() step in MegatronTrainRayActor
         * If a witness_id appears in the weight, it means the corresponding data is consumed at least once.
-    * TrainGroupStepEndEvent: after each train() step in RayTrainGroup
+    * TrainGroupStepEndEvent: after each train() step in TrainerController
 
     Check:
     1. For each (rollout_id, cell_index),

miles/utils/ft_utils/api_server/server.py

@@ -11,7 +11,7 @@
 
 from miles.ray.specs.inference import compute_engine_pool_ids
 from miles.ray.specs.train import compute_trainer_pool_id
-from miles.ray.train.group import RayTrainGroup
+from miles.ray.train.group import TrainerController
 from miles.utils.ft_utils.api_server.handles import _CellHandler
 from miles.utils.ft_utils.api_server.models import Cell, CellList, CellPatch, FaultInjection, K8sStatus, _OkResponse
 from miles.utils.ft_utils.api_server.registry import _CellRegistry

miles/utils/test_utils/ft_test_actions.py

@@ -9,7 +9,7 @@
 from miles.utils.workers.ray_worker_manager import RayWorkerManager
 
 if TYPE_CHECKING:
-    from miles.ray.train.controller import RayTrainGroup
+    from miles.ray.train.group import TrainerController
 
 logger = logging.getLogger(__name__)
 
@@ -47,12 +47,12 @@ def _load_actions(args: object, action_filter: set[str]) -> list[FTTestAction]:
 
 
 class FTTestActionControllerExecutor:
-    def __init__(self, *, actions: list[FTTestAction], controller: "RayTrainGroup") -> None:
+    def __init__(self, *, actions: list[FTTestAction], controller: "TrainerController") -> None:
         self._actions = actions
         self._controller = controller
 
     @staticmethod
-    def from_args(args: object, *, controller: "RayTrainGroup") -> "FTTestActionControllerExecutor":
+    def from_args(args: object, *, controller: "TrainerController") -> "FTTestActionControllerExecutor":
         return FTTestActionControllerExecutor(actions=_load_actions(args, _CONTROLLER_ACTIONS), controller=controller)
 
     async def run_after_step(self, rollout_id: int) -> None:

tests/e2e/ft/README.md

@@ -154,7 +154,7 @@ cell 0, alive back to N). Baseline and phase_a event dirs must contain zero reco
 events. This positively proves the crash -> shrink -> heal path executed; without it the
 comparison could silently degenerate to two fault-free runs.
 
-Fault injection via --ci-ft-test-actions JSON (data-driven, executed by RayTrainGroup).
+Fault injection via --ci-ft-test-actions JSON (data-driven, executed by TrainerController).
 The JSON `at_rollout` field specifies which rollout_id triggers the action.
 The `attempt` field (for actor-level actions like `crash_before_allreduce`) specifies which retry attempt to match.
 ```

tests/fast/ray/rollout/test_server_cell_dispose.py

@@ -7,11 +7,7 @@
 from tests.fast.ray.rollout.conftest import make_args
 
 from miles.ray.rollout import server_cell as server_cell_module
-from miles.ray.rollout.cell_state import (
-    CellAddrInfo,
-    StateDisposed,
-    StateServing,
-)
+from miles.ray.rollout.cell_state import CellAddrInfo, StateDisposed, StateServing
 from miles.ray.rollout.server_cell import ServerCell, ServerCellMetadata
 
 

tests/fast/ray/test_placement_group.py

@@ -2,11 +2,13 @@
 
 from argparse import Namespace
 from copy import deepcopy
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from miles.ray.placement_group import create_rollout_components
+from miles.ray.placement_group import create_rollout_components, create_training_models
+from miles.ray.train.group import TrainerController
 
 pytestmark = pytest.mark.asyncio
 
@@ -34,7 +36,7 @@ def __init__(self, handle: MagicMock) -> None:
     def options(self, **_kwargs):
         return self
 
-    def remote(self, args):
+    def remote(self, *, args):
         self.arg_snapshots.append(deepcopy(args))
         return self._handle
 
@@ -44,11 +46,14 @@ def fake_components():
     controller = MagicMock(name="inference_controller")
     controller.check_weights = AsyncMock()
     controller.offload = AsyncMock()
+    controller.eval_fleet = None
+    eval_fleet = object()
 
     def construct_controller(args):
         async def _init():
             args.sglang_router_ip = "10.0.0.1"
             args.sglang_router_port = 4321
+            controller.eval_fleet = eval_fleet
 
         controller.init = AsyncMock(side_effect=_init)
         return controller
@@ -62,7 +67,12 @@ async def _init():
     with patch("miles.ray.placement_group.InferenceController", controller_cls), patch(
         "miles.ray.placement_group.RolloutExecutor", executor_cls
     ), patch("miles.ray.placement_group.ray.get", return_value=5):
-        yield Namespace(controller=controller, executor_cls=executor_cls, executor_handle=executor_handle)
+        yield Namespace(
+            controller=controller,
+            executor_cls=executor_cls,
+            executor_handle=executor_handle,
+            eval_fleet=eval_fleet,
+        )
 
 
 class TestCreateRolloutComponents:
@@ -124,6 +134,109 @@ async def test_the_baseline_snapshot_is_not_taken_after_the_engines_are_up(self,
         actions = [call.kwargs["action"] for call in fake_components.controller.check_weights.await_args_list]
         assert "snapshot" not in actions
 
+    async def test_the_executor_is_handed_the_fleet_the_controller_just_built(self, fake_components):
+        """Checkpoint eval pins snapshots to these engines, so publishing a pre-init fleet evaluates nothing."""
+        args = _make_args(num_rollout=1)
+
+        await create_rollout_components(args)
+
+        fake_components.executor_handle.set_eval_fleet.remote.assert_awaited_once_with(fake_components.eval_fleet)
+
+
+class _FakeRolloutExecutorHandle:
+    def __init__(self) -> None:
+        self.loaded_rollout_ids: list[int] = []
+        self.load = SimpleNamespace(remote=self._load_remote)
+
+    async def _load_remote(self, rollout_id: int) -> None:
+        self.loaded_rollout_ids.append(rollout_id)
+
+
+_TRAINER_START_ROLLOUT_ID = 7
+
+
+@pytest.fixture
+def fake_trainer_controllers(monkeypatch: pytest.MonkeyPatch):
+    events: list[tuple[str, str]] = []
+
+    async def _fake_init(self: TrainerController) -> list[int]:
+        events.append(("init", self._role))
+        return [_TRAINER_START_ROLLOUT_ID]
+
+    async def _fake_set_rollout_executor(self: TrainerController) -> None:
+        events.append(("set_rollout_executor", self._role))
+
+    monkeypatch.setattr(TrainerController, "init", _fake_init)
+    monkeypatch.setattr(TrainerController, "set_rollout_executor", _fake_set_rollout_executor)
+    return SimpleNamespace(events=events)
+
+
+def _training_args(**overrides) -> Namespace:
+    defaults = dict(
+        actor_num_nodes=1,
+        actor_num_gpus_per_node=2,
+        critic_num_nodes=1,
+        critic_num_gpus_per_node=2,
+        use_critic=False,
+        kl_coef=0.0,
+        use_kl_loss=False,
+        use_opd=False,
+        opd_type=None,
+        disable_param_buffers_cpu_backup=True,
+        start_rollout_id=None,
+        rollout_global_dataset=False,
+        indep_dp=False,
+        enable_witness=False,
+    )
+    defaults.update(overrides)
+    return Namespace(**defaults)
+
+
+class TestCreateTrainingModels:
+    async def test_only_the_actor_is_wired_to_the_rollout_path(self, fake_trainer_controllers):
+        """The critic never broadcasts weights, so handing it the engines would let it publish over the actor's."""
+        inference_controller = object()
+        rollout_executor = _FakeRolloutExecutorHandle()
+
+        actor, critic = await create_training_models(
+            _training_args(use_critic=True, use_opd=True, opd_type="megatron"),
+            inference_controller,
+            rollout_executor,
+        )
+
+        assert actor._inference_controller is inference_controller
+        assert actor._rollout_executor is rollout_executor
+        assert critic._inference_controller is None
+        assert critic._rollout_executor is None
+        assert critic._with_opd_teacher is False
+
+    @pytest.mark.parametrize(
+        ("use_opd", "opd_type", "expected"),
+        [(True, "megatron", True), (True, "sglang", False), (False, "megatron", False)],
+    )
+    async def test_the_actor_hosts_the_teacher_only_for_megatron_opd(
+        self, fake_trainer_controllers, use_opd: bool, opd_type: str, expected: bool
+    ):
+        """Only the in-process Megatron teacher lives in the trainer; the sglang teacher is served by the engines."""
+        actor, _ = await create_training_models(
+            _training_args(use_opd=use_opd, opd_type=opd_type),
+            object(),
+            _FakeRolloutExecutorHandle(),
+        )
+
+        assert actor._with_opd_teacher is expected
+
+    async def test_the_executor_is_connected_and_rewound_once_the_trainers_are_up(self, fake_trainer_controllers):
+        """Cells accept the executor only after init, and the executor resumes from the checkpoint's rollout."""
+        args = _training_args(rollout_global_dataset=False)
+        rollout_executor = _FakeRolloutExecutorHandle()
+
+        await create_training_models(args, object(), rollout_executor)
+
+        assert fake_trainer_controllers.events == [("init", "actor"), ("set_rollout_executor", "actor")]
+        assert args.start_rollout_id == _TRAINER_START_ROLLOUT_ID
+        assert rollout_executor.loaded_rollout_ids == [_TRAINER_START_ROLLOUT_ID - 1]
+
 
 class TestCreatePlacementGroups:
     @staticmethod

tests/fast/ray/test_placement_group_shared_ppo.py

@@ -6,7 +6,7 @@
 
 from miles.ray import placement_group as placement_group_module
 from miles.ray.placement_group import _get_placement_group_layout
-from miles.ray.train.group import RayTrainGroup
+from miles.ray.train.group import TrainerController
 from miles.utils.workers.worker_info import WorkerInfo
 from miles.utils.workers.worker_provider.base import BaseWorkerProvider, ReconcileFn, StopWatchFn
 from miles.utils.workers.worker_provider.ray import RayWorkerProvider
@@ -71,22 +71,22 @@ async def watch_cells(self, reconcile: ReconcileFn) -> StopWatchFn:
         return _stop_watch
 
 
-async def _fake_init(self: RayTrainGroup) -> list[int]:
+async def _fake_init(self: TrainerController) -> list[int]:
     """Stand in for init(), keeping its cell-observation prologue and dropping the GPU work."""
     provider = RayWorkerProvider.create(pool_ids=[self._pool_id])
     self._watcher_disposer = await provider.watch_cells(self._reconcile)
     await self._wait_expected_num_cells()
     return [0]
 
 
-async def _fake_set_rollout_executor(self: RayTrainGroup) -> None:
+async def _fake_set_rollout_executor(self: TrainerController) -> None:
     return None
 
 
 _waited_roles: list[str] = []
 
 
-async def _fake_wait_expected_num_cells(self: RayTrainGroup) -> None:
+async def _fake_wait_expected_num_cells(self: TrainerController) -> None:
     """The startup barrier waits for the provider to report cells, and this provider reports none.
 
     Recording the role keeps init()'s call to the barrier under test: deleting that await
@@ -97,9 +97,9 @@ async def _fake_wait_expected_num_cells(self: RayTrainGroup) -> None:
 async def test_critic_role_disables_reward_kl_and_preserves_actor_args(monkeypatch):
     """Both training groups go through the real create(), and only the critic args are rewritten."""
     provider = _RecordingWorkerProvider()
-    monkeypatch.setattr(RayTrainGroup, "init", _fake_init)
-    monkeypatch.setattr(RayTrainGroup, "set_rollout_executor", _fake_set_rollout_executor)
-    monkeypatch.setattr(RayTrainGroup, "_wait_expected_num_cells", _fake_wait_expected_num_cells)
+    monkeypatch.setattr(TrainerController, "init", _fake_init)
+    monkeypatch.setattr(TrainerController, "set_rollout_executor", _fake_set_rollout_executor)
+    monkeypatch.setattr(TrainerController, "_wait_expected_num_cells", _fake_wait_expected_num_cells)
     _waited_roles.clear()
 
     args = Namespace(