| """Production K1 composition, including both DROP_MTP_BUFFER None paths.""" |
| from __future__ import annotations |
|
|
| import ast |
| import copy |
| import hashlib |
| import json |
| from collections.abc import Callable |
| from pathlib import Path |
| from types import SimpleNamespace |
| from typing import cast |
|
|
| import pytest |
|
|
| from offload.dspark.probabilistic.pp3 import prepare_verifier as pp3 |
| from offload.launch.windows import pp4_prepare as p |
|
|
| SOURCE = Path('/opt/dsv41/data/src-image') |
|
|
|
|
| def test_exact_k1_composition_and_removal_of_only_new_additions(tmp_path: Path) -> None: |
| out = tmp_path/'prepared' |
| receipt = p.prepare(SOURCE, out) |
| original = p.K1_RUNNER.read_text() |
| actual = (out/pp3.RUNNER).read_text() |
| assert receipt == json.loads((p.ROOT/'offload/dspark/probabilistic/pp4/verifier-k1-receipt.json').read_text()) |
| assert receipt['production_k1_sha256'] == p.K1_SHA256 |
| assert hashlib.sha256(original.encode()).hexdigest() == p.K1_SHA256 |
| |
| |
| actual = actual.removeprefix('from vllm.v1.worker.gpu import dsv41_pp4_round_trace as _pp4_round\n') |
| actual = actual.replace(' @_pp4_round.execution\n', '').replace(' @_pp4_round.sampling\n', '') |
| assert actual == pp3.runner(original) |
| prefix = 'REMOTE_DRAFT_PROBS_ABI = ' + repr(pp3.ABI) + '\n' |
| delta = pp3.runner(original).removeprefix(prefix) |
| start = delta.index(' # PP3 failed/missing-q rounds are target-only.') |
| end = delta.index(' sampler_output: SamplerOutput | None\n', start) |
| assert delta[:start]+delta[end:] == original |
| assert (out/pp3.UTILS).read_text() == pp3.utilities((SOURCE/pp3.UTILS).read_text()) |
|
|
|
|
| @pytest.mark.parametrize('site', [0, 1]) |
| @pytest.mark.parametrize('value', [None, 'mtp']) |
| def test_actual_none_return_branches_preserve_hidden_states(tmp_path: Path, site: int, value: str | None) -> None: |
| out = tmp_path/'prepared' |
| p.prepare(SOURCE, out) |
| source = (out/pp3.RUNNER).read_text() |
| blocks = [n for n in ast.walk(ast.parse(source)) if isinstance(n, ast.If) |
| and isinstance(n.test, ast.Call) and isinstance(n.test.func, ast.Name) |
| and n.test.func.id == 'hasattr' and len(n.test.args) > 1 |
| and isinstance(n.test.args[1], ast.Constant) and n.test.args[1].value == 'get_mtp_target_hidden_states'] |
| assert len(blocks) == 2 |
| fn = ast.parse('def probe(self, hidden_states, draft_hidden_states):\n spec_hidden_states = hidden_states\n return spec_hidden_states\n') |
| definition = cast(ast.FunctionDef, fn.body[0]) |
| definition.body.insert(1, copy.deepcopy(blocks[site])); ast.fix_missing_locations(fn) |
| ns: dict[str, object] = {} |
| exec(compile(fn, 'actual-K1-branch', 'exec'), ns) |
| class Hidden: |
| shape = (1, 1) |
| def size(self, dim: int) -> int: return 1 |
| def __getitem__(self, index: object) -> str: return 'mtp-selected' |
| hidden = Hidden() |
| model = SimpleNamespace(model=SimpleNamespace(get_mtp_target_hidden_states=lambda: None if value is None else hidden)) |
| result = cast(Callable[[object, object, object], object], ns['probe'])(model, hidden, hidden) |
| assert result is hidden if value is None else result == 'mtp-selected' |
|
|
|
|
| def test_k1_hash_drift_refuses_before_any_output(tmp_path: Path) -> None: |
| changed = tmp_path/'changed.py'; changed.write_bytes(p.K1_RUNNER.read_bytes()+b'\n') |
| with pytest.raises(ValueError, match='pinned source mismatch'): |
| p.prepare(SOURCE, tmp_path/'output', changed) |
| assert not (tmp_path/'output').exists() |
|
|
|
|
| def test_off_failure_mask_does_not_change_original_invalid_drafts() -> None: |
| original = p.K1_RUNNER.read_text() |
| generated = pp3.runner(original) |
| start = generated.index(' pp3_valid = ') |
| end = generated.index(' sampler_output: SamplerOutput | None', start) |
| block = generated[start:end] |
| for sentinel in (None, object()): |
| ns = {'self': SimpleNamespace(speculator=SimpleNamespace()), 'invalid_drafts': sentinel} |
| exec(compile(__import__('textwrap').dedent(block), 'actual-OFF-mask', 'exec'), ns) |
| assert ns['invalid_drafts'] is sentinel |
|
|
