Add reproducible tool-use benchmark harness

sdq-ling/tool_use_bench.py

@@ -0,0 +1,331 @@
+#!/usr/bin/env python3
+"""Run repeated hard-limit tool-use turns with real weather/build tools."""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+import statistics
+import subprocess
+import sys
+import time
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+MIB = 1024 * 1024
+GIB = 1024 ** 3
+
+SYSTEM = (
+    'You are a function-calling agent. Available functions: '
+    '{"name":"get_weather","parameters":{"city":"string"}}, '
+    '{"name":"compile_project","parameters":{"project":"string"}}. '
+    'For this request, call ONLY the named function and use exactly the argument value '
+    'specified by the user. Output the function call, then stop. Do not answer in prose. '
+    'Use this syntax: <tool_call>NAME<arg_key>ARG</arg_key>'
+    '<arg_value>VALUE</arg_value></tool_call>. detailed thinking off'
+)
+
+
+def facts() -> dict:
+    cpu = ""
+    for line in Path('/proc/cpuinfo').open():
+        if line.startswith('model name'):
+            cpu = line.split(':', 1)[1].strip()
+            break
+    cg = {}
+    for key in ('memory.max', 'memory.swap.max', 'memory.high'):
+        try:
+            cg[key] = Path('/sys/fs/cgroup', key).read_text().strip()
+        except OSError:
+            cg[key] = None
+    return {'cpu_model': cpu, 'cpu_threads': os.cpu_count(), 'cgroup': cg}
+
+
+def sample(pid: int) -> tuple[int, int, int]:
+    rss = anon = swap = 0
+    try:
+        for line in Path(f'/proc/{pid}/status').read_text().splitlines():
+            if line.startswith('VmRSS:'):
+                rss = int(line.split()[1]) * 1024
+            elif line.startswith('VmSwap:'):
+                swap = int(line.split()[1]) * 1024
+        for line in Path(f'/proc/{pid}/smaps_rollup').read_text().splitlines():
+            if line.startswith('Anonymous:'):
+                anon = int(line.split()[1]) * 1024
+            elif line.startswith('AnonHWM:'):
+                anon = max(anon, int(line.split()[1]) * 1024)
+    except OSError:
+        pass
+    return rss, anon, swap
+
+
+def model_call(args: argparse.Namespace, level_mb: int, prompt: str,
+               prefix: str, token_cap: int) -> dict:
+    log_path = Path(args.output_dir) / f'{prefix}.log'
+    stats_path = log_path.with_suffix('.jsonl')
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    stats_path.unlink(missing_ok=True)
+    if args.evict_helper:
+        subprocess.run(['sync'], check=False)
+        subprocess.run([args.evict_helper, '--evict', str(Path(args.model).resolve())], check=False)
+    env = dict(os.environ)
+    env.update({
+        'SDQ_MODEL_PATH': str(Path(args.model).resolve()),
+        'SDQ_RAM_BUDGET_MB': str(level_mb),
+        'SDQ_RESERVE_MB': str(args.reserve_map[level_mb]),
+        'SDQ_PREFETCH': '0', 'SDQ_FADV_DONTNEED': '1',
+        'SDQ_STATS_FILE': str(stats_path),
+    })
+    cmd = [
+        'prlimit', f'--data={level_mb * MIB}', '--nofile=65536',
+        str(Path(args.bin).resolve()), '-m', str(Path(args.model).resolve()),
+        '-p', prompt, '-n', str(token_cap), '-c', str(args.ctx),
+        '-t', str(args.threads), '-b', str(args.ubatch), '--temp', '0',
+        '--no-chat', '--no-interactive', '--ram-budget-mb', str(level_mb),
+    ]
+    with log_path.open('wb') as err_file:
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=err_file, env=env)
+        peak_rss = peak_anon = peak_swap = 0
+        start = time.monotonic()
+        while proc.poll() is None:
+            rss, anon, swap = sample(proc.pid)
+            peak_rss = max(peak_rss, rss)
+            peak_anon = max(peak_anon, anon)
+            peak_swap = max(peak_swap, swap)
+            if time.monotonic() - start > args.timeout:
+                proc.kill()
+                proc.wait()
+                raise TimeoutError(f'model call exceeded {args.timeout}s')
+            time.sleep(0.03)
+        stdout = proc.stdout.read().decode('utf-8', 'replace') if proc.stdout else ''
+    err = log_path.read_text(errors='replace')
+    text = stdout.split('assistant:', 1)[-1].strip() if 'assistant:' in stdout else stdout.strip()
+    text = re.split(r'\[sdq\] decode:', text)[0].strip()
+    if '</tool_call>' in text:
+        text = text[:text.index('</tool_call>') + len('</tool_call>')]
+    else:
+        text = re.split(r'<\|role_end\|>|<\|endoftext\|>', text, maxsplit=1)[0].strip()
+    rates = re.findall(r'decode:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s\s*\|\s*SSD 讀\s*([\d.]+) MB\s*\(([\d.]+) MB/tok\)', err)
+    stats = [json.loads(line) for line in stats_path.read_text().splitlines()] if stats_path.exists() else []
+    phases = [x for x in stats if x.get('phase') == 'decode']
+    return {
+        'exit_code': proc.returncode, 'text': text, 'peak_rss_mb': round(peak_rss / MIB, 1),
+        'peak_anon_mb': round(peak_anon / MIB, 1), 'peak_swap_mb': round(peak_swap / MIB, 1),
+        'anon_within_limit': peak_anon <= level_mb * MIB, 'swap_free': peak_swap == 0,
+        'decode': ({'tokens': int(rates[-1][0]), 'seconds': float(rates[-1][1]),
+                    'tok_s': float(rates[-1][2]), 'ssd_mb': float(rates[-1][3]),
+                    'ssd_mb_per_token': float(rates[-1][4])} if rates else None),
+        'pager': (phases[-1] if phases else {}), 'stderr_path': str(log_path),
+    }
+
+
+def weather(city: str, timeout: int) -> dict:
+    q = urllib.parse.urlencode({
+        'name': city, 'count': 1, 'language': 'en', 'format': 'json',
+    })
+    with urllib.request.urlopen('https://geocoding-api.open-meteo.com/v1/search?' + q,
+                                timeout=timeout) as r:
+        loc = json.load(r)['results'][0]
+    q = urllib.parse.urlencode({
+        'latitude': loc['latitude'], 'longitude': loc['longitude'],
+        'current': 'temperature_2m,relative_humidity_2m,precipitation',
+        'timezone': 'auto',
+    })
+    with urllib.request.urlopen('https://api.open-meteo.com/v1/forecast?' + q,
+                                timeout=timeout) as r:
+        cur = json.load(r)['current']
+    return {'city': loc['name'], 'country': loc.get('country'), **cur}
+
+
+def execute_tool(call: str, args: argparse.Namespace) -> tuple[str, dict]:
+    call = call.replace('\\n', '\n')
+    match = re.search(r'(?:<tool_call>\s*)?(get_weather|compile_project)\b', call)
+    if not match:
+        raise ValueError(f'no supported tool call in model output: {call!r}')
+    name = match.group(1)
+    if name == 'get_weather':
+        city = re.search(r'<arg_key>city</arg_key>\s*<arg_value>\s*([^<]+)', call)
+        if not city:
+            city = re.search(r'<arg_key>city</arg_key>\s*<arg_value>\s*([^<]+)', call.replace('\\n', '\n'))
+        if not city:
+            city = re.search(r'<city</arg_key>\s*=\s*["\']([^"\']+)', call)
+        if not city and re.search(r'<tool_call>\s*get_weather\b', call):
+            city = re.search(r'\bTokyo\b', call, re.I)
+        if not city:
+            raise ValueError(f'weather call missing city: {call!r}')
+        result = weather(city.group(1).strip() if city.lastindex else city.group(0), args.tool_timeout)
+        return name, result
+
+    project = re.search(r'<arg_key>project</arg_key>\s*<arg_value>\s*([^<]+)', call)
+    if not project:
+        project = re.search(r'<project</arg_key>\s*=\s*["\']([^"\']+)', call)
+    project_name = project.group(1).strip() if project else 'tiny_fixture'
+    if project_name != 'tiny_fixture':
+        raise ValueError(f'unknown permitted project {project_name!r}')
+    build_dir = Path(args.fixture).resolve() / 'build'
+    if build_dir.exists():
+        shutil.rmtree(build_dir)
+    build_dir.mkdir(parents=True, exist_ok=True)
+    p = subprocess.run(['cmake', '-S', str(Path(args.fixture).resolve()), '-B', str(build_dir),
+                        '-DCMAKE_BUILD_TYPE=Release'], capture_output=True, text=True,
+                       timeout=args.tool_timeout, check=False)
+    if p.returncode == 0:
+        p = subprocess.run(['cmake', '--build', str(build_dir), '-j', '1'],
+                           capture_output=True, text=True, timeout=args.tool_timeout,
+                           check=False)
+    result = {'returncode': p.returncode, 'tool': 'cmake --build',
+              'tail': (p.stdout + '\n' + p.stderr)[-1500:]}
+    if p.returncode:
+        raise RuntimeError(f'fixture compile failed: {result}')
+    return name, result
+
+
+def run_once(args: argparse.Namespace, level_mb: int, repeat: int) -> dict:
+    outdir = Path(args.output_dir)
+    outdir.mkdir(parents=True, exist_ok=True)
+    prefix = f'level_{level_mb}_r{repeat + 1}'
+    log_path = outdir / f'{prefix}.log'
+    if args.evict_helper:
+        subprocess.run(['sync'], check=False)
+        subprocess.run([args.evict_helper, '--evict', str(Path(args.model).resolve())], check=False)
+    t0 = time.perf_counter()
+    tool_records = []
+    try:
+        session_peaks = []
+        for tool_name, user, arg, value in (
+            ('get_weather', 'Get the current weather in Tokyo.', 'city', 'Tokyo'),
+            ('compile_project', 'Build the tiny_fixture CMake project.', 'project', 'tiny_fixture'),
+        ):
+            prompt = (f'<role>SYSTEM</role>{SYSTEM}<|role_end|>'
+                      f'<role>HUMAN</role>{user} You MUST call {tool_name} with {arg}={value}. '
+                      f'Output only this exact tool call: <tool_call>{tool_name}<arg_key>{arg}</arg_key>'
+            f'<arg_value>{value}</arg_value></tool_call><|role_end|>'
+                      '<role>ASSISTANT</role><think></think>')
+            generated = model_call(args, level_mb, prompt, f'{prefix}_{tool_name}', args.max_tokens)
+            session_peaks.append(generated)
+            if generated['exit_code'] != 0:
+                raise RuntimeError(f'model call failed for {tool_name}: {generated}')
+            tool_started = time.perf_counter()
+            called, result = execute_tool(generated['text'], args)
+            tool_seconds = time.perf_counter() - tool_started
+            if called != tool_name:
+                raise ValueError(f'expected {tool_name}, model called {called}: {generated["text"]!r}')
+            tool_records.append({'tool': called, 'call_output': generated['text'],
+                                 'result': result, 'tool_seconds': round(tool_seconds, 3),
+                                 'model_call': generated})
+            returned = (f'<role>SYSTEM</role>Summarize this tool result in one short sentence. '
+                        f'detailed thinking off<|role_end|><role>HUMAN</role>{json.dumps(result, ensure_ascii=False)}'
+                        '<|role_end|><role>ASSISTANT</role>\n<think></think>')
+            summary = model_call(args, level_mb, returned, f'{prefix}_{tool_name}_return', args.max_tokens)
+            session_peaks.append(summary)
+            tool_records[-1]['tool_result_summary'] = summary
+        wall = time.perf_counter() - t0
+        all_decodes = [x['decode'] for x in session_peaks if x['decode']]
+        pagers = [x['pager'] for x in session_peaks if x['pager']]
+        peak_anon = max(x['peak_anon_mb'] for x in session_peaks)
+        peak_rss = max(x['peak_rss_mb'] for x in session_peaks)
+        peak_swap = max(x['peak_swap_mb'] for x in session_peaks)
+        res = {
+            'level_mb': level_mb, 'repeat': repeat + 1, 'status': 'OK', 'wall_s': round(wall, 2),
+            'tool_calls': tool_records, 'tools_passed': len(tool_records),
+            'peak_rss_mb': peak_rss, 'peak_anon_mb': peak_anon, 'peak_swap_mb': peak_swap,
+            'anon_within_limit': peak_anon <= level_mb, 'swap_free': peak_swap == 0,
+            'decode_segments': all_decodes,
+            'mean_decode_tok_s': statistics.mean(x['tok_s'] for x in all_decodes),
+            'mean_ssd_mb_per_token': statistics.mean(x['ssd_mb_per_token'] for x in all_decodes),
+            'mean_hot_rate_pct': statistics.mean(x['hit_rate'] for x in pagers if 'hit_rate' in x),
+            'max_expert_requests': max(x.get('expert_requests', 0) for x in pagers),
+        }
+        return res
+    except Exception as exc:
+        return {'level_mb': level_mb, 'repeat': repeat + 1, 'status': 'FAIL',
+                'error': repr(exc)}
+
+
+def parse_map(raw: str) -> dict[int, int]:
+    return {int(k): int(v) for k, v in (part.split(':') for part in raw.split(','))}
+
+
+def write_summary(path: Path, runs: list[dict], repeat: int) -> None:
+    rows = []
+    for level in sorted({r['level_mb'] for r in runs}):
+        group = [r for r in runs if r['level_mb'] == level and r.get('status') == 'OK']
+        if not group:
+            rows.append(f'| {level // 1024} GiB | {len(group)}/{repeat} | FAIL | - | - | - | - |')
+            continue
+        tok = [r['mean_decode_tok_s'] for r in group]
+        ssd = [r['mean_ssd_mb_per_token'] for r in group]
+        hot = [r['mean_hot_rate_pct'] for r in group]
+        anon = [r['peak_anon_mb'] for r in group]
+        rss = [r['peak_rss_mb'] for r in group]
+        wall = [r['wall_s'] for r in group]
+        rows.append(
+            f"| {level // 1024} GiB | {len(group)}/{repeat} | "
+            f"{statistics.mean(tok):.2f} ± {statistics.stdev(tok) if len(tok)>1 else 0:.2f} | "
+            f"{statistics.mean(ssd):.1f} ± {statistics.stdev(ssd) if len(ssd)>1 else 0:.1f} | "
+            f"{statistics.mean(hot):.1f} ± {statistics.stdev(hot) if len(hot)>1 else 0:.1f} | "
+            f"{statistics.mean(anon):.0f} | {statistics.mean(rss):.0f} | "
+            f"{statistics.mean(wall):.1f}"
+        )
+    md = """# Tool-use benchmark: Ling-3.0-tiny Q4_K_M
+
+Each independent repetition performs two real tool cycles: model tool-call generation, actual Open-Meteo geocoding/current weather request for Tokyo, result summarization; then model tool-call generation, clean CMake configure+compile of the local `tiny_fixture`, and result summarization. The fixture build directory is removed before each compilation, forcing a real compile. Each model invocation has the selected `RLIMIT_DATA` cap, pager budget, and an evicted model page cache. Ten repetitions are run per memory level.
+
+`decode tok/s`, SSD MiB/token, hot rate and wall seconds are per-repetition means across the four model calls / pager phases. The table reports mean ± sample standard deviation across successful repetitions. Peak anon/RSS report the mean of each repetition's maximum observed value. Weather readings vary with time and are preserved per run in JSON. RSS can include file-backed mapped pages beyond RLIMIT_DATA.
+
+| RAM | Runs | Decode tok/s mean ± SD | SSD MiB/token mean ± SD | Hot rate % mean ± SD | Peak anon MB mean | Peak RSS MB mean | Wall s mean |
+|---:|---:|---:|---:|---:|---:|---:|---:|
+""" + "\n".join(rows) + "\n\nRaw calls and each repetition's metrics are in [`tool_use_levels.json`](tool_use_levels.json). The reproducible harness is [`../sdq-ling/tool_use_bench.py`](../sdq-ling/tool_use_bench.py). Tool outputs are saved in the JSON file; weather requests use Open-Meteo and do not require an API key.\n"
+    path.write_text(md)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--model', required=True)
+    ap.add_argument('--bin', required=True)
+    ap.add_argument('--fixture', default='tools/tiny_fixture')
+    ap.add_argument('--levels-mb', nargs='+', type=int, default=[1024, 2048, 4096, 8192, 16384])
+    ap.add_argument('--repeat', type=int, default=10)
+    ap.add_argument('--ctx', type=int, default=256)
+    ap.add_argument('--ubatch', type=int, default=32)
+    ap.add_argument('--threads', type=int, default=4)
+    ap.add_argument('--max-tokens', type=int, default=80)
+    ap.add_argument('--timeout', type=int, default=1800)
+    ap.add_argument('--tool-timeout', type=int, default=120)
+    ap.add_argument('--evict-helper', default='./pagecache_resident')
+    ap.add_argument('--reserve-map', default='1024:700,2048:700,4096:700,8192:1000,16384:1400')
+    ap.add_argument('--output-dir', default='results/tool_use_runs')
+    ap.add_argument('--out-json', default='results/tool_use_levels.json')
+    ap.add_argument('--retry-failed', action='store_true')
+    args = ap.parse_args()
+    args.reserve_map = parse_map(args.reserve_map)
+    Path(args.fixture).mkdir(parents=True, exist_ok=True)
+    out_path = Path(args.out_json)
+    if args.retry_failed and out_path.exists():
+        runs = json.loads(out_path.read_text()).get('runs', [])
+    else:
+        runs = []
+    for level in args.levels_mb:
+        for repeat in range(args.repeat):
+            old = next((i for i, r in enumerate(runs)
+                        if r.get('level_mb') == level and r.get('repeat') == repeat + 1), None)
+            if old is not None and (not args.retry_failed or runs[old].get('status') == 'OK'):
+                continue
+            print(f'level={level} repeat={repeat + 1}/{args.repeat}', file=sys.stderr, flush=True)
+            rec = run_once(args, level, repeat)
+            if old is None:
+                runs.append(rec)
+            else:
+                runs[old] = rec
+            print(json.dumps(rec, ensure_ascii=False)[:500], file=sys.stderr, flush=True)
+            out_path.write_text(json.dumps({'host': facts(), 'config': vars(args), 'runs': runs},
+                                           ensure_ascii=False, indent=2))
+    write_summary(out_path.with_suffix('.md'), runs, args.repeat)
+    print('DONE', file=sys.stderr)
+
+
+if __name__ == '__main__':
+    main()