File size: 4,864 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Experiment-only complete model-runner round timing; disabled without a numeric marker.
Installed as a read-only engine helper. No torch import/CUDA work while disabled. Profiled cells deliberately
synchronize completed rounds; their request rates NEVER enter the unprofiled headline throughput result.
"""
from __future__ import annotations
import hashlib
import importlib
import json
import os
import re
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from functools import wraps
from pathlib import Path
from typing import Concatenate, ParamSpec, Protocol, TypeVar, cast
P = ParamSpec('P')
R = TypeVar('R')
class Event(Protocol):
def record(self) -> None: ...
def synchronize(self) -> None: ...
def elapsed_time(self, other: Event) -> float: ...
class Scheduler(Protocol):
@property
def num_scheduled_tokens(self) -> Mapping[str, int]: ...
@property
def scheduled_spec_decode_tokens(self) -> Mapping[str, list[int]]: ...
@dataclass
class Span:
marker: str
root: Path
rank: int
batch: int
sequence: int
request_hash: str
start: Event
wall: float
_active: dict[int, Span] = {}
_sequences: dict[tuple[int, str], int] = {}
def begin(runner: object, scheduler: object, dummy: bool) -> None:
if dummy or not os.environ.get('VLLM_PP4_ROUND_DIR'):
return
root = Path(os.environ['VLLM_PP4_ROUND_DIR'])
marker_path = root / 'ACTIVE'
if not marker_path.exists():
return
marker = marker_path.read_text().strip()
if re.fullmatch('[A-Za-z0-9:_-]+', marker) is None:
raise ValueError('PP4 invalid profile cell marker')
scheduled_state = cast(Scheduler, scheduler)
scheduled = scheduled_state.num_scheduled_tokens
drafts = scheduled_state.scheduled_spec_decode_tokens
# Complete pure decode rounds only: no prefill or mixed batch can contaminate the p95.
if not scheduled or any(not drafts.get(k) or n != len(drafts[k]) + 1 for k, n in scheduled.items()):
return
key = id(runner)
if key in _active:
raise ValueError('PP4 prior execute_model round has no sample_tokens completion')
seq_key = key, marker
sequence = _sequences.get(seq_key, 0)
_sequences[seq_key] = sequence + 1
torch = importlib.import_module('torch')
rank = cast(Callable[[], int], torch.distributed.get_rank)()
if rank not in (0, 1):
raise ValueError('PP4 complete-round profile requires the reviewed two-rank target')
event_factory = cast(Callable[..., Event], torch.cuda.Event)
start = event_factory(enable_timing=True)
start.record()
_active[key] = Span(marker, root, rank, len(scheduled), sequence,
hashlib.sha256('\n'.join(sorted(scheduled)).encode()).hexdigest(), start, time.perf_counter())
def finish(runner: object, failed: bool = False) -> None:
span = _active.pop(id(runner), None)
if span is None:
return
torch = importlib.import_module('torch')
factory = cast(Callable[..., Event], torch.cuda.Event)
end = factory(enable_timing=True)
end.record()
end.synchronize() # PROFILED cells only; complete all target/verification/proposal stream work.
wall_ms = 1000 * (time.perf_counter() - span.wall)
cuda_ms = span.start.elapsed_time(end)
row = {'schema': 'pp4-complete-round-v1', 'cell': span.marker, 'rank': span.rank,
'round': span.sequence, 'batch': span.batch, 'request_hash': span.request_hash,
'complete_wall_ms': wall_ms, 'cuda_elapsed_ms': cuda_ms, 'failed': failed,
'scope': 'execute_model entry through sample_tokens completion plus end-event synchronization'}
with (span.root / f'rank{span.rank}.jsonl').open('a') as out:
out.write(json.dumps(row, allow_nan=False) + '\n')
def execution(fn: Callable[Concatenate[object, P], R]) -> Callable[Concatenate[object, P], R]:
@wraps(fn)
def call(runner: object, /, *args: P.args, **kwargs: P.kwargs) -> R:
scheduler = args[0] if args else kwargs.get('scheduler_output')
dummy = bool(kwargs.get('dummy_run', False)) or (len(args) > 2 and bool(args[2]))
begin(runner, scheduler, dummy)
try:
result = fn(runner, *args, **kwargs)
except BaseException:
finish(runner, True)
raise
if getattr(runner, 'execute_model_state', None) is None:
finish(runner, True)
return result
return call
def sampling(fn: Callable[Concatenate[object, P], R]) -> Callable[Concatenate[object, P], R]:
@wraps(fn)
def call(runner: object, /, *args: P.args, **kwargs: P.kwargs) -> R:
try:
result = fn(runner, *args, **kwargs)
except BaseException:
finish(runner, True)
raise
finish(runner)
return result
return call
|