scripts/inference/run.py

#!/usr/bin/env python3
"""ADSD inference with the HSD speculative decoding engine."""
import argparse
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys

ROOT = Path(__file__).resolve().parents[2]
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from native_hooks import atomic


def require(condition, message):
    if not condition:
        raise ValueError(message)


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as f:
        for block in iter(lambda: f.read(8 * 1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def module(path):
    spec = importlib.util.spec_from_file_location('workspace_tools', path)
    result = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(result)
    return result


def prepare(args):
    """Copy the installed package; never patch a shared Python installation."""
    import importlib.metadata
    installed = Path(importlib.metadata.distribution('transformers').locate_file('transformers'))
    runtime = args.runtime.resolve()
    runtime.mkdir(parents=True, exist_ok=False)
    overlay = runtime / 'overlay' / 'transformers'
    shutil.copytree(installed, overlay, ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
    helper = module(ROOT / 'scripts/eval/internal/prepare_workspace.py')
    source = ROOT / 'external/hsd'
    workspace = runtime / 'workspace'
    helper.copy_workspace(source, workspace)
    evaluator = workspace / 'chain-of-thought-hub/gsm8k/eval_speculative_decoding_llm.py'
    helper.patch_eval_script(evaluator)
    helper.patch_generation_utils(workspace / 'transformers/generation/utils.py')
    helper.ensure_outputs_dir(workspace)
    for src in (workspace / 'transformers').rglob('*.py'):
        dest = overlay / src.relative_to(workspace / 'transformers')
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text(src.read_text().expandtabs(4))
    text = evaluator.read_text()
    # A snapshot path's basename is a revision hash, not its model name.
    text = text.replace('model_ref = os.path.basename(model2_name.rstrip("/"))',
                        'model_ref = model2_name')
    text = text.replace("gsm8k = load_dataset('gsm8k', 'main')", '''from datasets import Dataset, DatasetDict
gsm8k = DatasetDict({split: Dataset.from_list([
    json.loads(line) for line in open(os.environ['ADSD_' + split.upper() + '_JSONL'])
    if line.strip()]) for split in ('train', 'test')})''')
    marker = '            end = time.time()\n'
    require(text.count(marker) == 1, 'Expected one evaluator insertion point')
    text = text.replace(marker, marker + '            from native_hooks import capture\n            capture(globals())\n')
    marker = '    print("start training")'
    require(text.count(marker) == 1, 'Expected one evaluator insertion point')
    text = text.replace(marker, '    from native_hooks import policy\n    policy(globals())\n' + marker)
    evaluator.write_text(text)
    atomic(runtime / 'manifest.json', {
        'transformers_version': importlib.metadata.version('transformers'),
        'base_package': str(installed),
        'evaluator_sha256': digest(evaluator),
        'overlay_hashes': {str(p.relative_to(overlay)): digest(p)
                           for p in overlay.rglob('*.py')},
        'note': 'Transformers runtime with HSD inference components.'})


def environment(args):
    env = os.environ.copy()
    cache_base = getattr(args, 'cache_root', None) or os.environ.get('ADSD_CACHE_ROOT')
    cache = Path(cache_base).expanduser().resolve() / args.output.name if cache_base else args.output.resolve() / 'cache'
    for key in ('HF_HOME', 'HUGGINGFACE_HUB_CACHE', 'TRANSFORMERS_CACHE',
                'HF_DATASETS_CACHE', 'TORCH_HOME', 'XDG_CACHE_HOME',
                'TRITON_CACHE_DIR', 'TORCH_EXTENSIONS_DIR', 'SHARED_CACHE_ROOT'):
        env[key] = str(cache / key.lower())
    for key in ('HF_HUB_OFFLINE', 'TRANSFORMERS_OFFLINE', 'HF_DATASETS_OFFLINE',
                'PYTHONNOUSERSITE', 'PYTHONDONTWRITEBYTECODE', 'PYTHONUNBUFFERED'):
        env[key] = '1'
    env['TOKENIZERS_PARALLELISM'] = 'false'
    env.pop('ADSD_ACCEPTANCE_FALLBACK_THRESHOLD', None)
    env['PYTHONPATH'] = os.pathsep.join([str(args.runtime.resolve() / 'overlay'),
                                       str(HERE), str(ROOT)])
    env['ADSD_TRAIN_JSONL'] = str(args.train.resolve())
    env['ADSD_TEST_JSONL'] = str(args.test.resolve())
    env['ADSD_RUN_OUTPUT'] = str(args.output.resolve())
    return env


def provenance(args):
    args.output = args.output.resolve()
    args.output.mkdir(parents=True, exist_ok=False)
    if not (args.runtime / 'manifest.json').is_file():
        raise ValueError('Run prepare first with this Python environment')
    manifest = {'command': sys.argv, 'args': {k: str(v) if isinstance(v, Path) else v
                                              for k, v in vars(args).items()},
                'runtime_manifest_sha256': digest(args.runtime / 'manifest.json'),
                'dataset_sha256': {split: digest(getattr(args, split)) for split in ('train', 'test')},
                'model_files': {}}
    for role in ('target', 'draft'):
        path = getattr(args, role).resolve()
        if not (path / 'config.json').is_file():
            raise ValueError(f'cache preload incomplete: {path}')
        required = list(path.glob('*.safetensors'))
        if not required:
            raise ValueError(f'cache preload incomplete: no safetensors in {path}')
        manifest['model_files'][role] = {
            p.name: {'bytes': p.stat().st_size, 'sha256': digest(p)}
            for p in sorted(path.iterdir()) if p.is_file() and not p.name.startswith('.')}
    source_paths = [HERE/'run.py', HERE/'native_hooks.py',
                    ROOT/'scripts/eval/internal/prepare_workspace.py']
    manifest['source_file_sha256'] = {str(p.relative_to(ROOT)):digest(p) for p in source_paths if p.exists()}
    manifest['source_commit_hint'] = os.environ.get('ADSD_SOURCE_COMMIT')
    if args.attack:
        manifest['attack_sha256'] = digest(args.attack)
    manifest['gpu'] = subprocess.check_output(['nvidia-smi', '--query-gpu=name,uuid,driver_version',
                                               '--format=csv,noheader'], text=True).strip()
    atomic(args.output / 'manifest.json', manifest)


def evaluate(args):
    provenance(args)
    base = args.runtime.resolve() / 'workspace/chain-of-thought-hub/gsm8k'
    work = args.output / 'native'
    work.mkdir()
    shutil.copytree(base / 'lib_prompt', work / 'lib_prompt')
    shutil.copy2(base / 'eval_speculative_decoding_llm.py', work)
    (work / 'outputs').mkdir()
    suffix = ''
    if args.attack:
        best = json.loads(args.attack.read_text())
        suffix = best.get('suffix_text')
        if not isinstance(suffix, str) or not suffix:
            raise ValueError('Attack artifact lacks nonempty suffix/prefix text')
    command = [sys.executable, str(work / 'eval_speculative_decoding_llm.py'),
               '--target-model', str(args.target.resolve()), '--draft-model', str(args.draft.resolve()),
               '--num-samples', str(args.samples), '--max-new-tokens', str(args.cap),
               '--sample-offset', str(args.sample_offset),
               '--seed', str(args.seed), '--gamma', '10', '--temperature', '1.0',
               '--prompt-tag', 'inference']
    if args.mode == 'tokenwise':
        command.append('--speculative')
    if suffix:
        command += ['--prompt-suffix', suffix]
    atomic(args.output / 'command.json', command)
    with (args.output / 'stdout.log').open('w') as log:
        completed = subprocess.run(command, cwd=work, env=environment(args), stdout=log,
                                   stderr=subprocess.STDOUT)
    rows = sorted(args.output.glob('row_*.json'))
    ok = completed.returncode == 0 and len(rows) == args.samples
    if not ok:
        atomic(args.output / 'status.json', {'status': 'failed',
            'exit_code': completed.returncode, 'rows': len(rows)})
        raise SystemExit(completed.returncode or 1)
    try:
        summarize(args.output)
    except Exception as exc:
        atomic(args.output / 'status.json', {'status': 'failed', 'stage': 'summarization',
            'error': str(exc), 'exit_code': 1, 'rows': len(rows)})
        raise
    atomic(args.output / 'progress.json', {'status': 'complete', 'completed': len(rows),
                                         'expected': args.samples})
    atomic(args.output / 'status.json', {'status': 'complete', 'exit_code': 0, 'rows': len(rows)})


def summarize(path):
    rows = [json.loads(p.read_text()) for p in sorted(path.glob('row_*.json'))]
    def total(row, field):
        return sum(row['counts'].get(field, []))
    for row in rows:
        require(row['elapsed_seconds'] > 0, 'Nonpositive generation time')
        if row['counts']:
            require(total(row, 'sample_length') == len(row['generated_ids']), 'Native token accounting mismatch')
            require(0 <= total(row, 'accepted_draft_tokens') <= total(row, 'draft_eval'), 'Invalid accepted-token accounting')
    seconds = sum(r['elapsed_seconds'] for r in rows)
    tokens = sum(len(r['generated_ids']) for r in rows)
    drafted = sum(total(r, 'draft_eval') for r in rows)
    accepted = sum(total(r, 'accepted_draft_tokens') for r in rows)
    calls = sum(total(r, 'total_step') for r in rows)
    block = [len(r['generated_ids']) / total(r, 'total_step') for r in rows if total(r, 'total_step')]
    acceptance = [total(r, 'accepted_draft_tokens') / total(r, 'draft_eval')
                  for r in rows if total(r, 'draft_eval')]
    summary = dict(samples=len(rows), correct=sum(r['correct_last_number'] for r in rows),
                   accuracy=sum(r['correct_last_number'] for r in rows)/len(rows),
                   latency_mean=seconds/len(rows), generated_tokens_mean=tokens/len(rows),
                   throughput_global=tokens/seconds, capped=sum(r['reached_cap'] for r in rows),
                   aggregate_acceptance_global=accepted/drafted if drafted else None,
                   aggregate_acceptance_mean=sum(acceptance)/len(acceptance) if acceptance else None,
                   block_efficiency_mean=sum(block)/len(block) if block else None,
                   block_efficiency_global=tokens/calls if calls else None,
                   scoring='last number of generated continuation; native full-text parser retained in stdout.log')
    import re
    native_score = re.findall(r'num_q (\d+) correct (\d+) ratio ([\d.]+)',
                             (path / 'stdout.log').read_text(errors='replace'))
    if native_score:
        n, correct, rounded = native_score[-1]
        summary['native_parser'] = dict(samples=int(n), correct=int(correct),
                                        accuracy=int(correct)/int(n), reported_rounded=float(rounded))
    atomic(path / 'summary.json', summary)
    return summary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='action', required=True)
    p = sub.add_parser('prepare')
    p.add_argument('--runtime', type=Path, required=True)
    p = sub.add_parser('evaluate')
    p.add_argument('--cache-root', type=Path, help='Cache directory; defaults to ADSD_CACHE_ROOT')
    for name in ('runtime', 'output', 'target', 'draft', 'train', 'test'):
        p.add_argument('--' + name, type=Path, required=True)
    p.add_argument('--attack', type=Path, help='Fixed suffix JSON; omit for benign inference')
    p.add_argument('--seed', type=int, default=2027)
    p.add_argument('--samples', type=int, default=263)
    p.add_argument('--sample-offset', type=int, default=0)
    p.add_argument('--cap', type=int, default=512)
    p.add_argument('--mode', choices=['tokenwise', 'target_only'], default='tokenwise')
    args = parser.parse_args()
    {'prepare': prepare, 'evaluate': evaluate}[args.action](args)


if __name__ == '__main__':
    main()