[Fix] vLLM Metrics Scrapper throughput calculation (NovaSky-AI#1794)

For more context checkout: 
[NovaSky-AI#1636 comment
1](https://github.com/NovaSky-AI/SkyRL/pull/1636/changes#r3411074771)
[NovaSky-AI#1636 comment
2](https://github.com/NovaSky-AI/SkyRL/pull/1636/changes#r3411079924)

The current vLLM Metrics Scrapper calculates rollout throughput with the
full step size time, and doesn't account for variances caused during
eval (different sampling params, eval batch sizes). This PR addresses
those issues by updating the timing logic and separating metrics for
training and eval rollout (for synchronous RL, in async the rollouts are
combined)

Author: zanderjiang

skyrl/train/evaluate.py

@@ -1,6 +1,7 @@
+import time
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import torch
 from loguru import logger
@@ -29,6 +30,9 @@
     validate_generator_output,
 )
 
+if TYPE_CHECKING:
+    from skyrl.train.utils.vllm_metrics_scraper import VLLMMetricsScraper
+
 
 @torch.no_grad()
 async def evaluate(
@@ -37,6 +41,7 @@ async def evaluate(
     cfg: SkyRLTrainConfig,
     global_step: int | None,
     tokenizer: AutoTokenizer,
+    vllm_metrics_scraper: Optional["VLLMMetricsScraper"] = None,
 ) -> Dict[str, float]:
     """Runs generation and evaluation of trajectories.
 
@@ -47,6 +52,9 @@ async def evaluate(
         global_step (int | None): current global step, or
             `None` to indicate a non-training context (e.g., eval-only)
         tokenizer (AutoTokenizer): tokenizer to use
+        vllm_metrics_scraper: when set, the open ``vllm/eval`` window is resumed
+            around each generation and paused after, so only generation time
+            counts toward eval throughput.
 
     Returns:
         Dict[str, float]: evaluation metrics
@@ -58,6 +66,7 @@ async def evaluate(
     concat_env_extras: List[Dict[str, Any]] = []
     concat_uids: List[str] = []
     sampling_params = cfg.generator.eval_sampling_params
+    eval_generate_time = 0.0
     pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress")
     for _, prompts in enumerate(eval_dataloader):
         pbar.update(1)
@@ -69,7 +78,13 @@ async def evaluate(
             "eval",
             global_step,
         )
+        gen_start = time.monotonic()
+        if vllm_metrics_scraper is not None:
+            vllm_metrics_scraper.resume()
         generator_output: GeneratorOutput = await generator.generate(generator_input)
+        if vllm_metrics_scraper is not None:
+            vllm_metrics_scraper.pause()
+        eval_generate_time += time.monotonic() - gen_start
         validate_generator_output(len(generator_input["prompts"]), generator_output)
         generator_outputs.append(generator_output)
         concat_all_envs.extend(generator_input["env_classes"])
@@ -127,6 +142,7 @@ async def evaluate(
                 eval_metrics,
             )
 
+    eval_metrics["timing/eval_generate"] = eval_generate_time
     return eval_metrics
 
 
@@ -137,6 +153,7 @@ async def evaluate_step_wise(
     cfg: SkyRLTrainConfig,
     global_step: int | None,
     tokenizer: AutoTokenizer,
+    vllm_metrics_scraper: Optional["VLLMMetricsScraper"] = None,
 ) -> Dict[str, float]:
     """Runs generation and evaluation of trajectories for step-wise training.
 
@@ -149,6 +166,9 @@ async def evaluate_step_wise(
         global_step (int | None): current global step, or
             `None` to indicate a non-training context (e.g., eval-only)
         tokenizer (AutoTokenizer): tokenizer to use
+        vllm_metrics_scraper: when set, the open ``vllm/eval`` window is resumed
+            around each generation and paused after, so only generation time
+            counts toward eval throughput.
 
     Returns:
         Dict[str, float]: evaluation metrics
@@ -160,6 +180,7 @@ async def evaluate_step_wise(
     concat_env_extras: List[Dict[str, Any]] = []
     concat_uids: List[str] = []
     sampling_params = cfg.generator.eval_sampling_params
+    eval_generate_time = 0.0
     pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress")
     for _, prompts in enumerate(eval_dataloader):
         pbar.update(1)
@@ -171,7 +192,13 @@ async def evaluate_step_wise(
             "eval",
             global_step,
         )
+        gen_start = time.monotonic()
+        if vllm_metrics_scraper is not None:
+            vllm_metrics_scraper.resume()
         generator_output: GeneratorOutput = await generator.generate(generator_input)
+        if vllm_metrics_scraper is not None:
+            vllm_metrics_scraper.pause()
+        eval_generate_time += time.monotonic() - gen_start
         traj_id_to_input = {
             traj_id.instance_id: {"env_class": env_class, "env_extras": env_extra}
             for traj_id, env_class, env_extra in zip(
@@ -244,4 +271,5 @@ async def evaluate_step_wise(
                 eval_metrics,
             )
 
+    eval_metrics["timing/eval_generate"] = eval_generate_time
     return eval_metrics

skyrl/train/trainer.py

@@ -204,13 +204,18 @@ def _build_train_dataloader_and_compute_training_steps(self):
                 self.total_training_steps = min(self.total_training_steps, self.cfg.trainer.max_training_steps)
 
     @torch.no_grad()
-    async def eval(self) -> Dict[str, float]:
+    async def eval(self, vllm_metrics_scraper: Optional[VLLMMetricsScraper] = None) -> Dict[str, float]:
         """
         Run generation and scoring on the evaluation dataset.
 
         The eval metrics are recorded after having finished training `self.global_step` steps.
         Metrics recorded in global_step 0 corresponds to evaluations before training.
 
+        Args:
+            vllm_metrics_scraper: when provided, the eval loop calls
+                ``resume()``/``pause()`` around each generation so the scraper
+                attributes only generation time to the open ``vllm/eval`` window.
+
         Returns:
             A dictionary of evaluation metrics.
         """
@@ -221,6 +226,7 @@ async def eval(self) -> Dict[str, float]:
                 cfg=self.cfg,
                 global_step=self.global_step,
                 tokenizer=self.tokenizer,
+                vllm_metrics_scraper=vllm_metrics_scraper,
             )
         else:
             eval_metrics = await evaluate(
@@ -229,6 +235,7 @@ async def eval(self) -> Dict[str, float]:
                 cfg=self.cfg,
                 global_step=self.global_step,
                 tokenizer=self.tokenizer,
+                vllm_metrics_scraper=vllm_metrics_scraper,
             )
         return eval_metrics
 
@@ -294,6 +301,13 @@ async def train(self):
                 if not step_started:
                     self._fire("on_step_start")
                     step_started = True
+                    # Open the train-rollout metrics window once per logical
+                    # step; paused so only the generation spans count toward the
+                    # throughput denominator (dynamic sampling may generate more
+                    # than once before the step completes).
+                    if self._vllm_metrics_scraper is not None:
+                        await self._vllm_metrics_scraper.start("vllm/train")
+                        self._vllm_metrics_scraper.pause()
                 with Timer("step", self.all_timings):
                     # for colocate_all=true, inference engine is always on GPU when starting the training step
 
@@ -311,8 +325,12 @@ async def train(self):
                     )
 
                     # 1.1. generation phase
+                    if self._vllm_metrics_scraper is not None:
+                        self._vllm_metrics_scraper.resume()
                     with Timer("generate", self.all_timings):
                         generator_output: GeneratorOutput = await self.generate(generator_input)
+                    if self._vllm_metrics_scraper is not None:
+                        self._vllm_metrics_scraper.pause()
 
                     if self.cfg.generator.step_wise_trajectories:
                         # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids`
@@ -331,6 +349,13 @@ async def train(self):
                         # if we are not continuing sampling, we sleep the inference engine
                         await self.inference_engine_client.sleep()
 
+                    # The train rollout for this step is done generating; close
+                    # its metrics window. ``vllm/eval/*`` is collected separately
+                    # around eval below.
+                    vllm_metrics: Dict[str, float] = {}
+                    if self._vllm_metrics_scraper is not None:
+                        vllm_metrics = await self._vllm_metrics_scraper.stop()
+
                     # 1.2 postprocess rewards (and merge step-wise turns if enabled)
                     with Timer("postprocess_generator_output", self.all_timings):
                         generator_output, uids = self.postprocess_generator_output(generator_output, uids)
@@ -434,18 +459,26 @@ async def train(self):
                     or self.global_step == self.total_training_steps
                 )
                 if force_eval or interval_eval:
+                    # Open the eval-rollout window; the scraper itself measures
+                    # the generation spans via resume()/pause() inside eval().
+                    if self._vllm_metrics_scraper is not None:
+                        await self._vllm_metrics_scraper.start("vllm/eval")
+                        self._vllm_metrics_scraper.pause()
                     self._fire("on_eval_start")
                     with Timer("eval", self.all_timings):
-                        eval_metrics = await self.eval()
+                        eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper)
                         self.all_metrics.update(eval_metrics)
                     self._fire("on_eval_end", metrics=eval_metrics)
+                    if self._vllm_metrics_scraper is not None:
+                        vllm_metrics.update(await self._vllm_metrics_scraper.stop())
 
                 log_payload = {
                     **self.all_metrics,
                     **{f"timing/{k}": v for k, v in self.all_timings.items()},
+                    # vllm/train/* = train rollout, vllm/eval/* = eval rollout,
+                    # each over its own generation time (owned by the scraper).
+                    **vllm_metrics,
                 }
-                if self._vllm_metrics_scraper is not None:
-                    log_payload.update(await self._vllm_metrics_scraper.sample())
 
                 if self._ray_gpu_monitor is not None:
                     log_payload.update(self._ray_gpu_monitor.flush())