"""Focused parser/accounting/scheduling tests; never contact the serving model.""" from __future__ import annotations import argparse import io import json import runpy import tempfile import threading import time import unittest from pathlib import Path from unittest.mock import patch import bench2 as b import report class NullTelemetry: def __init__(self, *args: object) -> None: self.samples: list[b.JSON] = [] self.errors: list[b.JSON] = [] def start(self) -> None: pass def finish(self) -> None: pass def sample(self) -> None: pass def event(delta: b.Object, finish: str | None = None) -> str: return 'data: ' + json.dumps({'choices': [{'delta': delta, 'finish_reason': finish}]}) + '\n\n' class BenchTests(unittest.TestCase): def test_original_workloads_and_record_builder_preserved(self) -> None: original = runpy.run_path(str(Path(__file__).resolve().parents[1] / 'bench.py')) self.assertEqual(b.WORKLOADS, original['WORKLOADS']) self.assertEqual(b.build_record(3000, 42), original['build_record'](3000, 42)) def test_split_simultaneous_fields_and_server_usage(self) -> None: wire = event({'reasoning_content': 'why', 'content': 'yes'}) wire += event({'content': ' indeed'}, 'stop') wire += 'data: ' + json.dumps({'usage': {'prompt_tokens': 11, 'completion_tokens': 7, 'completion_tokens_details': {'reasoning_tokens': 2}}}) + '\n\n' wire += 'data: [DONE]\n\n' with patch('urllib.request.urlopen', return_value=io.BytesIO(wire.encode())), patch.object(b, 'tok_count', return_value=2): row = b.stream_request({}) self.assertEqual(row['reasoning_text'], 'why') self.assertEqual(row['answer_text'], 'yes indeed') self.assertEqual(row['reasoning_tokens'], 2) self.assertEqual(row['answer_tokens'], 5) self.assertIsNotNone(row['first_visible_s']) self.assertEqual(row['completion_tokens'], 7) def test_channel_retokenization_is_explicit(self) -> None: wire = event({'reasoning': 'two chunks'}) + event({}, 'length') wire += 'data: {"usage":{"prompt_tokens":1,"completion_tokens":9}}\n\ndata: [DONE]\n' with patch('urllib.request.urlopen', return_value=io.BytesIO(wire.encode())), patch.object(b, 'tok_count', return_value=3): row = b.stream_request({}) self.assertEqual(row['reasoning_tokens'], 3) self.assertEqual(row['reasoning_token_source'], 'retokenized_text') self.assertIsNone(row['first_visible_s']) self.assertEqual(row['completion_tokens'], 9) def test_truncated_stream_is_not_success(self) -> None: with patch('urllib.request.urlopen', return_value=io.BytesIO(event({'content': 'partial'}).encode())), self.assertRaisesRegex(RuntimeError, 'incomplete stream'): b.stream_request({}) def test_metric_deltas_labels_and_reset(self) -> None: before = b.parse_metrics('vllm:spec_decode_num_drafts_total{engine="0"} 10\nvllm:spec_decode_num_draft_tokens_total 40\nvllm:spec_decode_num_accepted_tokens_total 20\n') after = b.parse_metrics('vllm:spec_decode_num_drafts_total{engine="0"} 12\nvllm:spec_decode_num_draft_tokens_total 48\nvllm:spec_decode_num_accepted_tokens_total 25\nvllm:spec_decode_num_accepted_tokens_per_pos_total{position="0",engine="0"} 2\n') row = b.spec_delta(before, after) self.assertEqual(row['mean_accepted_per_round'], 2.5) self.assertEqual(row['acceptance_per_position_per_round'], {'0': 1.0}) self.assertEqual(row['proposed_tokens'], 8) self.assertEqual(row['accepted_tokens'], 5) with self.assertRaisesRegex(RuntimeError, 'reset/disappeared'): b.spec_delta(after, before) self.assertIsNone(b.spec_delta({}, {})['rounds']) def test_idle_fail_closed(self) -> None: cases: tuple[dict[str, float], ...] = ({}, {'vllm:num_requests_running': 0}, {'vllm:num_requests_running': 1, 'vllm:num_requests_waiting': 0}) for metrics in cases: with self.assertRaises(RuntimeError): b.require_idle(metrics) b.require_idle({'vllm:num_requests_running': 0, 'vllm:num_requests_waiting': 0}) def test_prompt_counts_include_template_and_instruction(self) -> None: with patch.object(b, 'build_real', return_value='x' * 200), patch.object(b, 'prompt_tokens', side_effect=lambda p, t: len(p)): prompt, count = b.make_prompt(300, 'extract', 'real', 1, {}) self.assertEqual(count, len(prompt)) self.assertLessEqual(count, 300) self.assertIn(b.WORKLOADS['extract'], prompt) def test_settings_priority_and_salts(self) -> None: args = argparse.Namespace(max_tokens=10, temperature=1, top_p=.9, seed=42, thinking=True, reasoning_effort='high', priority=10) row = b.request_body('input', args, 'salt') self.assertEqual(row['cache_salt'], 'salt') self.assertEqual(row['priority'], 10) self.assertEqual(row['chat_template_kwargs'], {'thinking': True, 'reasoning_effort': 'high'}) def test_concurrent_and_mixed_overlap_and_jsonl(self) -> None: lock = threading.Lock() active = 0 peak = 0 def stream(body: b.Object, started: threading.Event | None = None) -> b.Object: nonlocal active, peak with lock: active += 1 peak = max(peak, active) if started: started.set() time.sleep(.1) with lock: active -= 1 return {'completion_tokens': 10, 'total_s': .1, 'decode_tps': 100, 'ttft_s': .01, 'first_visible_s': .02, 'finish': 'stop'} with tempfile.TemporaryDirectory() as tmp: args = argparse.Namespace(label='A', max_tokens=10, temperature=1, top_p=1, seed=42, thinking=True, reasoning_effort='high', priority=10, keep_text=False, spark_hosts=[], x86host_host='unused', process_name=[], extra_metrics=[], x86host_nic=[], out=str(Path(tmp)/'rows.jsonl'), mixed_timeout=1, mixed_delay=0, metrics_settle=0, cache_mode='cold') with patch.object(b, 'Telemetry', NullTelemetry), patch.object(b, 'scrape', return_value='vllm:num_requests_running 0\nvllm:num_requests_waiting 0\n'), patch.object(b, 'stream_request', side_effect=stream): for mixed in (False, True): row = b.run_group(args, {}, 32000, 'legal', 3, 0, [('input', 5)]*3, ['s']*3, mixed) self.assertEqual(row['delivered_tokens'], 30) self.assertEqual(row['request_count'], 3) self.assertIsInstance(row['aggregate_tps'], float) self.assertEqual(peak, 3) rows = [json.loads(x) for x in Path(args.out).read_text().splitlines()] self.assertEqual(sum(r['type']=='run' for r in rows), 2) self.assertEqual(sum(r['type']=='request' for r in rows), 6) rendered = report.table([Path(args.out)]) self.assertIn('1/1', rendered) self.assertIn('±', rendered) def test_dry_run_never_calls_inference_or_tokenize(self) -> None: with tempfile.TemporaryDirectory() as tmp: path = Path(tmp)/'dry.jsonl' with patch('sys.argv', ['bench2', '--dry-run', '--out', str(path)]), patch.object(b, 'models', return_value={'data':[{'id':'fixture'}]}), patch.object(b, 'scrape', return_value=''), patch.object(b, 'Telemetry', NullTelemetry), patch.object(b, 'post') as post, patch.object(b, 'stream_request') as stream: b.main() post.assert_not_called() stream.assert_not_called() self.assertEqual(json.loads(path.read_text())['generation_requests'], 0) def test_report_duplicate_detection_and_repeat_spread(self) -> None: self.assertEqual(report.spread([1,3]), '2.00 ± 1.41') with tempfile.TemporaryDirectory() as tmp: path = Path(tmp)/'rows.jsonl' row = {'type':'request', 'run_id':'same', 'request_index':0} path.write_text(json.dumps(row)+'\n'+json.dumps(row)+'\n') with self.assertRaisesRegex(ValueError, 'duplicate'): report.table([path]) if __name__ == '__main__': unittest.main()