| """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() |
|
|
