45 lines
1.7 KiB
Python
45 lines
1.7 KiB
Python
"""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
|