Add CollectiveSpec P0 phase trace harness

scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py

@@ -0,0 +1,44 @@
+"""P0 worker: record actual rank-side target execution phases."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.v1.worker.gpu_worker import Worker
+
+from .common import log_event, plan_summary
+
+
+class P0Worker(Worker):
+    """Log rank-local execution without modifying model or collective calls."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._p0_execute_epoch = 0
+
+    def _p0_rank_fields(self) -> dict[str, Any]:
+        config = self.vllm_config.parallel_config
+        return {
+            "global_rank": getattr(self, "rank", None),
+            "local_rank": getattr(self, "local_rank", None),
+            "data_parallel_rank": getattr(config, "data_parallel_rank", None),
+            "tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None),
+            "expert_parallel_size": getattr(config, "expert_parallel_size", None),
+            "tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
+            "data_parallel_size": getattr(config, "data_parallel_size", None),
+            "speculative_kmax": getattr(
+                getattr(self, "speculative_config", None), "num_speculative_tokens", None
+            ),
+        }
+
+    def execute_model(self, scheduler_output: Any):  # type: ignore[no-untyped-def]
+        self._p0_execute_epoch += 1
+        fields = {
+            "worker_phase_epoch": self._p0_execute_epoch,
+            **self._p0_rank_fields(),
+            **plan_summary(scheduler_output),
+        }
+        log_event("worker", "target_execute_begin", **fields)
+        output = super().execute_model(scheduler_output)
+        log_event("worker", "target_execute_end", **fields)
+        return output