| """Owned experiment seams; installed ONLY in separately prepared admission overlays. |
| |
| No torch import in the CPU proxy. Every seam needs the explicit admission env, |
| run-bound expiring marker, and private writable experiment mount. Never deployed |
| by ordinary PP4 preparation. Evidence contains counters/hashes, never text/logits. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import importlib |
| import json |
| import os |
| import struct |
| import time |
| from collections.abc import Callable |
| from pathlib import Path |
| from typing import TYPE_CHECKING, Protocol, cast |
|
|
| if TYPE_CHECKING: |
| import torch |
|
|
| _seen: set[tuple[int, str]] = set() |
|
|
|
|
| class Client(Protocol): |
| session: int |
| epoch: int |
| def new_session(self, reason: str) -> None: ... |
| def mark_failed(self, reason: str, *, backoff: bool = True) -> None: ... |
|
|
|
|
| class Speculator(Protocol): |
| sampled_valid: torch.Tensor |
|
|
|
|
| class Runner(Protocol): |
| speculator: Speculator |
|
|
|
|
| class Batch(Protocol): |
| idx_mapping: torch.Tensor |
|
|
|
|
| class Sample(Protocol): |
| sampled_token_ids: torch.Tensor |
| num_sampled: torch.Tensor |
|
|
|
|
| def marker() -> dict[str, object] | None: |
| run = os.environ.get('VLLM_PP4_ADMISSION_RUN') |
| if not run: |
| return None |
| root = Path(os.environ['VLLM_PP4_ROUND_DIR']) |
| path = root/'FAULT.json' |
| if not path.exists(): |
| return None |
| data = cast(dict[str, object], json.loads(path.read_text())) |
| if data.get('run') != run or not isinstance(data.get('id'), str) or float(str(data['expires'])) <= time.time(): |
| raise ValueError('PP4 admission marker owner/lease mismatch') |
| if data.get('mode') not in ('q-rank0', 'q-rank1', 'reset', 'reconnect', 'timeout', 'late-reply'): |
| raise ValueError('PP4 admission unknown fault') |
| return data |
|
|
|
|
| def event(kind: str, **values: object) -> None: |
| run = os.environ.get('VLLM_PP4_ADMISSION_RUN') |
| if not run: |
| raise ValueError('admission event outside owned experiment') |
| root = Path(os.environ['VLLM_PP4_ROUND_DIR']) |
| with (root/f'admission-{os.getpid()}.jsonl').open('a') as handle: |
| handle.write(json.dumps({'kind': kind, 'run': run, 'time': time.time(), **values}, allow_nan=False)+'\n') |
|
|
|
|
| def client_before(value: object) -> None: |
| data = marker() |
| if data is None or data['mode'] not in ('reset', 'reconnect'): |
| return |
| key = id(value), str(data['id']) |
| if key in _seen: |
| return |
| _seen.add(key) |
| client = cast(Client, value) |
| before = client.session |
| if data['mode'] == 'reset': |
| client.new_session('PP4 owned admission reset') |
| else: |
| client.mark_failed('PP4 owned admission reconnect', backoff=False) |
| event('client_fault', fault=data['mode'], id=data['id'], session_before=before, |
| session_after=client.session, epoch=client.epoch) |
|
|
|
|
| def client_after(value: object, reply: bytearray) -> None: |
| data = marker() |
| client = cast(Client, value) |
| if data is None: |
| return |
| key = id(value), str(data['id']) |
| if data['mode'] in ('timeout', 'late-reply') and key not in _seen: |
| _seen.add(key) |
| event('client_fault', fault=data['mode'], id=data['id'], session_before=client.session, |
| session_after=client.session, epoch=client.epoch) |
| if data['mode'] == 'timeout': |
| |
| |
| |
| module = importlib.import_module(type(value).__module__.rsplit('.', 1)[0]+'.transport') |
| error = cast(Callable[[str], Exception], module.TransportTimeout) |
| raise error('PP4 owned lost reply after worker application') |
| |
| |
| seq = struct.unpack_from('<Q', reply, 8)[0] |
| struct.pack_into('<Q', reply, 8, seq-1) |
| else: |
| event('client_recovered', id=data['id'], epoch=client.epoch, session=client.session) |
|
|
|
|
| def invalidate(runner: object, batch: object) -> None: |
| data = marker() |
| if data is None or data['mode'] not in ('q-rank0', 'q-rank1'): |
| return |
| import torch |
| rank = torch.distributed.get_rank() |
| valid = cast(Runner, runner).speculator.sampled_valid |
| mapping = cast(Batch, batch).idx_mapping.long() |
| flag = torch.tensor([int(rank == int(str(data['mode'])[-1]))], device=valid.device) |
| torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX) |
| if int(flag.item()) != 1: |
| raise ValueError('PP4 q failure not agreed by both ranks') |
| valid.index_fill_(0, mapping, False) |
| event('q_invalidated', id=data['id'], rank=rank, origin=int(str(data['mode'])[-1]), |
| slots=int(mapping.numel()), all_invalid=bool((~valid[mapping]).all().item())) |
|
|
|
|
| def verify(fn: Callable[..., tuple[torch.Tensor, torch.Tensor]]) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: |
| """Same real processed target logits/positions/seeds, independently target-only. |
| |
| Exact equality of every emitted token/length (not a histogram tolerance). |
| This proves coupled target-only behavior on each observed engine invocation; |
| it does not claim exhaustive equality for unobserved logits. |
| """ |
| def call(*args: object, **kwargs: object) -> tuple[torch.Tensor, torch.Tensor]: |
| data = marker() |
| actual = fn(*args, **kwargs) |
| if data is None or data['mode'] not in ('q-rank0', 'q-rank1'): |
| return actual |
| import torch |
| bound = __import__('inspect').signature(fn).bind(*args, **kwargs) |
| drafts = cast(torch.Tensor, bound.arguments['draft_sampled']) |
| local = cast(torch.Tensor, bound.arguments['expanded_local_pos']) |
| if not bool((local > 0).any().item()): |
| return actual |
| if not bool((drafts[local > 0] == -1).all().item()): |
| raise ValueError('PP4 failed q reached verifier as a proposal') |
| bound.arguments['draft_logits'] = None |
| previous = os.environ['VLLM_DSPARK_SAMPLED_TOP32'] |
| try: |
| os.environ['VLLM_DSPARK_SAMPLED_TOP32'] = '0' |
| expected = fn(*bound.args, **bound.kwargs) |
| finally: |
| os.environ['VLLM_DSPARK_SAMPLED_TOP32'] = previous |
| torch.testing.assert_close(actual[1], expected[1], rtol=0, atol=0) |
| if not bool((actual[1] == 1).all().item()): |
| raise ValueError('failed q accepted drafts') |
| torch.testing.assert_close(actual[0][:, 0], expected[0][:, 0], rtol=0, atol=0) |
| digest = hashlib.sha256(actual[0][:, 0].cpu().numpy().tobytes()).hexdigest() |
| event('target_only_exact', id=data['id'], rank=torch.distributed.get_rank(), |
| requests=int(actual[1].numel()), emitted_sha256=digest, lengths_one=True, exact=True) |
| return actual |
| return call |
|
|
|
|
| def observe_sample(output: object) -> None: |
| """Prove model EOS from actual sampler IDs even when API stop_reason is null.""" |
| run = os.environ.get('VLLM_PP4_ADMISSION_RUN') |
| if not run: |
| return |
| path = Path(os.environ['VLLM_PP4_ROUND_DIR'])/'EOS.json' |
| if not path.exists(): |
| return |
| data = cast(dict[str, object], json.loads(path.read_text())) |
| if data.get('run') != run or float(str(data['expires'])) <= time.time(): |
| raise ValueError('EOS observation owner/lease mismatch') |
| import torch |
| sampled = cast(Sample, output).sampled_token_ids |
| lengths = cast(Sample, output).num_sampled |
| valid = torch.arange(sampled.shape[1], device=sampled.device)[None, :] < lengths[:, None] |
| eos = int(str(data['eos'])) |
| seen = bool(((sampled == eos) & valid).any().item()) |
| if seen: |
| event('model_eos_sampled', rank=torch.distributed.get_rank(), eos=eos, observed=True) |
|
|
