fix validate in agent loop (#34)

Author: baymax591

recipe/transfer_queue/agent_loop.py

@@ -67,10 +67,10 @@ def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: Data
 
         return timing
 
-    def create_transferqueue_client(self, controller_infos, storage_infos):
+    def create_transferqueue_client(self, controller_infos, storage_infos, role):
         ray.get(
             [
-                worker._create_transferqueue_client.remote(controller_infos, storage_infos)
+                worker.create_transferqueue_client.remote(controller_infos, storage_infos, role)
                 for worker in self.agent_loop_workers
             ]
         )

recipe/transfer_queue/ray_trainer.py

@@ -748,24 +748,47 @@ def _validate(self):
             ground_truths = [item.get("ground_truth", None) for item in data.get("reward_model", {})]
             sample_gts.extend(ground_truths)
 
-            test_gen_meta = asyncio.run(
-                self.val_data_system_client.async_get_meta(
-                    data_fields=[
-                        "input_ids",
-                        "attention_mask",
-                        "position_ids",
-                        "index",
-                        "tools_kwargs",
-                        "interaction_kwargs",
-                        "ability",
-                        "raw_prompt_ids",
-                    ],
-                    batch_size=self.val_batch_size * self.config.actor_rollout_ref.rollout.val_kwargs.n,
-                    global_step=self.global_steps - 1,  # self.global_steps start from 1
-                    get_n_samples=False,
-                    task_name="generate_sequences",
+            if not self.async_rollout_mode:
+                test_gen_meta = asyncio.run(
+                    self.val_data_system_client.async_get_meta(
+                        data_fields=[
+                            "input_ids",
+                            "attention_mask",
+                            "position_ids",
+                            "index",
+                            "tools_kwargs",
+                            "interaction_kwargs",
+                            "ability",
+                            "raw_prompt_ids",
+                        ],
+                        batch_size=self.val_batch_size * self.config.actor_rollout_ref.rollout.val_kwargs.n,
+                        global_step=self.global_steps - 1,  # self.global_steps start from 1
+                        get_n_samples=False,
+                        task_name="generate_sequences",
+                    )
+                )
+            else:
+                test_gen_meta = asyncio.run(
+                    self.val_data_system_client.async_get_meta(
+                        data_fields=[
+                            "input_ids",
+                            "attention_mask",
+                            "position_ids",
+                            "index",
+                            "tools_kwargs",
+                            "interaction_kwargs",
+                            "ability",
+                            "raw_prompt_ids",
+                            "raw_prompt",
+                            "reward_model",
+                            "data_source",
+                        ],
+                        batch_size=self.val_batch_size * self.config.actor_rollout_ref.rollout.val_kwargs.n,
+                        global_step=self.global_steps - 1,  # self.global_steps start from 1
+                        get_n_samples=False,
+                        task_name="async_generate_sequences",
+                    )
                 )
-            )
 
             test_gen_meta.extra_info = {
                 "eos_token_id": self.tokenizer.eos_token_id,
@@ -1028,8 +1051,12 @@ def init_workers(self):
             self.async_rollout_manager = AgentLoopManager(
                 config=self.config, worker_group=self.actor_rollout_wg, rm_wg=self.rm_wg
             )
+
+            self.async_rollout_manager.create_transferqueue_client(
+                self.data_system_controller_infos, self.data_system_storage_unit_infos, role="train"
+            )
             self.async_rollout_manager.create_transferqueue_client(
-                self.data_system_controller_infos, self.data_system_storage_unit_infos
+                self.val_data_system_controller_infos, self.val_data_system_storage_unit_infos, role="val"
             )
 
     def _save_checkpoint(self):

verl/experimental/agent_loop/agent_loop.py

@@ -725,13 +725,13 @@ def _postprocess(self, inputs: list[_InternalAgentLoopOutput]) -> DataProto:
             meta_info={"metrics": metrics, "reward_extra_keys": reward_extra_keys},
         )
 
-    def _create_transferqueue_client(self, controller_infos, storage_infos):
+    def create_transferqueue_client(self, controller_infos, storage_infos, role):
         from verl.single_controller.ray.base import get_random_string
         from verl.utils.transferqueue_utils import create_transferqueue_client
 
         client_name = get_random_string(length=6)
         create_transferqueue_client(
-            client_id=f"worker_{client_name}",
+            client_id=f"{role}_worker_{client_name}",
             controller_infos=controller_infos,
             storage_infos=storage_infos,
         )