File size: 8,153 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """Diagnostic host wrappers only. Never add synchronization between peer launches."""
from __future__ import annotations
import functools
import importlib
import json
import os
import threading
import time
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from pathlib import Path
from typing import Any
OUT = Path(os.environ.get('PP4_SERVER_OUT', '.'))
STAGE = 'startup'
EVENTS: list[dict[str, Any]] = []
LOCK = threading.Lock()
EPOCHS: dict[int, list[int]] = {}
DEFER_EVENTS = False
PENDING: list[dict[str, Any]] = []
def emit(event: str, **fields: object) -> None:
row = {'event': event, 'stage': STAGE, 'host_ns': time.perf_counter_ns(),
'pid': os.getpid(), 'python_thread': threading.get_ident(), 'native_thread': threading.get_native_id(), **fields}
EVENTS.append(row)
if DEFER_EVENTS:
# No lock, file operation or print while a paired launch can be in
# flight. The serialized worker owns this interval; append is under GIL.
PENDING.append(row)
return
with LOCK:
rows = [*PENDING, row]
PENDING.clear()
with (OUT/'phases.jsonl').open('a') as handle:
for item in rows:
handle.write(json.dumps(item, allow_nan=False)+'\n')
print('\n'.join('PP4 SERVER '+json.dumps(item, allow_nan=False) for item in rows), flush=True)
@contextmanager
def phase(name: str) -> Iterator[None]:
global DEFER_EVENTS
emit('BEGIN', phase=name)
start = time.perf_counter_ns()
try:
yield
except BaseException as exc:
if name in ('model._replay', 'selector.select'):
DEFER_EVENTS = False
emit('FAIL', phase=name, error=type(exc).__name__, elapsed_ns=time.perf_counter_ns()-start)
raise
else:
if name in ('model._replay', 'selector.select'):
DEFER_EVENTS = False
emit('END', phase=name, elapsed_ns=time.perf_counter_ns()-start)
def patch(target: Any, name: str, label: str | None = None) -> None:
original = getattr(target, name)
@functools.wraps(original)
def call(*args: Any, **kwargs: Any) -> Any:
with phase(label or name):
return original(*args, **kwargs)
setattr(target, name, call)
def snapshot(ctx: Any) -> None:
"""Called ONLY after original result(), hence both submissions and drain.
Do not infer device overlap from these host records. Exact advancing peer
state proves completed exchange, not a CUDA activity trace.
"""
import torch
states = {r: {k: ctx.state[r][k].detach().cpu() for k in ('sticky', 'error', 'done', 'epoch', 'ids', 'q')}
for r in (0, 1)}
healthy = all(not bool(states[r][k].any().item()) for r in (0, 1) for k in ('sticky', 'error', 'done'))
exact = all(torch.equal(states[0][k], states[1][k]) for k in ('epoch', 'ids', 'q'))
outputs = getattr(ctx, '_pp4_probe_outputs', {})
proposed_exact = (set(outputs) == {0, 1} and torch.equal(outputs[0].detach().cpu(), outputs[1].detach().cpu()))
epochs = [int(x) for x in states[0]['epoch'].tolist()]
previous = EPOCHS.get(id(ctx))
advancing = previous is None or all(a > b for a, b in zip(epochs, previous, strict=True))
# Keep the context alive through the model graph dictionary. Eager contexts
# can reuse object IDs, so monotonicity is enforced only for graph records.
if ctx in [row['context'] for row in getattr(ctx.model, 'sampled_graphs', {}).values()]:
EPOCHS[id(ctx)] = epochs
else:
advancing = True
emit('peer-state-after-drain', healthy=healthy, both_rank_exact=exact, proposed_ids_exact=proposed_exact, advancing=advancing,
epochs=epochs, batch=ctx.batch, depth=ctx.depth)
if not (healthy and exact and proposed_exact and advancing):
raise RuntimeError('peer sticky/error/done/epoch or exact q/ID failure after drain')
def install(namespace: dict[str, Any]) -> None:
"""Called at entry to REAL worker.server.main, after definitions, before build."""
import torch
sampled_fused: Any = importlib.import_module('worker.sampled_fused')
tp2 = importlib.import_module('worker.tp2')
for name in ('capture', '_capture_one', 'draft', '_draft', 'heads', 'draft_sampled', '_replay', 'insert_eager',
'insert_taps_graphed', 'insert_mx_eager'):
patch(tp2.TPDrafter, name, 'model.'+name)
for name in ('reset', 'stage', 'flush', 'discard', 'draft_sampled'):
patch(tp2.TP2Backend, name, 'backend.'+name)
patch(namespace['DraftCore'], 'handle', 'DraftCore.handle')
patch(namespace['DraftCore'], '_step', 'DraftCore._step')
patch(namespace['P'], 'encode_step_reply', 'reply.serialization')
patch(namespace['P'], 'reply', 'frame.serialization')
original_select = sampled_fused.FusedContext.select
def select(ctx: Any, ranks: Any, logits: Any, previous: Any, outputs: Any, position: int) -> None:
# Retain existing tensor references only; inspect AFTER both submissions
# and result drain, including the physical rank1 proposal IDs.
if not hasattr(ctx, '_pp4_probe_outputs'):
ctx._pp4_probe_outputs = {}
ctx._pp4_probe_outputs.update(outputs)
with phase('selector.select'):
original_select(ctx, ranks, logits, previous, outputs, position)
sampled_fused.FusedContext.select = select
original_result = sampled_fused.FusedContext.result
def result(ctx: Any) -> Any:
with phase('selector.result'):
value = original_result(ctx)
snapshot(ctx)
return value
sampled_fused.FusedContext.result = result
original_replay = torch.cuda.CUDAGraph.replay
def replay(graph: Any) -> None:
global DEFER_EVENTS
DEFER_EVENTS = True
# Only host device/stream queries; no event, item, D2H or synchronize.
with phase('graph.rank-submission'):
emit('graph-stream', device=torch.cuda.current_device(),
stream=int(torch.cuda.current_stream().cuda_stream))
original_replay(graph)
graph_type: Any = torch.cuda.CUDAGraph # diagnostic monkeypatch of the untyped CUDA extension boundary
graph_type.replay = replay
# This interception sees actual lazy handle initialization anywhere in the
# model; no implementation or kernel source is copied into the probe.
compiler = importlib.import_module('triton.compiler.compiler')
patch(compiler.CompiledKernel, '_init_handles', 'triton.handle-load')
jit = importlib.import_module('triton.runtime.jit')
patch(jit.JITFunction, 'warmup', 'triton.warmup')
original_kernel = sampled_fused._kernel
class Compiled:
def __init__(self, value: Any):
self.value = value
@property
def run(self) -> Callable[..., Any]:
with phase('selector.handle-resolution'):
value: Callable[..., Any] = self.value.run
return value
class Kernel:
def warmup(self, *args: Any, **kwargs: Any) -> Compiled:
with phase('selector.warmup'):
return Compiled(original_kernel.warmup(*args, **kwargs))
def __getitem__(self, grid: Any) -> Callable[..., Any]:
def launch(*args: Any, **kwargs: Any) -> Any:
global DEFER_EVENTS
DEFER_EVENTS = True
with phase('selector.rank-submission'):
emit('selector-stream', device=torch.cuda.current_device(),
stream=int(torch.cuda.current_stream().cuda_stream), position=int(args[43]))
return original_kernel[grid](*args, **kwargs)
return launch
sampled_fused._kernel = Kernel()
ctypes_module = namespace['ctypes']
class CtypesProxy:
def __getattr__(self, name: str) -> Any:
if name != 'memmove':
return getattr(ctypes_module, name)
def memmove(*args: Any) -> Any:
with phase('native.reply-copy'):
return ctypes_module.memmove(*args)
return memmove
namespace['ctypes'] = CtypesProxy()
|