Blame for code/offload/launch/windows/dist_gate.py

"""WIN2-R3 distributional gate for candidates (F-METHOD §3.4): is the candidate's next-token distribution the same model?

Probe: 64 public-corpus prompts (32 short ~1K tokens, 24 at 32K, 8 at 150K; offsets registered below), each tokenized
once (/tokenize, raw text), then the top-20 next-token logprobs at 8 FIXED positions: the prompt's last 8 tokens are the
fixed continuation, and position k asks for the next token after ids[:L - 8 + k] (/v1/completions, max_tokens 1,
temperature 0, logprobs 20, token-id prompts, one cache salt per run so only the first probe of a prompt prefills). No
generation, so no generated text exists; only token ids and logprob vectors are recorded, on Controller (never git).

Compare (control || candidate) per position with greedy_kl.top20_kl (shared top-20 + OTHER; common mass >= 0.99 is
reported per position) and top-1 agreement. Gate against an A/A band registered BEFORE use from two production starts:
- kl_ceiling = 1.5 x the 99th percentile of the A/A MEAN KL (seeded bootstrap over the A/A positions); LOSS when the
  candidate's mean KL exceeds it (simulated no-effect false LOSS 0.0-0.2 %);
- top-1: F-METHOD's literal rule (LOSS when agreement < A/A minimum - 1) rejects a no-effect candidate 13-34 % of the
  time once the A/A itself shows >= 1 expected near-tie flip (one A/A run is a noisy minimum). The gate therefore uses an
  exact conditional test of the two flip counts instead: with D_c candidate and D_a A/A disagreements over k registered
  A/A comparisons, D_c | D_c + D_a ~ Binomial(D_c + D_a, 1 / (k + 1)) under equal flip rates; LOSS when the one-sided
  p < 0.01 (false LOSS <= 1 % by construction). The literal rule is still reported (`fmethod_top1_rule`).
- common mass: positions whose shared top-20 mass is < 0.99 (where the coarse KL is not trustworthy) are counted; LOSS
  when the candidate has significantly more of them than the A/A (same exact test). Greedy identity stays the restore
gate (W8), unchanged. COST (ESTIMATED from MEASURED TTFTs): one cold prefill per prompt per arm, 32 x ~0.5 s + 24 x ~15 s
+ 8 x ~69 s ~ 15.5 min per arm; `--subset short,32k` drops the 150K prompts (~6.5 min).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import random
import statistics
import sys
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol, cast

from offload.gb10_budget import window_cells as wc
from offload.launch.windows import greedy_kl
from offload.prefill_maxlen.evidence import Json, number, record

SCHEMA = 'dsv41-dist-gate-v1'
POSITIONS = 8
KL_FACTOR = 1.5
BOOTSTRAP = 2000
BOOTSTRAP_SEED = 20261027
TOP1_ALPHA = 0.01
COORDINATOR = 'Deepseek/coordinator'


@dataclass(frozen=True)
class DistPrompt:
    key: str
    target: int
    offset: int


def _prompts() -> tuple[DistPrompt, ...]:
    out = [DistPrompt(f'dist-short-{i:02d}', 1024, 200_000 + 10_000 * i) for i in range(32)]
    out += [DistPrompt(f'dist-32k-{i:02d}', 32000, 1_000_000 + 120_000 * i) for i in range(24)]
    out += [DistPrompt(f'dist-150k-{i}', 150000, 4_200_000 + 560_000 * i) for i in range(8)]
    return tuple(out)


DIST_PROMPTS = _prompts()
SUBSETS = {'short': 1024, '32k': 32000, '150k': 150000}


def selected(subset: Sequence[str] | None) -> tuple[DistPrompt, ...]:
    if not subset:
        return DIST_PROMPTS
    unknown = set(subset) - set(SUBSETS)
    if unknown:
        raise ValueError(f'unknown prompt subsets {sorted(unknown)}; choose from {sorted(SUBSETS)}')
    targets = {SUBSETS[s] for s in subset}
    return tuple(p for p in DIST_PROMPTS if p.target in targets)


class DistApi(Protocol):
    def model(self) -> dict[str, Json]: ...
    def token_ids(self, text: str) -> list[int]: ...
    def top20(self, prefix: Sequence[int], salt: str) -> dict[str, Json]: ...


class Http(wc.Http):
    def token_ids(self, text: str) -> list[int]:
        r = self.post('/tokenize', {'model': self.model()['id'], 'prompt': text, 'add_special_tokens': True})
        ids = r.get('tokens')
        if not isinstance(ids, list) or not all(isinstance(t, int) and not isinstance(t, bool) for t in ids):
            raise ValueError('tokenize returned no token ids')
        return cast(list[int], ids)

    def top20(self, prefix: Sequence[int], salt: str) -> dict[str, Json]:
        r = self.post('/v1/completions', {'model': self.model()['id'], 'prompt': cast(Json, list(prefix)), 'max_tokens': 1,
                                          'temperature': 0, 'logprobs': 20, 'return_tokens_as_token_ids': True,
                                          'priority': 10, 'cache_salt': salt})
        choices = r.get('choices')
        if not isinstance(choices, list) or len(choices) != 1:
            raise ValueError('one completion choice expected')
        tops = record(record(choices[0])['logprobs'])['top_logprobs']
        if not isinstance(tops, list) or len(tops) != 1:
            raise ValueError('one top-logprob position expected')
        return record(tops[0])


def cuts(api: DistApi, materials: str, hashes: Mapping[str, Json], prompts: Sequence[DistPrompt]) -> dict[str, Json]:
    """Numeric cuts (characters, tokens, sha256) of the registered prompts; built once, reused by every arm."""
    rows: dict[str, Json] = {}
    for p in prompts:
        chars, tokens = wc.cut(materials, wc.Prompt(p.key, p.target, p.offset), lambda t: len(api.token_ids(t)))
        text = wc.render(materials, p.offset, chars)
        rows[p.key] = {'target': p.target, 'offset': p.offset, 'chars': chars, 'tokens': tokens,
                       'sha256': hashlib.sha256(text.encode()).hexdigest()}
    return {'schema': SCHEMA, 'type': 'cuts', 'corpus_sha256': cast(Json, dict(hashes)), 'prompts': rows}


def probe(api: DistApi, materials: str, cut_rows: Mapping[str, Json], salt_prefix: str) -> dict[str, Json]:
    """Top-20 at the 8 fixed positions of every prompt in `cut_rows` on the server under test."""
    if not salt_prefix.strip():
        raise ValueError('a cache salt prefix per run is required')
    positions: list[Json] = []
    for key, value in record(cast(Json, dict(cut_rows))).items():
        row = record(value)
        text = wc.render(materials, int(number(row['offset'])), int(number(row['chars'])))
        if hashlib.sha256(text.encode()).hexdigest() != row.get('sha256'):
            raise ValueError(f'{key}: rebuilt prompt differs from its registered cut')
        ids = api.token_ids(text)
        if len(ids) <= POSITIONS:
            raise ValueError(f'{key}: prompt shorter than the fixed continuation')
        salt = f'{salt_prefix}:dist:{key}'
        for k in range(POSITIONS):
            at = len(ids) - POSITIONS + k
            positions.append({'prompt': key, 'position': at, 'k': k,
                              'prefix_sha256': hashlib.sha256(json.dumps(ids[:at]).encode()).hexdigest(),
                              'top20': api.top20(ids[:at], salt)})
    return {'schema': SCHEMA, 'type': 'probe', 'positions': positions, 'utc': time.strftime('%Y-%m-%dT%H:%M:%SZ')}


def _top1(top: Mapping[str, Json]) -> str:
    return max(top, key=lambda k: number(top[k]))


def compare(control: Mapping[str, Json], candidate: Mapping[str, Json]) -> dict[str, Json]:
    a, b = [record(p) for p in cast(list[Json], control['positions'])], [record(p) for p in cast(list[Json], candidate['positions'])]
    if [(p['prompt'], p['position'], p['prefix_sha256']) for p in a] != [(p['prompt'], p['position'], p['prefix_sha256']) for p in b]:
        raise ValueError('control and candidate probed different prompts or positions')
    per: list[Json] = []
    kls: list[float] = []
    agree = low = 0
    for x, y in zip(a, b):
        r = greedy_kl.top20_kl(record(x['top20']), record(y['top20']))
        same = _top1(record(x['top20'])) == _top1(record(y['top20']))
        agree += same
        low += r['common_mass'] < greedy_kl.MIN_COMMON_MASS
        kls.append(r['kl_nats'])
        per.append({'prompt': x['prompt'], 'k': x['k'], 'kl_nats': r['kl_nats'], 'common_mass': r['common_mass'],
                    'top1_agree': same})
    return {'schema': SCHEMA, 'type': 'compare', 'positions': len(kls), 'mean_kl': statistics.fmean(kls),
            'max_kl': max(kls), 'top1_agree': agree, 'low_common_mass_positions': low, 'per_position': per}


def _p99(values: Sequence[float]) -> float:
    ordered = sorted(values)
    return ordered[min(len(ordered) - 1, math.ceil(0.99 * len(ordered)) - 1)]


def register_band(aa: Sequence[Mapping[str, Json]], *, factor: float = KL_FACTOR) -> dict[str, Json]:
    """A/A band from compare() results of production vs production (two starts). Register it before any candidate use."""
    if not aa:
        raise ValueError('at least one A/A comparison is required')
    counts = {int(number(c['positions'])) for c in aa}
    if len(counts) != 1:
        raise ValueError('A/A comparisons cover different position sets')
    rng = random.Random(BOOTSTRAP_SEED)
    means: list[float] = []
    for c in aa:
        kls = [number(record(p)['kl_nats']) for p in cast(list[Json], c['per_position'])]
        means += [statistics.fmean(rng.choices(kls, k=len(kls))) for _ in range(BOOTSTRAP)]
    band: dict[str, Json] = {'schema': SCHEMA, 'type': 'band', 'positions': counts.pop(), 'factor': factor,
                             'aa_mean_kl_p99': _p99(means), 'kl_ceiling': factor * _p99(means),
                             'top1_min': min(int(number(c['top1_agree'])) for c in aa), 'aa_comparisons': len(aa),
                             'aa_top1_disagreements': sum(int(number(c['positions']) - number(c['top1_agree'])) for c in aa),
                             'aa_low_common_mass': sum(int(number(c['low_common_mass_positions'])) for c in aa),
                             'bootstrap': {'resamples': BOOTSTRAP, 'seed': BOOTSTRAP_SEED},
                             'registered_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}
    band['sha256'] = hashlib.sha256(json.dumps(band, sort_keys=True).encode()).hexdigest()
    return band


def gate(result: Mapping[str, Json], band: Mapping[str, Json]) -> dict[str, Json]:
    body = {k: v for k, v in band.items() if k != 'sha256'}
    if band.get('sha256') != hashlib.sha256(json.dumps(body, sort_keys=True).encode()).hexdigest():
        raise ValueError('A/A band changed after registration')
    if result.get('positions') != band.get('positions'):
        raise ValueError('candidate comparison and A/A band cover different position sets')
    reasons = []
    if number(result['mean_kl']) > number(band['kl_ceiling']):
        reasons.append(f'mean KL {number(result["mean_kl"]):.3g} > ceiling {number(band["kl_ceiling"]):.3g}')
    flips = int(number(result['positions']) - number(result['top1_agree']))
    p = top1_p(flips, int(number(band['aa_top1_disagreements'])), int(number(band['aa_comparisons'])))
    if p < TOP1_ALPHA:
        reasons.append(f'top-1 disagreements {flips} vs A/A {band["aa_top1_disagreements"]} over '
                       f'{band["aa_comparisons"]} comparison(s): one-sided p {p:.4f} < {TOP1_ALPHA}')
    low = int(number(result['low_common_mass_positions']))
    p_low = top1_p(low, int(number(band['aa_low_common_mass'])), int(number(band['aa_comparisons'])))
    if p_low < TOP1_ALPHA:
        reasons.append(f'positions with top-20 common mass < {greedy_kl.MIN_COMMON_MASS} {low} vs A/A '
                       f'{band["aa_low_common_mass"]}: one-sided p {p_low:.4f} < {TOP1_ALPHA}')
    return {'verdict': 'LOSS' if reasons else 'PASS', 'reasons': cast(Json, reasons), 'band_sha256': band['sha256'],
            'mean_kl': result['mean_kl'], 'top1_agree': result['top1_agree'], 'top1_disagreements': flips,
            'top1_p': p, 'fmethod_top1_rule': number(result['top1_agree']) < number(band['top1_min']) - 1,
            'low_common_mass_positions': result['low_common_mass_positions']}


def top1_p(candidate_flips: int, aa_flips: int, aa_comparisons: int) -> float:
    """P(X >= candidate_flips), X ~ Binomial(candidate_flips + aa_flips, 1 / (aa_comparisons + 1)): the exact
    conditional test that the candidate's top-1 flip rate exceeds the A/A rate (Poisson counts, same positions)."""
    if candidate_flips < 0 or aa_flips < 0 or aa_comparisons < 1:
        raise ValueError('invalid flip counts')
    n, q = candidate_flips + aa_flips, 1 / (aa_comparisons + 1)
    return sum(math.comb(n, i) * q ** i * (1 - q) ** (n - i) for i in range(candidate_flips, n + 1))


def _read(path: Path) -> dict[str, Json]:
    return record(cast(Json, json.loads(path.read_text())))


def _write(path: Path, data: Mapping[str, Json]) -> None:
    with path.open('x') as handle:
        handle.write(json.dumps(data, allow_nan=False) + '\n')


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    sub = ap.add_subparsers(dest='cmd', required=True)
    for name in ('cuts', 'probe'):
        s = sub.add_parser(name)
        s.add_argument('--execute', action='store_true')
        s.add_argument('--url', default=wc.URL)
        s.add_argument('--data-dir', type=Path, default=Path('/opt/dsv41/data/bench-data'))
        s.add_argument('--output', type=Path, required=True)
        s.add_argument('--subset', help='comma list of short,32k,150k (default: all 64 prompts)')
        if name == 'probe':
            s.add_argument('--cuts', type=Path, required=True)
            s.add_argument('--salt', required=True)
    c = sub.add_parser('compare')
    c.add_argument('control', type=Path)
    c.add_argument('candidate', type=Path)
    c.add_argument('--output', type=Path, required=True)
    b = sub.add_parser('band')
    b.add_argument('aa', type=Path, nargs='+', help='compare() outputs of production vs production')
    b.add_argument('--output', type=Path, required=True)
    g = sub.add_parser('gate')
    g.add_argument('comparison', type=Path)
    g.add_argument('--band', type=Path, required=True)
    args = ap.parse_args(argv)
    if args.cmd in ('cuts', 'probe'):
        prompts = selected(args.subset.split(',') if args.subset else None)
        if not args.execute:
            print(json.dumps({'executed': False, 'prompts': len(prompts), 'positions': POSITIONS * len(prompts)}))
            return 0
        if os.environ.get('WINDOW_COORDINATOR') != COORDINATOR:
            print(f'STOP: probes are coordinator-only (WINDOW_COORDINATOR={COORDINATOR})', file=sys.stderr)
            return 2
        from offload.prefill_trace.window_requests import corpus
        materials, hashes = corpus(args.data_dir)
        api = Http(args.url)
        if args.cmd == 'cuts':
            _write(args.output, cuts(api, materials, hashes, prompts))
        else:
            registered = _read(args.cuts)
            if registered.get('corpus_sha256') != hashes:
                raise ValueError('public corpus changed since the cuts were built')
            rows = record(registered['prompts'])
            _write(args.output, probe(api, materials, {p.key: rows[p.key] for p in prompts}, args.salt))
        return 0
    if args.cmd == 'compare':
        _write(args.output, compare(_read(args.control), _read(args.candidate)))
        return 0
    if args.cmd == 'band':
        _write(args.output, register_band([_read(p) for p in args.aa]))
        return 0
    verdict = gate(_read(args.comparison), _read(args.band))
    print(json.dumps(verdict, allow_nan=False))
    return 0 if verdict['verdict'] == 'PASS' else 2


if __name__ == '__main__':
    raise SystemExit(main())