Trace P0 batch-level DP coordination

scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py

@@ -15,6 +15,38 @@ class P0Worker(Worker):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._p0_execute_epoch = 0
+        self._p0_batch_epoch = 0
+
+    def init_device(self):  # type: ignore[no-untyped-def]
+        result = super().init_device()
+        model_runner = self.model_runner
+        if getattr(model_runner, "_collectivespec_p0_wrapped", False):
+            return result
+        original = model_runner._determine_batch_execution_and_padding
+
+        def traced_determine(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
+            output = original(*args, **kwargs)
+            self._p0_batch_epoch += 1
+            scheduler_output = kwargs.get("scheduler_output")
+            if scheduler_output is None and args:
+                scheduler_output = args[0]
+            cudagraph_mode, batch_desc, should_ubatch, rows_across_dp, _ = output
+            log_event(
+                "worker",
+                "batch_execution_plan",
+                batch_phase_epoch=self._p0_batch_epoch,
+                **self._p0_rank_fields(),
+                **plan_summary(scheduler_output),
+                cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
+                physical_batch_rows=getattr(batch_desc, "num_tokens", None),
+                should_ubatch=should_ubatch,
+                rows_across_dp=rows_across_dp,
+            )
+            return output
+
+        model_runner._determine_batch_execution_and_padding = traced_determine
+        model_runner._collectivespec_p0_wrapped = True
+        return result
 
     def _p0_rank_fields(self) -> dict[str, Any]:
         config = self.vllm_config.parallel_config