[Weight Sync Orchestrator 2/N] Ｒoute sync_weights through the weight sync coordinator

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())
 

tunix/experimental/orchestrator/distributed_rl_engine.py

@@ -104,13 +104,16 @@ 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(
         self._rollout_workers
     )
     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

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(