profiler bug fix for agent loop

Author: tangmengcheng

verl/experimental/agent_loop/agent_loop.py

@@ -837,6 +837,14 @@ def sleep(self):
         """Sleep all rollout replica instances."""
         self._run_all([replica.sleep() for replica in self.rollout_replicas])
 
+    def start_profile(self, **kwargs):
+        """Start profiling on all replicas."""
+        self._run_all([replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+    def stop_profile(self):
+        """Stop profiling on all replicas."""
+        self._run_all([replica.stop_profile() for replica in self.rollout_replicas])
+
     def _run_all(self, tasks: list[asyncio.Task]):
         async def run_all():
             await asyncio.gather(*tasks)

verl/trainer/ppo/ray_trainer.py

@@ -904,6 +904,8 @@ def _start_profiling(self, do_profile: bool) -> None:
                 self.critic_wg.start_profile(profile_step=self.global_steps)
             if self.use_rm:
                 self.rm_wg.start_profile(profile_step=self.global_steps)
+            if hasattr(self, "async_rollout_manager") and self.async_rollout_manager is not None:
+                self.async_rollout_manager.start_profile(role="rollout", profile_step=self.global_steps)
 
     def _stop_profiling(self, do_profile: bool) -> None:
         """Stop profiling for all worker groups if profiling is enabled."""
@@ -915,6 +917,8 @@ def _stop_profiling(self, do_profile: bool) -> None:
                 self.critic_wg.stop_profile()
             if self.use_rm:
                 self.rm_wg.stop_profile()
+            if hasattr(self, "async_rollout_manager") and self.async_rollout_manager is not None:
+                self.async_rollout_manager.stop_profile()
 
     def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen", keep_minibatch=False):
         """Reorder the data on single controller such that each dp rank gets similar total tokens"""

verl/workers/rollout/vllm_rollout/vllm_async_server.py

@@ -440,6 +440,16 @@ async def sleep(self):
     async def wait_for_requests_to_drain(self):
         await self.engine.wait_for_requests_to_drain()
 
+    def start_profile(self, **kwargs):
+        """Start profiling on all workers."""
+        if self.workers:
+            ray.get([worker.start_profile.remote(**kwargs) for worker in self.workers])
+
+    def stop_profile(self):
+        """Stop profiling on all workers."""
+        if self.workers:
+            ray.get([worker.stop_profile.remote() for worker in self.workers])
+
 
 @ray.remote(num_cpus=1)
 class vLLMHttpServer(vLLMHttpServerBase):
@@ -559,6 +569,14 @@ async def sleep(self):
         await self.servers[0].wait_for_requests_to_drain.remote()
         await asyncio.gather(*[server.sleep.remote() for server in self.servers])
 
+    async def start_profile(self, **kwargs):
+        """Start profiling on all servers."""
+        await asyncio.gather(*[server.start_profile.remote(**kwargs) for server in self.servers])
+
+    async def stop_profile(self):
+        """Stop profiling on all servers."""
+        await asyncio.gather(*[server.stop_profile.remote() for server in self.servers])
+
 
 def _qwen2_5_vl_dedup_image_tokens(prompt_ids: list[int], processor):
     """Deduplicate consecutive image tokens in prompt_ids for Qwen2.5-VL, since vLLM will replicate the