diff --git a/tests/experimental/orchestrator/distributed_rl_engine_test.py b/tests/experimental/orchestrator/distributed_rl_engine_test.py index 93aef8af7..8fb59bfa6 100644 --- a/tests/experimental/orchestrator/distributed_rl_engine_test.py +++ b/tests/experimental/orchestrator/distributed_rl_engine_test.py @@ -40,6 +40,10 @@ def __init__(self, *args, **kwargs): self.fwd_bwd = mock.AsyncMock() self.update = mock.AsyncMock() self.prepare_weight_sync = mock.AsyncMock() + self.release_weight_sync = mock.AsyncMock() + self.pre_weight_sync = mock.AsyncMock() + self.post_weight_sync = mock.AsyncMock() + self.abort_weight_sync = mock.AsyncMock() self.score = mock.AsyncMock() self.per_token_logps = mock.AsyncMock() @@ -219,31 +223,38 @@ async def _run(): asyncio.run(_run()) - def test_sync_weights_coordination(self): + def test_sync_weights_delegates_to_coordinator(self): async def _run(): - mock_meta = datatypes.WeightSyncMetadata( - new_policy_version=42, - transfer_mode="p2p", - ) - self.mock_actor.prepare_weight_sync.return_value = mock_meta - self.mock_rollout_1.weight_sync.return_value = None - self.mock_rollout_2.weight_sync.return_value = None + class _FakeResult: + policy_version = 7 + + class _FakeCoordinator: - ver = await self.engine.sync_weights(role=datatypes.Role.ACTOR) - self.assertEqual(ver, 42) + def __init__(self): + self.calls = [] - self.mock_actor.prepare_weight_sync.assert_called_once() - self.mock_rollout_1.weight_sync.assert_called_once_with(metadata=mock_meta) - self.mock_rollout_2.weight_sync.assert_called_once_with(metadata=mock_meta) + async def sync(self, policy_version=0, **kwargs): + self.calls.append(policy_version) + _FakeResult.policy_version = policy_version + return _FakeResult + + coordinator = _FakeCoordinator() + engine = distributed_rl_engine.DistributedRLEngine( + rollout_workers=[self.mock_rollout_1, self.mock_rollout_2], + trainer_workers={datatypes.Role.ACTOR: self.mock_actor}, + inference_workers={datatypes.Role.REFERENCE: self.mock_ref}, + weight_sync_coordinator=coordinator, + ) + self.assertEqual(await engine.sync_weights(), 1) + self.assertEqual(await engine.sync_weights(), 2) + self.assertEqual(coordinator.calls, [1, 2]) asyncio.run(_run()) - def test_sync_weights_requires_weight_sync_metadata(self): + def test_sync_weights_requires_a_coordinator(self): async def _run(): - self.mock_actor.prepare_weight_sync.return_value = datatypes.Response() - - with self.assertRaisesRegex(RuntimeError, "WeightSyncMetadata"): - await self.engine.sync_weights(role=datatypes.Role.ACTOR) + with self.assertRaises(RuntimeError): + await self.engine.sync_weights() asyncio.run(_run()) diff --git a/tunix/experimental/orchestrator/distributed_rl_engine.py b/tunix/experimental/orchestrator/distributed_rl_engine.py index 1891482cb..cd84e1205 100644 --- a/tunix/experimental/orchestrator/distributed_rl_engine.py +++ b/tunix/experimental/orchestrator/distributed_rl_engine.py @@ -104,6 +104,7 @@ def __init__( inference_workers: ( Mapping[datatypes.Role, remote_execution.ActorHandle] | None ) = None, + weight_sync_coordinator: Any = None, ): self._rollout_workers = list(rollout_workers) self._rollout_pool = remote_execution.RoutingActorPool( @@ -111,6 +112,8 @@ def __init__( ) self._trainer_workers = dict(trainer_workers) self._inference_workers = dict(inference_workers or {}) + self._policy_version = 0 + self._weight_sync_coordinator = weight_sync_coordinator async def _invoke_worker( self, @@ -344,23 +347,15 @@ async def sync_weights( # pyrefly: ignore[bad-override] role: datatypes.Role = datatypes.Role.ACTOR, target_roles: Sequence[datatypes.Role] | None = None, ) -> int: - """Executes accelerator-to-accelerator collective weight broadcast.""" - # TODO: integrate with raiden controller instead - del target_roles - trainer = self._trainer_workers.get(role) - if trainer is None: - return 0 - sync_metadata = await self._invoke_worker(trainer, "prepare_weight_sync") - if not isinstance(sync_metadata, datatypes.WeightSyncMetadata): + """Runs one weight sync round through the coordinator.""" + del role, target_roles + if self._weight_sync_coordinator is None: raise RuntimeError( - "prepare_weight_sync must return WeightSyncMetadata; got " - f"{type(sync_metadata).__name__}." + "sync_weights needs a coordinator; construct the engine with" + " weight_sync_coordinator." ) - tasks = [ - self._invoke_worker(w, "weight_sync", metadata=sync_metadata) - for w in self._rollout_workers - if hasattr(w, "weight_sync") or hasattr(w, "asubmit") - ] - if tasks: - await asyncio.gather(*tasks) - return getattr(sync_metadata, "new_policy_version", 1) + result = await self._weight_sync_coordinator.sync( + policy_version=self._policy_version + 1 + ) + self._policy_version = result.policy_version + return result.policy_version diff --git a/tunix/experimental/orchestrator/orchestrator.py b/tunix/experimental/orchestrator/orchestrator.py index 6ae6dbdc2..7c3b1fb5c 100644 --- a/tunix/experimental/orchestrator/orchestrator.py +++ b/tunix/experimental/orchestrator/orchestrator.py @@ -49,6 +49,7 @@ def __init__( registry: worker_registry.WorkerRegistry | None = None, lifecycle_driver: lifecycle.LifecycleDriver | None = None, monitor: health_monitor.HealthMonitor | None = None, + weight_sync_coordinator: Any = None, ): """Initializes ClusterOrchestrator.""" self.config = config @@ -65,6 +66,7 @@ def __init__( ] = {} self._remote_worker_infos: dict[str, datatypes.WorkerInfo] = {} self.engine: distributed_rl_engine.DistributedRLEngine | None = None + self._weight_sync_coordinator = weight_sync_coordinator def __enter__(self) -> "ClusterOrchestrator": """Interactive context manager bring-up.""" @@ -230,6 +232,7 @@ def _create_engine(self) -> distributed_rl_engine.DistributedRLEngine: rollout_workers=rollout_workers, trainer_workers=trainer_workers, inference_workers=inference_workers, + weight_sync_coordinator=self._weight_sync_coordinator, ) def run_program(