Update rollout topology and weight update tests

Author: YanhuiDua

tests/rl/test_multi_task_agent_loop_manager.py

@@ -167,7 +167,7 @@ def _fake_agent_loop():
     rollout_ctl = MagicMock()
     rollout_ctl.continue_generation.remote = AsyncMock()
     rollout_ctl.pause_generation.remote = AsyncMock()
-    rollout_ctl.get_rollout_metadata.remote = AsyncMock(return_value={"server_url_dict": {}})
+    rollout_ctl.get_weight_update_targets.remote = AsyncMock(return_value=())
     agent_loop = MagicMock()
     agent_loop.rollout_ctl = rollout_ctl
     return agent_loop
@@ -177,7 +177,7 @@ def _fake_rollout_controller():
     rollout_controller = MagicMock()
     rollout_controller.continue_generation.remote = AsyncMock()
     rollout_controller.pause_generation.remote = AsyncMock()
-    rollout_controller.get_rollout_metadata.remote = AsyncMock(return_value={"server_url_dict": {}})
+    rollout_controller.get_weight_update_targets.remote = AsyncMock(return_value=())
     return rollout_controller
 
 

tests/rl/test_producer.py

@@ -100,7 +100,7 @@ def _build_agent_loop(self, sleep_by_id: dict[int, float] | None = None):
         mock_agent_loop = MagicMock()
         mock_agent_loop.rollout_ctl.continue_generation.remote = AsyncMock(return_value=None)
         mock_agent_loop.rollout_ctl.pause_generation.remote = AsyncMock(return_value=None)
-        mock_agent_loop.rollout_ctl.get_rollout_metadata.remote = AsyncMock(return_value={"server_url_dict": {}})
+        mock_agent_loop.rollout_ctl.get_weight_update_targets.remote = AsyncMock(return_value=())
 
         async def mock_pause():
             await mock_agent_loop.rollout_ctl.pause_generation.remote()

tests/rl/test_rl_colocate_trainer.py

@@ -75,7 +75,7 @@ def _build_fake_rollout_controller():
     rollout_ctl = MagicMock()
     rollout_ctl.continue_generation.remote = AsyncMock(return_value=None)
     rollout_ctl.pause_generation.remote = AsyncMock(return_value=None)
-    rollout_ctl.get_rollout_metadata.remote = AsyncMock(return_value={"server_url_dict": {}})
+    rollout_ctl.get_weight_update_targets.remote = AsyncMock(return_value=())
     return rollout_ctl
 
 

tests/rl/test_rl_trainer_checkpoint.py

@@ -104,7 +104,7 @@ def __init__(self):
         self.restart_inactive_workers = _RemoteMethod(return_value="rollout_restarted")
         self.onload_weights = _RemoteMethod(return_value="weights_loaded")
         self.onload_kvcache = _RemoteMethod(return_value="kvcache_loaded")
-        self.get_rollout_metadata = _RemoteMethod(return_value={"server_url_dict": {}})
+        self.get_weight_update_targets = _RemoteMethod(return_value=())
         self.set_enable_partial_rollout = _RemoteMethod(return_value=None)
         self.validate_registered_workers_to_proxy = _RemoteMethod(return_value=_AwaitableValue(None))
 
@@ -127,14 +127,19 @@ def __init__(self):
         self.update_weights_count = 0
         self.rollout_info = None
 
-    def update_rollout_info(
-            self,
-            info,
-            train_rollout_mode,
-            weight_update_host,
-            weight_update_port
-        ):
-        self.rollout_info = info
+    def bind_rollout_weight_update(
+        self,
+        *,
+        targets,
+        rollout_config,
+        train_rollout_mode,
+        weight_update_host=None,
+        weight_update_port=None,
+    ):
+        self.rollout_info = {
+            "targets": targets,
+            "rollout_config": rollout_config,
+        }
         self.train_rollout_mode = train_rollout_mode
         self.weight_update_host = weight_update_host
         self.weight_update_port = weight_update_port