add metric first_micro_batch_rollout_time in fully diagg mode.

PiperOrigin-RevId: 848314187

Author: jiangyangmu

tests/perf/export_test.py

@@ -260,6 +260,7 @@ def test_export_grpo_metrics_fully_disaggregated(self, mock_perf_counter):
         "perf/global_step_time": 1.3,
         "perf/weight_sync_time": 0.1,
         "perf/rollout_idle_time": 0.69,
+        "perf/first_micro_batch_rollout_time": 0.41,
         "perf/sum/rollout_time": 0.11,
         "perf/sum/refer_inference_time": 0.11,
         "perf/sum/actor_train_time": 0.11,

tunix/perf/export.py

@@ -288,6 +288,10 @@ def _grpo_metrics_fully_disaggregated(
         span.duration for span in actor_train_step_spans
     ]
 
+    first_micro_batch_rollout_time: float = (
+        rollout_spans[0].end - global_step_group.begin
+    )
+
     # append [0.0] to make size equal to micro batch
     actor_gap_time: list[float] = [
         b.end - a.begin
@@ -299,6 +303,7 @@ def _grpo_metrics_fully_disaggregated(
         "perf/global_step_time": (global_step_time, None),
         "perf/weight_sync_time": (weight_sync_time, None),
         "perf/rollout_idle_time": (rollout_idle_time, None),
+        "perf/first_micro_batch_rollout_time": (first_micro_batch_rollout_time, None),
         "perf/sum/rollout_time": (np.sum(rollout_time), None),
         "perf/sum/refer_inference_time": (np.sum(refer_inference_time), None),
         "perf/sum/refer_gap_time": (np.sum(refer_gap_time), None),