File size: 3,888 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 | """Offline PP4 overlay preparation, composed on PP3's exact pinned probability verifier."""
from __future__ import annotations
import argparse
import ast
import hashlib
import json
from pathlib import Path
from offload.dspark.probabilistic.pp3 import prepare_verifier as pp3
HELPER = 'vllm/v1/worker/gpu/dsv41_pp4_round_trace.py'
K1_SHA256 = '682eb895b8fb3d68b9580954b6cda255a49f882bddb4b06342d7d96258eb6e62'
ROOT = Path(__file__).resolve().parents[3]
K1_RUNNER = ROOT / 'offload/kvpack/overlay' / pp3.RUNNER
def instrument(text: str) -> str:
tree = ast.parse(text)
lines = text.splitlines(keepends=True)
hits: set[str] = set()
edits: list[tuple[int, str]] = []
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name in ('execute_model', 'sample_tokens'):
if node.name in hits:
raise ValueError('PP4 duplicate model-runner timing anchor')
hits.add(node.name)
decorator = 'execution' if node.name == 'execute_model' else 'sampling'
at = min([node.lineno, *(x.lineno for x in node.decorator_list)]) - 1
edits.append((at, ' ' * node.col_offset + '@_pp4_round.' + decorator + '\n'))
if hits != {'execute_model', 'sample_tokens'}:
raise ValueError('PP4 model-runner timing anchors changed')
for at, insert in sorted(edits, reverse=True):
lines.insert(at, insert)
result = 'from vllm.v1.worker.gpu import dsv41_pp4_round_trace as _pp4_round\n' + ''.join(lines)
compile(result, '<pp4-round-overlay>', 'exec')
return result
def prepare(source: Path, output: Path, k1_runner: Path = K1_RUNNER) -> dict[str, object]:
if output.exists():
raise ValueError('output must not exist; refuse overwriting any prepared overlay')
manifest = json.loads((ROOT / 'offload/kvpack/MANIFEST.json').read_text())
entries = [entry for entry in manifest['files'] if entry['container_path'].endswith('/' + pp3.RUNNER)]
if len(entries) != 1 or entries[0]['overlay_sha256'] != K1_SHA256:
raise ValueError('reviewed K1 manifest runner pin changed')
runner_data = k1_runner.read_bytes()
utils_data = (source / pp3.UTILS).read_bytes()
pins = {pp3.RUNNER: K1_SHA256, pp3.UTILS: pp3.PINS[pp3.UTILS]}
for name, data in ((pp3.RUNNER, runner_data), (pp3.UTILS, utils_data)):
if hashlib.sha256(data).hexdigest() != pins[name]:
raise ValueError('PP4 pinned source mismatch: ' + name)
# Compose the narrow probability/failure mask and timing additions onto
# the FULL production K1 source. Never replace it with the old image runner.
texts = {pp3.RUNNER: instrument(pp3.runner(runner_data.decode())),
pp3.UTILS: pp3.utilities(utils_data.decode()),
HELPER: Path(__file__).with_name('pp4_round_trace.py').read_text()}
for name, text in texts.items():
compile(text, name, 'exec')
output.mkdir(parents=True)
for name, text in texts.items():
path = output / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(text)
outputs = {name: hashlib.sha256((output / name).read_bytes()).hexdigest() for name in texts}
receipt: dict[str, object] = {'abi': pp3.ABI, 'source_sha256': pins, 'production_k1_sha256': K1_SHA256,
'output_sha256': outputs, 'round_scope': 'execute_model through sample_tokens completed wall/CUDA',
'profiled_only': True}
(output / 'receipt.json').write_text(json.dumps(receipt, indent=2) + '\n')
return receipt
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--output', type=Path, required=True)
args = parser.parse_args()
print(json.dumps(prepare(args.source, args.output), indent=2))
if __name__ == '__main__':
main()
|