From 1f1c175a3faef837dde1ffdf0dfe2005026c0585 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:39:19 -0500 Subject: [PATCH 01/26] test: pin an archived pre-usage run against post-run consumers Co-Authored-By: Claude Haiku 5.5 --- .../echo_char_streaming/README.md | 28 ++ .../echo_char_streaming/dataset.jsonl | 5 + .../echo_char_streaming/events.jsonl | 47 ++ .../metrics/final_snapshot.json | 442 +++++++++++++++++ .../performance/result_summary.json | 469 ++++++++++++++++++ .../echo_char_streaming/sample_idx_map.json | 19 + .../unit/metrics/test_archived_run_compat.py | 200 ++++++++ 7 files changed, 1210 insertions(+) create mode 100644 tests/assets/archived_runs/echo_char_streaming/README.md create mode 100644 tests/assets/archived_runs/echo_char_streaming/dataset.jsonl create mode 100644 tests/assets/archived_runs/echo_char_streaming/events.jsonl create mode 100644 tests/assets/archived_runs/echo_char_streaming/metrics/final_snapshot.json create mode 100644 tests/assets/archived_runs/echo_char_streaming/performance/result_summary.json create mode 100644 tests/assets/archived_runs/echo_char_streaming/sample_idx_map.json create mode 100644 tests/unit/metrics/test_archived_run_compat.py diff --git a/tests/assets/archived_runs/echo_char_streaming/README.md b/tests/assets/archived_runs/echo_char_streaming/README.md new file mode 100644 index 000000000..0f83f37e0 --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/README.md @@ -0,0 +1,28 @@ +# Archived run: echo server, `char` tokenizer, streaming + +Frozen report directory from one real run of commit `93ec321e`, before the +benchmark read server-reported token counts. Every token count in it comes from +the client tokenizer, and no record carries `usage`. +`tests/unit/metrics/test_archived_run_compat.py` checks that current code still +decodes these files and recomputes the numbers this run reported. Do not +regenerate or edit them: the test pins values from these exact files. + +Kept from the report directory: `events.jsonl`, `sample_idx_map.json`, +`metrics/final_snapshot.json` and `performance/result_summary.json`. +`dataset.jsonl` is the input (five prompts, the last one empty, so three samples +have empty output). + +Produced from the repository root, with the echo server on a free local port: + +```bash +python -m inference_endpoint.testing.echo_server --port 18765 & +inference-endpoint benchmark online \ + --endpoints http://127.0.0.1:18765 \ + --model tests/assets/tokenizers/char \ + --dataset tests/assets/archived_runs/echo_char_streaming/dataset.jsonl \ + --load-pattern concurrency --concurrency 2 --num-samples 15 \ + --streaming on --steady-state \ + --workers 1 --no-cpu-affinity --max-connections 32 \ + --report-dir "$OUT" +kill %1 +``` diff --git a/tests/assets/archived_runs/echo_char_streaming/dataset.jsonl b/tests/assets/archived_runs/echo_char_streaming/dataset.jsonl new file mode 100644 index 000000000..ae1913b43 --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/dataset.jsonl @@ -0,0 +1,5 @@ +{"prompt": "Name three primary colors"} +{"prompt": "Explain why the sky looks blue during the day"} +{"prompt": "Count from one to five"} +{"prompt": "Summarize a short story about a lighthouse keeper"} +{"prompt": ""} diff --git a/tests/assets/archived_runs/echo_char_streaming/events.jsonl b/tests/assets/archived_runs/echo_char_streaming/events.jsonl new file mode 100644 index 000000000..ccc1b2cf6 --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/events.jsonl @@ -0,0 +1,47 @@ +{"event_type":"session.started","timestamp_ns":1984175435233708,"sample_uuid":"","conversation_id":"","turn":null,"data":null} +{"event_type":"session.phase_start","timestamp_ns":1984175439782125,"sample_uuid":"","conversation_id":"","turn":null,"data":["PhaseData","performance",true,5,0]} +{"event_type":"session.start_performance_tracking","timestamp_ns":1984175439813375,"sample_uuid":"","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.issued","timestamp_ns":1984175440037041,"sample_uuid":"d34749f3dac14301b3b510b87ad6bb66","conversation_id":"","turn":null,"data":["PromptData","Summarize a short story about a lighthouse keeper",null,null,null,null,null,null]} +{"event_type":"sample.issued","timestamp_ns":1984175440764083,"sample_uuid":"5f9fbd6c315545f9a9768ceb990bc623","conversation_id":"","turn":null,"data":["PromptData","Explain why the sky looks blue during the day",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175443295458,"sample_uuid":"d34749f3dac14301b3b510b87ad6bb66","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175443350583,"sample_uuid":"d34749f3dac14301b3b510b87ad6bb66","conversation_id":"","turn":null,"data":["TextModelOutput",["Summarize"," a short story about a lighthouse keeper"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175443617583,"sample_uuid":"fa36867a47fe40679d072e2cc9901dff","conversation_id":"","turn":null,"data":["PromptData","Count from one to five",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175443631041,"sample_uuid":"5f9fbd6c315545f9a9768ceb990bc623","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175443810041,"sample_uuid":"5f9fbd6c315545f9a9768ceb990bc623","conversation_id":"","turn":null,"data":["TextModelOutput",["Explain"," why the sky looks blue during the day"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175443871666,"sample_uuid":"a40baa4de8f5438ba44e8c1413ae6a9b","conversation_id":"","turn":null,"data":["PromptData","",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175444239458,"sample_uuid":"fa36867a47fe40679d072e2cc9901dff","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175444311333,"sample_uuid":"fa36867a47fe40679d072e2cc9901dff","conversation_id":"","turn":null,"data":["TextModelOutput",["Count"," from one to five"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175444374791,"sample_uuid":"ba2041fe3d77426ebf231083159ea9b1","conversation_id":"","turn":null,"data":["PromptData","Name three primary colors",null,null,null,null,null,null]} +{"event_type":"sample.complete","timestamp_ns":1984175444430375,"sample_uuid":"a40baa4de8f5438ba44e8c1413ae6a9b","conversation_id":"","turn":null,"data":["TextModelOutput",[],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175444484000,"sample_uuid":"6c60e58013aa41c4bd3ccb4abf198acb","conversation_id":"","turn":null,"data":["PromptData","",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175444757916,"sample_uuid":"ba2041fe3d77426ebf231083159ea9b1","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175444872541,"sample_uuid":"ba2041fe3d77426ebf231083159ea9b1","conversation_id":"","turn":null,"data":["TextModelOutput",["Name"," three primary colors"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175444933250,"sample_uuid":"6feaea2825914b7ba055447894a3be67","conversation_id":"","turn":null,"data":["PromptData","Count from one to five",null,null,null,null,null,null]} +{"event_type":"sample.complete","timestamp_ns":1984175445032000,"sample_uuid":"6c60e58013aa41c4bd3ccb4abf198acb","conversation_id":"","turn":null,"data":["TextModelOutput",[],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175445099000,"sample_uuid":"4e71a95113ca4feba36fadaa453f9217","conversation_id":"","turn":null,"data":["PromptData","Summarize a short story about a lighthouse keeper",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175445276250,"sample_uuid":"6feaea2825914b7ba055447894a3be67","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175445427916,"sample_uuid":"6feaea2825914b7ba055447894a3be67","conversation_id":"","turn":null,"data":["TextModelOutput",["Count"," from one to five"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175445477416,"sample_uuid":"66eb38a736b74b71af2257009f139e82","conversation_id":"","turn":null,"data":["PromptData","Name three primary colors",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175445529041,"sample_uuid":"4e71a95113ca4feba36fadaa453f9217","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175445775083,"sample_uuid":"4e71a95113ca4feba36fadaa453f9217","conversation_id":"","turn":null,"data":["TextModelOutput",["Summarize"," a short story about a lighthouse keeper"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175445836625,"sample_uuid":"fb098c4dfb0e447da3a00c43844c333f","conversation_id":"","turn":null,"data":["PromptData","Explain why the sky looks blue during the day",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175445998416,"sample_uuid":"66eb38a736b74b71af2257009f139e82","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175446101458,"sample_uuid":"66eb38a736b74b71af2257009f139e82","conversation_id":"","turn":null,"data":["TextModelOutput",["Name"," three primary colors"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175446150500,"sample_uuid":"123093183f044a978476c7fb29564702","conversation_id":"","turn":null,"data":["PromptData","Name three primary colors",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175446228041,"sample_uuid":"fb098c4dfb0e447da3a00c43844c333f","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175446373416,"sample_uuid":"fb098c4dfb0e447da3a00c43844c333f","conversation_id":"","turn":null,"data":["TextModelOutput",["Explain"," why the sky looks blue during the day"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175446430625,"sample_uuid":"ca0da7923b944b78be04e56ff5dbd5d0","conversation_id":"","turn":null,"data":["PromptData","Count from one to five",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175446558291,"sample_uuid":"123093183f044a978476c7fb29564702","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175446596875,"sample_uuid":"123093183f044a978476c7fb29564702","conversation_id":"","turn":null,"data":["TextModelOutput",["Name"," three primary colors"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175446647583,"sample_uuid":"1482678cb4fc4fdc845688081eabcbd7","conversation_id":"","turn":null,"data":["PromptData","Summarize a short story about a lighthouse keeper",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175446818041,"sample_uuid":"ca0da7923b944b78be04e56ff5dbd5d0","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175446871291,"sample_uuid":"ca0da7923b944b78be04e56ff5dbd5d0","conversation_id":"","turn":null,"data":["TextModelOutput",["Count"," from one to five"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175446899416,"sample_uuid":"1eaa615af3f14b92bae05bda2321220e","conversation_id":"","turn":null,"data":["PromptData","",null,null,null,null,null,null]} +{"event_type":"sample.recv_first","timestamp_ns":1984175447067541,"sample_uuid":"1482678cb4fc4fdc845688081eabcbd7","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175447139333,"sample_uuid":"1482678cb4fc4fdc845688081eabcbd7","conversation_id":"","turn":null,"data":["TextModelOutput",["Summarize"," a short story about a lighthouse keeper"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.issued","timestamp_ns":1984175447185083,"sample_uuid":"8b00dd93572c4178bebdab64214359e6","conversation_id":"","turn":null,"data":["PromptData","Explain why the sky looks blue during the day",null,null,null,null,null,null]} +{"event_type":"sample.complete","timestamp_ns":1984175447298125,"sample_uuid":"1eaa615af3f14b92bae05bda2321220e","conversation_id":"","turn":null,"data":["TextModelOutput",[],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"sample.recv_first","timestamp_ns":1984175447475000,"sample_uuid":"8b00dd93572c4178bebdab64214359e6","conversation_id":"","turn":null,"data":null} +{"event_type":"sample.complete","timestamp_ns":1984175447651041,"sample_uuid":"8b00dd93572c4178bebdab64214359e6","conversation_id":"","turn":null,"data":["TextModelOutput",["Explain"," why the sky looks blue during the day"],null,null],"finish_reason":"stop","worker_id":0} +{"event_type":"session.stop_performance_tracking","timestamp_ns":1984175447716250,"sample_uuid":"","conversation_id":"","turn":null,"data":null} +{"event_type":"session.ended","timestamp_ns":1984175447840875,"sample_uuid":"","conversation_id":"","turn":null,"data":null} diff --git a/tests/assets/archived_runs/echo_char_streaming/metrics/final_snapshot.json b/tests/assets/archived_runs/echo_char_streaming/metrics/final_snapshot.json new file mode 100644 index 000000000..9833c6a0a --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/metrics/final_snapshot.json @@ -0,0 +1,442 @@ +{ + "counter": 1, + "timestamp_ns": 1984175454722166, + "state": "complete", + "n_pending_tasks": 0, + "metrics": [ + { + "type": "counter", + "name": "total_samples_issued", + "value": 15 + }, + { + "type": "counter", + "name": "total_samples_completed", + "value": 15 + }, + { + "type": "counter", + "name": "total_samples_failed", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_samples_issued", + "value": 15 + }, + { + "type": "counter", + "name": "tracked_samples_completed", + "value": 15 + }, + { + "type": "counter", + "name": "tracked_samples_failed", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_finish_reason_stop", + "value": 15 + }, + { + "type": "counter", + "name": "tracked_finish_reason_length", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_finish_reason_tool_calls", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_finish_reason_content_filter", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_finish_reason_function_call", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_finish_reason_other", + "value": 0 + }, + { + "type": "counter", + "name": "tracked_duration_ns", + "value": 7837666 + }, + { + "type": "counter", + "name": "legacy_loadgen_window_duration_ns", + "value": 7614000 + }, + { + "type": "counter", + "name": "total_duration_ns", + "value": 12607167 + }, + { + "type": "series", + "name": "sample_latency_ns", + "count": 15, + "total": 13232749, + "min": 398709, + "max": 3313542, + "sum_sq": 23989621216225.0, + "percentiles": { + "99.9": 3045958.0, + "99.0": 3045958.0, + "97.0": 3045958.0, + "95.0": 3045958.0, + "90.0": 693750.0, + "80.0": 676083.0, + "75.0": 624042.0, + "50.0": 536791.0, + "25.0": 465958.0, + "10.0": 440666.0, + "5.0": 398709.0, + "1.0": 398709.0 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 3045958.0 + }, + "histogram": [ + [[398708.9999999999, 430516.07519651874], 1], + [[430516.07519651874, 464860.5649800096], 2], + [[464860.5649800096, 501944.8920110408], 4], + [[501944.8920110408, 541987.6272507855], 1], + [[541987.6272507855, 585224.7781943273], 2], + [[585224.7781943273, 631911.1798729778], 1], + [[631911.1798729778, 682321.9968240393], 1], + [[682321.9968240393, 736754.3448804442], 1], + [[736754.3448804442, 795529.0423389273], 0], + [[795529.0423389273, 858992.5008279228], 0], + [[858992.5008279228, 927518.7670197559], 0], + [[927518.7670197559, 1001511.7272207525], 0], + [[1001511.7272207525, 1081407.4878328932], 0], + [[1081407.4878328932, 1167676.9457171634], 0], + [[1167676.9457171634, 1260828.5636080746], 0], + [[1260828.5636080746, 1361411.366937322], 0], + [[1361411.366937322, 1470018.1797295364], 0], + [[1470018.1797295364, 1587289.118642145], 0], + [[1587289.118642145, 1713915.3657428266], 0], + [[1713915.3657428266, 1850643.242260917], 0], + [[1850643.242260917, 1998278.607322961], 0], + [[1998278.607322961, 2157691.6075981404], 0], + [[2157691.6075981404, 2329821.80584742], 0], + [[2329821.80584742, 2515683.718603536], 0], + [[2515683.718603536, 2716372.79562031], 0], + [[2716372.79562031, 2933071.8763334905], 0], + [[2933071.8763334905, 3167058.161386841], 1], + [[3167058.161386841, 3419710.7403127686], 1], + [[3419710.7403127686, 3692518.719735033], 0], + [[3692518.719735033, 3987090.000000003], 0] + ] + }, + { + "type": "series", + "name": "isl", + "count": 12, + "total": 423, + "min": 22, + "max": 49, + "sum_sq": 16605.0, + "percentiles": { + "99.9": 49.0, + "99.0": 49.0, + "97.0": 49.0, + "95.0": 49.0, + "90.0": 49.0, + "80.0": 45.0, + "75.0": 45.0, + "50.0": 25.0, + "25.0": 22.0, + "10.0": 22.0, + "5.0": 22.0, + "1.0": 22.0 + }, + "histogram": [ + [[22.000000000000004, 23.755053571209615], 3], + [[23.755053571209615, 25.650116825956307], 3], + [[25.650116825956307, 27.696359059471682], 0], + [[27.696359059471682, 29.905840599327576], 0], + [[29.905840599327576, 32.29158388768554], 0], + [[32.29158388768554, 34.8676502341445], 0], + [[34.8676502341445, 37.64922269155919], 0], + [[37.64922269155919, 40.652695543290406], 0], + [[40.652695543290406, 43.89577092931537], 0], + [[43.89577092931537, 47.397563180701454], 3], + [[47.397563180701454, 51.1787114773798], 3], + [[51.1787114773798, 55.26150149321076], 0], + [[55.26150149321076, 59.66999674530464], 0], + [[59.66999674530464, 64.43018042175524], 0], + [[64.43018042175524, 69.5701085237044], 0], + [[69.5701085237044, 75.12007522433925], 0], + [[75.12007522433925, 81.11279141943069], 0], + [[81.11279141943069, 87.58357752176946], 0], + [[87.58357752176946, 94.5705716358101], 0], + [[94.5705716358101, 102.11495433948117], 0], + [[102.11495433948117, 110.26119139799997], 0], + [[110.26119139799997, 119.05729584022211], 0], + [[119.05729584022211, 128.5551109421739], 0], + [[128.5551109421739, 138.8106157856426], 0], + [[138.8106157856426, 149.88425519275148], 0], + [[149.88425519275148, 161.84129598112114], 0], + [[161.84129598112114, 174.75221163934202], 0], + [[174.75221163934202, 188.69309768999685], 0], + [[188.69309768999685, 203.74612018833477], 0], + [[203.74612018833477, 220.00000000000009], 0] + ] + }, + { + "type": "series", + "name": "osl", + "count": 12, + "total": 423, + "min": 22, + "max": 49, + "sum_sq": 16605.0, + "percentiles": { + "99.9": 49.0, + "99.0": 49.0, + "97.0": 49.0, + "95.0": 49.0, + "90.0": 49.0, + "80.0": 45.0, + "75.0": 45.0, + "50.0": 25.0, + "25.0": 22.0, + "10.0": 22.0, + "5.0": 22.0, + "1.0": 22.0 + }, + "histogram": [ + [[22.000000000000004, 23.755053571209615], 3], + [[23.755053571209615, 25.650116825956307], 3], + [[25.650116825956307, 27.696359059471682], 0], + [[27.696359059471682, 29.905840599327576], 0], + [[29.905840599327576, 32.29158388768554], 0], + [[32.29158388768554, 34.8676502341445], 0], + [[34.8676502341445, 37.64922269155919], 0], + [[37.64922269155919, 40.652695543290406], 0], + [[40.652695543290406, 43.89577092931537], 0], + [[43.89577092931537, 47.397563180701454], 3], + [[47.397563180701454, 51.1787114773798], 3], + [[51.1787114773798, 55.26150149321076], 0], + [[55.26150149321076, 59.66999674530464], 0], + [[59.66999674530464, 64.43018042175524], 0], + [[64.43018042175524, 69.5701085237044], 0], + [[69.5701085237044, 75.12007522433925], 0], + [[75.12007522433925, 81.11279141943069], 0], + [[81.11279141943069, 87.58357752176946], 0], + [[87.58357752176946, 94.5705716358101], 0], + [[94.5705716358101, 102.11495433948117], 0], + [[102.11495433948117, 110.26119139799997], 0], + [[110.26119139799997, 119.05729584022211], 0], + [[119.05729584022211, 128.5551109421739], 0], + [[128.5551109421739, 138.8106157856426], 0], + [[138.8106157856426, 149.88425519275148], 0], + [[149.88425519275148, 161.84129598112114], 0], + [[161.84129598112114, 174.75221163934202], 0], + [[174.75221163934202, 188.69309768999685], 0], + [[188.69309768999685, 203.74612018833477], 0], + [[203.74612018833477, 220.00000000000009], 0] + ] + }, + { + "type": "series", + "name": "ttft_ns", + "count": 12, + "total": 10320914, + "min": 289917, + "max": 3258417, + "sum_sq": 20674275793030.0, + "percentiles": { + "99.9": 2866958.0, + "99.0": 2866958.0, + "97.0": 2866958.0, + "95.0": 2866958.0, + "90.0": 621875.0, + "80.0": 521000.0, + "75.0": 521000.0, + "50.0": 407791.0, + "25.0": 383125.0, + "10.0": 343000.0, + "5.0": 289917.0, + "1.0": 289917.0 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 3258417.0 + }, + "histogram": [ + [[289917.00000000006, 314266.5187543599], 1], + [[314266.5187543599, 340661.1023499292], 0], + [[340661.1023499292, 369272.5114792686], 1], + [[369272.5114792686, 400286.932654067], 3], + [[400286.932654067, 433906.1897993375], 3], + [[433906.1897993375, 470349.05760685407], 0], + [[470349.05760685407, 509852.68519438285], 0], + [[509852.68519438285, 552674.1393350563], 1], + [[552674.1393350563, 599092.0772993311], 0], + [[599092.0772993311, 649408.5601954305], 1], + [[649408.5601954305, 703951.0186084262], 0], + [[703951.0186084262, 763074.383329214], 0], + [[763074.383329214, 827163.3950389333], 0], + [[827163.3950389333, 896635.1079789143], 0], + [[896635.1079789143, 971941.6038986086], 0], + [[971941.6038986086, 1053572.9339422823], 0], + [[1053572.9339422823, 1142060.3076185903], 0], + [[1142060.3076185903, 1237979.5496049824], 0], + [[1237979.5496049824, 1341954.8468818597], 0], + [[1341954.8468818597, 1454662.8105806217], 0], + [[1454662.8105806217, 1576836.8789777928], 0], + [[1576836.8789777928, 1709272.0902873608], 0], + [[1709272.0902873608, 1852830.2563099014], 0], + [[1852830.2563099014, 2008445.5706055905], 0], + [[2008445.5706055905, 2177130.687685899], 0], + [[2177130.687685899, 2359983.3127837726], 0], + [[2359983.3127837726, 2558193.345084762], 0], + [[2558193.345084762, 2773050.6209030827], 0], + [[2773050.6209030827, 3005953.307190775], 1], + [[3005953.307190775, 3258416.999999999], 1] + ] + }, + { + "type": "series", + "name": "chunk_delta_ns", + "count": 0, + "total": 0, + "min": 0, + "max": 0, + "sum_sq": 0, + "percentiles": { + "99.9": 0.0, + "99.0": 0.0, + "97.0": 0.0, + "95.0": 0.0, + "90.0": 0.0, + "80.0": 0.0, + "75.0": 0.0, + "50.0": 0.0, + "25.0": 0.0, + "10.0": 0.0, + "5.0": 0.0, + "1.0": 0.0 + }, + "histogram": [] + }, + { + "type": "series", + "name": "tpot_ns", + "count": 12, + "total": 50977.06920610349, + "min": 1378.125, + "max": 8921.529411764706, + "sum_sq": 265768504.85745338, + "percentiles": { + "99.9": 6151.05, + "99.0": 6151.05, + "97.0": 6151.05, + "95.0": 6151.05, + "90.0": 5458.333333333333, + "80.0": 4906.761904761905, + "75.0": 4906.761904761905, + "50.0": 4227.941176470588, + "25.0": 1837.3333333333333, + "10.0": 1794.8, + "5.0": 1378.125, + "1.0": 1378.125 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 8921.529411764706 + }, + "histogram": [ + [[1378.125, 1488.0651455828754], 1], + [[1488.0651455828754, 1606.7757841259552], 0], + [[1606.7757841259552, 1734.9565831288367], 0], + [[1734.9565831288367, 1873.3630261794685], 2], + [[1873.3630261794685, 2022.8108656916656], 0], + [[2022.8108656916656, 2184.180930860472], 0], + [[2184.180930860472, 2358.4243191729547], 0], + [[2358.4243191729547, 2546.568002072596], 0], + [[2546.568002072596, 2749.720877816487], 0], + [[2749.720877816487, 2969.08030720019], 0], + [[2969.08030720019, 3205.939170671092], 1], + [[3205.939170671092, 3461.6934884241423], 0], + [[3461.6934884241423, 3737.850648391955], 0], + [[3737.850648391955, 4036.0382906241516], 1], + [[4036.0382906241516, 4358.013900419546], 1], + [[4358.013900419546, 4705.675166751932], 1], + [[4705.675166751932, 5081.071167041041], 2], + [[5081.071167041041, 5486.414444190387], 1], + [[5486.414444190387, 5924.094047072767], 0], + [[5924.094047072767, 6396.689611322618], 1], + [[6396.689611322618, 6906.986563425843], 0], + [[6906.986563425843, 7457.992537718451], 0], + [[7457.992537718451, 8052.9551030537905], 0], + [[8052.9551030537905, 8695.380903617666], 0], + [[8695.380903617666, 9389.056326705035], 1], + [[9389.056326705035, 10138.069819271941], 0], + [[10138.069819271941, 10946.835984794017], 0], + [[10946.835984794017, 11820.121602455756], 0], + [[11820.121602455756, 12763.073722024936], 0], + [[12763.073722024936, 13781.250000000004], 0] + ] + } + ], + "steady_state": { + "found": false, + "reason": "no admissible steady plateau", + "superpass_size": 5, + "n_super_passes": 3, + "warmup": 1, + "window": null, + "ttft": null, + "tpot": null, + "osl": null, + "latency": null, + "tps": null, + "cov": {}, + "cov_basis": null, + "anomaly": { + "detected": false, + "change_point_sp": null, + "delta_pct": 0.0, + "pettitt": null, + "plateaus": [] + }, + "short_window": null, + "global_trend": { + "ttft_p50": "insufficient", + "ttft_p90": "insufficient", + "tpot_p50": "insufficient", + "tpot_p90": "insufficient" + }, + "drifting_up": [] + } +} diff --git a/tests/assets/archived_runs/echo_char_streaming/performance/result_summary.json b/tests/assets/archived_runs/echo_char_streaming/performance/result_summary.json new file mode 100644 index 000000000..34fd5538c --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/performance/result_summary.json @@ -0,0 +1,469 @@ +{ + "version": "0.1.0", + "git_sha": "93ec321", + "test_started_at": 0, + "n_samples_issued": 15, + "n_samples_completed": 15, + "n_samples_failed": 0, + "duration_ns": 7837666, + "state": "complete", + "complete": true, + "ttft": { + "total": 10320914, + "min": 289917, + "max": 3258417, + "median": 407791.0, + "avg": 860076.1666666666, + "std_dev": 1035615.9187746316, + "percentiles": { + "99.9": 2866958.0, + "99.0": 2866958.0, + "97.0": 2866958.0, + "95.0": 2866958.0, + "90.0": 621875.0, + "80.0": 521000.0, + "75.0": 521000.0, + "50.0": 407791.0, + "25.0": 383125.0, + "10.0": 343000.0, + "5.0": 289917.0, + "1.0": 289917.0 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 3258417.0 + }, + "histogram": { + "buckets": [ + [289917.00000000006, 314266.5187543599], + [314266.5187543599, 340661.1023499292], + [340661.1023499292, 369272.5114792686], + [369272.5114792686, 400286.932654067], + [400286.932654067, 433906.1897993375], + [433906.1897993375, 470349.05760685407], + [470349.05760685407, 509852.68519438285], + [509852.68519438285, 552674.1393350563], + [552674.1393350563, 599092.0772993311], + [599092.0772993311, 649408.5601954305], + [649408.5601954305, 703951.0186084262], + [703951.0186084262, 763074.383329214], + [763074.383329214, 827163.3950389333], + [827163.3950389333, 896635.1079789143], + [896635.1079789143, 971941.6038986086], + [971941.6038986086, 1053572.9339422823], + [1053572.9339422823, 1142060.3076185903], + [1142060.3076185903, 1237979.5496049824], + [1237979.5496049824, 1341954.8468818597], + [1341954.8468818597, 1454662.8105806217], + [1454662.8105806217, 1576836.8789777928], + [1576836.8789777928, 1709272.0902873608], + [1709272.0902873608, 1852830.2563099014], + [1852830.2563099014, 2008445.5706055905], + [2008445.5706055905, 2177130.687685899], + [2177130.687685899, 2359983.3127837726], + [2359983.3127837726, 2558193.345084762], + [2558193.345084762, 2773050.6209030827], + [2773050.6209030827, 3005953.307190775], + [3005953.307190775, 3258416.999999999] + ], + "counts": [ + 1, 0, 1, 3, 3, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 1 + ] + } + }, + "tpot": { + "total": 50977.06920610349, + "min": 1378.125, + "max": 8921.529411764706, + "median": 4227.941176470588, + "avg": 4248.089100508624, + "std_dev": 2115.1697285674973, + "percentiles": { + "99.9": 6151.05, + "99.0": 6151.05, + "97.0": 6151.05, + "95.0": 6151.05, + "90.0": 5458.333333333333, + "80.0": 4906.761904761905, + "75.0": 4906.761904761905, + "50.0": 4227.941176470588, + "25.0": 1837.3333333333333, + "10.0": 1794.8, + "5.0": 1378.125, + "1.0": 1378.125 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 8921.529411764706 + }, + "histogram": { + "buckets": [ + [1378.125, 1488.0651455828754], + [1488.0651455828754, 1606.7757841259552], + [1606.7757841259552, 1734.9565831288367], + [1734.9565831288367, 1873.3630261794685], + [1873.3630261794685, 2022.8108656916656], + [2022.8108656916656, 2184.180930860472], + [2184.180930860472, 2358.4243191729547], + [2358.4243191729547, 2546.568002072596], + [2546.568002072596, 2749.720877816487], + [2749.720877816487, 2969.08030720019], + [2969.08030720019, 3205.939170671092], + [3205.939170671092, 3461.6934884241423], + [3461.6934884241423, 3737.850648391955], + [3737.850648391955, 4036.0382906241516], + [4036.0382906241516, 4358.013900419546], + [4358.013900419546, 4705.675166751932], + [4705.675166751932, 5081.071167041041], + [5081.071167041041, 5486.414444190387], + [5486.414444190387, 5924.094047072767], + [5924.094047072767, 6396.689611322618], + [6396.689611322618, 6906.986563425843], + [6906.986563425843, 7457.992537718451], + [7457.992537718451, 8052.9551030537905], + [8052.9551030537905, 8695.380903617666], + [8695.380903617666, 9389.056326705035], + [9389.056326705035, 10138.069819271941], + [10138.069819271941, 10946.835984794017], + [10946.835984794017, 11820.121602455756], + [11820.121602455756, 12763.073722024936], + [12763.073722024936, 13781.250000000004] + ], + "counts": [ + 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 2, 1, 0, 1, 0, 0, 0, 0, + 1, 0, 0, 0, 0, 0 + ] + } + }, + "latency": { + "total": 13232749, + "min": 398709, + "max": 3313542, + "median": 536791.0, + "avg": 882183.2666666667, + "std_dev": 937927.482592221, + "percentiles": { + "99.9": 3045958.0, + "99.0": 3045958.0, + "97.0": 3045958.0, + "95.0": 3045958.0, + "90.0": 693750.0, + "80.0": 676083.0, + "75.0": 624042.0, + "50.0": 536791.0, + "25.0": 465958.0, + "10.0": 440666.0, + "5.0": 398709.0, + "1.0": 398709.0 + }, + "early_stopping_percentiles": { + "99.9": null, + "99.0": null, + "97.0": null, + "95.0": null, + "90.0": null, + "80.0": null, + "75.0": null, + "50.0": 3045958.0 + }, + "histogram": { + "buckets": [ + [398708.9999999999, 430516.07519651874], + [430516.07519651874, 464860.5649800096], + [464860.5649800096, 501944.8920110408], + [501944.8920110408, 541987.6272507855], + [541987.6272507855, 585224.7781943273], + [585224.7781943273, 631911.1798729778], + [631911.1798729778, 682321.9968240393], + [682321.9968240393, 736754.3448804442], + [736754.3448804442, 795529.0423389273], + [795529.0423389273, 858992.5008279228], + [858992.5008279228, 927518.7670197559], + [927518.7670197559, 1001511.7272207525], + [1001511.7272207525, 1081407.4878328932], + [1081407.4878328932, 1167676.9457171634], + [1167676.9457171634, 1260828.5636080746], + [1260828.5636080746, 1361411.366937322], + [1361411.366937322, 1470018.1797295364], + [1470018.1797295364, 1587289.118642145], + [1587289.118642145, 1713915.3657428266], + [1713915.3657428266, 1850643.242260917], + [1850643.242260917, 1998278.607322961], + [1998278.607322961, 2157691.6075981404], + [2157691.6075981404, 2329821.80584742], + [2329821.80584742, 2515683.718603536], + [2515683.718603536, 2716372.79562031], + [2716372.79562031, 2933071.8763334905], + [2933071.8763334905, 3167058.161386841], + [3167058.161386841, 3419710.7403127686], + [3419710.7403127686, 3692518.719735033], + [3692518.719735033, 3987090.000000003] + ], + "counts": [ + 1, 2, 4, 1, 2, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 1, 1, 0, 0 + ] + } + }, + "output_sequence_lengths": { + "total": 423, + "min": 22, + "max": 49, + "median": 25.0, + "avg": 35.25, + "std_dev": 12.410589320122043, + "percentiles": { + "99.9": 49.0, + "99.0": 49.0, + "97.0": 49.0, + "95.0": 49.0, + "90.0": 49.0, + "80.0": 45.0, + "75.0": 45.0, + "50.0": 25.0, + "25.0": 22.0, + "10.0": 22.0, + "5.0": 22.0, + "1.0": 22.0 + }, + "histogram": { + "buckets": [ + [22.000000000000004, 23.755053571209615], + [23.755053571209615, 25.650116825956307], + [25.650116825956307, 27.696359059471682], + [27.696359059471682, 29.905840599327576], + [29.905840599327576, 32.29158388768554], + [32.29158388768554, 34.8676502341445], + [34.8676502341445, 37.64922269155919], + [37.64922269155919, 40.652695543290406], + [40.652695543290406, 43.89577092931537], + [43.89577092931537, 47.397563180701454], + [47.397563180701454, 51.1787114773798], + [51.1787114773798, 55.26150149321076], + [55.26150149321076, 59.66999674530464], + [59.66999674530464, 64.43018042175524], + [64.43018042175524, 69.5701085237044], + [69.5701085237044, 75.12007522433925], + [75.12007522433925, 81.11279141943069], + [81.11279141943069, 87.58357752176946], + [87.58357752176946, 94.5705716358101], + [94.5705716358101, 102.11495433948117], + [102.11495433948117, 110.26119139799997], + [110.26119139799997, 119.05729584022211], + [119.05729584022211, 128.5551109421739], + [128.5551109421739, 138.8106157856426], + [138.8106157856426, 149.88425519275148], + [149.88425519275148, 161.84129598112114], + [161.84129598112114, 174.75221163934202], + [174.75221163934202, 188.69309768999685], + [188.69309768999685, 203.74612018833477], + [203.74612018833477, 220.00000000000009] + ], + "counts": [ + 3, 3, 0, 0, 0, 0, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0 + ] + } + }, + "input_sequence_lengths": { + "total": 423, + "min": 22, + "max": 49, + "median": 25.0, + "avg": 35.25, + "std_dev": 12.410589320122043, + "percentiles": { + "99.9": 49.0, + "99.0": 49.0, + "97.0": 49.0, + "95.0": 49.0, + "90.0": 49.0, + "80.0": 45.0, + "75.0": 45.0, + "50.0": 25.0, + "25.0": 22.0, + "10.0": 22.0, + "5.0": 22.0, + "1.0": 22.0 + }, + "histogram": { + "buckets": [ + [22.000000000000004, 23.755053571209615], + [23.755053571209615, 25.650116825956307], + [25.650116825956307, 27.696359059471682], + [27.696359059471682, 29.905840599327576], + [29.905840599327576, 32.29158388768554], + [32.29158388768554, 34.8676502341445], + [34.8676502341445, 37.64922269155919], + [37.64922269155919, 40.652695543290406], + [40.652695543290406, 43.89577092931537], + [43.89577092931537, 47.397563180701454], + [47.397563180701454, 51.1787114773798], + [51.1787114773798, 55.26150149321076], + [55.26150149321076, 59.66999674530464], + [59.66999674530464, 64.43018042175524], + [64.43018042175524, 69.5701085237044], + [69.5701085237044, 75.12007522433925], + [75.12007522433925, 81.11279141943069], + [81.11279141943069, 87.58357752176946], + [87.58357752176946, 94.5705716358101], + [94.5705716358101, 102.11495433948117], + [102.11495433948117, 110.26119139799997], + [110.26119139799997, 119.05729584022211], + [119.05729584022211, 128.5551109421739], + [128.5551109421739, 138.8106157856426], + [138.8106157856426, 149.88425519275148], + [149.88425519275148, 161.84129598112114], + [161.84129598112114, 174.75221163934202], + [174.75221163934202, 188.69309768999685], + [188.69309768999685, 203.74612018833477], + [203.74612018833477, 220.00000000000009] + ], + "counts": [ + 3, 3, 0, 0, 0, 0, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0 + ] + } + }, + "legacy_loadgen_window_duration_ns": null, + "qps": 1913.8350626321662, + "tps": 53970.14876622709, + "e2e_avg_interactivity": 31966.147019035878, + "finish_reason_counts": { + "stop": 15, + "length": 0, + "tool_calls": 0, + "content_filter": 0, + "function_call": 0, + "other": 0 + }, + "run_config": { + "load_pattern": { + "type": "concurrency", + "target_qps": null, + "target_concurrency": 2 + }, + "warmup": { + "enabled": false, + "n_requests": null, + "salt": false, + "drain": false, + "warmup_random_seed": 42 + }, + "scheduler_random_seed": 42, + "dataloader_random_seed": 42 + }, + "output_sequence_lengths_full_run": { + "output_sequence_lengths": { + "total": 423, + "min": 22, + "max": 49, + "median": 25.0, + "avg": 35.25, + "std_dev": 12.410589320122043, + "percentiles": { + "99.9": 49.0, + "99.0": 49.0, + "97.0": 49.0, + "95.0": 49.0, + "90.0": 49.0, + "80.0": 45.0, + "75.0": 45.0, + "50.0": 25.0, + "25.0": 22.0, + "10.0": 22.0, + "5.0": 22.0, + "1.0": 22.0 + }, + "histogram": { + "buckets": [ + [22.000000000000004, 23.755053571209615], + [23.755053571209615, 25.650116825956307], + [25.650116825956307, 27.696359059471682], + [27.696359059471682, 29.905840599327576], + [29.905840599327576, 32.29158388768554], + [32.29158388768554, 34.8676502341445], + [34.8676502341445, 37.64922269155919], + [37.64922269155919, 40.652695543290406], + [40.652695543290406, 43.89577092931537], + [43.89577092931537, 47.397563180701454], + [47.397563180701454, 51.1787114773798], + [51.1787114773798, 55.26150149321076], + [55.26150149321076, 59.66999674530464], + [59.66999674530464, 64.43018042175524], + [64.43018042175524, 69.5701085237044], + [69.5701085237044, 75.12007522433925], + [75.12007522433925, 81.11279141943069], + [81.11279141943069, 87.58357752176946], + [87.58357752176946, 94.5705716358101], + [94.5705716358101, 102.11495433948117], + [102.11495433948117, 110.26119139799997], + [110.26119139799997, 119.05729584022211], + [119.05729584022211, 128.5551109421739], + [128.5551109421739, 138.8106157856426], + [138.8106157856426, 149.88425519275148], + [149.88425519275148, 161.84129598112114], + [161.84129598112114, 174.75221163934202], + [174.75221163934202, 188.69309768999685], + [188.69309768999685, 203.74612018833477], + [203.74612018833477, 220.00000000000009] + ], + "counts": [ + 3, 3, 0, 0, 0, 0, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0 + ] + } + }, + "n_turns_counted": 12, + "n_empty": 3, + "n_errors": 0, + "n_undecodable": 0, + "n_missing": 0, + "partial": false + }, + "steady_state": { + "found": false, + "reason": "no admissible steady plateau", + "superpass_size": 5, + "n_super_passes": 3, + "warmup": 1, + "window": null, + "ttft": null, + "tpot": null, + "osl": null, + "latency": null, + "tps": null, + "cov": {}, + "cov_basis": null, + "anomaly": { + "detected": false, + "change_point_sp": null, + "delta_pct": 0.0, + "pettitt": null, + "plateaus": [] + }, + "short_window": null, + "global_trend": { + "ttft_p50": "insufficient", + "ttft_p90": "insufficient", + "tpot_p50": "insufficient", + "tpot_p90": "insufficient" + }, + "drifting_up": [] + }, + "n_samples_succeeded": 15, + "n_samples_dropped": 0 +} diff --git a/tests/assets/archived_runs/echo_char_streaming/sample_idx_map.json b/tests/assets/archived_runs/echo_char_streaming/sample_idx_map.json new file mode 100644 index 000000000..706807b29 --- /dev/null +++ b/tests/assets/archived_runs/echo_char_streaming/sample_idx_map.json @@ -0,0 +1,19 @@ +{ + "performance": { + "d34749f3dac14301b3b510b87ad6bb66": 3, + "5f9fbd6c315545f9a9768ceb990bc623": 1, + "fa36867a47fe40679d072e2cc9901dff": 2, + "a40baa4de8f5438ba44e8c1413ae6a9b": 4, + "ba2041fe3d77426ebf231083159ea9b1": 0, + "6c60e58013aa41c4bd3ccb4abf198acb": 4, + "6feaea2825914b7ba055447894a3be67": 2, + "4e71a95113ca4feba36fadaa453f9217": 3, + "66eb38a736b74b71af2257009f139e82": 0, + "fb098c4dfb0e447da3a00c43844c333f": 1, + "123093183f044a978476c7fb29564702": 0, + "ca0da7923b944b78be04e56ff5dbd5d0": 2, + "1482678cb4fc4fdc845688081eabcbd7": 3, + "1eaa615af3f14b92bae05bda2321220e": 4, + "8b00dd93572c4178bebdab64214359e6": 1 + } +} diff --git a/tests/unit/metrics/test_archived_run_compat.py b/tests/unit/metrics/test_archived_run_compat.py new file mode 100644 index 000000000..8392fe0ac --- /dev/null +++ b/tests/unit/metrics/test_archived_run_compat.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A run archived before server token counts existed keeps its numbers. + +``tests/assets/archived_runs/echo_char_streaming`` is a real report directory +(client tokenizer counts only, no ``usage`` anywhere). Every post-run consumer +must still decode it and recompute exactly what that run reported. A failure +here means an archived ``events.jsonl`` or ``final_snapshot.json`` changed +meaning, which is never acceptable; fix the code, not the pinned values. +""" + +import importlib.util +import json +from collections import Counter +from pathlib import Path + +import msgspec +import pytest +from inference_endpoint.commands.benchmark.full_run_osl import full_run_osl_for_report +from inference_endpoint.core.record import EventRecord, EventType +from inference_endpoint.metrics import steady_state_diagnostics as ssd +from inference_endpoint.metrics.report import Report + +pytestmark = pytest.mark.unit + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_RUN_DIR = _REPO_ROOT / "tests/assets/archived_runs/echo_char_streaming" +_EVENTS = _RUN_DIR / "events.jsonl" +_SUMMARY = _RUN_DIR / "performance/result_summary.json" +_CHAR_TOKENIZER_DIR = _REPO_ROOT / "tests/assets/tokenizers/char" +_SUMMARY_FIELDS = ( + "n_samples_issued", + "n_samples_completed", + "n_samples_failed", + "duration_ns", + "state", + "complete", + "ttft", + "tpot", + "latency", + "output_sequence_lengths", + "input_sequence_lengths", + "legacy_loadgen_window_duration_ns", + "qps", + "tps", + "e2e_avg_interactivity", + "finish_reason_counts", + "steady_state", +) + + +def _json(path): + return json.loads(Path(path).read_bytes()) + + +def _char_counts(texts): + # The char tokenizer emits one token per ASCII character. + return [len(t) for t in texts] + + +def _load_es_script(): + spec = importlib.util.spec_from_file_location( + "early_stopping_estimate_from_events_archived", + _REPO_ROOT / "scripts/early_stopping_estimate_from_events.py", + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_every_line_decodes_as_event_record(): + decoder = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) + records = [decoder.decode(line) for line in _EVENTS.read_bytes().splitlines()] + assert Counter(r.event_type.topic for r in records) == { + "session.started": 1, + "session.phase_start": 1, + "session.start_performance_tracking": 1, + "sample.issued": 15, + "sample.recv_first": 12, + "sample.complete": 15, + "session.stop_performance_tracking": 1, + "session.ended": 1, + } + for record in records: + assert getattr(record, "prompt_tokens", msgspec.UNSET) is msgspec.UNSET + assert getattr(record, "completion_tokens", msgspec.UNSET) is msgspec.UNSET + + +def test_report_from_archived_snapshot_reproduces_archived_summary(): + summary = _json(_SUMMARY) + report = Report.from_snapshot( + _json(_RUN_DIR / "metrics/final_snapshot.json"), + run_config=summary["run_config"], + use_legacy_loadgen_qps_metrics=False, # concurrency run + ) + rebuilt = json.loads(report.to_json()) + for field in _SUMMARY_FIELDS: + assert rebuilt[field] == summary[field], field + assert (report.n_samples_issued, report.n_samples_completed) == (15, 15) + assert (report.n_samples_failed, report.duration_ns) == (0, 7837666) + assert (report.state, report.complete) == ("complete", True) + osl = report.output_sequence_lengths + assert (osl["total"], osl["min"], osl["max"], osl["avg"]) == (423, 22, 49, 35.25) + assert report.input_sequence_lengths["total"] == 423 + assert report.tpot["total"] == pytest.approx(50977.06920610349, rel=1e-12) + assert report.tpot["avg"] == pytest.approx(4248.089100508624, rel=1e-12) + assert (report.ttft["total"], report.latency["total"]) == (10320914, 13232749) + assert report.qps == pytest.approx(1913.8350626321662, rel=1e-12) + assert report.tps == pytest.approx(53970.14876622709, rel=1e-12) + assert report.finish_reason_counts["stop"] == 15 + assert report.steady_state is not None + assert report.steady_state.found is False + assert report.steady_state.superpass_size == 5 + assert report.steady_state.n_super_passes == 3 + + +def test_full_run_osl_recomputes_archived_block(): + block = full_run_osl_for_report(_RUN_DIR, str(_CHAR_TOKENIZER_DIR)) + archived = _json(_SUMMARY)["output_sequence_lengths_full_run"] + assert block is not None + for key in ( + "output_sequence_lengths", + "n_turns_counted", + "n_empty", + "n_errors", + "n_undecodable", + "n_missing", + "partial", + ): + assert json.loads(json.dumps(block[key])) == archived[key], key + assert (block["n_turns_counted"], block["n_empty"]) == (12, 3) + assert block["partial"] is False + osl = block["output_sequence_lengths"] + assert (osl["total"], osl["min"], osl["max"], osl["avg"]) == (423, 22, 49, 35.25) + + +def test_steady_state_replay_rebuilds_archived_series(): + size = ssd.superpass_size_from_events(str(_EVENTS)) + assert size == 5 + series = ssd.build_super_pass_series(str(_EVENTS), size, _char_counts) + assert [ + (sp.n_issued, len(sp.ttft_ns), len(sp.tpot_ns), len(sp.latency_ns)) + for sp in series + ] == [(5, 4, 4, 5)] * 3 + assert [sp.out_tokens for sp in series] == [116, 116, 116] + snap = { + m["name"]: m for m in _json(_RUN_DIR / "metrics/final_snapshot.json")["metrics"] + } + assert sum(sum(sp.tpot_ns) for sp in series) == pytest.approx( + snap["tpot_ns"]["total"], rel=1e-12 + ) + assert sum(sum(sp.ttft_ns) for sp in series) == snap["ttft_ns"]["total"] + assert snap["ttft_ns"]["total"] == 10320914 + assert sum(sum(sp.latency_ns) for sp in series) == 13232749 + + +def test_steady_state_cli_replays_archived_verdict(tmp_path): + out = tmp_path / "steady_state.json" + argv = [str(_RUN_DIR), "--tokenizer", str(_CHAR_TOKENIZER_DIR), "--json", str(out)] + assert ssd.main(argv) == 0 + replayed = _json(out)["steady_state"] + assert replayed == _json(_SUMMARY)["steady_state"] + assert (replayed["found"], replayed["n_super_passes"]) == (False, 3) + + +def test_early_stopping_script_reproduces_archived_estimates(tmp_path): + es = _load_es_script() + series = es.compute_series(str(_EVENTS), count_tokens=_char_counts) + assert {k: len(v) for k, v in series.items()} == { + "ttft_ns": 12, + "tpot_ns": 12, + "sample_latency_ns": 15, + } + assert sum(series["ttft_ns"]) == 10320914 + assert sum(series["sample_latency_ns"]) == 13232749 + assert sum(series["tpot_ns"]) == pytest.approx(50977.06920610349, rel=1e-12) + + out = tmp_path / "augmented.json" + es.main( + [ + str(_EVENTS), + "--summary", + str(_SUMMARY), + "--tokenizer", + str(_CHAR_TOKENIZER_DIR), + "--json", + str(out), + ] + ) + augmented, summary = _json(out), _json(_SUMMARY) + for field in ("ttft", "tpot", "latency"): + esp = augmented[field]["early_stopping_percentiles"] + assert esp == summary[field]["early_stopping_percentiles"], field + assert esp["90.0"] is None + assert augmented["ttft"]["early_stopping_percentiles"]["50.0"] == 3258417 + assert augmented["latency"]["early_stopping_percentiles"]["50.0"] == 3045958 + assert augmented["tpot"]["early_stopping_percentiles"]["50.0"] == pytest.approx( + 8921.529411764706, rel=1e-12 + ) From 88901387cbcfdc129102d760a9edbc75697b3371 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:01 -0500 Subject: [PATCH 02/26] feat(core): token-count basis type, naming rule and server-usage helpers TokenCountSource (client | server) says where a token count comes from. basis_key(base, basis, *, primary) is the single naming rule: the primary basis keeps the base name and the other basis gets a _ suffix (osl_server beside a client primary, osl_client after a flip). Server counts ride QueryResult.metadata under PROMPT_TOKENS_KEY and COMPLETION_TOKENS_KEY. well_formed_count and server_usage_count validate a raw JSON value and warn on a malformed one; put_server_usage writes only well-formed counts, so a bad server value never fails a frame. EventRecord gains prompt_tokens and completion_tokens (UNSET when absent, as in archived logs) for the COMPLETE event. core/record.py adds is_complete_line, the byte-level filter that post-run readers use to find COMPLETE lines in events.jsonl, and EventLogFormatError for a log that is not in the compact form. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- src/inference_endpoint/core/record.py | 27 +++ src/inference_endpoint/core/types.py | 137 ++++++++++++++- .../services/event_logger/test_file_writer.py | 24 +++ tests/unit/core/test_record.py | 55 ++++++ tests/unit/core/test_token_count_source.py | 165 ++++++++++++++++++ 5 files changed, 401 insertions(+), 7 deletions(-) create mode 100644 tests/unit/core/test_token_count_source.py diff --git a/src/inference_endpoint/core/record.py b/src/inference_endpoint/core/record.py index 66fa430d4..70b3a0934 100644 --- a/src/inference_endpoint/core/record.py +++ b/src/inference_endpoint/core/record.py @@ -168,6 +168,33 @@ class EventRecord(msgspec.Struct, kw_only=True, frozen=True, gc=False): # type: data: OUTPUT_TYPE | PromptData | ErrorData | PhaseData | None = None finish_reason: str | msgspec.UnsetType = msgspec.UNSET worker_id: int | msgspec.UnsetType = msgspec.UNSET + prompt_tokens: int | msgspec.UnsetType = msgspec.UNSET + completion_tokens: int | msgspec.UnsetType = msgspec.UNSET + + +# ``event_type`` is EventRecord's first field, so in an ``events.jsonl`` line +# (msgspec JSON) the marker sits in the leading bytes. A JSON string value +# escapes its quotes, so the marker cannot occur inside one. +_COMPLETE_LINE_MARKER: Final = b'"event_type":"sample.complete"' +_MARKER_WINDOW: Final = 256 + + +def is_complete_line(line: bytes) -> bool: + """Whether one ``events.jsonl`` line holds a COMPLETE record. + + Searches only the line's leading bytes, so a reader skips every other + record without decoding it or scanning its prompt and response payload. + + Precondition: the line is compact msgspec JSON with ``event_type`` first, + as ``JSONLWriter`` writes it. Any other encoding (e.g. ``json.dumps`` + with its default ``": "`` separator) never matches. + """ + return line.find(_COMPLETE_LINE_MARKER, 0, _MARKER_WINDOW) != -1 + + +class EventLogFormatError(ValueError): + """An event log a reader cannot read without losing samples, such as one + not in the compact form ``is_complete_line`` requires.""" EVENT_LINE_LEAD: Final[bytes] = b'{"event_type":"' diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index 03cc29942..b71a3b360 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -19,13 +19,17 @@ This module defines the basic data structures used throughout the system. """ +import logging +import reprlib import time import uuid from enum import Enum -from typing import Any, cast +from typing import Any, Final, cast import msgspec +logger = logging.getLogger(__name__) + class APIType(str, Enum): """Supported inference API protocols. @@ -54,6 +58,43 @@ def default_route(self) -> str: raise ValueError(f"Invalid API type: {self}") +class TokenCountSource(str, Enum): + """The basis a token count is measured on. + + - SERVER: the endpoint's per-request ``usage`` + - CLIENT: the local tokenizer + + A run has one primary basis, whose series and report keys keep their base + name; ``other`` names the sibling basis, whose keys are suffixed (see + ``basis_key``). + """ + + SERVER = "server" + CLIENT = "client" + + @property + def other(self) -> "TokenCountSource": + """The opposite basis.""" + if self is TokenCountSource.SERVER: + return TokenCountSource.CLIENT + return TokenCountSource.SERVER + + +def basis_key(base: str, basis: TokenCountSource, *, primary: TokenCountSource) -> str: + """Name ``base`` on ``basis``: unchanged on the primary basis, else suffixed. + + Setup, registration and report time only; never call this per sample. + """ + return base if basis is primary else f"{base}_{basis.value}" + + +# Bounds of the HDR histogram behind every token-count series. A count above +# ``TOKEN_HDR_HIGH`` cannot be recorded, so it is also the largest well-formed +# count (see ``well_formed_count``). +TOKEN_HDR_LOW: Final[int] = 1 +TOKEN_HDR_HIGH: Final[int] = 10_000_000 # 10M tokens + + class QueryStatus(Enum): """Status of a query in its lifecycle. @@ -388,8 +429,8 @@ class Query( instead of {"id": ..., "data": ..., ...}). Provides ~6-50% size reduction and ~6-29% ser/des speedup for ZMQ transport depending on payload size. - omit_defaults=True: Fields with default values are omitted during encoding, - further reducing message size for queries with empty headers. + omit_defaults=True: Has no effect together with array_like=True, which always + encodes every field; kept for intent only. """ id: str = msgspec.field(default_factory=lambda: str(uuid.uuid4())) @@ -437,8 +478,9 @@ class QueryResult( gc=False: Safe because metadata contains only scalar key-value pairs. Do NOT store cyclic references in metadata or response_output fields. - omit_defaults=True: Fields with static defaults (ie. those NOT using default_factory) - are omitted if value equals default. + omit_defaults=True: Has no effect together with array_like=True, which always + encodes every field; kept for intent only. Do NOT add fields: carry extra + per-response data (e.g. server token counts) in metadata. array_like=True: Encodes as array instead of object (e.g. ["id", "chunk", false, {}] instead of {"id": ..., "response_chunk": ..., ...}). Reduces payload size. @@ -477,6 +519,86 @@ def get_response_output_string(self) -> str: return "" +# ``QueryResult.metadata`` keys carrying the server's per-request token usage. +PROMPT_TOKENS_KEY: Final = "prompt_tokens" +COMPLETION_TOKENS_KEY: Final = "completion_tokens" + + +def well_formed_count(value: object) -> int | None: + """Return ``value`` if it is a usable token count, else ``None``. + + Only an exact ``int`` in ``[0, TOKEN_HDR_HIGH]`` qualifies: ``bool`` (an + ``int`` subclass), floats, strings, negatives and out-of-range counts do not. + """ + # Exact type check: ``isinstance`` would accept ``bool``. + if type(value) is int and 0 <= value <= TOKEN_HDR_HIGH: # noqa: E721 + return value + return None + + +# A malformed ``usage`` comes from the endpoint, so it recurs on every response: +# warn on the first one only, once per process (each HTTP worker is its own +# process), never per frame. +_malformed_usage_warned = False + + +def _warn_malformed_usage(what: str, value: object) -> None: + global _malformed_usage_warned + if _malformed_usage_warned: + return + _malformed_usage_warned = True + logger.warning( + "Ignoring malformed server usage %s: %s (type '%s'); token counts must be " + "ints in [0, %d]. Further malformed usage values are not logged.", + what, + reprlib.repr(value), + type(value).__name__, + TOKEN_HDR_HIGH, + ) + + +def server_usage_count(value: object, key: str) -> int | None: + """``well_formed_count`` for a count the server reported under ``key``. + + A non-null value that is not well formed triggers the once-per-process + malformed-usage warning; null (absent) does not. + """ + count = well_formed_count(value) + if count is None and value is not None: + _warn_malformed_usage(f"'{key}'", value) + return count + + +def server_usage_counts(usage: object) -> tuple[int | None, int | None]: + """Read ``(prompt_tokens, completion_tokens)`` from a JSON ``usage`` object. + + Each count goes through ``server_usage_count``. A non-null ``usage`` that is + not a ``dict`` yields ``(None, None)`` and triggers the once-per-process + malformed-usage warning. + """ + if usage is None: + return None, None + if not isinstance(usage, dict): + _warn_malformed_usage("object", usage) + return None, None + return ( + server_usage_count(usage.get(PROMPT_TOKENS_KEY), PROMPT_TOKENS_KEY), + server_usage_count(usage.get(COMPLETION_TOKENS_KEY), COMPLETION_TOKENS_KEY), + ) + + +def put_server_usage( + metadata: dict[str, Any], + prompt_tokens: int | None, + completion_tokens: int | None, +) -> None: + """Write each non-``None`` count into ``metadata`` under its usage key.""" + if prompt_tokens is not None: + metadata[PROMPT_TOKENS_KEY] = prompt_tokens + if completion_tokens is not None: + metadata[COMPLETION_TOKENS_KEY] = completion_tokens + + class StreamChunk( msgspec.Struct, tag="stream_chunk", @@ -504,8 +626,9 @@ class StreamChunk( gc=False: Safe because metadata contains only scalar key-value pairs. Do NOT store cyclic references in metadata field. - omit_defaults=True: Fields with static defaults (ie. those NOT using default_factory) - are omitted if value equals default. + omit_defaults=True: Has no effect together with array_like=True, which always + encodes every field; kept for intent only. Do NOT add fields: carry extra + data in metadata. array_like=True: Encodes as array instead of object (e.g. ["id", "chunk", {}] instead of {"id": ..., "response_chunk": ..., ...}). Reduces payload size. diff --git a/tests/unit/async_utils/services/event_logger/test_file_writer.py b/tests/unit/async_utils/services/event_logger/test_file_writer.py index 9e9ae6f96..a413d8c50 100644 --- a/tests/unit/async_utils/services/event_logger/test_file_writer.py +++ b/tests/unit/async_utils/services/event_logger/test_file_writer.py @@ -73,6 +73,30 @@ def test_writes_valid_jsonl(self, tmp_path): assert "finish_reason" not in records[0] assert records[1]["finish_reason"] == "tool_calls" + def test_writes_server_usage_only_when_set(self, tmp_path): + writer = JSONLWriter(tmp_path / "events", flush_interval=1) + try: + writer.write(_record(SampleEventType.COMPLETE, uuid="s0", ts=1000)) + writer.write( + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=2000, + sample_uuid="s1", + prompt_tokens=11, + completion_tokens=3, + ) + ) + finally: + writer.close() + + lines = (tmp_path / "events.jsonl").read_text().strip().split("\n") + records = [json.loads(line) for line in lines] + assert "prompt_tokens" not in records[0] + assert "completion_tokens" not in records[0] + assert (records[1]["prompt_tokens"], records[1]["completion_tokens"]) == (11, 3) + # Consumers scan the line prefix for the event type, so it stays first. + assert next(iter(records[1])) == "event_type" + def test_record_roundtrip_fields(self, tmp_path): writer = JSONLWriter(tmp_path / "events", flush_interval=1) try: diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index ecade65a6..895c7e3d1 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -347,3 +347,58 @@ def test_every_event_type_a_writer_emits_is_found_at_the_lead(self, event_type): ) def test_any_other_layout_has_no_leading_topic(self, line): assert leading_event_topic(line) is None + + +@pytest.mark.unit +class TestEventRecordServerTokenCounts: + _JSON_ENC = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + _JSON_DEC = msgspec.json.Decoder(EventRecord, dec_hook=EventType.decode_hook) + # A COMPLETE line in the archived shape, without the count fields. + _ARCHIVED_LINE = ( + b'{"event_type":"sample.complete","timestamp_ns":123,"sample_uuid":"u",' + b'"conversation_id":"","turn":null,"data":["TextModelOutput","ab",null,null],' + b'"finish_reason":"stop","worker_id":2}' + ) + _KW = { + "event_type": SampleEventType.COMPLETE, + "timestamp_ns": 123, + "sample_uuid": "u", + "data": TextModelOutput(output="ab"), + "finish_reason": "stop", + "worker_id": 2, + } + + def test_counts_default_to_unset(self): + r = EventRecord(event_type=SampleEventType.COMPLETE) + assert r.prompt_tokens is msgspec.UNSET + assert r.completion_tokens is msgspec.UNSET + + def test_record_without_counts_encodes_byte_identical_to_archived_line(self): + assert self._JSON_ENC.encode(EventRecord(**self._KW)) == self._ARCHIVED_LINE + _, payload = _codec.encode(EventRecord(**self._KW)) + assert list(msgspec.msgpack.decode(payload)) == [ + "event_type", + "timestamp_ns", + "sample_uuid", + "conversation_id", + "turn", + "data", + "finish_reason", + "worker_id", + ] + + def test_counts_round_trip_and_append_after_worker_id(self): + record = EventRecord(**self._KW, prompt_tokens=5, completion_tokens=7) + line = self._JSON_ENC.encode(record) + assert ( + line + == self._ARCHIVED_LINE[:-1] + b',"prompt_tokens":5,"completion_tokens":7}' + ) + for d in (self._JSON_DEC.decode(line), _codec.decode(_codec.encode(record)[1])): + assert (d.prompt_tokens, d.completion_tokens) == (5, 7) + + def test_archived_line_decodes_with_unset_counts(self): + d = self._JSON_DEC.decode(self._ARCHIVED_LINE) + assert (d.finish_reason, d.worker_id) == ("stop", 2) + assert d.prompt_tokens is msgspec.UNSET + assert d.completion_tokens is msgspec.UNSET diff --git a/tests/unit/core/test_token_count_source.py b/tests/unit/core/test_token_count_source.py new file mode 100644 index 000000000..ace3ed113 --- /dev/null +++ b/tests/unit/core/test_token_count_source.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the token-count basis vocabulary in ``core.types``.""" + +import logging +from typing import Any + +import pytest +from inference_endpoint.core.types import ( + COMPLETION_TOKENS_KEY, + PROMPT_TOKENS_KEY, + TOKEN_HDR_HIGH, + TokenCountSource, + basis_key, + put_server_usage, + server_usage_count, + server_usage_counts, + well_formed_count, +) + + +@pytest.mark.unit +def test_other_basis(): + assert TokenCountSource.CLIENT.other is TokenCountSource.SERVER + assert TokenCountSource.SERVER.other is TokenCountSource.CLIENT + + +@pytest.mark.unit +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_basis_key_suffixes_only_the_other_basis(primary): + other = primary.other + assert basis_key("osl", primary, primary=primary) == "osl" + assert basis_key("osl", other, primary=primary) == f"osl_{other.value}" + assert basis_key("tpot_ns", other, primary=primary) == f"tpot_ns_{other.value}" + + +@pytest.mark.unit +@pytest.mark.parametrize( + "value, expected", + [ + (0, 0), + (7, 7), + (TOKEN_HDR_HIGH, TOKEN_HDR_HIGH), + (TOKEN_HDR_HIGH + 1, None), + (-1, None), + (True, None), + (False, None), + (3.0, None), + ("3", None), + (None, None), + ([], None), + ], +) +def test_well_formed_count(value, expected): + assert well_formed_count(value) == expected + + +@pytest.mark.unit +@pytest.mark.parametrize( + "usage, expected", + [ + ({"prompt_tokens": 3, "completion_tokens": 5}, (3, 5)), + ({"prompt_tokens": 3, "completion_tokens": "x"}, (3, None)), + ({"prompt_tokens": 12.0}, (None, None)), + ({}, (None, None)), + (None, (None, None)), + ([1, 2], (None, None)), + ("usage", (None, None)), + ], +) +def test_server_usage_counts(usage, expected): + assert server_usage_counts(usage) == expected + + +@pytest.mark.unit +@pytest.mark.parametrize( + "usage, shown", + [ + ("x", "'x' (type 'str')"), + ([], "[] (type 'list')"), + ({"prompt_tokens": 12.0}, "'prompt_tokens': 12.0 (type 'float')"), + ({"completion_tokens": True}, "'completion_tokens': True (type 'bool')"), + ( + {"prompt_tokens": 3, "completion_tokens": TOKEN_HDR_HIGH + 1}, + f"'completion_tokens': {TOKEN_HDR_HIGH + 1} (type 'int')", + ), + ], +) +def test_first_malformed_usage_warns_once_showing_value_and_type( + usage_warning_unlatched, usage, shown +): + server_usage_counts(usage) + server_usage_counts(usage) + server_usage_counts("another") + (record,) = usage_warning_unlatched.records + assert record.levelno == logging.WARNING + assert shown in record.getMessage() + + +@pytest.mark.unit +def test_warning_shows_a_short_repr(usage_warning_unlatched): + server_usage_counts({"prompt_tokens": "9" * 1000}) + (record,) = usage_warning_unlatched.records + assert len(record.getMessage()) < 200 + + +@pytest.mark.unit +@pytest.mark.parametrize( + "value, expected, warns", + [(5, 5, False), (None, None, False), (-1, None, True), (5.0, None, True)], +) +def test_server_usage_count_warns_only_on_non_null_rejects( + usage_warning_unlatched, value, expected, warns +): + assert server_usage_count(value, PROMPT_TOKENS_KEY) == expected + assert bool(usage_warning_unlatched.records) is warns + + +@pytest.mark.unit +@pytest.mark.parametrize( + "usage", + [ + None, + {}, + {"prompt_tokens": None, "completion_tokens": None}, + {"prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8}, + ], + ids=["absent", "empty", "null-counts", "well-formed"], +) +def test_absent_or_well_formed_usage_never_warns(usage_warning_unlatched, usage): + server_usage_counts(usage) + assert usage_warning_unlatched.records == [] + + +@pytest.mark.unit +def test_put_server_usage_writes_only_present_counts(): + md: dict[str, Any] = {} + put_server_usage(md, 4, None) + assert md == {PROMPT_TOKENS_KEY: 4} + + +@pytest.mark.unit +def test_put_server_usage_writes_both_counts(): + md: dict[str, Any] = {"other": 1} + put_server_usage(md, 4, 9) + assert md == {"other": 1, PROMPT_TOKENS_KEY: 4, COMPLETION_TOKENS_KEY: 9} + + +@pytest.mark.unit +def test_metadata_keys_are_the_openai_usage_field_names(): + assert PROMPT_TOKENS_KEY == "prompt_tokens" + assert COMPLETION_TOKENS_KEY == "completion_tokens" From db3bfd714250c2e13abb4cd68e3299ccb12ca935 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:01 -0500 Subject: [PATCH 03/26] feat(config): token_count_source and stream_include_usage settings settings.token_count_source (--token-count-source) picks the primary basis for ISL/OSL/TPS/TPOT. It defaults to client, so every unsuffixed number matches main; flipping the default to server is a one-line change here plus the regenerated templates. settings.client.stream_include_usage (--no-stream-include-usage to opt out) asks streaming OpenAI-compatible endpoints for a usage frame. It is on by default so the server basis is always reported beside the client one; turning it off restores main's request bytes. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- src/inference_endpoint/config/schema.py | 22 +++++++++++++-- .../templates/concurrency_template_full.yaml | 4 ++- .../templates/offline_template_full.yaml | 4 ++- .../templates/online_template_full.yaml | 4 ++- .../endpoint_client/config.py | 17 +++++++++++ tests/unit/config/test_schema.py | 24 ++++++++++++++++ tests/unit/config/test_yaml_loader.py | 28 ++++++++++++++++++- 7 files changed, 97 insertions(+), 6 deletions(-) diff --git a/src/inference_endpoint/config/schema.py b/src/inference_endpoint/config/schema.py index 222d31057..4231ab617 100644 --- a/src/inference_endpoint/config/schema.py +++ b/src/inference_endpoint/config/schema.py @@ -44,7 +44,7 @@ model_validator, ) -from ..core.types import APIType +from ..core.types import APIType, TokenCountSource from ..endpoint_client.config import HTTPClientConfig from ..exceptions import CLIError from ..utils import WithUpdatesMixin @@ -318,7 +318,7 @@ class ModelParams(BaseModel): str | None, cyclopts.Parameter( alias="--tokenizer", - help="HF repo ID or local path for the tokenizer. Overrides model name for client-side token metrics (ISL/OSL/TPOT).", + help="HF repo ID or local path for the tokenizer. Overrides model name for client-side token counting: every token metric when the primary basis is client, the client-basis ISL/OSL/TPOT blocks beside them otherwise.", ), ] = None @@ -1016,6 +1016,24 @@ class Settings(WithUpdatesMixin, BaseModel): True, description="Collect input sequence length (ISL) metrics (default: true).", ) + token_count_source: Annotated[ + TokenCountSource, + cyclopts.Parameter( + alias="--token-count-source", + help=( + "Primary basis for ISL/OSL/TPS/TPOT: " + "server (endpoint usage) or client (tokenizer). " + "The other basis is reported beside it" + ), + ), + ] = Field( + TokenCountSource.CLIENT, + description=( + "Primary token-count basis for ISL/OSL/TPS/TPOT: server = the " + "endpoint's per-request usage, client = the local tokenizer. The " + "other basis is reported beside it under a `_` suffix." + ), + ) @model_validator(mode="after") def _min_issue_duration_requires_qps(self) -> Self: diff --git a/src/inference_endpoint/config/templates/concurrency_template_full.yaml b/src/inference_endpoint/config/templates/concurrency_template_full.yaml index 7ed842b91..869467c49 100644 --- a/src/inference_endpoint/config/templates/concurrency_template_full.yaml +++ b/src/inference_endpoint/config/templates/concurrency_template_full.yaml @@ -18,7 +18,7 @@ model_params: skip_special_tokens: true # Whether OpenAI text-completions servers omit special tokens from decoded output osl_distribution: null # Output sequence length distribution streaming: 'on' # Streaming mode: auto/on/off | options: auto, on, off - tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token metrics (ISL/OSL/TPOT). + tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token counting: every token metric when the primary basis is client, the client-basis ISL/OSL/TPOT blocks beside them otherwise. datasets: # Dataset configs - name: perf type: performance # Dataset purpose: performance or accuracy | options: performance, accuracy @@ -77,6 +77,7 @@ settings: linger: -1 # ZMQ linger on close (-1=block until sent) immediate: 1 # ZMQ IMMEDIATE (1=only enqueue on ready) stream_all_chunks: false # Stream all chunks to main thread (caution: perf overhead) + stream_include_usage: true # Request stream_options.include_usage on streaming OpenAI requests (disable if the endpoint rejects it) worker_initialization_timeout: 60.0 # Worker init timeout (seconds) worker_graceful_shutdown_wait: 0.5 # Post-run graceful shutdown wait (seconds) worker_force_kill_timeout: 0.5 # Force kill timeout after graceful wait (seconds) @@ -108,6 +109,7 @@ settings: enabled: false # Steady-state window detection (default off) metrics_tokenizer_workers: 4 # In-process tokenizer threads for live (mid-run) ISL/OSL/TPOT (default: 4; 0 = defer everything to the end-of-run drain). metrics_isl: true # Collect input sequence length (ISL) metrics (default: true). + token_count_source: client # Primary token-count basis for ISL/OSL/TPS/TPOT: server = the endpoint's per-request usage, client = the local tokenizer. The other basis is reported beside it under a `_` suffix. | options: server, client endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/config/templates/offline_template_full.yaml b/src/inference_endpoint/config/templates/offline_template_full.yaml index 6ed1aebb8..4bc457573 100644 --- a/src/inference_endpoint/config/templates/offline_template_full.yaml +++ b/src/inference_endpoint/config/templates/offline_template_full.yaml @@ -18,7 +18,7 @@ model_params: skip_special_tokens: true # Whether OpenAI text-completions servers omit special tokens from decoded output osl_distribution: null # Output sequence length distribution streaming: 'off' # Streaming mode: auto/on/off | options: auto, on, off - tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token metrics (ISL/OSL/TPOT). + tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token counting: every token metric when the primary basis is client, the client-basis ISL/OSL/TPOT blocks beside them otherwise. datasets: # Dataset configs - name: perf type: performance # Dataset purpose: performance or accuracy | options: performance, accuracy @@ -77,6 +77,7 @@ settings: linger: -1 # ZMQ linger on close (-1=block until sent) immediate: 1 # ZMQ IMMEDIATE (1=only enqueue on ready) stream_all_chunks: false # Stream all chunks to main thread (caution: perf overhead) + stream_include_usage: true # Request stream_options.include_usage on streaming OpenAI requests (disable if the endpoint rejects it) worker_initialization_timeout: 60.0 # Worker init timeout (seconds) worker_graceful_shutdown_wait: 0.5 # Post-run graceful shutdown wait (seconds) worker_force_kill_timeout: 0.5 # Force kill timeout after graceful wait (seconds) @@ -108,6 +109,7 @@ settings: enabled: false # Steady-state window detection (default off) metrics_tokenizer_workers: 4 # In-process tokenizer threads for live (mid-run) ISL/OSL/TPOT (default: 4; 0 = defer everything to the end-of-run drain). metrics_isl: true # Collect input sequence length (ISL) metrics (default: true). + token_count_source: client # Primary token-count basis for ISL/OSL/TPS/TPOT: server = the endpoint's per-request usage, client = the local tokenizer. The other basis is reported beside it under a `_` suffix. | options: server, client endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/config/templates/online_template_full.yaml b/src/inference_endpoint/config/templates/online_template_full.yaml index 36a844354..27e6b9a19 100644 --- a/src/inference_endpoint/config/templates/online_template_full.yaml +++ b/src/inference_endpoint/config/templates/online_template_full.yaml @@ -18,7 +18,7 @@ model_params: skip_special_tokens: true # Whether OpenAI text-completions servers omit special tokens from decoded output osl_distribution: null # Output sequence length distribution streaming: 'on' # Streaming mode: auto/on/off | options: auto, on, off - tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token metrics (ISL/OSL/TPOT). + tokenizer_name: null # HF repo ID or local path for the tokenizer. Overrides model name for client-side token counting: every token metric when the primary basis is client, the client-basis ISL/OSL/TPOT blocks beside them otherwise. datasets: # Dataset configs - name: perf type: performance # Dataset purpose: performance or accuracy | options: performance, accuracy @@ -78,6 +78,7 @@ settings: linger: -1 # ZMQ linger on close (-1=block until sent) immediate: 1 # ZMQ IMMEDIATE (1=only enqueue on ready) stream_all_chunks: false # Stream all chunks to main thread (caution: perf overhead) + stream_include_usage: true # Request stream_options.include_usage on streaming OpenAI requests (disable if the endpoint rejects it) worker_initialization_timeout: 60.0 # Worker init timeout (seconds) worker_graceful_shutdown_wait: 0.5 # Post-run graceful shutdown wait (seconds) worker_force_kill_timeout: 0.5 # Force kill timeout after graceful wait (seconds) @@ -109,6 +110,7 @@ settings: enabled: false # Steady-state window detection (default off) metrics_tokenizer_workers: 4 # In-process tokenizer threads for live (mid-run) ISL/OSL/TPOT (default: 4; 0 = defer everything to the end-of-run drain). metrics_isl: true # Collect input sequence length (ISL) metrics (default: true). + token_count_source: client # Primary token-count basis for ISL/OSL/TPS/TPOT: server = the endpoint's per-request usage, client = the local tokenizer. The other basis is reported beside it under a `_` suffix. | options: server, client endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/endpoint_client/config.py b/src/inference_endpoint/endpoint_client/config.py index b2839996d..89c5dae3c 100644 --- a/src/inference_endpoint/endpoint_client/config.py +++ b/src/inference_endpoint/endpoint_client/config.py @@ -187,6 +187,23 @@ class HTTPClientConfig(WithUpdatesMixin, BaseModel): False, description="Stream all chunks to main thread (caution: perf overhead)" ) + # Streaming OpenAI chat/completions requests ask for the final usage frame + # (stream_options.include_usage) so server-reported token counts can be + # recorded. Turn off for endpoints that reject stream_options. + stream_include_usage: Annotated[ + bool, + cyclopts.Parameter( + alias="--stream-include-usage", + help="Ask streaming OpenAI requests for server token usage", + ), + ] = Field( + True, + description=( + "Request stream_options.include_usage on streaming OpenAI requests " + "(disable if the endpoint rejects it)" + ), + ) + # Worker lifecycle timeouts worker_initialization_timeout: float = Field( 60.0, description="Worker init timeout (seconds)" diff --git a/tests/unit/config/test_schema.py b/tests/unit/config/test_schema.py index 9a881c184..d1ecd181c 100644 --- a/tests/unit/config/test_schema.py +++ b/tests/unit/config/test_schema.py @@ -33,19 +33,43 @@ LoadPattern, LoadPatternType, ModelParams, + OfflineSettings, + OnlineSettings, OSLDistribution, OSLDistributionType, ProfilerEngine, ProfilingConfig, RuntimeConfig, + Settings, StreamingMode, SubmissionReference, TestType, ) +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.exceptions import CLIError from pydantic import ValidationError +class TestTokenCountSource: + @pytest.mark.unit + def test_values(self): + assert [m.value for m in TokenCountSource] == ["server", "client"] + + @pytest.mark.unit + def test_default_token_count_source_is_client(self): + for cls in (Settings, OfflineSettings, OnlineSettings): + assert cls().token_count_source is TokenCountSource.CLIENT + + @pytest.mark.unit + def test_rejects_unknown_value(self): + with pytest.raises(ValidationError, match="token_count_source"): + Settings(token_count_source="tokenizer") + + @pytest.mark.unit + def test_is_run_level_not_per_dataset(self): + assert "token_count_source" not in ModelParams.model_fields + + class TestOSLDistribution: @pytest.mark.unit def test_fixed_distribution(self): diff --git a/tests/unit/config/test_yaml_loader.py b/tests/unit/config/test_yaml_loader.py index ffde68610..4c0bd160c 100644 --- a/tests/unit/config/test_yaml_loader.py +++ b/tests/unit/config/test_yaml_loader.py @@ -18,8 +18,12 @@ from pathlib import Path import pytest -from inference_endpoint.config.schema import BenchmarkConfig, LoadPatternType +from inference_endpoint.config.schema import ( + BenchmarkConfig, + LoadPatternType, +) from inference_endpoint.config.schema import TestType as BenchmarkTestType +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.exceptions import CLIError from pydantic import ValidationError @@ -123,6 +127,28 @@ def test_load_metrics_isl_disabled(self, tmp_path): assert config.settings.metrics_isl is False + @pytest.mark.unit + def test_load_token_count_source_server(self, tmp_path): + config_file = tmp_path / "token_count_source_server.yaml" + config_file.write_text( + """ +type: offline +model_params: + name: test-model +datasets: + - path: test.jsonl +settings: + token_count_source: server +endpoint_config: + endpoints: + - http://localhost:8000 +""" + ) + + config = BenchmarkConfig.from_yaml_file(config_file) + + assert config.settings.token_count_source is TokenCountSource.SERVER + def test_load_nonexistent_file(self): """Test error when file doesn't exist.""" with pytest.raises(FileNotFoundError, match="not found"): From 58485960925996c17f1e7c080dd27660e469dbbb Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:01 -0500 Subject: [PATCH 04/26] feat(endpoint_client): request and capture server token usage Request side: every adapter takes encode_query(query, *, include_usage). The worker passes the run's stream_include_usage. The OpenAI chat and /v1/completions adapters send stream_options: {"include_usage": true} only on streaming requests with at most one choice; SGLang (meta_info) and videogen ignore the flag. Query.data keeps main's shape. Response side: usage is read from any SSE frame (the last well-formed value per field wins) and from non-streaming bodies. Count fields decode broadly and are validated at capture, so a malformed count warns and never drops the frame's content or finish_reason. Each SSE frame decodes in its own try, so one bad frame no longer drops the frames after it in the same read. SGLang accepts the current meta_info shape (null finish_reason, num_retractions, optional e2e_latency/output_ids) and validates prompt_tokens once per changed value rather than once per frame. session.py lifts the counts from QueryResult.metadata onto the COMPLETE EventRecord. It trusts put_server_usage, the single writer, and adds no validation to the per-request path. Hot path: the OpenAI SSE fast path keeps main's cost plus one usage test per frame (decode micro-benchmark within +0.7..+2.4% of main). Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- .../endpoint_client/adapter_protocol.py | 46 ++- .../endpoint_client/worker.py | 36 ++- .../load_generator/session.py | 10 + src/inference_endpoint/openai/accumulator.py | 54 +++- .../openai/completions_adapter.py | 65 +++- .../openai/openai_adapter.py | 4 +- .../openai/openai_msgspec_adapter.py | 49 ++- src/inference_endpoint/openai/types.py | 50 ++- src/inference_endpoint/sglang/accumulator.py | 51 +++- src/inference_endpoint/sglang/adapter.py | 38 ++- src/inference_endpoint/sglang/types.py | 28 +- .../utils/benchmark_httpclient.py | 29 +- src/inference_endpoint/videogen/adapter.py | 5 +- .../endpoint_client/test_sglang_adapter.py | 30 +- tests/integration/videogen/test_adapter.py | 12 +- tests/performance/openai/test_adapter.py | 2 +- .../openai/test_msgspec_adapter.py | 2 +- tests/performance/openai/test_types.py | 7 +- .../performance/sglang/test_sglang_adapter.py | 2 +- .../unit/dataset_manager/test_data_loader.py | 54 ++++ .../endpoint_client/test_adapter_protocol.py | 100 +++++- tests/unit/endpoint_client/test_worker.py | 145 +++++++++ .../unit/load_generator/test_async_session.py | 30 ++ tests/unit/openai/test_accumulator.py | 288 +++++++++++++++++- .../unit/openai/test_accumulator_contract.py | 12 +- tests/unit/openai/test_completions_adapter.py | 229 +++++++++++++- tests/unit/openai/test_msgspec_adapter.py | 225 +++++++++++++- tests/unit/openai/test_openai_adapter.py | 6 +- tests/unit/sglang/test_sglang_usage.py | 246 +++++++++++++++ tests/unit/videogen/test_adapter.py | 22 +- 30 files changed, 1718 insertions(+), 159 deletions(-) create mode 100644 tests/unit/endpoint_client/test_worker.py create mode 100644 tests/unit/sglang/test_sglang_usage.py diff --git a/src/inference_endpoint/endpoint_client/adapter_protocol.py b/src/inference_endpoint/endpoint_client/adapter_protocol.py index a96faaa0a..a7efd0d7d 100644 --- a/src/inference_endpoint/endpoint_client/adapter_protocol.py +++ b/src/inference_endpoint/endpoint_client/adapter_protocol.py @@ -69,12 +69,16 @@ def dataset_transforms( @classmethod @abstractmethod - def encode_query(cls, query: Query) -> bytes: + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: """ Encode a Query to bytes for HTTP transmission. Args: query: Input query with prompt and parameters + include_usage: The run's ``stream_include_usage`` setting. Adapters + whose API has ``stream_options`` ask for the final usage frame + on streaming requests with at most one choice; the others + accept and ignore it. Returns: Encoded request bytes ready for HTTP POST @@ -102,11 +106,19 @@ def decode_sse_message(cls, json_bytes: bytes) -> Any: """ Decode SSE message and return adapter-specific chunk object. + The OpenAI adapters (chat and completions) return + ``SSEChoice | SSEMessage | None``: the first ``SSEChoice`` for a frame + without ``usage`` (the common case), ``None`` for a frame with neither + usage nor choices, and the whole ``SSEMessage`` for a frame that + carries usage. A malformed ``usage`` value never fails the decode; + counts are validated when the accumulator captures them. + Args: json_bytes: Raw JSON bytes from SSE stream Returns: - Adapter-specific chunk object passed to accumulator.add_chunk() + Adapter-specific chunk object passed to accumulator.add_chunk(), + or None for a frame with nothing to accumulate """ raise NotImplementedError("decode_sse_message not implemented") @@ -114,9 +126,11 @@ def decode_sse_message(cls, json_bytes: bytes) -> Any: def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]: """Parse SSE chunk and extract all chunk objects. - Extracts JSON documents from SSE stream and decodes them to chunk objects. - Filters None returns from decode_sse_message and skips frames that fail - to decode (e.g. role-only or finish_reason-only frames). + Extracts JSON documents from SSE stream and decodes each to a chunk + object. Filters None returns from decode_sse_message. A frame that fails + to decode is skipped on its own, so later frames in the same read (e.g. + a trailing usage frame) are still returned. At most one warning is + logged per call, after all frames are processed. Args: buffer: Byte buffer containing SSE data @@ -127,12 +141,20 @@ def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]: """ json_docs = cls.SSE_DATA_PATTERN.findall(buffer[:end_pos]) parsed: list[Any] = [] - # Note: if one frame is malformed, remaining frames are skipped - try: - for json_doc in json_docs: + n_bad = 0 + first_exc_name = "" + for json_doc in json_docs: + try: content = cls.decode_sse_message(json_doc) - if content is not None: - parsed.append(content) - except (msgspec.DecodeError, msgspec.ValidationError) as exc: - logger.warning("skipping malformed SSE batch (%s)", type(exc).__name__) + except (msgspec.DecodeError, msgspec.ValidationError) as exc: + if not n_bad: + first_exc_name = type(exc).__name__ + n_bad += 1 + continue + if content is not None: + parsed.append(content) + if n_bad: + logger.warning( + "skipping %d malformed SSE frame(s) (%s)", n_bad, first_exc_name + ) return parsed diff --git a/src/inference_endpoint/endpoint_client/worker.py b/src/inference_endpoint/endpoint_client/worker.py index 321528baa..5278c22e9 100644 --- a/src/inference_endpoint/endpoint_client/worker.py +++ b/src/inference_endpoint/endpoint_client/worker.py @@ -63,6 +63,30 @@ # Already set, which is fine (likely in tests or when importing multiple times) pass +# Streaming single-choice requests in a run all carry the same stream_options, +# so an endpoint that rejects it rejects them all; one hint per worker process +# is enough. A module flag because each worker process holds exactly one Worker. +_stream_options_hint_logged = False + + +def _is_stream_options_rejection( + status_code: int, body: bytes, streaming: bool +) -> bool: + """Whether an error response reads as the endpoint refusing ``stream_options``.""" + return streaming and 400 <= status_code < 500 and b"stream_options" in body + + +def _log_stream_options_hint(status_code: int) -> None: + global _stream_options_hint_logged + _stream_options_hint_logged = True + logger.warning( + "Endpoint rejected a streaming request that names stream_options " + "(HTTP %d). If it cannot report usage on streams, rerun with " + "--no-stream-include-usage (YAML: settings.client.stream_include_usage: " + "false).", + status_code, + ) + def worker_main( worker_id: int, @@ -335,7 +359,9 @@ async def _run_main_loop(self) -> None: def _prepare_request(self, query: Query) -> InFlightRequest: """Build InFlightRequest with serialized HTTP bytes.""" # Encode Query into HTTP payload bytes using adapter - body_bytes = self._adapter.encode_query(query) + body_bytes = self._adapter.encode_query( + query, include_usage=self.http_config.stream_include_usage + ) is_streaming = query.data.get("stream", False) # Build complete HTTP request bytes @@ -395,6 +421,14 @@ async def _process_response(self, req: InFlightRequest) -> None: if status_code != 200: error_body = await conn.protocol.read_body() self._pool.release(conn) + if ( + not _stream_options_hint_logged + and self.http_config.stream_include_usage + and _is_stream_options_rejection( + status_code, error_body, req.is_streaming + ) + ): + _log_stream_options_hint(status_code) await self._handle_error( req.query_id, f"HTTP {status_code}: {error_body.decode('utf-8', errors='replace')}", diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index c4cd615ba..9be343bca 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -39,6 +39,8 @@ SessionEventType, ) from ..core.types import ( + COMPLETION_TOKENS_KEY, + PROMPT_TOKENS_KEY, PhaseData, PhaseType, PromptData, @@ -824,6 +826,14 @@ def _handle_response(self, resp: QueryResult | StreamChunk, now_ns: int) -> None worker_id=( worker_id if isinstance(worker_id, int) else msgspec.UNSET ), + # put_server_usage, the only writer of these keys, + # stores validated counts. + prompt_tokens=resp.metadata.get( + PROMPT_TOKENS_KEY, msgspec.UNSET + ), + completion_tokens=resp.metadata.get( + COMPLETION_TOKENS_KEY, msgspec.UNSET + ), ) ) diff --git a/src/inference_endpoint/openai/accumulator.py b/src/inference_endpoint/openai/accumulator.py index 5d1d6b311..5bbf5fdd1 100644 --- a/src/inference_endpoint/openai/accumulator.py +++ b/src/inference_endpoint/openai/accumulator.py @@ -17,11 +17,17 @@ from typing import Any -from inference_endpoint.core.types import QueryResult, StreamChunk, TextModelOutput +from inference_endpoint.core.types import ( + QueryResult, + StreamChunk, + TextModelOutput, + put_server_usage, + server_usage_counts, +) from inference_endpoint.endpoint_client.accumulator_protocol import ( SSEAccumulatorProtocol, ) -from inference_endpoint.openai.types import SSEChoice, SSEDelta +from inference_endpoint.openai.types import SSEChoice, SSEDelta, SSEMessage class OpenAISSEAccumulator(SSEAccumulatorProtocol): @@ -32,14 +38,48 @@ def __init__(self, query_id: str, stream_all_chunks: bool): self.reasoning_chunks: list[str] = [] self.tool_call_chunks: list[tuple[dict[str, Any], ...]] = [] self._finish_reason: str | None = None + self._prompt_tokens: int | None = None + self._completion_tokens: int | None = None + # Usage sums the tokens of every choice while only the first choice of + # each frame is accumulated, so a stream whose usage-bearing frame holds + # several choices reports no counts. + self._multi_choice = False self._first_delta: SSEDelta | None = None self.first_chunk_sent = False self.query_id = query_id self.stream_all_chunks = stream_all_chunks - def add_chunk(self, choice: SSEChoice | None) -> StreamChunk | None: - if not isinstance(choice, SSEChoice): + def _take_message(self, msg: SSEMessage) -> SSEChoice | None: + """Record a usage-bearing frame's usage and choice count. + + Returns the frame's first choice to accumulate, or ``None``. + """ + if msg.usage is not None: + prompt_tokens, completion_tokens = server_usage_counts(msg.usage) + # Per-field merge: a later frame with a null, missing or malformed + # count (e.g. a proxy ending with "usage": {}) must not erase a + # count seen earlier. + if prompt_tokens is not None: + self._prompt_tokens = prompt_tokens + if completion_tokens is not None: + self._completion_tokens = completion_tokens + choices = msg.choices + if not choices: + return None + if len(choices) > 1: + self._multi_choice = True + return choices[0] + + def add_chunk(self, chunk: SSEChoice | SSEMessage | None) -> StreamChunk | None: + if isinstance(chunk, SSEChoice): + choice = chunk + elif isinstance(chunk, SSEMessage): + message_choice = self._take_message(chunk) + if message_choice is None: + return None + choice = message_choice + else: return None if choice.finish_reason: @@ -76,7 +116,7 @@ def add_chunk(self, choice: SSEChoice | None) -> StreamChunk | None: if content is not None and ( self.stream_all_chunks or not self.first_chunk_sent ): - chunk = StreamChunk( + stream_chunk = StreamChunk( id=self.query_id, response_chunk=content, metadata={ @@ -84,7 +124,7 @@ def add_chunk(self, choice: SSEChoice | None) -> StreamChunk | None: }, ) self.first_chunk_sent = True - return chunk + return stream_chunk else: return None @@ -112,6 +152,8 @@ def get_final_output(self) -> QueryResult: } if self._finish_reason: metadata["finish_reason"] = self._finish_reason + if not self._multi_choice: + put_server_usage(metadata, self._prompt_tokens, self._completion_tokens) _content, _reasoning, merged_tool_calls = text_output.as_message_parts() if merged_tool_calls: metadata["tool_calls"] = list(merged_tool_calls) diff --git a/src/inference_endpoint/openai/completions_adapter.py b/src/inference_endpoint/openai/completions_adapter.py index 6c0124ee3..09d4f1fc0 100644 --- a/src/inference_endpoint/openai/completions_adapter.py +++ b/src/inference_endpoint/openai/completions_adapter.py @@ -15,9 +15,17 @@ """OpenAI /v1/completions adapter for pre-tokenized prompts (bypasses server chat template).""" +from typing import Any + import msgspec from inference_endpoint.config.schema import ModelParams, StreamingMode -from inference_endpoint.core.types import Query, QueryResult, TextModelOutput +from inference_endpoint.core.types import ( + Query, + QueryResult, + TextModelOutput, + put_server_usage, + server_usage_counts, +) from inference_endpoint.dataset_manager.transforms import ( AddStaticColumns, ColumnFilter, @@ -27,11 +35,14 @@ from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter from .types import ( + STREAM_OPTIONS_INCLUDE_USAGE, SSEChoice, SSEDelta, + SSEMessage, TextCompletionRequest, TextCompletionResponse, TextCompletionSSEMessage, + requests_multiple_choices, ) @@ -75,16 +86,23 @@ def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: ] @classmethod - def encode_query(cls, query: Query) -> bytes: + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: if "input_tokens" not in query.data: raise KeyError( f"input_tokens not found in query.data: {list(query.data.keys())}" ) + stream = query.data.get("stream") + n = query.data.get("n") return cls._request_encoder.encode( TextCompletionRequest( model=query.data.get("model", "no-model-name"), prompt=query.data["input_tokens"], - stream=query.data.get("stream"), + stream=stream, + stream_options=( + STREAM_OPTIONS_INCLUDE_USAGE + if include_usage and stream and not requests_multiple_choices(n) + else None + ), max_tokens=query.data.get("max_tokens"), min_tokens=query.data.get("min_tokens"), skip_special_tokens=query.data.get("skip_special_tokens"), @@ -92,7 +110,7 @@ def encode_query(cls, query: Query) -> bytes: top_p=query.data.get("top_p"), top_k=query.data.get("top_k"), repetition_penalty=query.data.get("repetition_penalty"), - n=query.data.get("n"), + n=n, presence_penalty=query.data.get("presence_penalty"), frequency_penalty=query.data.get("frequency_penalty"), stop=query.data.get("stop"), @@ -105,21 +123,40 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> QueryResult: if not resp.choices: raise ValueError("Response must contain at least one choice") choice = resp.choices[0] + metadata: dict[str, Any] = {} + if choice.finish_reason: + metadata["finish_reason"] = choice.finish_reason + # Usage sums every choice while only the first is measured. + if len(resp.choices) == 1: + put_server_usage(metadata, *server_usage_counts(resp.usage)) return QueryResult( id=query_id, response_output=TextModelOutput(output=choice.text), - metadata=( - {"finish_reason": choice.finish_reason} if choice.finish_reason else {} - ), + metadata=metadata, ) @classmethod - def decode_sse_message(cls, json_bytes: bytes) -> SSEChoice: + def decode_sse_message(cls, json_bytes: bytes) -> SSEChoice | SSEMessage | None: + """Decode an SSE frame to the chat-shaped part the accumulator consumes. + + Returns an ``SSEChoice`` built from the first choice for a frame without + usage (the common case), ``None`` for a frame with neither usage nor + choices, and an ``SSEMessage`` carrying every choice and the usage for + a frame that carries usage. + """ msg = cls._sse_decoder.decode(json_bytes) - if not msg.choices: - return SSEChoice() - choice = msg.choices[0] - return SSEChoice( - delta=SSEDelta(content=choice.text), - finish_reason=choice.finish_reason, + if msg.usage is None: + if not msg.choices: + return None + choice = msg.choices[0] + return SSEChoice( + delta=SSEDelta(content=choice.text), + finish_reason=choice.finish_reason, + ) + return SSEMessage( + choices=tuple( + SSEChoice(delta=SSEDelta(content=c.text), finish_reason=c.finish_reason) + for c in msg.choices + ), + usage=msg.usage, ) diff --git a/src/inference_endpoint/openai/openai_adapter.py b/src/inference_endpoint/openai/openai_adapter.py index db75e91ed..11d2c590e 100644 --- a/src/inference_endpoint/openai/openai_adapter.py +++ b/src/inference_endpoint/openai/openai_adapter.py @@ -67,8 +67,8 @@ def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: ] @classmethod - def encode_query(cls, query: Query) -> bytes: - """Encode a Query to bytes for HTTP transmission.""" + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: + """Encode a Query to bytes for HTTP transmission; ignores ``include_usage``.""" request = cls.to_endpoint_request(query) return cls.encode_request(request) diff --git a/src/inference_endpoint/openai/openai_msgspec_adapter.py b/src/inference_endpoint/openai/openai_msgspec_adapter.py index 63318c88e..d6428589a 100644 --- a/src/inference_endpoint/openai/openai_msgspec_adapter.py +++ b/src/inference_endpoint/openai/openai_msgspec_adapter.py @@ -22,7 +22,13 @@ import msgspec from inference_endpoint.config.schema import ModelParams, StreamingMode -from inference_endpoint.core.types import Query, QueryResult, TextModelOutput +from inference_endpoint.core.types import ( + Query, + QueryResult, + TextModelOutput, + put_server_usage, + server_usage_counts, +) from inference_endpoint.dataset_manager.transforms import ( AddStaticColumns, ColumnFilter, @@ -33,6 +39,7 @@ from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter from .types import ( + STREAM_OPTIONS_INCLUDE_USAGE, ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, @@ -40,6 +47,7 @@ ChatMessage, SSEChoice, SSEMessage, + requests_multiple_choices, ) # ============================================================================ @@ -110,9 +118,9 @@ def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: ] @classmethod - def encode_query(cls, query: Query) -> bytes: + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: """Encode a Query directly to bytes for HTTP transmission.""" - request = cls.to_endpoint_request(query) + request = cls.to_endpoint_request(query, include_usage=include_usage) return cls.encode_request(request) @classmethod @@ -122,19 +130,26 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> QueryResult: return cls.from_endpoint_response(openai_response, result_id=query_id) @classmethod - def decode_sse_message(cls, json_bytes: bytes) -> SSEChoice | None: - """Decode SSE message and return the SSEChoice (delta + finish_reason).""" + def decode_sse_message(cls, json_bytes: bytes) -> SSEChoice | SSEMessage | None: + """Decode an SSE frame to the part of it the accumulator consumes. + + Returns the first ``SSEChoice`` for a frame without usage (the common + case), ``None`` for a frame with neither usage nor choices, and the + whole ``SSEMessage`` for a frame that carries usage. + """ msg = cls._sse_decoder.decode(json_bytes) - if not msg.choices: - return None - return msg.choices[0] + if msg.usage is None: + return msg.choices[0] if msg.choices else None + return msg # ======================================================================== # Internal APIs # ======================================================================== @classmethod - def to_endpoint_request(cls, query: Query) -> ChatCompletionRequest: + def to_endpoint_request( + cls, query: Query, *, include_usage: bool + ) -> ChatCompletionRequest: """ Convert a Query to an OpenAI request struct. @@ -144,6 +159,8 @@ def to_endpoint_request(cls, query: Query) -> ChatCompletionRequest: Args: query: Input query with prompt and parameters + include_usage: Ask for the final usage frame on a streaming request + with at most one choice (see ``HttpRequestAdapter.encode_query``). Returns: msgspec.Struct ChatCompletionRequest @@ -170,17 +187,24 @@ def to_endpoint_request(cls, query: Query) -> ChatCompletionRequest: ), ) + stream = query.data.get("stream") + n = query.data.get("n") return ChatCompletionRequest( model=query.data.get("model", "no-model-name"), messages=messages, - stream=query.data.get("stream"), + stream=stream, + stream_options=( + STREAM_OPTIONS_INCLUDE_USAGE + if include_usage and stream and not requests_multiple_choices(n) + else None + ), max_completion_tokens=query.data.get("max_completion_tokens"), temperature=query.data.get("temperature"), seed=query.data.get("seed"), top_p=query.data.get("top_p"), top_k=query.data.get("top_k"), repetition_penalty=query.data.get("repetition_penalty"), - n=query.data.get("n"), + n=n, presence_penalty=query.data.get("presence_penalty"), frequency_penalty=query.data.get("frequency_penalty"), stop=query.data.get("stop"), @@ -216,6 +240,9 @@ def from_endpoint_response( metadata["finish_reason"] = response.choices[0].finish_reason if message.tool_calls: metadata["tool_calls"] = message.tool_calls + # Usage sums every choice while only the first is measured. + if len(response.choices) == 1: + put_server_usage(metadata, *server_usage_counts(response.usage)) assert not ( message.reasoning_content and message.reasoning diff --git a/src/inference_endpoint/openai/types.py b/src/inference_endpoint/openai/types.py index 400baf704..7eb0508c3 100644 --- a/src/inference_endpoint/openai/types.py +++ b/src/inference_endpoint/openai/types.py @@ -17,6 +17,7 @@ msgspec types for OpenAI API serialization/deserialization. """ +import builtins from typing import Any import msgspec @@ -66,12 +67,18 @@ class SSEChoice( finish_reason: str | None = None +# Server ``usage`` fields below are typed ``object``: msgspec decodes any JSON +# value into builtins without raising, so a malformed ``usage`` never fails the +# frame or body it rides in. Counts are validated at capture by +# ``server_usage_counts``. gc=False stays safe: a JSON-decoded builtin cannot +# reference the struct that holds it, so no cycle can form. class SSEMessage( msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False ): # type: ignore[call-arg] """SSE message structure for OpenAI streaming responses.""" choices: tuple[SSEChoice, ...] = () + usage: object = None # validated by server_usage_counts # ============================================================================ @@ -103,6 +110,30 @@ class ChatMessage( reasoning_content: str | None = None +# gc=False: single bool field, no containers +class StreamOptions( + msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False +): # type: ignore[call-arg] + """OpenAI ``stream_options`` request object.""" + + include_usage: bool | None = None + + +# Shared by every request that asks for the final usage frame; the struct is +# frozen, so reusing one instance avoids a per-request allocation. +STREAM_OPTIONS_INCLUDE_USAGE = StreamOptions(include_usage=True) + + +def requests_multiple_choices(n: object) -> bool: + """Whether a request's ``n`` asks for more than one choice. + + Server usage then sums the tokens of every choice while only ``choices[0]`` + is measured. ``n`` is whatever the dataset row holds: a float when its + column has gaps (NaN compares False), never a bool. + """ + return isinstance(n, int | float) and not isinstance(n, bool) and n > 1 + + # gc=False: audit 2026-05: request containers set at construction; frozen=True blocks field reassignment. class ChatCompletionRequest( msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False @@ -120,6 +151,7 @@ class ChatCompletionRequest( temperature: float | None = None max_completion_tokens: int | None = None stream: bool | None = None + stream_options: StreamOptions | None = None seed: int | None = None top_p: float | None = None top_k: int | None = None @@ -169,16 +201,6 @@ class ChatCompletionChoice( finish_reason: str | None = None -class CompletionUsage( - msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False -): # type: ignore[call-arg] - """Token usage statistics.""" - - prompt_tokens: int - completion_tokens: int - total_tokens: int - - # gc=False: audit 2026-05: choices set at construction; frozen=True blocks field reassignment. class ChatCompletionResponse( msgspec.Struct, @@ -201,7 +223,8 @@ class ChatCompletionResponse( created: int = 0 model: str = "" choices: list[ChatCompletionChoice] - usage: CompletionUsage | None = None + # builtins.object: the ``object`` field above shadows the builtin here. + usage: builtins.object = None # validated by server_usage_counts system_fingerprint: str | None = None @@ -223,6 +246,7 @@ class TextCompletionRequest( min_tokens: int | None = None skip_special_tokens: bool | None = None stream: bool | None = None + stream_options: StreamOptions | None = None top_p: float | None = None top_k: int | None = None repetition_penalty: float | None = None @@ -252,7 +276,8 @@ class TextCompletionResponse( created: int model: str choices: list[TextCompletionChoice] - usage: CompletionUsage | None + # builtins.object: the ``object`` field above shadows the builtin here. + usage: builtins.object = None # validated by server_usage_counts class TextCompletionSSEChoice( @@ -270,3 +295,4 @@ class TextCompletionSSEMessage( """SSE message structure for /v1/completions streaming responses.""" choices: tuple[TextCompletionSSEChoice, ...] = () + usage: object = None # validated by server_usage_counts diff --git a/src/inference_endpoint/sglang/accumulator.py b/src/inference_endpoint/sglang/accumulator.py index 081106ebd..e9792762e 100644 --- a/src/inference_endpoint/sglang/accumulator.py +++ b/src/inference_endpoint/sglang/accumulator.py @@ -15,7 +15,17 @@ """SGLang SSE stream accumulator implementation.""" -from inference_endpoint.core.types import QueryResult, StreamChunk, TextModelOutput +from typing import Any + +from inference_endpoint.core.types import ( + COMPLETION_TOKENS_KEY, + PROMPT_TOKENS_KEY, + QueryResult, + StreamChunk, + TextModelOutput, + put_server_usage, + server_usage_count, +) from inference_endpoint.endpoint_client.accumulator_protocol import ( SSEAccumulatorProtocol, ) @@ -29,6 +39,12 @@ def __init__(self, query_id: str, stream_all_chunks: bool): self.text = "" self.token_ids: list[int] = [] self.total_tokens = 0 + self.prompt_tokens: int | None = None + # The last raw prompt_tokens validated; frames repeat one value. + self._raw_prompt_tokens: object = None + # Set by the first decoded frame (each carries meta_info counts); until + # then the initial total_tokens of 0 is not a server-reported count. + self._usage_seen = False self.retraction_occurred = False self.first_chunk_sent = False @@ -39,6 +55,18 @@ def add_chunk(self, delta: SGLangSSEDelta) -> StreamChunk | None: if not isinstance(delta, SGLangSSEDelta): return None + # Captured before the unchanged-count early return so a stream whose + # completion count never moves (e.g. an abort) still reports its counts. + # Only a changed raw value is validated; a frame without a usable + # prompt_tokens keeps the earlier count. + self._usage_seen = True + raw = delta.prompt_tokens + if raw is not None and raw != self._raw_prompt_tokens: + self._raw_prompt_tokens = raw + count = server_usage_count(raw, PROMPT_TOKENS_KEY) + if count is not None: + self.prompt_tokens = count + if delta.total_completion_tokens == self.total_tokens: return None @@ -60,7 +88,6 @@ def add_chunk(self, delta: SGLangSSEDelta) -> StreamChunk | None: "first_chunk": not self.first_chunk_sent, "final_chunk": False, "retraction_occurred": delta.has_retractions, - "n_tokens": len(delta.token_delta), } chunk = StreamChunk( id=self.query_id, @@ -73,14 +100,20 @@ def add_chunk(self, delta: SGLangSSEDelta) -> StreamChunk | None: return None def get_final_output(self) -> QueryResult: + metadata: dict[str, Any] = { + "first_chunk": not self.first_chunk_sent, + "final_chunk": True, + "retraction_occurred": self.retraction_occurred, + "token_ids": self.token_ids, + } + if self._usage_seen: + put_server_usage( + metadata, + self.prompt_tokens, + server_usage_count(self.total_tokens, COMPLETION_TOKENS_KEY), + ) return QueryResult( id=self.query_id, response_output=TextModelOutput(output=self.text), - metadata={ - "first_chunk": not self.first_chunk_sent, - "final_chunk": True, - "retraction_occurred": self.retraction_occurred, - "n_tokens": self.total_tokens, - "token_ids": self.token_ids, - }, + metadata=metadata, ) diff --git a/src/inference_endpoint/sglang/adapter.py b/src/inference_endpoint/sglang/adapter.py index a6614fa35..9fe0a3527 100644 --- a/src/inference_endpoint/sglang/adapter.py +++ b/src/inference_endpoint/sglang/adapter.py @@ -13,10 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + import msgspec from ..config.schema import ModelParams, StreamingMode -from ..core.types import Query, QueryResult, TextModelOutput +from ..core.types import ( + COMPLETION_TOKENS_KEY, + PROMPT_TOKENS_KEY, + Query, + QueryResult, + TextModelOutput, + put_server_usage, + server_usage_count, +) from ..dataset_manager.transforms import ( AddStaticColumns, ColumnFilter, @@ -67,8 +77,12 @@ def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: ] @classmethod - def encode_query(cls, query: Query) -> bytes: - """Encode a Query to bytes for HTTP transmission.""" + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: + """Encode a Query to bytes for HTTP transmission. + + ``include_usage`` is ignored: ``/generate`` reports token counts in + ``meta_info`` on every response. + """ # Get sampling parameters superset = query.data @@ -100,13 +114,17 @@ def encode_query(cls, query: Query) -> bytes: @classmethod def decode_response(cls, response_bytes: bytes, query_id: str) -> QueryResult: resp = cls._response_decoder.decode(response_bytes) + mi = resp.meta_info + metadata: dict[str, Any] = {"token_ids": resp.output_ids} + put_server_usage( + metadata, + server_usage_count(mi.prompt_tokens, PROMPT_TOKENS_KEY), + server_usage_count(mi.completion_tokens, COMPLETION_TOKENS_KEY), + ) return QueryResult( id=query_id, response_output=TextModelOutput(output=resp.text), - metadata={ - "token_ids": resp.output_ids, - "n_tokens": resp.meta_info.completion_tokens, - }, + metadata=metadata, ) @classmethod @@ -122,10 +140,12 @@ def decode_sse_message(cls, json_bytes: bytes) -> SGLangSSEDelta: # the output_tokens is the delta from the previous chunk token_delta = resp.output_ids - has_retractions = resp.meta_info.total_retractions > 0 + mi = resp.meta_info + has_retractions = bool(mi.num_retractions or mi.total_retractions) return SGLangSSEDelta( text=total_text, token_delta=token_delta, - total_completion_tokens=resp.meta_info.completion_tokens, + total_completion_tokens=mi.completion_tokens, has_retractions=has_retractions, + prompt_tokens=mi.prompt_tokens, ) diff --git a/src/inference_endpoint/sglang/types.py b/src/inference_endpoint/sglang/types.py index 9e75b3a4d..5e35dabf0 100644 --- a/src/inference_endpoint/sglang/types.py +++ b/src/inference_endpoint/sglang/types.py @@ -51,21 +51,29 @@ class SGLangGenerateRequest( class MetaInfo(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False): # type: ignore[call-arg] - id: str - finish_reason: dict[str, Any] - prompt_tokens: int - weight_version: str - total_retractions: int + # Only completion_tokens is required (it drives the stream): SGLang sends a + # null finish_reason until the last frame, e2e_latency only on it, and names + # the retraction counter num_retractions (older servers: total_retractions). + id: str | None = None + finish_reason: dict[str, Any] | None = None + # Any JSON value decodes without raising, so a malformed count never fails + # the frame; server_usage_count validates it (the accumulator on a stream, + # once per changed value; the adapter on a non-streaming response). gc=False + # stays safe: a JSON-decoded builtin cannot reference this struct. + prompt_tokens: object = None + weight_version: str | None = None + total_retractions: int | None = None + num_retractions: int | None = None completion_tokens: int - cached_tokens: int - e2e_latency: float + cached_tokens: int | None = None + e2e_latency: float | None = None class SGLangGenerateResponse( msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False ): # type: ignore[call-arg] text: str - output_ids: list[int] + output_ids: list[int] = msgspec.field(default_factory=list) meta_info: MetaInfo @@ -76,3 +84,7 @@ class SGLangSSEDelta( token_delta: list[int] = msgspec.field(default_factory=list) total_completion_tokens: int = 0 has_retractions: bool = False + # The raw meta_info value: every frame repeats it, so the accumulator + # validates it only when it changes. gc=False stays safe: a JSON-decoded + # builtin cannot reference this struct. + prompt_tokens: object = None diff --git a/src/inference_endpoint/utils/benchmark_httpclient.py b/src/inference_endpoint/utils/benchmark_httpclient.py index 91e2d5c68..35c9e7ceb 100644 --- a/src/inference_endpoint/utils/benchmark_httpclient.py +++ b/src/inference_endpoint/utils/benchmark_httpclient.py @@ -47,10 +47,7 @@ from inference_endpoint.endpoint_client.http_client import ( HTTPEndpointClient, ) -from inference_endpoint.testing.max_throughput_server import ( - MaxThroughputServer, - build_response, -) +from inference_endpoint.testing.max_throughput_server import MaxThroughputServer try: import matplotlib @@ -405,6 +402,10 @@ def _get_memory_info(self) -> str: _SEND_BATCH = 32 # Max requests issued per event-loop yield +# Whether the benchmark client asks for the usage frame on streaming requests; +# the stub then sends it, which changes the SSE events per response. +_STREAM_INCLUDE_USAGE = True + def build_prompt(length: int) -> str: """Build a prompt of specified length using repeated 'hello world '.""" @@ -445,6 +446,7 @@ def _create_client( num_workers=num_workers if num_workers > 0 else -1, max_connections=max_connections if max_connections > 0 else -1, warmup_connections=0, + stream_include_usage=_STREAM_INCLUDE_USAGE, worker_gc_mode="relaxed", log_level="CRITICAL", cpu_affinity=cpu_affinity_plan, @@ -693,13 +695,16 @@ def _resolve_stream_intervals(output_length: int, fractions: list[float]) -> lis return sorted(set(vals)) -def _sse_events_per_response(output_length: int, stream_interval: int = 1) -> int: +def _sse_events_per_response( + output_length: int, stream_interval: int = 1, *, include_usage: bool +) -> int: """Compute the number of SSE events the server sends per streaming response. ``ceil(output_length / stream_interval)`` content events, plus 1 finish - event and 1 [DONE] sentinel. + event, 1 [DONE] sentinel and, when the request asks for it + (``include_usage``), 1 usage event. """ - return -(-output_length // stream_interval) + 2 # ceil division + 2 + return -(-output_length // stream_interval) + 2 + int(include_usage) def _wait_for_port_free(host: str, port: int, timeout: float = 5.0) -> None: @@ -732,7 +737,7 @@ def _restart_server( if server._port is not None: _wait_for_port_free(server.host, server._port) server.stream = stream - server._response = build_response(output_length, stream, stream_interval) + server.set_response(output_length, stream_interval) # Reset shutdown flags and worker lists so _start_workers can re-use them server._shutdown.clear() server._workers.clear() @@ -764,7 +769,9 @@ def run_single( _restart_server(server, prompt_length, args.streaming, stream_interval) epr = ( - _sse_events_per_response(prompt_length, stream_interval) + _sse_events_per_response( + prompt_length, stream_interval, include_usage=_STREAM_INCLUDE_USAGE + ) if args.streaming else 1 ) @@ -941,7 +948,9 @@ def run_sweep( prompt = build_prompt(prompt_length) epr = ( - _sse_events_per_response(prompt_length, stream_interval) + _sse_events_per_response( + prompt_length, stream_interval, include_usage=_STREAM_INCLUDE_USAGE + ) if args.streaming else 1 ) diff --git a/src/inference_endpoint/videogen/adapter.py b/src/inference_endpoint/videogen/adapter.py index ff4f9a704..d8bf3cc04 100644 --- a/src/inference_endpoint/videogen/adapter.py +++ b/src/inference_endpoint/videogen/adapter.py @@ -61,12 +61,13 @@ def dataset_transforms(cls, model_params: "ModelParams") -> "list[Transform]": ] @classmethod - def encode_query(cls, query: Query) -> bytes: + def encode_query(cls, query: Query, *, include_usage: bool) -> bytes: """Serialise query.data to VideoPathRequest JSON bytes. Only `prompt` is required. All other fields fall back to defaults on VideoPathRequest but can be overridden via query.data. Streaming is - not supported — `stream=True` raises. + not supported — `stream=True` raises. `include_usage` is ignored: the + endpoint reports no token usage. """ data = query.data if "prompt" not in data: diff --git a/tests/integration/endpoint_client/test_sglang_adapter.py b/tests/integration/endpoint_client/test_sglang_adapter.py index 78aee685b..7c2de11e9 100644 --- a/tests/integration/endpoint_client/test_sglang_adapter.py +++ b/tests/integration/endpoint_client/test_sglang_adapter.py @@ -93,9 +93,11 @@ async def test_sglang_non_streaming_request(self, sglang_futures_client): # Verify metadata assert result.metadata is not None assert "token_ids" in result.metadata - assert "n_tokens" in result.metadata + assert "completion_tokens" in result.metadata assert isinstance(result.metadata["token_ids"], list) - assert isinstance(result.metadata["n_tokens"], int) + assert isinstance(result.metadata["completion_tokens"], int) + assert result.metadata["prompt_tokens"] == len(input_tokens) + assert "n_tokens" not in result.metadata @pytest.mark.skip( reason="Running this test requires a running GPT-OSS server at localhost:30000." @@ -136,16 +138,28 @@ async def test_sglang_streaming_request(self, sglang_futures_client): assert result.metadata is not None assert "token_ids" in result.metadata - assert "n_tokens" in result.metadata + assert "completion_tokens" in result.metadata assert isinstance(result.metadata["token_ids"], list) - assert isinstance(result.metadata["n_tokens"], int) + assert isinstance(result.metadata["completion_tokens"], int) + assert result.metadata["prompt_tokens"] == len(input_tokens) + assert "n_tokens" not in result.metadata # Check that something was generated, but no more than max_new_tokens - assert 0 < result.metadata["n_tokens"] and result.metadata["n_tokens"] <= 100 + assert ( + 0 < result.metadata["completion_tokens"] + and result.metadata["completion_tokens"] <= 100 + ) - # The token IDs in the result should be at most n_tokens because of retractions + # The token IDs in the result should be at most completion_tokens because of + # retractions if result.metadata["retraction_occurred"]: - assert len(result.metadata["token_ids"]) <= result.metadata["n_tokens"] + assert ( + len(result.metadata["token_ids"]) + <= result.metadata["completion_tokens"] + ) else: # STOP token is not included in the response, but counts towards generated - assert len(result.metadata["token_ids"]) + 1 == result.metadata["n_tokens"] + assert ( + len(result.metadata["token_ids"]) + 1 + == result.metadata["completion_tokens"] + ) diff --git a/tests/integration/videogen/test_adapter.py b/tests/integration/videogen/test_adapter.py index 2bf34a1b1..7c88b853b 100644 --- a/tests/integration/videogen/test_adapter.py +++ b/tests/integration/videogen/test_adapter.py @@ -56,7 +56,7 @@ def test_perf_mode_round_trip_returns_video_path( ) -> None: """Default response_format=video_path → server returns the saved-video path.""" query = Query(id="q1", data={"prompt": "a golden retriever running on a beach"}) - request_bytes = VideoGenAdapter.encode_query(query) + request_bytes = VideoGenAdapter.encode_query(query, include_usage=False) status, content = _post( f"{mock_trtllm_serve.url}/v1/videos/generations", request_bytes @@ -83,7 +83,7 @@ def test_accuracy_mode_round_trip_returns_video_bytes( "response_format": "video_bytes", }, ) - request_bytes = VideoGenAdapter.encode_query(query) + request_bytes = VideoGenAdapter.encode_query(query, include_usage=False) status, content = _post( f"{mock_trtllm_serve.url}/v1/videos/generations", request_bytes @@ -102,7 +102,7 @@ def test_request_default_response_format_is_video_path( self, mock_trtllm_serve: MockTrtllmServe ) -> None: query = Query(id="q2", data={"prompt": "ocean waves at sunset"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["response_format"] == "video_path" def test_request_accuracy_mode_asks_for_video_bytes( @@ -112,14 +112,14 @@ def test_request_accuracy_mode_asks_for_video_bytes( id="q2b", data={"prompt": "ocean waves at sunset", "response_format": "video_bytes"}, ) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["response_format"] == "video_bytes" def test_request_carries_mlperf_defaults( self, mock_trtllm_serve: MockTrtllmServe ) -> None: query = Query(id="q3", data={"prompt": "ocean waves at sunset"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["fps"] == 16 assert payload["seconds"] == pytest.approx(5.0) assert payload["size"] == "720x1280" @@ -136,7 +136,7 @@ def test_http_500_response_raises_on_decode( query = Query(id="q4", data={"prompt": "a cat"}) status, content = _post( f"{mock_trtllm_serve_error.url}/v1/videos/generations", - VideoGenAdapter.encode_query(query), + VideoGenAdapter.encode_query(query, include_usage=False), ) assert status == 500 with pytest.raises((ValidationError, json.JSONDecodeError)): diff --git a/tests/performance/openai/test_adapter.py b/tests/performance/openai/test_adapter.py index 6ed12ab28..94f21498a 100644 --- a/tests/performance/openai/test_adapter.py +++ b/tests/performance/openai/test_adapter.py @@ -82,7 +82,7 @@ def test_encode_query(benchmark, size_name, text): """Benchmark encode_query (Query -> HTTP bytes).""" query = make_query(text) benchmark.group = "openai_adapter_encode_query" - benchmark(OpenAIAdapter.encode_query, query) + benchmark(OpenAIAdapter.encode_query, query, include_usage=False) @pytest.mark.performance diff --git a/tests/performance/openai/test_msgspec_adapter.py b/tests/performance/openai/test_msgspec_adapter.py index 21bedd252..b5f04cf1e 100644 --- a/tests/performance/openai/test_msgspec_adapter.py +++ b/tests/performance/openai/test_msgspec_adapter.py @@ -82,7 +82,7 @@ def test_encode_query(benchmark, size_name, text): """Benchmark encode_query (Query -> HTTP bytes).""" query = make_query(text) benchmark.group = "msgspec_adapter_encode_query" - benchmark(OpenAIMsgspecAdapter.encode_query, query) + benchmark(OpenAIMsgspecAdapter.encode_query, query, include_usage=True) @pytest.mark.performance diff --git a/tests/performance/openai/test_types.py b/tests/performance/openai/test_types.py index bdf8aa6bb..838c08345 100644 --- a/tests/performance/openai/test_types.py +++ b/tests/performance/openai/test_types.py @@ -36,7 +36,6 @@ ChatCompletionResponse, ChatCompletionResponseMessage, ChatMessage, - CompletionUsage, SSEChoice, SSEDelta, SSEMessage, @@ -98,11 +97,7 @@ def make_chat_response(text: str) -> ChatCompletionResponse: finish_reason="stop", ) ], - usage=CompletionUsage( - prompt_tokens=10, - completion_tokens=20, - total_tokens=30, - ), + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, system_fingerprint="fp_test", ) diff --git a/tests/performance/sglang/test_sglang_adapter.py b/tests/performance/sglang/test_sglang_adapter.py index f614cc0d7..e205567f3 100644 --- a/tests/performance/sglang/test_sglang_adapter.py +++ b/tests/performance/sglang/test_sglang_adapter.py @@ -92,7 +92,7 @@ def test_encode_query(benchmark, size_name, input_tokens): """Benchmark encode_query (Query -> HTTP bytes).""" query = make_query(input_tokens) benchmark.group = "sglang_adapter_encode_query" - benchmark(SGLangGenerateAdapter.encode_query, query) + benchmark(SGLangGenerateAdapter.encode_query, query, include_usage=False) @pytest.mark.performance diff --git a/tests/unit/dataset_manager/test_data_loader.py b/tests/unit/dataset_manager/test_data_loader.py index 2d2b0fbd3..2410d4fba 100644 --- a/tests/unit/dataset_manager/test_data_loader.py +++ b/tests/unit/dataset_manager/test_data_loader.py @@ -13,9 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import pandas as pd import pytest +from inference_endpoint.config.schema import APIType, ModelParams, StreamingMode +from inference_endpoint.core.types import Query +from inference_endpoint.dataset_manager.agentic_inference_dataset import ( + AgenticInferenceDataset, +) +from inference_endpoint.dataset_manager.dataset import Dataset from inference_endpoint.dataset_manager.predefined.random import RandomDataset +from inference_endpoint.dataset_manager.transforms import ColumnFilter +from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter + +_STREAMING = ModelParams(name="m", streaming=StreamingMode.ON) + + +def _plain_sample() -> dict: + # A preset ColumnFilter that keeps the column replaces the adapter's filter. + ds = Dataset( + pd.DataFrame([{"prompt": "hi", "include_usage": True}]), + transforms=[ + ColumnFilter( + required_columns=["prompt"], optional_columns=["include_usage"] + ) + ], + ) + ds.load(api_type=APIType.OPENAI, model_params=_STREAMING) + return ds.load_sample(0) + + +def _agentic_sample() -> dict: + row = { + "conversation_id": "c1", + "turn": 1, + "role": "user", + "content": "hi", + "stream": True, + "include_usage": True, + } + ds = AgenticInferenceDataset(pd.DataFrame([row])) + ds.load(api_type=APIType.OPENAI, model_params=_STREAMING) + return ds.load_sample(0) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "load", [_plain_sample, _agentic_sample], ids=["dataset", "agentic"] +) +def test_dataset_column_named_include_usage_passes_through(load): + sample = load() + assert sample["include_usage"] is True + # The run setting, not the column, decides whether usage is requested. + body = json.loads( + OpenAIMsgspecAdapter.encode_query(Query(data=sample), include_usage=False) + ) + assert "stream_options" not in body and "include_usage" not in body @pytest.mark.unit diff --git a/tests/unit/endpoint_client/test_adapter_protocol.py b/tests/unit/endpoint_client/test_adapter_protocol.py index 591dc4236..dec9f0630 100644 --- a/tests/unit/endpoint_client/test_adapter_protocol.py +++ b/tests/unit/endpoint_client/test_adapter_protocol.py @@ -13,15 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for HttpRequestAdapter.parse_sse_chunk.""" +"""Tests for HttpRequestAdapter: the encode_query contract and parse_sse_chunk.""" +import inspect +import logging from typing import Any import msgspec import pytest from inference_endpoint.config.schema import ModelParams +from inference_endpoint.core.types import Query from inference_endpoint.dataset_manager.transforms import Transform from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter +from inference_endpoint.endpoint_client.config import ADAPTER_MAP +from inference_endpoint.openai.completions_adapter import OpenAITextCompletionsAdapter +from inference_endpoint.openai.openai_adapter import OpenAIAdapter +from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter +from inference_endpoint.sglang.adapter import SGLangGenerateAdapter +from inference_endpoint.videogen.adapter import VideoGenAdapter class _SimpleAdapter(HttpRequestAdapter): @@ -32,7 +41,7 @@ def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: return [] @classmethod - def encode_query(cls, query: Any) -> bytes: + def encode_query(cls, query: Any, *, include_usage: bool) -> bytes: return b"" @classmethod @@ -46,6 +55,48 @@ def decode_sse_message(cls, json_bytes: bytes) -> Any: return {"x": 1} +_ADAPTERS = [ + OpenAIMsgspecAdapter, + OpenAITextCompletionsAdapter, + SGLangGenerateAdapter, + VideoGenAdapter, + OpenAIAdapter, +] + + +@pytest.mark.unit +def test_adapter_list_covers_every_api_type(): + names = {f"{a.__module__}.{a.__qualname__}" for a in _ADAPTERS} + assert set(ADAPTER_MAP.values()) <= names + + +@pytest.mark.unit +@pytest.mark.parametrize( + "adapter", [HttpRequestAdapter, *_ADAPTERS], ids=lambda a: a.__name__ +) +def test_encode_query_requires_keyword_only_include_usage(adapter): + param = inspect.signature(adapter.encode_query).parameters["include_usage"] + assert param.kind is inspect.Parameter.KEYWORD_ONLY + assert param.default is inspect.Parameter.empty + + +@pytest.mark.unit +@pytest.mark.parametrize( + "adapter, data", + [ + (OpenAIAdapter, {"model": "m", "prompt": "hi", "stream": True}), + (SGLangGenerateAdapter, {"input_tokens": [1, 2], "stream": True}), + (VideoGenAdapter, {"prompt": "a cat"}), + ], + ids=["openai-test-only", "sglang", "videogen"], +) +def test_adapters_without_usage_requests_ignore_include_usage(adapter, data): + query = Query(id="q", data=data) + assert adapter.encode_query(query, include_usage=True) == adapter.encode_query( + query, include_usage=False + ) + + @pytest.mark.unit def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(): """A malformed SSE frame is skipped; surrounding valid frames are preserved.""" @@ -53,3 +104,48 @@ def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(): end_pos = len(buffer) result = _SimpleAdapter.parse_sse_chunk(buffer, end_pos) assert result == [{"x": 1}, {"x": 1}] + + +_LOGGER = "inference_endpoint.endpoint_client.adapter_protocol" + + +class _StrictAdapter(_SimpleAdapter): + """Decodes frames as dict[str, int] so bad frames raise real msgspec errors.""" + + @classmethod + def decode_sse_message(cls, json_bytes: bytes) -> Any: + return msgspec.json.decode(json_bytes, type=dict[str, int]) + + +def _sse(*frames: bytes) -> bytes: + return b"".join(b"data: " + f + b"\n\n" for f in frames) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "bad", [b"{not json}", b'{"x":"s"}'], ids=["decode-error", "validation-error"] +) +def test_parse_sse_chunk_keeps_frames_after_bad_frame(bad, caplog): + caplog.set_level(logging.WARNING, logger=_LOGGER) + buffer = _sse(b'{"x":1}', bad, b'{"x":2}', b'{"x":3}') + result = _StrictAdapter.parse_sse_chunk(buffer, len(buffer)) + assert result == [{"x": 1}, {"x": 2}, {"x": 3}] + assert len(caplog.records) == 1 + + +@pytest.mark.unit +def test_parse_sse_chunk_warns_once_per_call(caplog): + caplog.set_level(logging.WARNING, logger=_LOGGER) + buffer = _sse(b"{not json}", b'{"x":1}', b'{"x":"s"}') + assert _StrictAdapter.parse_sse_chunk(buffer, len(buffer)) == [{"x": 1}] + assert [r.getMessage() for r in caplog.records] == [ + "skipping 2 malformed SSE frame(s) (DecodeError)" + ] + + +@pytest.mark.unit +def test_parse_sse_chunk_clean_batch_logs_nothing(caplog): + caplog.set_level(logging.WARNING, logger=_LOGGER) + buffer = _sse(b'{"x":1}', b'{"x":2}') + assert _StrictAdapter.parse_sse_chunk(buffer, len(buffer)) == [{"x": 1}, {"x": 2}] + assert caplog.records == [] diff --git a/tests/unit/endpoint_client/test_worker.py b/tests/unit/endpoint_client/test_worker.py new file mode 100644 index 000000000..79b6f78c3 --- /dev/null +++ b/tests/unit/endpoint_client/test_worker.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Worker request encoding and error handling, without sockets.""" + +import logging +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from inference_endpoint.core.types import Query +from inference_endpoint.endpoint_client import worker as worker_mod +from inference_endpoint.endpoint_client.config import HTTPClientConfig +from inference_endpoint.endpoint_client.http import InFlightRequest +from inference_endpoint.endpoint_client.worker import ( + Worker, + _is_stream_options_rejection, +) +from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter + +_REJECTION = ( + b'{"object":"error","message":"Unrecognized request argument supplied: ' + b'stream_options","code":400}' +) + + +def _worker(**config: Any) -> Worker: + http_config = HTTPClientConfig( + endpoint_urls=["http://localhost:8000"], + num_workers=1, + warmup_connections=0, + max_connections=1, + **config, + ) + worker = Worker(worker_id=0, connector=MagicMock(), http_config=http_config) + worker._pool = MagicMock() + worker._responses = MagicMock() + worker._http_template = MagicMock() + return worker + + +@pytest.mark.unit +@pytest.mark.parametrize("stream_include_usage", [True, False]) +def test_prepare_request_passes_the_run_usage_setting( + monkeypatch, stream_include_usage +): + encode = MagicMock(return_value=b"{}") + monkeypatch.setattr(OpenAIMsgspecAdapter, "encode_query", encode) + worker = _worker(stream_include_usage=stream_include_usage) + query = Query(id="q", data={"prompt": "hi", "stream": True}) + + worker._prepare_request(query) + + encode.assert_called_once_with(query, include_usage=stream_include_usage) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "status_code, body, streaming, expected", + [ + (400, _REJECTION, True, True), + (422, b"extra field stream_options not permitted", True, True), + (400, _REJECTION, False, False), + (400, b'{"message":"max_tokens too large"}', True, False), + (500, _REJECTION, True, False), + (399, _REJECTION, True, False), + ], +) +def test_is_stream_options_rejection(status_code, body, streaming, expected): + assert _is_stream_options_rejection(status_code, body, streaming) is expected + + +def _failing_request(query_id: str, status_code: int, body: bytes) -> InFlightRequest: + conn = MagicMock() + conn.protocol.read_headers = AsyncMock(return_value=(status_code, [])) + conn.protocol.read_body = AsyncMock(return_value=body) + req = InFlightRequest(query_id=query_id, http_bytes=b"", is_streaming=True) + req.connection = conn + return req + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_stream_options_rejection_logs_opt_out_hint_once(monkeypatch, caplog): + monkeypatch.setattr(worker_mod, "_stream_options_hint_logged", False) + caplog.set_level(logging.WARNING, logger=worker_mod.__name__) + worker = _worker() + + for query_id in ("q0", "q1"): + await worker._process_response(_failing_request(query_id, 400, _REJECTION)) + + hints = [ + r.getMessage() + for r in caplog.records + if "--no-stream-include-usage" in r.getMessage() + ] + assert len(hints) == 1 + assert "HTTP 400" in hints[0] + assert "settings.client.stream_include_usage" in hints[0] + # The hint is advice only: both requests still fail. + errors = [c.args[0] for c in worker._responses.send.call_args_list] + assert [e.id for e in errors] == ["q0", "q1"] + assert all(e.error is not None for e in errors) + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_other_errors_log_no_opt_out_hint(monkeypatch, caplog): + monkeypatch.setattr(worker_mod, "_stream_options_hint_logged", False) + caplog.set_level(logging.WARNING, logger=worker_mod.__name__) + worker = _worker() + + await worker._process_response(_failing_request("q0", 400, b"bad max_tokens")) + + assert not any( + "--no-stream-include-usage" in r.getMessage() for r in caplog.records + ) + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_opted_out_worker_never_logs_opt_out_hint(monkeypatch, caplog): + monkeypatch.setattr(worker_mod, "_stream_options_hint_logged", False) + caplog.set_level(logging.WARNING, logger=worker_mod.__name__) + worker = _worker(stream_include_usage=False) + + await worker._process_response(_failing_request("q0", 400, _REJECTION)) + + assert not any( + "--no-stream-include-usage" in r.getMessage() for r in caplog.records + ) + (error,) = (c.args[0] for c in worker._responses.send.call_args_list) + assert error.id == "q0" and error.error is not None diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 87443a4ee..c8acd6a19 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -21,7 +21,9 @@ import random import time from types import SimpleNamespace +from typing import Any +import msgspec import pytest from inference_endpoint.config.runtime_settings import RuntimeSettings from inference_endpoint.config.schema import LoadPattern, LoadPatternType @@ -32,12 +34,14 @@ SessionEventType, ) from inference_endpoint.core.types import ( + TOKEN_HDR_HIGH, ErrorData, PhaseData, PhaseType, Query, QueryResult, StreamChunk, + put_server_usage, ) from inference_endpoint.dataset_manager.dataset import Dataset from inference_endpoint.load_generator.session import ( @@ -1920,6 +1924,32 @@ async def test_drops_late_response_after_timeout(self): assert publisher.events == [] assert phase_issuer.inflight == 1 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "counts,expected", + [ + ((11, 3), (11, 3)), + ((None, None), (msgspec.UNSET, msgspec.UNSET)), + ((None, 3), (msgspec.UNSET, 3)), + ((11, None), (11, msgspec.UNSET)), + ((0, 0), (0, 0)), + ((TOKEN_HDR_HIGH, TOKEN_HDR_HIGH), (TOKEN_HDR_HIGH, TOKEN_HDR_HIGH)), + ], + ids=["ints", "missing", "prompt_missing", "completion_missing", "zero", "max"], + ) + async def test_session_lifts_the_counts_put_server_usage_wrote( + self, counts, expected + ): + metadata: dict[str, Any] = {} + put_server_usage(metadata, *counts) + publisher = FakePublisher() + session = BenchmarkSession(FakeIssuer(), publisher, asyncio.get_running_loop()) + session._handle_response( + QueryResult(id="q", response_output="ab", metadata=metadata), 1_000 + ) + (complete,) = publisher.events_of_type(SampleEventType.COMPLETE) + assert (complete.prompt_tokens, complete.completion_tokens) == expected + class NoopStrategy: """Issues nothing, so ``_run_phase`` returns without any in-flight work.""" diff --git a/tests/unit/openai/test_accumulator.py b/tests/unit/openai/test_accumulator.py index bd2f16182..022383907 100644 --- a/tests/unit/openai/test_accumulator.py +++ b/tests/unit/openai/test_accumulator.py @@ -15,10 +15,13 @@ """Unit tests for OpenAISSEAccumulator tool-call handling.""" +import msgspec import pytest -from inference_endpoint.core.types import StreamChunk, TextModelOutput +from inference_endpoint.core.types import QueryResult, StreamChunk, TextModelOutput from inference_endpoint.openai.accumulator import OpenAISSEAccumulator -from inference_endpoint.openai.types import SSEChoice, SSEDelta +from inference_endpoint.openai.completions_adapter import OpenAITextCompletionsAdapter +from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter +from inference_endpoint.openai.types import SSEChoice, SSEDelta, SSEMessage def _make_tc_partial(idx: int, tc_id: str, name: str, args: str) -> SSEChoice: @@ -223,3 +226,284 @@ def test_mixed_reasoning_and_content_are_both_preserved( assert output.output == ("42", "!") assert "".join(output.reasoning) == "Think: done" assert output.tool_calls == ((tuple(tool_calls),) if tool_calls else None) + + +def _usage_frame(**counts: object) -> SSEMessage: + """A usage-only trailer as the adapters decode it.""" + return SSEMessage(usage=counts) + + +def _counts(metadata: dict) -> dict: + return { + k: metadata[k] for k in ("prompt_tokens", "completion_tokens") if k in metadata + } + + +@pytest.mark.unit +class TestAccumulatorUsage: + def test_usage_frame_sets_metadata_without_stream_chunk(self): + acc = OpenAISSEAccumulator("qid", stream_all_chunks=True) + assert acc.add_chunk(_usage_frame(prompt_tokens=7, completion_tokens=0)) is None + assert acc.first_chunk_sent is False + assert acc.add_chunk(_make_content_choice("Hi")).metadata["first_chunk"] is True + assert acc.add_chunk(_usage_frame(prompt_tokens=7, completion_tokens=3)) is None + result = acc.get_final_output() + assert result.response_output.output == ("Hi",) + assert _counts(result.metadata) == {"prompt_tokens": 7, "completion_tokens": 3} + + def test_usage_only_stream_reports_zero(self): + acc = OpenAISSEAccumulator("qid", stream_all_chunks=False) + acc.add_chunk(_usage_frame(prompt_tokens=7, completion_tokens=0)) + result = acc.get_final_output() + assert result.response_output.output == () + assert result.metadata == { + "first_chunk": True, + "final_chunk": True, + "prompt_tokens": 7, + "completion_tokens": 0, + } + + def test_message_applies_choice_and_usage(self): # LiteLLM shape + acc = OpenAISSEAccumulator("qid", stream_all_chunks=False) + chunk = acc.add_chunk( + SSEMessage( + choices=( + SSEChoice(delta=SSEDelta(content="Hi"), finish_reason="stop"), + ), + usage={"prompt_tokens": 7, "completion_tokens": 1}, + ) + ) + assert (chunk.response_chunk, chunk.metadata["first_chunk"]) == ("Hi", True) + md = acc.get_final_output().metadata + assert (md["finish_reason"], md["prompt_tokens"], md["completion_tokens"]) == ( + "stop", + 7, + 1, + ) + + def test_usage_on_every_frame_last_wins(self): + acc = OpenAISSEAccumulator("qid", stream_all_chunks=True) + for n, text in enumerate("abc", start=1): + acc.add_chunk( + SSEMessage( + choices=(_make_content_choice(text),), + usage={"prompt_tokens": 5, "completion_tokens": n}, + ) + ) + result = acc.get_final_output() + assert result.response_output.output == ("a", "bc") + assert result.metadata["completion_tokens"] == 3 + + def test_null_usage_fields_do_not_erase_earlier_counts( + self, + ): # proxy ends with "usage": {} + acc = OpenAISSEAccumulator("qid", stream_all_chunks=False) + acc.add_chunk( + SSEMessage( + choices=(_make_content_choice("Hi"),), + usage={"prompt_tokens": 7, "completion_tokens": 3}, + ) + ) + assert acc.add_chunk(_usage_frame()) is None + assert _counts(acc.get_final_output().metadata) == { + "prompt_tokens": 7, + "completion_tokens": 3, + } + acc.add_chunk( + _usage_frame(prompt_tokens=7, completion_tokens=5) + ) # valid ints still override + assert acc.get_final_output().metadata["completion_tokens"] == 5 + + @pytest.mark.parametrize( + "chunk, expected", + [ + ( + _usage_frame(prompt_tokens=None, completion_tokens=2), + {"completion_tokens": 2}, + ), + (None, {}), + ], + ids=["null-prompt-tokens", "no-usage"], + ) + def test_missing_counts_omit_keys(self, chunk, expected): + acc = OpenAISSEAccumulator("qid", stream_all_chunks=False) + acc.add_chunk(_make_content_choice("Hi")) + assert acc.add_chunk(chunk) is None + assert _counts(acc.get_final_output().metadata) == expected + + +def _chat_choice(text: str, **fields: object) -> dict: + return {"delta": {"content": text}, **fields} + + +def _completions_choice(text: str, **fields: object) -> dict: + return {"text": text, **fields} + + +_WIRE_ADAPTERS = pytest.mark.parametrize( + "adapter, choice", + [ + (OpenAIMsgspecAdapter, _chat_choice), + (OpenAITextCompletionsAdapter, _completions_choice), + ], + ids=["chat", "completions"], +) + + +def _accumulate_wire(adapter, frames: list[dict]) -> QueryResult: + buffer = ( + b"".join(b"data: " + msgspec.json.encode(f) + b"\n\n" for f in frames) + + b"data: [DONE]\n\n" + ) + acc = OpenAISSEAccumulator("qid", stream_all_chunks=True) + for decoded in adapter.parse_sse_chunk(buffer, len(buffer)): + acc.add_chunk(decoded) + return acc.get_final_output() + + +@pytest.mark.unit +@_WIRE_ADAPTERS +@pytest.mark.parametrize( + "usage", + [{"prompt_tokens": 12.0}, [], "x", {"completion_tokens": -1}], + ids=["float-count", "list", "string", "negative-count"], +) +def test_malformed_usage_keeps_frame_content(adapter, choice, usage): + result = _accumulate_wire( + adapter, [{"choices": [choice("Hi", finish_reason="stop")], "usage": usage}] + ) + assert result.response_output.output == ("Hi",) + assert result.metadata["finish_reason"] == "stop" + assert _counts(result.metadata) == {} + + +@pytest.mark.unit +@_WIRE_ADAPTERS +def test_usage_on_every_frame_last_wellformed_wins(adapter, choice): + result = _accumulate_wire( + adapter, + [ + { + "choices": [choice("a")], + "usage": {"prompt_tokens": 5, "completion_tokens": 1}, + }, + { + "choices": [choice("b")], + "usage": {"prompt_tokens": 5, "completion_tokens": 2}, + }, + { + "choices": [choice("c", finish_reason="stop")], + "usage": {"completion_tokens": "bad"}, + }, + ], + ) + assert result.response_output.output == ("a", "bc") + assert _counts(result.metadata) == {"prompt_tokens": 5, "completion_tokens": 2} + + +@pytest.mark.unit +@_WIRE_ADAPTERS +def test_null_usage_frames_record_nothing(adapter, choice): + result = _accumulate_wire( + adapter, + [ + {"choices": [choice("a")], "usage": None}, + {"choices": [choice("b", finish_reason="stop")], "usage": None}, + {"choices": [], "usage": None}, + ], + ) + assert result.response_output.output == ("a", "b") + assert _counts(result.metadata) == {} + + +@pytest.mark.unit +@_WIRE_ADAPTERS +def test_malformed_usage_on_the_wire_warns_once( + adapter, choice, usage_warning_unlatched +): + _accumulate_wire( + adapter, + [ + {"choices": [choice("a")], "usage": "x"}, + {"choices": [choice("b")], "usage": {"completion_tokens": -1}}, + {"choices": [], "usage": [1]}, + ], + ) + (record,) = usage_warning_unlatched.records + assert "'x' (type 'str')" in record.getMessage() + + +@pytest.mark.unit +@_WIRE_ADAPTERS +def test_well_formed_usage_on_the_wire_never_warns( + adapter, choice, usage_warning_unlatched +): + _accumulate_wire( + adapter, + [ + {"choices": [choice("a")], "usage": None}, + {"choices": [], "usage": {"prompt_tokens": 5, "completion_tokens": 1}}, + ], + ) + assert usage_warning_unlatched.records == [] + + +@pytest.mark.unit +@_WIRE_ADAPTERS +def test_usage_bearing_frame_with_two_choices_drops_counts(adapter, choice): + result = _accumulate_wire( + adapter, + [ + {"choices": [choice("a")]}, + { + "choices": [choice("b"), choice("x", index=1)], + "usage": {"prompt_tokens": 5, "completion_tokens": 4}, + }, + ], + ) + # Each frame contributes only its first choice to the output. + assert result.response_output.output == ("a", "b") + assert _counts(result.metadata) == {} + + +@pytest.mark.unit +@pytest.mark.parametrize( + "adapter, frames", + [ + ( + OpenAIMsgspecAdapter, + [ + b'{"choices":[{"delta":{"role":"assistant"}}],"usage":null}', + b'{"choices":[{"delta":{"content":"Hi"}}],"usage":null}', + b'{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":null}', + ], + ), + ( + OpenAITextCompletionsAdapter, + [ + b'{"choices":[{"index":0,"text":"Hi"}]}', + b'{"choices":[{"index":0,"text":"","finish_reason":"stop"}]}', + ], + ), + ], + ids=["chat", "completions"], +) +def test_wire_usage_trailer_reaches_final_metadata(adapter, frames): + trailer = b'{"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9}}' + buffer = ( + b"".join(b"data: " + f + b"\n\n" for f in [*frames, trailer]) + + b"data: [DONE]\n\n" + ) + acc = OpenAISSEAccumulator("qid", stream_all_chunks=False) + chunks = [ + c + for d in adapter.parse_sse_chunk(buffer, len(buffer)) + if (c := acc.add_chunk(d)) + ] + assert [c.response_chunk for c in chunks] == ["Hi"] + md = acc.get_final_output().metadata + assert (md["finish_reason"], md["prompt_tokens"], md["completion_tokens"]) == ( + "stop", + 7, + 2, + ) diff --git a/tests/unit/openai/test_accumulator_contract.py b/tests/unit/openai/test_accumulator_contract.py index 5fc6d4f84..9b454d314 100644 --- a/tests/unit/openai/test_accumulator_contract.py +++ b/tests/unit/openai/test_accumulator_contract.py @@ -22,7 +22,7 @@ import pytest from hypothesis import given, settings from hypothesis import strategies as st -from inference_endpoint.core.types import QueryResult +from inference_endpoint.core.types import QueryResult, TextModelOutput from inference_endpoint.metrics.steady_state_diagnostics import text_after_first_chunk from inference_endpoint.openai.accumulator import OpenAISSEAccumulator from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter @@ -130,6 +130,13 @@ def add_wire(choices): ) assert add_wire([{"delta": {}, "finish_reason": "stop"}]) is None assert add_wire([]) is None + trailer = {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 5}} + assert ( + acc.add_chunk( + OpenAIMsgspecAdapter.decode_sse_message(msgspec.json.encode(trailer)) + ) + is None + ) return acc.get_final_output(), emitted @@ -150,6 +157,7 @@ def _assert_contract( assert output.text_after_first_chunk() == expected_text assert result.metadata["first_chunk"] is (first is None) assert result.metadata["finish_reason"] == "stop" + assert result.metadata["completion_tokens"] == 5 assert result.metadata.get("tool_calls") == (list(full[2]) if full[2] else None) expected_chunks = [] @@ -167,7 +175,7 @@ def _assert_contract( assert decoded.metadata == result.metadata assert decoded.response_output.as_message_parts_after_first_chunk() == tail # Offline diagnostics intentionally omit tools, but share the text boundary. - wire_output = msgspec.json.decode(msgspec.json.encode(output)) + wire_output = msgspec.json.decode(msgspec.json.encode(output), type=TextModelOutput) assert text_after_first_chunk(wire_output) == (tail[1] or "") + tail[0] return result diff --git a/tests/unit/openai/test_completions_adapter.py b/tests/unit/openai/test_completions_adapter.py index 391ab88db..84660f19e 100644 --- a/tests/unit/openai/test_completions_adapter.py +++ b/tests/unit/openai/test_completions_adapter.py @@ -34,15 +34,33 @@ from inference_endpoint.dataset_manager.predefined.livecodebench import ( presets as livecodebench_presets, ) -from inference_endpoint.dataset_manager.transforms import AddStaticColumns, Harmonize +from inference_endpoint.dataset_manager.transforms import ( + AddStaticColumns, + Harmonize, +) from inference_endpoint.openai.accumulator import OpenAISSEAccumulator from inference_endpoint.openai.completions_adapter import OpenAITextCompletionsAdapter from inference_endpoint.openai.types import ( SSEChoice, + SSEDelta, + SSEMessage, TextCompletionSSEChoice, TextCompletionSSEMessage, ) +_COMPLETION = { + "model": "m", + "input_tokens": [1, 2, 3], + "max_tokens": 8, + "min_tokens": 1, +} +# The exact body sent when usage is not requested. +_MAIN_COMPLETION_BODY = { + True: b'{"model":"m","prompt":[1,2,3],"max_tokens":8,"min_tokens":1,"stream":true}', + False: b'{"model":"m","prompt":[1,2,3],"max_tokens":8,"min_tokens":1,' + b'"stream":false}', +} + class TestOpenAITextCompletionsAdapterDatasetTransforms: @pytest.mark.unit @@ -74,7 +92,7 @@ def test_generation_controls_flow_from_model_params_to_request(self): payload = json.loads( OpenAITextCompletionsAdapter.encode_query( - Query(data=data.to_dict(orient="records")[0]) + Query(data=data.to_dict(orient="records")[0]), include_usage=False ) ) assert payload["min_tokens"] == 0 @@ -89,7 +107,7 @@ def test_generation_controls_use_defaults_in_request(self): payload = json.loads( OpenAITextCompletionsAdapter.encode_query( - Query(data=data.to_dict(orient="records")[0]) + Query(data=data.to_dict(orient="records")[0]), include_usage=False ) ) assert payload["min_tokens"] == 1 @@ -100,7 +118,9 @@ class TestOpenAITextCompletionsAdapterEncodeQuery: @pytest.mark.unit def test_encode_with_token_ids(self): query = Query(data={"input_tokens": [1, 2, 3], "model": "m", "max_tokens": 100}) - payload = json.loads(OpenAITextCompletionsAdapter.encode_query(query)) + payload = json.loads( + OpenAITextCompletionsAdapter.encode_query(query, include_usage=False) + ) assert payload["prompt"] == [1, 2, 3] assert payload["model"] == "m" assert payload["max_tokens"] == 100 @@ -109,17 +129,59 @@ def test_encode_with_token_ids(self): def test_encode_missing_input_tokens_raises(self): query = Query(data={"model": "m"}) with pytest.raises(KeyError, match="input_tokens"): - OpenAITextCompletionsAdapter.encode_query(query) + OpenAITextCompletionsAdapter.encode_query(query, include_usage=False) @pytest.mark.unit def test_encode_omits_none_fields(self): query = Query(data={"input_tokens": [10, 20], "model": "m"}) - payload = json.loads(OpenAITextCompletionsAdapter.encode_query(query)) + payload = json.loads( + OpenAITextCompletionsAdapter.encode_query(query, include_usage=False) + ) assert "temperature" not in payload assert "top_p" not in payload assert "min_tokens" not in payload assert "skip_special_tokens" not in payload + @staticmethod + def _body(include_usage: bool, **data) -> dict: + query = Query(data={**_COMPLETION, **data}) + return json.loads( + OpenAITextCompletionsAdapter.encode_query( + query, include_usage=include_usage + ) + ) + + @pytest.mark.unit + def test_streaming_body_requests_usage_when_enabled(self): + body = self._body(True, stream=True) + assert body["stream_options"] == {"include_usage": True} + + @pytest.mark.unit + @pytest.mark.parametrize("stream", [True, False]) + def test_body_matches_main_when_usage_disabled(self, stream): + query = Query(data={**_COMPLETION, "stream": stream}) + assert ( + OpenAITextCompletionsAdapter.encode_query(query, include_usage=False) + == _MAIN_COMPLETION_BODY[stream] + ) + + @pytest.mark.unit + @pytest.mark.parametrize( + "data", [{"stream": False}, {}], ids=["stream-false", "unset"] + ) + def test_non_streaming_body_never_requests_usage(self, data): + assert "stream_options" not in self._body(True, **data) + + @pytest.mark.unit + @pytest.mark.parametrize( + "n, expected", + [(2, "absent"), (2.0, "absent"), (1, {"include_usage": True})], + ) + def test_n_gt_1_row_omits_stream_options(self, n, expected): + body = self._body(True, stream=True, n=n) + assert body["n"] == n + assert body.get("stream_options", "absent") == expected + class TestOpenAITextCompletionsAdapterDecodeResponse: @pytest.mark.unit @@ -157,6 +219,75 @@ def test_decode_empty_choices_raises(self): with pytest.raises(ValueError, match="at least one choice"): OpenAITextCompletionsAdapter.decode_response(response_bytes, "qid-1") + @pytest.mark.unit + def test_decode_non_streaming_copies_usage(self): + body = { + "id": "c", + "created": 1, + "model": "m", + "choices": [{"index": 0, "text": "hi", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, + } + result = OpenAITextCompletionsAdapter.decode_response( + json.dumps(body).encode(), "q" + ) + assert result.metadata == { + "finish_reason": "stop", + "prompt_tokens": 4, + "completion_tokens": 2, + } + + @pytest.mark.unit + def test_completions_body_without_usage_key_decodes(self): + body = { + "id": "c", + "created": 1, + "model": "m", + "choices": [{"index": 0, "text": "hi", "finish_reason": "stop"}], + } + result = OpenAITextCompletionsAdapter.decode_response( + json.dumps(body).encode(), "q" + ) + assert result.response_output.output == "hi" + assert result.metadata == {"finish_reason": "stop"} + + @pytest.mark.unit + @pytest.mark.parametrize( + "usage", [{}, "x", [], {"prompt_tokens": 4.0, "completion_tokens": -2}] + ) + def test_non_streaming_malformed_usage_succeeds(self, usage): + body = { + "id": "c", + "created": 1, + "model": "m", + "choices": [{"index": 0, "text": "hi", "finish_reason": "stop"}], + "usage": usage, + } + result = OpenAITextCompletionsAdapter.decode_response( + json.dumps(body).encode(), "q" + ) + assert result.error is None + assert result.response_output.output == "hi" + assert result.metadata == {"finish_reason": "stop"} + + @pytest.mark.unit + def test_non_streaming_n_gt_1_has_no_counts(self): + body = { + "id": "c", + "created": 1, + "model": "m", + "choices": [ + {"index": 0, "text": "hi", "finish_reason": "stop"}, + {"index": 1, "text": "yo", "finish_reason": "stop"}, + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 4, "total_tokens": 8}, + } + result = OpenAITextCompletionsAdapter.decode_response( + json.dumps(body).encode(), "q" + ) + assert result.response_output.output == "hi" + assert result.metadata == {"finish_reason": "stop"} + class TestOpenAITextCompletionsAdapterDecodeSSE: @pytest.mark.unit @@ -169,12 +300,86 @@ def test_decode_sse_message_returns_sse_choice(self): assert choice.delta.content == "tok" @pytest.mark.unit - def test_decode_sse_empty_choices_returns_empty_choice(self): - msg = TextCompletionSSEMessage(choices=()) - json_bytes = msgspec.json.encode(msg) - choice = OpenAITextCompletionsAdapter.decode_sse_message(json_bytes) - assert isinstance(choice, SSEChoice) - assert choice.delta is None + @pytest.mark.parametrize( + "frame, expected", + [ + pytest.param({"choices": []}, None, id="empty-choices-no-usage"), + pytest.param( + {"choices": [{"index": 0, "text": "tok"}], "usage": None}, + SSEChoice(delta=SSEDelta(content="tok")), + id="null-usage", + ), + pytest.param( + { + "choices": [], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 9, + "total_tokens": 14, + }, + }, + SSEMessage( + usage={ + "prompt_tokens": 5, + "completion_tokens": 9, + "total_tokens": 14, + } + ), + id="usage-trailer", + ), + pytest.param( + { + "choices": [{"index": 0, "text": "tok", "finish_reason": "length"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 9}, + }, + SSEMessage( + choices=( + SSEChoice( + delta=SSEDelta(content="tok"), finish_reason="length" + ), + ), + usage={"prompt_tokens": 5, "completion_tokens": 9}, + ), + id="choice-and-usage", + ), + pytest.param( + {"choices": [{"index": 0, "text": "tok"}], "usage": "x"}, + SSEMessage( + choices=(SSEChoice(delta=SSEDelta(content="tok")),), usage="x" + ), + id="malformed-usage", + ), + pytest.param( + {"choices": [{"index": 1, "text": "tok"}]}, + SSEChoice(delta=SSEDelta(content="tok")), + id="second-choice-index-no-usage", + ), + pytest.param( + {"choices": [{"index": 0, "text": "a"}, {"index": 1, "text": "b"}]}, + SSEChoice(delta=SSEDelta(content="a")), + id="two-choices-no-usage", + ), + pytest.param( + { + "choices": [{"index": 0, "text": "a"}, {"index": 1, "text": "b"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 9}, + }, + SSEMessage( + choices=( + SSEChoice(delta=SSEDelta(content="a")), + SSEChoice(delta=SSEDelta(content="b")), + ), + usage={"prompt_tokens": 5, "completion_tokens": 9}, + ), + id="two-choices-with-usage", + ), + ], + ) + def test_decode_sse_usage_carrier(self, frame, expected): + assert ( + OpenAITextCompletionsAdapter.decode_sse_message(msgspec.json.encode(frame)) + == expected + ) @pytest.mark.unit def test_decode_sse_empty_text_returns_choice_with_empty_content(self): diff --git a/tests/unit/openai/test_msgspec_adapter.py b/tests/unit/openai/test_msgspec_adapter.py index 72e8127f2..2e684252e 100644 --- a/tests/unit/openai/test_msgspec_adapter.py +++ b/tests/unit/openai/test_msgspec_adapter.py @@ -25,7 +25,14 @@ OpenAIMsgspecAdapter, _chat_message_from_dict, ) -from inference_endpoint.openai.types import ChatMessage +from inference_endpoint.openai.types import ( + STREAM_OPTIONS_INCLUDE_USAGE, + ChatMessage, + SSEChoice, + SSEDelta, + SSEMessage, + requests_multiple_choices, +) def _make_reasoning_response( @@ -105,7 +112,7 @@ def test_to_endpoint_request_preserves_tool_calls(): "messages": messages, }, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) encoded = msgspec.json.encode(request) payload = json.loads(encoded) @@ -134,7 +141,7 @@ def test_backward_compat_plain_messages_unchanged(): id="q2", data={"model": "m", "messages": messages}, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) encoded = msgspec.json.encode(request) payload = json.loads(encoded) @@ -263,7 +270,7 @@ def test_seed_in_request_payload(): "seed": 99, }, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) encoded = msgspec.json.encode(request) payload = json.loads(encoded) assert payload.get("seed") == 99 @@ -307,7 +314,7 @@ def test_chat_template_kwargs_forwarded_when_set(): "chat_template_kwargs": {"thinking": True, "preserve_thinking": True}, }, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) payload = json.loads(msgspec.json.encode(request)) assert payload["chat_template_kwargs"] == { "thinking": True, @@ -321,7 +328,7 @@ def test_chat_template_kwargs_omitted_when_unset(): id="q4", data={"model": "m", "messages": [{"role": "user", "content": "hi"}]}, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) payload = json.loads(msgspec.json.encode(request)) assert "chat_template_kwargs" not in payload @@ -356,7 +363,9 @@ def test_dataset_transforms_preserve_chat_template_kwargs_dict(): row = df.to_dict(orient="records")[0] assert row["chat_template_kwargs"] == chat_template_kwargs - request = OpenAIMsgspecAdapter.to_endpoint_request(Query(id="q5", data=row)) + request = OpenAIMsgspecAdapter.to_endpoint_request( + Query(id="q5", data=row), include_usage=False + ) payload = json.loads(msgspec.json.encode(request)) assert payload["chat_template_kwargs"] == chat_template_kwargs @@ -378,8 +387,208 @@ def test_messages_shadow_prompt_in_request(): "prompt": "[deadbeefdeadbeef] shadowed", }, ) - request = OpenAIMsgspecAdapter.to_endpoint_request(query) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=False) payload = json.loads(msgspec.json.encode(request)) assert [m["content"] for m in payload["messages"]] == ["authoritative"] assert "shadowed" not in json.dumps(payload) + + +_CHAT = { + "model": "m", + "system": "be brief", + "prompt": "hi", + "max_completion_tokens": 8, + "temperature": 0.0, +} + + +def _chat_body(include_usage: bool, **data) -> dict: + query = Query(id="q", data={**_CHAT, **data}) + return json.loads( + OpenAIMsgspecAdapter.encode_query(query, include_usage=include_usage) + ) + + +@pytest.mark.unit +def test_streaming_body_requests_usage_when_enabled(): + query = Query(id="q", data={**_CHAT, "stream": True}) + request = OpenAIMsgspecAdapter.to_endpoint_request(query, include_usage=True) + assert request.stream_options is STREAM_OPTIONS_INCLUDE_USAGE + assert _chat_body(True, stream=True)["stream_options"] == {"include_usage": True} + + +# The exact body sent when usage is not requested. +_MAIN_CHAT_BODY = { + True: b'{"model":"m","messages":[{"role":"system","content":"be brief"},' + b'{"role":"user","content":"hi"}],"temperature":0.0,' + b'"max_completion_tokens":8,"stream":true}', + False: b'{"model":"m","messages":[{"role":"system","content":"be brief"},' + b'{"role":"user","content":"hi"}],"temperature":0.0,' + b'"max_completion_tokens":8,"stream":false}', +} + + +@pytest.mark.unit +@pytest.mark.parametrize("stream", [True, False]) +def test_body_matches_main_when_usage_disabled(stream): + query = Query(id="q", data={**_CHAT, "stream": stream}) + assert ( + OpenAIMsgspecAdapter.encode_query(query, include_usage=False) + == _MAIN_CHAT_BODY[stream] + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("data", [{"stream": False}, {}], ids=["stream-false", "unset"]) +def test_non_streaming_body_never_requests_usage(data): + assert "stream_options" not in _chat_body(True, **data) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "n, expected", + [(2, "absent"), (2.0, "absent"), (1, {"include_usage": True})], +) +def test_n_gt_1_row_omits_stream_options(n, expected): + body = _chat_body(True, stream=True, n=n) + assert body["n"] == n + assert body.get("stream_options", "absent") == expected + + +@pytest.mark.unit +@pytest.mark.parametrize( + "n, expected", + [ + (2, True), + (2.5, True), + (1, False), + (1.0, False), + (0, False), + (None, False), + (True, False), + ("2", False), + (float("nan"), False), + ], +) +def test_requests_multiple_choices(n, expected): + assert requests_multiple_choices(n) is expected + + +_USAGE = {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} +_HI = SSEChoice(delta=SSEDelta(content="Hi")) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "frame, expected", + [ + pytest.param({"choices": [{"delta": {"content": "Hi"}}]}, _HI, id="choice"), + pytest.param( + {"choices": [{"delta": {"content": "Hi"}}], "usage": None}, + _HI, + id="openai-null-usage", + ), + pytest.param( + {"choices": [], "usage": _USAGE}, + SSEMessage(usage=_USAGE), + id="usage-trailer", + ), + pytest.param( + {"choices": [{"delta": {"content": "Hi"}}], "usage": _USAGE}, + SSEMessage(choices=(_HI,), usage=_USAGE), + id="litellm", + ), + pytest.param( + {"choices": [], "usage": {"prompt_tokens": None, "completion_tokens": 3}}, + SSEMessage(usage={"prompt_tokens": None, "completion_tokens": 3}), + id="null-prompt-tokens", + ), + pytest.param( + {"choices": [], "usage": {}}, SSEMessage(usage={}), id="empty-usage-object" + ), + pytest.param( + {"choices": [{"delta": {"content": "Hi"}}], "usage": "x"}, + SSEMessage(choices=(_HI,), usage="x"), + id="malformed-usage", + ), + pytest.param( + {"choices": [{"index": 1, "delta": {"content": "Hi"}}]}, + _HI, + id="second-choice-index-no-usage", + ), + pytest.param( + {"choices": [{"delta": {"content": "Hi"}}, {"index": 1, "delta": {}}]}, + _HI, + id="two-choices-no-usage", + ), + pytest.param( + { + "choices": [{"delta": {"content": "Hi"}}, {"index": 1, "delta": {}}], + "usage": _USAGE, + }, + SSEMessage(choices=(_HI, SSEChoice(delta=SSEDelta())), usage=_USAGE), + id="two-choices-with-usage", + ), + pytest.param({"choices": []}, None, id="sglang-empty-choices"), + pytest.param({"choices": [], "usage": None}, None, id="empty-null-usage"), + ], +) +def test_decode_sse_message_usage_carrier(frame, expected): + assert ( + OpenAIMsgspecAdapter.decode_sse_message(msgspec.json.encode(frame)) == expected + ) + + +@pytest.mark.unit +def test_decode_sse_message_common_frame_is_the_bare_choice(): + decoded = OpenAIMsgspecAdapter.decode_sse_message( + b'{"choices":[{"index":0,"delta":{"content":"Hi"}}],"usage":null}' + ) + assert type(decoded) is SSEChoice + + +def _chat_body_with(usage: object, n_choices: int = 1) -> bytes: + body = json.loads(_make_reasoning_response("reasoning", "r")) + body["choices"] = [{**body["choices"][0], "index": i} for i in range(n_choices)] + body["usage"] = usage + return json.dumps(body).encode() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "usage", [{}, "x", [], {"prompt_tokens": 7.0, "completion_tokens": -3}] +) +def test_non_streaming_malformed_usage_succeeds(usage): + result = OpenAIMsgspecAdapter.decode_response(_chat_body_with(usage), "q1") + assert result.error is None + assert result.response_output.output == "answer" + assert result.metadata["finish_reason"] == "stop" + assert "prompt_tokens" not in result.metadata + assert "completion_tokens" not in result.metadata + + +@pytest.mark.unit +def test_non_streaming_n_gt_1_has_no_counts(): + result = OpenAIMsgspecAdapter.decode_response( + _chat_body_with(_USAGE, n_choices=2), "q1" + ) + assert result.response_output.output == "answer" + assert "prompt_tokens" not in result.metadata + assert "completion_tokens" not in result.metadata + + +@pytest.mark.unit +@pytest.mark.parametrize( + "usage, expected", + [(_USAGE, {"prompt_tokens": 7, "completion_tokens": 3}), (None, {})], + ids=["usage", "no-usage"], +) +def test_decode_response_copies_usage(usage, expected): + body = json.loads(_make_reasoning_response("reasoning", "r")) + if usage is not None: + body["usage"] = usage + md = OpenAIMsgspecAdapter.decode_response(json.dumps(body).encode(), "q1").metadata + assert { + k: md[k] for k in ("prompt_tokens", "completion_tokens") if k in md + } == expected diff --git a/tests/unit/openai/test_openai_adapter.py b/tests/unit/openai/test_openai_adapter.py index 826727634..5b2414946 100644 --- a/tests/unit/openai/test_openai_adapter.py +++ b/tests/unit/openai/test_openai_adapter.py @@ -60,7 +60,7 @@ def test_tool_definitions_forwarded(): "stream": False, }, ) - encoded = OpenAIAdapter.encode_query(query) + encoded = OpenAIAdapter.encode_query(query, include_usage=False) payload = json.loads(encoded) assert "tools" in payload @@ -88,7 +88,7 @@ def test_tool_use_messages_roundtrip(): "stream": False, }, ) - encoded = OpenAIAdapter.encode_query(query) + encoded = OpenAIAdapter.encode_query(query, include_usage=False) payload = json.loads(encoded) msgs = payload["messages"] @@ -140,7 +140,7 @@ def test_no_tools_key_when_absent(): "stream": False, }, ) - encoded = OpenAIAdapter.encode_query(query) + encoded = OpenAIAdapter.encode_query(query, include_usage=False) payload = json.loads(encoded) # Pydantic model_dump includes None fields; tools must be None when not supplied diff --git a/tests/unit/sglang/test_sglang_usage.py b/tests/unit/sglang/test_sglang_usage.py new file mode 100644 index 000000000..83f4e1348 --- /dev/null +++ b/tests/unit/sglang/test_sglang_usage.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for SGLang meta_info decoding and server-reported token counts.""" + +import json + +import msgspec +import pytest +from inference_endpoint.core.types import ( + PROMPT_TOKENS_KEY, + TOKEN_HDR_HIGH, + server_usage_count, +) +from inference_endpoint.sglang import accumulator as sglang_accumulator +from inference_endpoint.sglang.accumulator import SGLangSSEAccumulator +from inference_endpoint.sglang.adapter import SGLangGenerateAdapter as A +from inference_endpoint.sglang.types import SGLangSSEDelta + +_STOP = {"type": "stop", "matched": 200002} + + +def _frame(text, ids, ct, *, fr=None, e2e=None, retr=("num_retractions", 0)) -> bytes: + meta = { + "id": "a1b2", + "finish_reason": fr, + "prompt_tokens": 7, + "weight_version": "default", + retr[0]: retr[1], + "completion_tokens": ct, + "cached_tokens": 0, + } + if e2e is not None: + meta["e2e_latency"] = e2e + body = {"text": text, "meta_info": meta} | ( + {} if ids is None else {"output_ids": ids} + ) + return json.dumps(body).encode() + + +# Current SGLang: finish_reason null until the last frame; e2e_latency only on it; +# the stop token counts but adds no text. +_STREAM = [ + _frame("Hello", [9906], 1), + _frame("Hello world", [1917], 2), + _frame("Hello world", [], 3, fr=_STOP, e2e=0.042), +] + + +@pytest.mark.unit +class TestSGLangMetaInfo: + def test_unfinished_frame_with_null_finish_reason_decodes(self): + assert A.decode_sse_message(_STREAM[0]) == SGLangSSEDelta( + text="Hello", + token_delta=[9906], + total_completion_tokens=1, + has_retractions=False, + prompt_tokens=7, + ) + + def test_minimal_meta_info_without_output_ids_decodes(self): + d = A.decode_sse_message( + b'{"text":"Hi","meta_info":{"prompt_tokens":7,"completion_tokens":1}}' + ) + assert (d.token_delta, d.prompt_tokens, d.has_retractions) == ([], 7, False) + + @pytest.mark.parametrize( + "retr,expected", + [ + (("num_retractions", 2), True), + (("num_retractions", 0), False), + (("total_retractions", 1), True), + (("total_retractions", 0), False), + ], + ) + def test_either_retraction_name_is_accepted(self, retr, expected): + frame = _frame("x", [1], 1, fr=_STOP, e2e=0.1, retr=retr) + assert A.decode_sse_message(frame).has_retractions is expected + + def test_completion_tokens_stays_required(self): + body = json.loads(_STREAM[0]) + del body["meta_info"]["completion_tokens"] + with pytest.raises(msgspec.ValidationError, match="completion_tokens"): + A.decode_sse_message(json.dumps(body).encode()) + + @pytest.mark.parametrize( + "prompt_tokens", ["absent", None, 7.0, "7", True, -1], ids=repr + ) + def test_unusable_prompt_tokens_decode_as_absent(self, prompt_tokens): + body = json.loads(_STREAM[0]) + if prompt_tokens == "absent": + del body["meta_info"]["prompt_tokens"] + else: + body["meta_info"]["prompt_tokens"] = prompt_tokens + frame = json.dumps(body).encode() + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + acc.add_chunk(A.decode_sse_message(frame)) + for md in ( + acc.get_final_output().metadata, + A.decode_response(frame, "q").metadata, + ): + assert "prompt_tokens" not in md and md["completion_tokens"] == 1 + + +@pytest.mark.unit +class TestSGLangServerTokenCounts: + def test_stream_final_output_carries_server_usage(self): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + chunks = [c for f in _STREAM if (c := acc.add_chunk(A.decode_sse_message(f)))] + r = acc.get_final_output() + assert r.response_output.output == "Hello world" + assert (r.metadata["prompt_tokens"], r.metadata["completion_tokens"]) == (7, 3) + assert r.metadata["token_ids"] == [9906, 1917] and "n_tokens" not in r.metadata + assert len(chunks) == 2 and all("n_tokens" not in c.metadata for c in chunks) + + def test_prompt_tokens_captured_when_completion_count_never_changes(self): + acc = SGLangSSEAccumulator("q", stream_all_chunks=False) + frame = _frame("", [], 0, fr={"type": "abort"}, e2e=0.01) + assert acc.add_chunk(A.decode_sse_message(frame)) is None + meta = acc.get_final_output().metadata + assert (meta["prompt_tokens"], meta["completion_tokens"]) == (7, 0) + + def test_sglang_frame_without_prompt_tokens_keeps_earlier_value(self): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + later = json.loads(_STREAM[2]) + del later["meta_info"]["prompt_tokens"] + for frame in (_STREAM[0], _STREAM[1], json.dumps(later).encode()): + acc.add_chunk(A.decode_sse_message(frame)) + meta = acc.get_final_output().metadata + assert (meta["prompt_tokens"], meta["completion_tokens"]) == (7, 3) + + @pytest.mark.parametrize("later", [7.0, "8", -1, None], ids=repr) + def test_later_unusable_prompt_tokens_keep_the_earlier_count(self, later): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + acc.add_chunk(A.decode_sse_message(_STREAM[0])) + for raw in _STREAM[1:]: + body = json.loads(raw) + body["meta_info"]["prompt_tokens"] = later + acc.add_chunk(A.decode_sse_message(json.dumps(body).encode())) + meta = acc.get_final_output().metadata + assert (meta["prompt_tokens"], meta["completion_tokens"]) == (7, 3) + + @pytest.mark.parametrize( + "sent, expected", [((7, 7, 8), 8), (("7", 7, 7), 7)], ids=["changed", "fixed"] + ) + def test_a_changed_usable_prompt_tokens_is_taken(self, sent, expected): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + for raw, prompt_tokens in zip(_STREAM, sent, strict=True): + body = json.loads(raw) + body["meta_info"]["prompt_tokens"] = prompt_tokens + acc.add_chunk(A.decode_sse_message(json.dumps(body).encode())) + assert acc.get_final_output().metadata["prompt_tokens"] == expected + + def test_prompt_tokens_validated_once_per_distinct_value(self, monkeypatch): + seen: list[object] = [] + + def counting(value: object, key: str) -> int | None: + if key == PROMPT_TOKENS_KEY: + seen.append(value) + return server_usage_count(value, key) + + monkeypatch.setattr(sglang_accumulator, "server_usage_count", counting) + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + for frame in (*_STREAM, _frame("Hello world!", [0], 4)): + acc.add_chunk(A.decode_sse_message(frame)) + assert acc.get_final_output().metadata["prompt_tokens"] == 7 + assert seen == [7] + + def test_stream_without_prompt_tokens_still_reports_completion(self): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + for raw in _STREAM: + body = json.loads(raw) + del body["meta_info"]["prompt_tokens"] + acc.add_chunk(A.decode_sse_message(json.dumps(body).encode())) + meta = acc.get_final_output().metadata + assert "prompt_tokens" not in meta and meta["completion_tokens"] == 3 + + def test_no_decoded_frame_reports_no_counts(self): + meta = ( + SGLangSSEAccumulator("q", stream_all_chunks=True) + .get_final_output() + .metadata + ) + assert "prompt_tokens" not in meta and "completion_tokens" not in meta + + def test_non_streaming_response_carries_server_usage(self): + r = A.decode_response( + _frame("Hello world", [9906, 1917], 3, fr=_STOP, e2e=0.042), "q" + ) + assert r.metadata == { + "token_ids": [9906, 1917], + "prompt_tokens": 7, + "completion_tokens": 3, + } + + +@pytest.mark.unit +class TestSGLangMalformedCountWarning: + def test_streaming_malformed_prompt_tokens_warns_once( + self, usage_warning_unlatched + ): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + for raw, prompt_tokens in zip(_STREAM, (7.0, 7.0, "7"), strict=True): + body = json.loads(raw) + body["meta_info"]["prompt_tokens"] = prompt_tokens + acc.add_chunk(A.decode_sse_message(json.dumps(body).encode())) + assert "prompt_tokens" not in acc.get_final_output().metadata + (record,) = usage_warning_unlatched.records + assert "'prompt_tokens': 7.0 (type 'float')" in record.getMessage() + + def test_non_streaming_malformed_completion_tokens_warns( + self, usage_warning_unlatched + ): + md = A.decode_response(_frame("Hi", [1], -1, fr=_STOP, e2e=0.1), "q").metadata + assert md["prompt_tokens"] == 7 and "completion_tokens" not in md + (record,) = usage_warning_unlatched.records + assert "'completion_tokens': -1 (type 'int')" in record.getMessage() + + def test_stream_with_out_of_range_completion_tokens_warns( + self, usage_warning_unlatched + ): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + acc.add_chunk(A.decode_sse_message(_frame("Hi", [1], TOKEN_HDR_HIGH + 1))) + meta = acc.get_final_output().metadata + assert meta["prompt_tokens"] == 7 and "completion_tokens" not in meta + (record,) = usage_warning_unlatched.records + assert f"'completion_tokens': {TOKEN_HDR_HIGH + 1}" in record.getMessage() + + def test_well_formed_stream_never_warns(self, usage_warning_unlatched): + acc = SGLangSSEAccumulator("q", stream_all_chunks=True) + for frame in _STREAM: + acc.add_chunk(A.decode_sse_message(frame)) + acc.get_final_output() + assert usage_warning_unlatched.records == [] diff --git a/tests/unit/videogen/test_adapter.py b/tests/unit/videogen/test_adapter.py index 810015042..8c68d1c16 100644 --- a/tests/unit/videogen/test_adapter.py +++ b/tests/unit/videogen/test_adapter.py @@ -31,13 +31,13 @@ class TestVideoGenAdapter: def test_encode_query_produces_valid_json(self): query = Query(id="q1", data={"prompt": "a golden retriever running"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["prompt"] == "a golden retriever running" def test_encode_query_default_response_format_is_video_path(self): """Perf mode default — server saves to Lustre, returns path only.""" query = Query(id="q1", data={"prompt": "ocean waves"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["response_format"] == "video_path" def test_encode_query_accuracy_mode_requests_video_bytes(self): @@ -45,12 +45,12 @@ def test_encode_query_accuracy_mode_requests_video_bytes(self): query = Query( id="q1", data={"prompt": "ocean waves", "response_format": "video_bytes"} ) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["response_format"] == "video_bytes" def test_encode_query_uses_mlperf_defaults(self): query = Query(id="q1", data={"prompt": "ocean waves"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["fps"] == 16 assert payload["seconds"] == pytest.approx(5.0) assert payload["size"] == "720x1280" @@ -64,43 +64,43 @@ def test_encode_query_allows_override_via_data(self): id="q1", data={"prompt": "test", "seed": 99, "fps": 24, "guidance_scale": 6.0}, ) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["seed"] == 99 assert payload["fps"] == 24 assert payload["guidance_scale"] == pytest.approx(6.0) def test_encode_query_includes_negative_prompt(self): query = Query(id="q1", data={"prompt": "test", "negative_prompt": "blurry"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["negative_prompt"] == "blurry" def test_encode_query_omits_negative_prompt_when_absent(self): """exclude_none=True so server can apply its own default.""" query = Query(id="q1", data={"prompt": "test"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert "negative_prompt" not in payload def test_encode_query_includes_latent_path(self): query = Query( id="q1", data={"prompt": "test", "latent_path": "/lustre/fixed_latent.pt"} ) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert payload["latent_path"] == "/lustre/fixed_latent.pt" def test_encode_query_omits_latent_path_when_absent(self): query = Query(id="q1", data={"prompt": "test"}) - payload = json.loads(VideoGenAdapter.encode_query(query)) + payload = json.loads(VideoGenAdapter.encode_query(query, include_usage=False)) assert "latent_path" not in payload def test_encode_query_missing_prompt_raises(self): query = Query(id="q1", data={"seed": 42}) with pytest.raises(KeyError): - VideoGenAdapter.encode_query(query) + VideoGenAdapter.encode_query(query, include_usage=False) def test_encode_query_rejects_stream_true(self): query = Query(id="q1", data={"prompt": "test", "stream": True}) with pytest.raises(ValueError, match="non-streaming"): - VideoGenAdapter.encode_query(query) + VideoGenAdapter.encode_query(query, include_usage=False) def test_decode_response_returns_video_bytes_in_metadata(self): resp = VideoPayloadResponse( From 90e7a6faf819770b3b51305fd4afa91de3190411 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:02 -0500 Subject: [PATCH 05/26] test(testing): stub servers report token usage when the request asks The echo, max-throughput and variable-throughput servers follow vLLM: - a streaming response carries a usage frame only when the request sets stream_options.include_usage, after the finish frame and before [DONE]; - a non-streaming body always carries usage. Roofline numbers are compared across versions, so the stream layout is pinned byte for byte by tests. The echo server also serves /v1/completions and rejects stream_options without stream. Fixtures: mock_http_echo_server_factory(report_usage=False) models a server that sends no usage, and usage_warning_unlatched resets the once-per-process malformed-usage warning. tests/test_helpers.py adds parse_sse_frames, make_report and DEFAULT_SOURCE, which the later layers' tests use. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- src/inference_endpoint/testing/echo_server.py | 212 ++++++++--- .../testing/max_throughput_server.py | 151 +++++++- .../testing/variable_throughput_server.py | 24 +- tests/conftest.py | 25 +- .../integration/test_variable_server_usage.py | 64 ++++ tests/test_helpers.py | 100 +++++ tests/unit/testing/__init__.py | 14 + tests/unit/testing/test_echo_server_usage.py | 170 +++++++++ tests/unit/testing/test_stub_servers.py | 346 ++++++++++++++++++ 9 files changed, 1035 insertions(+), 71 deletions(-) create mode 100644 tests/integration/test_variable_server_usage.py create mode 100644 tests/unit/testing/__init__.py create mode 100644 tests/unit/testing/test_echo_server_usage.py create mode 100644 tests/unit/testing/test_stub_servers.py diff --git a/src/inference_endpoint/testing/echo_server.py b/src/inference_endpoint/testing/echo_server.py index 9cb7847b0..981e1ed78 100644 --- a/src/inference_endpoint/testing/echo_server.py +++ b/src/inference_endpoint/testing/echo_server.py @@ -25,6 +25,7 @@ import uuid from abc import abstractmethod from collections.abc import Awaitable, Callable +from typing import Any from aiohttp import web @@ -36,6 +37,21 @@ RequestHandler = Callable[[web.Request], web.Response | Awaitable[web.Response]] +def _include_usage(payload: dict[str, Any]) -> bool: + """Whether a request asks for a trailing usage frame in its stream. + + Follows vLLM's truthiness rule: any truthy ``stream_options`` on a + non-streaming request is rejected, so ``stream_options: {}`` is accepted. + + Raises: + ValueError: If ``stream_options`` is set without ``stream``. + """ + stream_options = payload.get("stream_options") + if stream_options and not payload.get("stream"): + raise ValueError("Stream options can only be defined when `stream=True`.") + return bool((stream_options or {}).get("include_usage")) + + class HTTPServer: @property @abstractmethod @@ -59,11 +75,14 @@ def __init__( port: int = 0, max_osl: int | None = None, request_handler: RequestHandler | None = None, + report_usage: bool = True, ): self.host = host self.port = port # If 0, will auto-assign available port self.max_osl = max_osl self._request_handler = request_handler + # False models a server that never reports token usage. + self.report_usage = report_usage self._actual_port = None # Store the actual port after binding self.app = None @@ -93,6 +112,34 @@ def get_max_osl(self): """ return self.max_osl + def _usage( + self, prompt_tokens: int, completion_tokens: int + ) -> dict[str, int] | None: + """Build an OpenAI usage object, or None when usage is not reported.""" + if not self.report_usage: + return None + return { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + + def _apply_max_osl(self, text: str) -> str: + """Truncate or repeat ``text`` to exactly ``max_osl`` characters.""" + if self.max_osl is not None and len(text) > 0: + if len(text) > self.max_osl: + text = text[: self.max_osl] + else: + text = text * (self.max_osl // len(text) + 1) + text = text[: self.max_osl] + return text + + @staticmethod + async def _read_json_payload(request: web.Request) -> dict[str, Any]: + if request.content_type == "application/json": + return await request.json() + return json.loads(await request.text()) + def get_response(self, request: str) -> str: """ Return the input request string as the response. @@ -177,19 +224,26 @@ async def _handle_streaming_response( self, id: str, request: web.Request, - completion_request: CreateChatCompletionRequest, + model: str, content: str, + *, + prompt_tokens: int, + include_usage: bool, + text_completion: bool = False, ) -> web.StreamResponse: """ - Handle a streaming chat completion response using Server-Sent Events (SSE). + Handle a streaming completion response using Server-Sent Events (SSE). - Streams the response word by word, simulating an OpenAI-compatible chat completion endpoint. Creates an SSE-formatted stream with individual word chunks and a final completion marker. + Streams the response word by word, simulating an OpenAI-compatible endpoint. Creates an SSE-formatted stream with individual word chunks, a final finish chunk, an optional usage chunk (vLLM order) and the ``[DONE]`` marker. Args: id (str): Unique identifier for the streaming response. request (web.Request): The original web request. - completion_request (CreateChatCompletionRequest): The parsed chat completion request. + model (str): Model name echoed in every frame. content (str): The content to be streamed. + prompt_tokens (int): Prompt token count reported in the usage frame. + include_usage (bool): Whether the request asked for a usage frame. The frame is also suppressed when ``report_usage`` is False. + text_completion (bool): Frame as ``/v1/completions`` (``text_completion`` objects with ``text`` choices) instead of chat completion chunks. Returns: web.StreamResponse: A streaming HTTP response with chunked SSE data. @@ -197,6 +251,25 @@ async def _handle_streaming_response( Raises: Any underlying exceptions that occur during streaming, which will be logged. """ + object_type = "text_completion" if text_completion else "chat.completion.chunk" + + def frame(choices: list[dict[str, Any]], **extra: Any) -> bytes: + chunk = { + "id": id, + "object": object_type, + "created": int(time.time()), + "model": model, + "choices": choices, + **extra, + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + def choice(text: str, finish_reason: str | None) -> dict[str, Any]: + if text_completion: + return {"index": 0, "text": text, "finish_reason": finish_reason} + delta = {"content": text} if text else {} + return {"index": 0, "delta": delta, "finish_reason": finish_reason} + try: response = web.StreamResponse( status=200, @@ -213,36 +286,20 @@ async def _handle_streaming_response( words = raw_response.split() if raw_response else [] # Send chunks + completion_tokens = 0 for i, word in enumerate(words): # Add space before word (except first) chunk_content = f" {word}" if i > 0 else word - chunk_data = { - "id": id, - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": str(completion_request.model.root), - "choices": [ - { - "index": 0, - "delta": {"content": chunk_content}, - "finish_reason": None, - } - ], - } - - await response.write(f"data: {json.dumps(chunk_data)}\n\n".encode()) + await response.write(frame([choice(chunk_content, None)])) + completion_tokens += len(chunk_content) # Send final chunk with finish_reason - final_chunk = { - "id": id, - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": str(completion_request.model.root), - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - } + await response.write(frame([choice("", "stop")])) - await response.write(f"data: {json.dumps(final_chunk)}\n\n".encode()) + usage = self._usage(prompt_tokens, completion_tokens) + if include_usage and usage is not None: + await response.write(frame([], usage=usage)) await response.write(b"data: [DONE]\n\n") return response except Exception as e: @@ -270,11 +327,8 @@ async def _handle_echo_chat_completions_request( # Get request body try: - if request.content_type == "application/json": - json_payload = await request.json() - else: - raw_payload = await request.text() - json_payload = json.loads(raw_payload) + json_payload = await self._read_json_payload(request) + include_usage = _include_usage(json_payload) completion_request = CreateChatCompletionRequest(**json_payload) raw_request = "" if completion_request.messages and len(completion_request.messages) > 0: @@ -291,24 +345,19 @@ async def _handle_echo_chat_completions_request( self.logger.debug( f"Content of request: {raw_request} - response : {raw_response}" ) - if self.max_osl is not None and len(raw_response) > 0: - # if max_osl is specified, it can be either larger or smaller than the length of the prompt - # if max_osl is larger, we can repeate the prompt until we reach the max_osl - if len(raw_response) > self.max_osl: - raw_response = raw_response[: self.max_osl] - # if max_osl is smaller, we can truncate the prompt - else: - raw_response = raw_response * ( - self.max_osl // len(raw_response) + 1 - ) - raw_response = raw_response[: self.max_osl] + raw_response = self._apply_max_osl(raw_response) # Check if this is a streaming request self.logger.debug(f"Streaming response: {completion_request.stream}\n") if completion_request.stream is True: # Return SSE (Server-Sent Events) format for streaming return await self._handle_streaming_response( - id, request, completion_request, raw_response + id, + request, + str(completion_request.model.root), + raw_response, + prompt_tokens=len(raw_request), + include_usage=include_usage, ) else: # Non-streaming: return QueryResult @@ -320,6 +369,9 @@ async def _handle_echo_chat_completions_request( mode="json" ) echo_response["id"] = id + echo_response["usage"] = self._usage( + len(raw_request), len(raw_response) + ) self.logger.debug(f"Echo response (non-streaming): {echo_response}") return web.json_response(echo_response, status=200) @@ -331,6 +383,68 @@ async def _handle_echo_chat_completions_request( status=400, ) + async def _handle_echo_completions_request( + self, request: web.Request + ) -> web.StreamResponse: + """ + Handle an incoming HTTP request to the OpenAI text completions endpoint (``/v1/completions``). + + Echoes the prompt back as the completion text: a ``str`` prompt is returned as-is and a token-ID list is joined with spaces. Returns either an SSE stream or a JSON ``text_completion`` body, with the same usage and ``stream_options`` behavior as the chat completions handler. + + Raises: + ValueError: If the request payload has no prompt. + """ + custom = await self._dispatch(request) + if custom is not None: + return custom + + try: + json_payload = await self._read_json_payload(request) + include_usage = _include_usage(json_payload) + prompt = json_payload.get("prompt") + if prompt is None: + raise ValueError("Request must contain a prompt") + raw_request = ( + prompt + if isinstance(prompt, str) + else " ".join(str(token) for token in prompt) + ) + id = json_payload.get("id", str(uuid.uuid4())) + model = str(json_payload.get("model", "")) + raw_response = self._apply_max_osl(self.get_response(raw_request)) + + if json_payload.get("stream") is True: + return await self._handle_streaming_response( + id, + request, + model, + raw_response, + prompt_tokens=len(prompt), + include_usage=include_usage, + text_completion=True, + ) + return web.json_response( + { + "id": id, + "object": "text_completion", + "created": int(time.time()), + "model": model, + "choices": [ + {"index": 0, "text": raw_response, "finish_reason": "stop"} + ], + "usage": self._usage(len(prompt), len(raw_response)), + }, + status=200, + ) + + except Exception as e: + # A catch-all exception handler to help debug the issue without bringing down the server + self.logger.error(f"Error handling completions request: {e}") + return web.json_response( + {"error": f"error encountered : {str(e)}"}, + status=400, + ) + def _run_server(self): """Run the server in a separate thread.""" self._loop = asyncio.new_event_loop() @@ -346,6 +460,7 @@ def _register_routes(self, app: "web.Application") -> None: app.router.add_post( "/v1/chat/completions", self._handle_echo_chat_completions_request ) + app.router.add_post("/v1/completions", self._handle_echo_completions_request) app.router.add_post("/echo", self._handle_echo_request) async def _start_server(self): @@ -437,6 +552,11 @@ def create_parser() -> argparse.ArgumentParser: "--host", type=str, help="hostname/address to bind to", default="127.0.0.1" ) parser.add_argument("--port", type=int, help="port to bind to", default=12345) + parser.add_argument( + "--no-usage", + action="store_true", + help="never report token usage (models a server without usage)", + ) return parser @@ -467,7 +587,9 @@ def main(): server = None try: - server = EchoServer(host=args.host, port=args.port) + server = EchoServer( + host=args.host, port=args.port, report_usage=not args.no_usage + ) server.start() # Wait for the server thread to finish diff --git a/src/inference_endpoint/testing/max_throughput_server.py b/src/inference_endpoint/testing/max_throughput_server.py index 12f3267d6..53459a73d 100644 --- a/src/inference_endpoint/testing/max_throughput_server.py +++ b/src/inference_endpoint/testing/max_throughput_server.py @@ -35,6 +35,7 @@ import multiprocessing.sharedctypes import multiprocessing.synchronize import os +import re import signal import socket import sys @@ -68,6 +69,10 @@ def _counter_read(counter: _SyncInt | None) -> int: return 0 +# Prompt tokens every stub response reports in its usage object. +STUB_PROMPT_TOKENS = 10 + + def _chunked(data: bytes) -> bytes: """HTTP chunked transfer encoding.""" return f"{len(data):x}\r\n".encode() + data + b"\r\n" if data else b"0\r\n\r\n" @@ -94,19 +99,94 @@ def _sse_event( ) +def _sse_usage_event(model: str, created: int, completion_tokens: int) -> bytes: + """Encode the OpenAI ``stream_options.include_usage`` frame (empty ``choices``).""" + return ( + b"data: " + + msgspec.json.encode( + { + "id": "r", + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [], + "usage": { + "prompt_tokens": STUB_PROMPT_TOKENS, + "completion_tokens": completion_tokens, + "total_tokens": completion_tokens + STUB_PROMPT_TOKENS, + }, + } + ) + + b"\n\n" + ) + + +def build_stream_tail( + model: str, created: int, completion_tokens: int, *, include_usage: bool +) -> bytes: + """Build the end of an SSE stream, ready to write after the content chunks. + + The finish frame, then (only when ``include_usage``) the usage frame, then + the ``[DONE]`` sentinel (vLLM's order), then the zero-length HTTP chunk that + terminates the chunked body. With the usage frame each frame is its own HTTP + chunk; without it the finish frame and ``[DONE]`` share one chunk, the + layout of the usage-free stream. + + Args: + model: Model name in the response JSON. + created: ``created`` timestamp shared by every frame of the response. + completion_tokens: Completion tokens reported in the usage frame. The + stubs emit one character per token, so this is the content length. + include_usage: Whether the request asked for the usage frame + (``stream_options.include_usage``). + """ + finish = _sse_event(model, created, {}, "stop") + done = b"data: [DONE]\n\n" + if include_usage: + usage = _sse_usage_event(model, created, completion_tokens) + chunks = _chunked(finish) + _chunked(usage) + _chunked(done) + else: + chunks = _chunked(finish + done) + return chunks + b"0\r\n\r\n" + + +_INCLUDE_USAGE_KEY = b'"include_usage"' +_COLON_TRUE = re.compile(rb"\s*:\s*true") + + +def requests_usage(body: bytes) -> bool: + """Whether a request body sets ``stream_options.include_usage`` to true. + + A byte search, not a JSON decode: this runs once per request on the + roofline path. The key is searched from the end because the in-repo + encoders emit ``stream_options`` after the (large) messages or prompt. A + quoted ``"include_usage"`` inside a prompt is escaped in the JSON body, so + it never matches. + """ + i = body.rfind(_INCLUDE_USAGE_KEY) + return i >= 0 and _COLON_TRUE.match(body, i + len(_INCLUDE_USAGE_KEY)) is not None + + def build_streaming_response( - output_length: int = 160, stream_interval: int = 1, model: str = "max-tp" + output_length: int = 160, + stream_interval: int = 1, + model: str = "max-tp", + *, + include_usage: bool, ) -> bytes: """Build complete SSE streaming response (chunked encoding). *stream_interval* controls how many characters are packed into each SSE event. Total SSE content events = ``ceil(output_length / stream_interval)``. - Each event is sent in its own HTTP chunked frame. + Each event is sent in its own HTTP chunked frame, followed by the finish + frame, the usage frame when ``include_usage`` and ``[DONE]`` (see + :func:`build_stream_tail`). Args: output_length: Total characters in the response. stream_interval: Characters per SSE event (1 = finest granularity). model: Model name in the response JSON. + include_usage: Whether to send the usage frame. """ headers = ( b"HTTP/1.1 200 OK\r\n" @@ -124,10 +204,9 @@ def build_streaming_response( event = _sse_event(model, created, {"content": "x" * chars}) chunks.append(_chunked(event)) - # Finish event + [DONE] sentinel in one final chunk - finish = _sse_event(model, created, {}, "stop") - chunks.append(_chunked(finish + b"data: [DONE]\n\n")) - chunks.append(b"0\r\n\r\n") + chunks.append( + build_stream_tail(model, created, output_length, include_usage=include_usage) + ) return headers + b"".join(chunks) @@ -159,9 +238,9 @@ def build_non_streaming_response( } ], "usage": { - "prompt_tokens": 10, + "prompt_tokens": STUB_PROMPT_TOKENS, "completion_tokens": output_length, - "total_tokens": output_length + 10, + "total_tokens": output_length + STUB_PROMPT_TOKENS, }, "system_fingerprint": None, } @@ -179,6 +258,8 @@ def build_response( stream: bool = False, stream_interval: int = 1, model: str = "max-tp", + *, + include_usage: bool, ) -> bytes: """Build a pre-compiled HTTP response. @@ -187,9 +268,13 @@ def build_response( stream: SSE streaming (True) or single JSON (False). stream_interval: Characters per SSE event (streaming only). model: Model name in the response JSON. + include_usage: Whether to send the usage frame (streaming only; a JSON + response always carries ``usage``). """ if stream: - return build_streaming_response(output_length, stream_interval, model) + return build_streaming_response( + output_length, stream_interval, model, include_usage=include_usage + ) return build_non_streaming_response(output_length, model) @@ -232,15 +317,24 @@ class ServerProtocol(asyncio.Protocol): __slots__ = ( "transport", "response", - "response_size", + "response_no_usage", "_parser", "_request", ) - def __init__(self, response: bytes): + def __init__(self, response: bytes, response_no_usage: bytes | None): + """Create the protocol for one connection. + + Args: + response: Sent to every request, or only to requests that set + ``stream_options.include_usage``. + response_no_usage: Sent to a request that does not. ``None`` when + the response does not depend on the request (non-streaming), + so the request body is never inspected. + """ self.transport = None self.response = response - self.response_size = len(response) + self.response_no_usage = response_no_usage self._request: RequestParser | None = None self._parser: httptools.HttpRequestParser | None = None @@ -270,11 +364,16 @@ def data_received(self, data): if request.done(): _counter_add(_req_counter) + response = self.response + no_usage = self.response_no_usage + if no_usage is not None and not requests_usage(request.body): + response = no_usage + transport = self.transport if transport and not transport.is_closing(): - transport.write(self.response) # type: ignore[union-attr] + transport.write(response) # type: ignore[union-attr] _counter_add(_resp_counter) - _counter_add(_byte_counter, self.response_size) + _counter_add(_byte_counter, len(response)) # Reset for next request (HTTP/1.1 keep-alive) self._request = RequestParser() @@ -291,6 +390,7 @@ def _worker( host: str, port: int, response: bytes, + response_no_usage: bytes | None, ready: multiprocessing.synchronize.Event, shutdown: multiprocessing.synchronize.Event, counters: tuple[_SyncInt | None, _SyncInt | None, _SyncInt | None] = ( @@ -310,7 +410,7 @@ async def run(): loop = asyncio.get_running_loop() def protocol_factory(): - return ServerProtocol(response) + return ServerProtocol(response, response_no_usage) server = await loop.create_server( protocol_factory, @@ -382,7 +482,7 @@ def __init__( self._worker_shutdown: multiprocessing.synchronize.Event | None = None self._stats_thread: threading.Thread | None = None - self._response = build_response(output_length, stream, stream_interval) + self.set_response(output_length, stream_interval) global _req_counter, _resp_counter, _byte_counter if stats: @@ -390,6 +490,24 @@ def __init__( _resp_counter = multiprocessing.Value("L", 0) # type: ignore[assignment] _byte_counter = multiprocessing.Value("L", 0) # type: ignore[assignment] + def set_response(self, output_length: int, stream_interval: int = 1) -> None: + """(Re)build the pre-compiled responses for the current ``stream`` mode. + + A streaming server precompiles a second response without the usage + frame for requests that do not set ``stream_options.include_usage``; + serving the right one costs a body probe and one branch per request. + """ + self._response = build_response( + output_length, self.stream, stream_interval, include_usage=True + ) + self._response_no_usage = ( + build_response( + output_length, self.stream, stream_interval, include_usage=False + ) + if self.stream + else None + ) + @property def url(self) -> str: return f"http://{self.host}:{self._port or self.port}" @@ -430,6 +548,7 @@ def _start_workers(self): self.host, self._port, self._response, + self._response_no_usage, ready, self._worker_shutdown, (_req_counter, _resp_counter, _byte_counter), diff --git a/src/inference_endpoint/testing/variable_throughput_server.py b/src/inference_endpoint/testing/variable_throughput_server.py index b640b6b42..936ea5b1c 100644 --- a/src/inference_endpoint/testing/variable_throughput_server.py +++ b/src/inference_endpoint/testing/variable_throughput_server.py @@ -74,6 +74,8 @@ _counter_read, _sse_event, build_non_streaming_response, + build_stream_tail, + requests_usage, ) except ImportError: from max_throughput_server import ( # type: ignore[no-redef] @@ -83,6 +85,8 @@ _counter_read, _sse_event, build_non_streaming_response, + build_stream_tail, + requests_usage, ) # Shared counters (multiprocessing.Value for cross-process atomicity) @@ -265,8 +269,12 @@ def data_received(self, data): osl = int(rng.lognormvariate(self._osl_mu, self._osl_sigma)) osl = max(self._osl_min, min(osl, self._osl_max)) + # Streaming responses end with a usage frame only when the request + # sets stream_options.include_usage. + include_usage = self._stream and requests_usage(request.body) + # Dispatch async handler (prevent GC before completion) - task = self._loop.create_task(self._handle_request(osl)) + task = self._loop.create_task(self._handle_request(osl, include_usage)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) @@ -274,7 +282,7 @@ def data_received(self, data): self._request = RequestParser() self._parser = httptools.HttpRequestParser(self._request) - async def _handle_request(self, osl: int): + async def _handle_request(self, osl: int, include_usage: bool): """Handle a single request with sampled timing.""" transport = self.transport if transport is None or transport.is_closing(): @@ -319,7 +327,7 @@ async def _handle_request(self, osl: int): ) if self._stream: - await self._handle_streaming(transport, osl, ttft, tpot) + await self._handle_streaming(transport, osl, ttft, tpot, include_usage) else: await self._handle_non_streaming(transport, osl, ttft) except (RuntimeError, ConnectionError, OSError): @@ -342,12 +350,15 @@ async def _handle_non_streaming(self, transport, osl: int, delay: float): _counter_add(_resp_counter) _counter_add(_byte_counter, len(response)) - async def _handle_streaming(self, transport, osl: int, ttft: float, tpot: float): + async def _handle_streaming( + self, transport, osl: int, ttft: float, tpot: float, include_usage: bool + ): """Stream chunks with TTFT + per-token TPOT delays. tpot is the simulated per-token generation time (seconds). stream_interval is chars per SSE event (≈ tokens). Actual inter-event delay = tpot × stream_interval. + include_usage adds the usage frame between the finish frame and [DONE]. """ interval = self._stream_interval num_events = math.ceil(osl / interval) @@ -363,8 +374,7 @@ async def _handle_streaming(self, transport, osl: int, ttft: float, tpot: float) event = _sse_event(model, created, {"content": "x" * chars}) chunks.append(_chunked(event)) - finish = _sse_event(model, created, {}, "stop") - tail = _chunked(finish + b"data: [DONE]\n\n") + b"0\r\n\r\n" + tail = build_stream_tail(model, created, osl, include_usage=include_usage) # TTFT delay before first data if ttft > 0: @@ -403,7 +413,7 @@ async def _handle_streaming(self, transport, osl: int, ttft: float, tpot: float) transport.write(chunk) total_bytes += len(chunk) - # Finish event + [DONE] sentinel + terminator + # Finish event, optional usage frame, [DONE] sentinel and terminator if not transport.is_closing(): transport.write(tail) total_bytes += len(tail) diff --git a/tests/conftest.py b/tests/conftest.py index cc54197ee..d0e59b7e5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -27,6 +27,7 @@ from pathlib import Path from typing import Any +import inference_endpoint.core.types as core_types import pytest from inference_endpoint import metrics from inference_endpoint.config.runtime_settings import RuntimeSettings @@ -117,11 +118,23 @@ def temp_dir(tmp_path_factory: pytest.TempPathFactory) -> Path: return tmp_path_factory.mktemp("test_artifacts") +@pytest.fixture +def usage_warning_unlatched(monkeypatch, caplog) -> pytest.LogCaptureFixture: + """A fresh worker process: the malformed-usage warning has not fired yet. + + Returns ``caplog`` capturing WARNING records from ``core.types``. + """ + monkeypatch.setattr(core_types, "_malformed_usage_warned", False) + caplog.set_level(logging.WARNING, logger=core_types.__name__) + return caplog + + @pytest.fixture(scope="function") def mock_http_echo_server_factory(): """Factory fixture for EchoServer with an optional custom request handler. - Yields a factory callable: ``factory(handler=None) -> EchoServer``. + Yields a factory callable: + ``factory(handler=None, *, report_usage=True) -> EchoServer``. When called with no arguments the server uses the default echo behaviour. When called with a ``handler`` the handler is invoked for every incoming @@ -129,6 +142,10 @@ def mock_http_echo_server_factory(): (sync) callable or an async coroutine function; in both cases it receives the ``aiohttp.web.Request`` and must return an ``aiohttp.web.Response``. + ``report_usage=False`` models a server that never reports token usage: + non-streaming responses carry ``"usage": null`` and streams send no usage + frame. + All servers created by the factory are stopped automatically at the end of the test. @@ -142,8 +159,10 @@ def test_custom(mock_http_echo_server_factory): """ servers: list[EchoServer] = [] - def factory(handler: RequestHandler | None = None) -> EchoServer: - server = EchoServer(port=0, request_handler=handler) + def factory( + handler: RequestHandler | None = None, *, report_usage: bool = True + ) -> EchoServer: + server = EchoServer(port=0, request_handler=handler, report_usage=report_usage) logging.info("Starting mock HTTP echo server") server.start() servers.append(server) diff --git a/tests/integration/test_variable_server_usage.py b/tests/integration/test_variable_server_usage.py new file mode 100644 index 000000000..4d5211842 --- /dev/null +++ b/tests/integration/test_variable_server_usage.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test: the variable-response stub streams usage only when asked.""" + +import json +import urllib.request + +import pytest +from inference_endpoint.testing.variable_throughput_server import ( + VariableResponseServer, +) + +from tests.test_helpers import parse_sse_frames + + +@pytest.mark.integration +@pytest.mark.parametrize("include_usage", [True, False]) +def test_variable_server_stream_usage_follows_request(include_usage): + payload: dict = {"stream": True} + if include_usage: + payload["stream_options"] = {"include_usage": True} + with VariableResponseServer( + host="127.0.0.1", + port=0, + output_len_mean=40, + output_len_spread=0.0, + stream=True, + stream_interval=10, + num_workers=1, + quiet=True, + ) as srv: + req = urllib.request.Request( + f"{srv.url}/v1/chat/completions", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=10) as resp: + frames = parse_sse_frames(resp.read()) + # Content frames, then the finish frame, then usage (only when asked), [DONE]. + tail_len = 3 if include_usage else 2 + n = sum(len(f["choices"][0]["delta"]["content"]) for f in frames[:-tail_len]) + assert n > 0 + assert frames[-tail_len]["choices"][0]["finish_reason"] == "stop" + if include_usage: + assert frames[-2]["choices"] == [] + assert frames[-2]["usage"] == { + "prompt_tokens": 10, + "completion_tokens": n, + "total_tokens": n + 10, + } + assert frames[-1] == b"[DONE]" diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 6e5b99db5..5440c360f 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -23,15 +23,72 @@ import hashlib import json import random +import re import string import uuid from pathlib import Path +from typing import Any +import msgspec import zmq +from inference_endpoint.config.schema import Settings from inference_endpoint.core.types import ( + TOKEN_HDR_HIGH, Query, + TextModelOutput, ) from inference_endpoint.dataset_manager.dataset import Dataset +from inference_endpoint.metrics.report import OtherBasisMetrics, Report +from inference_endpoint.metrics.token_count_basis import ServerUsage + +# The token-count basis a run uses when none is configured. Basis-dependent tests +# take it from here instead of pinning a literal. +DEFAULT_SOURCE = Settings.model_fields["token_count_source"].default + +# One table pins the usable-server-usage rule for every consumer: +# (id, prompt_tokens, completion_tokens, output, expected classification). +SERVER_USAGE_CASES: list[tuple[str, object, object, TextModelOutput, ServerUsage]] = [ + ("usable", 10, 5, TextModelOutput(output="hi"), ServerUsage.USABLE), + ( + "missing_both", + msgspec.UNSET, + msgspec.UNSET, + TextModelOutput(output="hi"), + ServerUsage.MISSING, + ), + ( + "missing_prompt", + msgspec.UNSET, + 5, + TextModelOutput(output="hi"), + ServerUsage.MISSING, + ), + ("zero_beside_output", 10, 0, TextModelOutput(output="hi"), ServerUsage.MISSING), + ("true_zero", 10, 0, TextModelOutput(output=""), ServerUsage.TRUE_ZERO), + ("zero_prompt", 0, 5, TextModelOutput(output="hi"), ServerUsage.MISSING), + ("negative", 10, -1, TextModelOutput(output="hi"), ServerUsage.MISSING), + ("bool", True, 5, TextModelOutput(output="hi"), ServerUsage.MISSING), + ( + "above_bound", + 10, + TOKEN_HDR_HIGH + 1, + TextModelOutput(output="hi"), + ServerUsage.MISSING, + ), + ("float", 10.0, 5, TextModelOutput(output="hi"), ServerUsage.MISSING), +] + +# The cases an event log can carry. It stores counts as integers, so a bool or +# float count reaches a post-run reader as an undecodable line, not a usage case. +LOGGED_SERVER_USAGE_CASES = [ + case + for case in SERVER_USAGE_CASES + if all( + value is msgspec.UNSET + or (isinstance(value, int) and not isinstance(value, bool)) + for value in case[1:3] + ) +] def write_large_prompt_dataset(directory: Path, n_samples: int) -> Path: @@ -187,3 +244,46 @@ def get_test_socket_path(tmp_path: Path, test_name: str, suffix: str = "") -> st len(socket_path) <= zmq.IPC_PATH_MAX_LEN ), "socket path is too long for ZMQ IPC" return socket_path + + +def parse_sse_frames(body: bytes) -> list[dict[str, Any] | bytes]: + """Parse SSE ``data:`` frames from a raw or decoded streaming body. + + JSON frames decode to dicts and the ``[DONE]`` sentinel stays ``b"[DONE]"``. + HTTP chunk-size lines are ignored, so this accepts raw chunked bytes as well + as bodies already decoded by an HTTP client. + """ + return [ + m if m == b"[DONE]" else json.loads(m) + for m in re.findall(rb"data: (\{[^\n]+\}|\[DONE\])", body) + ] + + +_BASE_REPORT: dict[str, Any] = { + "version": "t", + "git_sha": None, + "test_started_at": 0, + "n_samples_issued": 10, + "n_samples_completed": 10, + "n_samples_failed": 0, + "duration_ns": 1, + "state": "complete", + "complete": True, + "ttft": {}, + "tpot": {}, + "latency": {}, + "input_sequence_lengths": {}, + "output_sequence_lengths": {}, + "token_count_source": DEFAULT_SOURCE, + "other_basis": OtherBasisMetrics( + input_sequence_lengths=None, + output_sequence_lengths=None, + tpot=None, + tps=None, + ), +} + + +def make_report(**over: Any) -> Report: + """A complete 10-sample ``Report`` with ``over`` applied to its fields.""" + return Report(**(_BASE_REPORT | over)) diff --git a/tests/unit/testing/__init__.py b/tests/unit/testing/__init__.py new file mode 100644 index 000000000..467079831 --- /dev/null +++ b/tests/unit/testing/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit/testing/test_echo_server_usage.py b/tests/unit/testing/test_echo_server_usage.py new file mode 100644 index 000000000..665a08ae8 --- /dev/null +++ b/tests/unit/testing/test_echo_server_usage.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the echo server's token usage reporting.""" + +import json + +import aiohttp +import pytest +from inference_endpoint.openai.completions_adapter import OpenAITextCompletionsAdapter + +from tests.test_helpers import parse_sse_frames + +PROMPT = "hello big world" # 18 chars; streamed as "hello big world" (15) +ERR = b"Stream options can only be defined when `stream=True`." +INCLUDE = {"stream": True, "stream_options": {"include_usage": True}} + + +async def _post(url, payload) -> tuple[int, bytes]: + async with aiohttp.ClientSession() as s, s.post(url, json=payload) as r: + return r.status, await r.read() + + +def _chat(content=PROMPT, **extra): + return {"model": "m", "messages": [{"role": "user", "content": content}], **extra} + + +def _url(server, path="/v1/chat/completions"): + return f"{server.url}{path}" + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_chat_non_streaming_reports_usage(mock_http_echo_server): + status, body = await _post(_url(mock_http_echo_server), _chat()) + assert ( + status == 200 and json.loads(body)["choices"][0]["message"]["content"] == PROMPT + ) + assert json.loads(body)["usage"] == { + "prompt_tokens": 18, + "completion_tokens": 18, + "total_tokens": 36, + } + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_chat_streaming_usage_counts_emitted_text(mock_http_echo_server): + frames = parse_sse_frames( + (await _post(_url(mock_http_echo_server), _chat(**INCLUDE)))[1] + ) + assert ( + "".join(f["choices"][0]["delta"]["content"] for f in frames[:-3]) + == "hello big world" + ) + assert frames[-3]["choices"][0]["finish_reason"] == "stop" + assert frames[-2]["choices"] == [] + assert frames[-2]["usage"] == { + "prompt_tokens": 18, + "completion_tokens": 15, + "total_tokens": 33, + } + assert frames[-1] == b"[DONE]" + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_max_osl_usage_counts_text_actually_sent(mock_http_echo_server): + mock_http_echo_server.set_max_osl(8) # "ab cd" -> "ab cdab " -> streamed "ab cdab" + frames = parse_sse_frames( + (await _post(_url(mock_http_echo_server), _chat("ab cd", **INCLUDE)))[1] + ) + assert frames[-2]["usage"] == { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + } + + +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra", + [{}, {"stream_options": None}, {"stream_options": {"include_usage": False}}], +) +async def test_streaming_without_include_usage_sends_no_usage( + mock_http_echo_server, extra +): + frames = parse_sse_frames( + (await _post(_url(mock_http_echo_server), _chat(stream=True, **extra)))[1] + ) + assert frames[-1] == b"[DONE]" and not any("usage" in f for f in frames[:-1]) + + +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "payload", "expected"), + [ + ("/v1/chat/completions", _chat(stream_options={"include_usage": True}), 400), + ( + "/v1/chat/completions", + _chat(stream=False, stream_options={"include_usage": True}), + 400, + ), + ("/v1/chat/completions", _chat(stream=False, stream_options={}), 200), + ( + "/v1/completions", + {"model": "m", "prompt": "hi", "stream_options": {"include_usage": True}}, + 400, + ), + ], +) +async def test_stream_options_without_stream_is_rejected( + mock_http_echo_server, path, payload, expected +): + status, body = await _post(_url(mock_http_echo_server, path), payload) + assert status == expected and (ERR in body) is (expected == 400) + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_report_usage_false_never_sends_usage(mock_http_echo_server_factory): + server = mock_http_echo_server_factory(report_usage=False) + assert json.loads((await _post(_url(server), _chat()))[1])["usage"] is None + server.set_max_osl(0) # every output becomes "" + frames = parse_sse_frames((await _post(_url(server), _chat(**INCLUDE)))[1]) + assert [f["choices"][0]["finish_reason"] for f in frames[:-1]] == ["stop"] + assert frames[-1] == b"[DONE]" + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_text_completions_route_reports_usage(mock_http_echo_server): + url = _url(mock_http_echo_server, "/v1/completions") + body = (await _post(url, {"model": "m", "prompt": "hi there"}))[1] + assert ( + OpenAITextCompletionsAdapter.decode_response( + body, "q" + ).get_response_output_string() + == "hi there" + ) + assert json.loads(body)["usage"] == { + "prompt_tokens": 8, + "completion_tokens": 8, + "total_tokens": 16, + } + frames = parse_sse_frames( + (await _post(url, {"model": "m", "prompt": [11, 22, 333], **INCLUDE}))[1] + ) + assert [f["object"] for f in frames[:-1]] == ["text_completion"] * 5 + assert "".join(f["choices"][0]["text"] for f in frames[:3]) == "11 22 333" + assert frames[3]["choices"][0]["finish_reason"] == "stop" + assert frames[4]["usage"] == { + "prompt_tokens": 3, + "completion_tokens": 9, + "total_tokens": 12, + } diff --git a/tests/unit/testing/test_stub_servers.py b/tests/unit/testing/test_stub_servers.py new file mode 100644 index 000000000..9306aa1b7 --- /dev/null +++ b/tests/unit/testing/test_stub_servers.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the built-in throughput stub servers' streaming responses.""" + +import asyncio +import random + +import msgspec +import pytest +from inference_endpoint.testing import max_throughput_server +from inference_endpoint.testing.max_throughput_server import ( + STUB_PROMPT_TOKENS, + MaxThroughputServer, + ServerProtocol, + build_response, + build_stream_tail, + build_streaming_response, + requests_usage, +) +from inference_endpoint.testing.variable_throughput_server import ( + VariableResponseProtocol, +) +from inference_endpoint.utils.benchmark_httpclient import _sse_events_per_response + +from tests.test_helpers import parse_sse_frames + +OUTPUT_LENGTH = 40 +STREAM_INTERVAL = 10 + + +class _RecordingTransport: + """Stands in for an asyncio transport; records what the stub writes.""" + + def __init__(self): + self.writes: list[bytes] = [] + + def write(self, data: bytes) -> None: + self.writes.append(data) + + def is_closing(self) -> bool: + return False + + def get_extra_info(self, name: str): + return None + + def close(self) -> None: + pass + + +def _stream_request(*, include_usage: bool) -> bytes: + """A raw HTTP streaming chat request, with or without ``include_usage``.""" + payload: dict = { + "model": "m", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + } + if include_usage: + payload["stream_options"] = {"include_usage": True} + body = msgspec.json.encode(payload) + head = ( + b"POST /v1/chat/completions HTTP/1.1\r\n" + b"Host: stub\r\n" + b"Content-Type: application/json\r\n" + b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + ) + return head + body + + +async def _max_throughput_reply(request: bytes) -> bytes: + server = MaxThroughputServer( + stream=True, + output_length=OUTPUT_LENGTH, + stream_interval=STREAM_INTERVAL, + quiet=True, + ) + protocol = ServerProtocol(server._response, server._response_no_usage) + transport = _RecordingTransport() + protocol.connection_made(transport) + protocol.data_received(request) + return b"".join(transport.writes) + + +async def _variable_reply(request: bytes) -> bytes: + # Zero means/spreads: no TTFT or inter-token sleeps, and the output length + # is clamped to OUTPUT_LENGTH. + protocol = VariableResponseProtocol( + asyncio.get_running_loop(), + random.Random(0), + None, + None, + stream=True, + stream_interval=STREAM_INTERVAL, + osl_mu=0.0, + osl_sigma=0.0, + osl_min=OUTPUT_LENGTH, + osl_max=OUTPUT_LENGTH, + osl_mean=OUTPUT_LENGTH, + mode="icl", + fcl_mu=0.0, + fcl_sigma=0.0, + fcl_mean=0.0, + rate_mu=0.0, + rate_sigma=0.0, + rate_mean=0.0, + icl_mu=0.0, + icl_sigma=0.0, + icl_mean=0.0, + ) + transport = _RecordingTransport() + protocol.connection_made(transport) + protocol.data_received(request) + await asyncio.gather(*protocol._tasks) + return b"".join(transport.writes) + + +STUB_REPLIES = { + "max_throughput": _max_throughput_reply, + "variable": _variable_reply, +} + +_CREATED = 1_700_000_000 +_CONTENT_FRAME = ( + b'a2\r\ndata: {"id":"r","object":"chat.completion.chunk","created":1700000000,' + b'"model":"max-tp","choices":[{"index":0,"delta":{"content":"xxxxx"},' + b'"finish_reason":null}]}\n\n\r\n' +) +# The usage-free stream: the finish frame and [DONE] share the last chunk. +# Roofline numbers are compared across versions, so this layout is pinned byte +# for byte. +_USAGE_FREE_STREAM = ( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: text/event-stream\r\n" + b"Cache-Control: no-cache\r\n" + b"Transfer-Encoding: chunked\r\n" + b"Connection: keep-alive\r\n\r\n" + + _CONTENT_FRAME + * 2 + + b'a1\r\ndata: {"id":"r","object":"chat.completion.chunk","created":1700000000,' + b'"model":"max-tp","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + b"\n\ndata: [DONE]\n\n\r\n" + b"0\r\n\r\n" +) + + +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize("stub", STUB_REPLIES) +@pytest.mark.parametrize("include_usage", [True, False]) +async def test_stub_sends_usage_only_when_requested(stub, include_usage): + frames = parse_sse_frames( + await STUB_REPLIES[stub](_stream_request(include_usage=include_usage)) + ) + usage_frames = [f for f in frames if isinstance(f, dict) and "usage" in f] + assert len(usage_frames) == int(include_usage) + if include_usage: + assert usage_frames[0]["usage"] == { + "prompt_tokens": STUB_PROMPT_TOKENS, + "completion_tokens": OUTPUT_LENGTH, + "total_tokens": OUTPUT_LENGTH + STUB_PROMPT_TOKENS, + } + content = "".join( + f["choices"][0]["delta"].get("content", "") + for f in frames + if isinstance(f, dict) and f["choices"] + ) + assert content == "x" * OUTPUT_LENGTH + assert frames[-1] == b"[DONE]" + + +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize("stub", STUB_REPLIES) +async def test_stub_usage_follows_finish_frame(stub): + frames = parse_sse_frames( + await STUB_REPLIES[stub](_stream_request(include_usage=True)) + ) + assert all( + f["choices"][0]["finish_reason"] is None + for f in frames[:-3] + if isinstance(f, dict) + ) + assert frames[-3]["choices"] == [{"index": 0, "delta": {}, "finish_reason": "stop"}] + assert frames[-2]["object"] == "chat.completion.chunk" + assert frames[-2]["choices"] == [] + assert "usage" in frames[-2] + assert frames[-1] == b"[DONE]" + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("body", "expected"), + [ + (b'{"stream":true,"stream_options":{"include_usage":true}}', True), + (b'{"stream": true, "stream_options": {"include_usage": true}}', True), + (b'{"stream_options":{"include_usage":false}}', False), + (b'{"stream_options":{"include_usage":null}}', False), + (b'{"stream_options":{}}', False), + (b'{"stream_options":null}', False), + (b"{}", False), + # A prompt quoting the key carries it escaped, so it never matches. + (rb'{"messages":[{"content":"say \"include_usage\":true"}]}', False), + ], +) +def test_requests_usage_reads_include_usage_from_body(body, expected): + assert requests_usage(body) is expected + + +@pytest.mark.unit +def test_streaming_response_without_usage_is_byte_stable(monkeypatch): + monkeypatch.setattr(max_throughput_server.time, "time", lambda: _CREATED) + assert ( + build_streaming_response(10, 5, "max-tp", include_usage=False) + == _USAGE_FREE_STREAM + ) + assert ( + build_response(10, True, 5, "max-tp", include_usage=False) == _USAGE_FREE_STREAM + ) + + +@pytest.mark.unit +def test_non_streaming_response_never_probes_the_request_body(monkeypatch): + server = MaxThroughputServer(stream=False, output_length=OUTPUT_LENGTH, quiet=True) + assert server._response_no_usage is None + + def probe(body: bytes) -> bool: + raise AssertionError("a non-streaming stub must not inspect the body") + + monkeypatch.setattr(max_throughput_server, "requests_usage", probe) + protocol = ServerProtocol(server._response, server._response_no_usage) + transport = _RecordingTransport() + protocol.connection_made(transport) + protocol.data_received(_stream_request(include_usage=False)) + assert transport.writes == [server._response] + assert b'"usage"' in server._response + + +@pytest.mark.unit +def test_set_response_follows_the_stream_mode(): + def usage_frames(response: bytes) -> int: + return sum( + isinstance(f, dict) and "usage" in f for f in parse_sse_frames(response) + ) + + server = MaxThroughputServer(stream=False, output_length=OUTPUT_LENGTH, quiet=True) + assert server._response_no_usage is None + assert b"text/event-stream" not in server._response + + server.stream = True + server.set_response(OUTPUT_LENGTH, STREAM_INTERVAL) + assert server._response_no_usage is not None + assert usage_frames(server._response) == 1 + assert usage_frames(server._response_no_usage) == 0 + + server.stream = False + server.set_response(OUTPUT_LENGTH) + assert server._response_no_usage is None + assert b"text/event-stream" not in server._response + + +@pytest.mark.unit +def test_streaming_response_usage_follows_finish_frame(): + frames = parse_sse_frames( + build_streaming_response( + output_length=160, + stream_interval=10, + model="max-tp", + include_usage=True, + ) + ) + assert len(frames) == 19 + assert ( + "".join(f["choices"][0]["delta"]["content"] for f in frames[:16]) == "x" * 160 + ) + assert frames[16]["choices"] == [{"index": 0, "delta": {}, "finish_reason": "stop"}] + assert frames[17]["object"] == "chat.completion.chunk" + assert frames[17]["choices"] == [] + assert frames[17]["usage"] == { + "prompt_tokens": 10, + "completion_tokens": 160, + "total_tokens": 170, + } + assert frames[18] == b"[DONE]" + + +@pytest.mark.unit +@pytest.mark.parametrize("include_usage", [True, False]) +def test_stream_tail_is_finish_then_optional_usage_then_done(include_usage): + tail = build_stream_tail("m", 123, completion_tokens=7, include_usage=include_usage) + frames = parse_sse_frames(tail) + assert frames[0]["choices"][0]["finish_reason"] == "stop" + if include_usage: + assert frames[1]["usage"] == { + "prompt_tokens": STUB_PROMPT_TOKENS, + "completion_tokens": 7, + "total_tokens": 7 + STUB_PROMPT_TOKENS, + } + assert frames[1 + include_usage :] == [b"[DONE]"] + assert tail.endswith(b"data: [DONE]\n\n\r\n0\r\n\r\n") + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("output_length", "stream_interval", "include_usage", "expected"), + [ + (160, 1, True, 163), + (160, 10, True, 19), + (160, 160, True, 4), + (5, 2, True, 6), + (160, 1, False, 162), + (160, 10, False, 18), + (160, 160, False, 3), + (5, 2, False, 5), + ], +) +def test_sse_events_per_response_follows_include_usage( + output_length, stream_interval, include_usage, expected +): + assert ( + _sse_events_per_response( + output_length, stream_interval, include_usage=include_usage + ) + == expected + ) + assert ( + len( + parse_sse_frames( + build_streaming_response( + output_length, stream_interval, include_usage=include_usage + ) + ) + ) + == expected + ) From 260d975113c2ac4ed683b0c1f218bcd6f7e7e586 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:02 -0500 Subject: [PATCH 06/26] feat(metrics): record server token counts beside the tokenizer's metrics/token_count_basis.py holds the one usage rule. classify_server_usage decides, once per sample, whether the server's counts are usable: usage is missing when there is none, or when completion_tokens is 0 beside non-empty output. read_server_usage returns the usable counts. The live aggregator and every post-run tool use these two functions. The module also reads a finished run's basis back (token_count_source_for_run, replay_basis). Aggregator: - Token series are registered from one table (TOKEN_SERIES) and named per basis with basis_key. Primary series keep main's names and registration order, so a default snapshot is a strict superset of main's. - ServerUsageTrigger classifies each COMPLETE once and records server ISL/OSL/TPOT synchronously, with no tokenizer and nothing added to n_pending_tasks. Server TPOT uses LoadGen's rule, (complete - first chunk) / (completion_tokens - 1), and skips completion_tokens <= 1 and non-positive windows. - Missing usage is counted on either primary (tracked_samples_missing_usage[_empty_output]). osl_paired_* counters total both OSLs over the samples that were counted on both bases. - Early stopping and steady state attach to the primary basis only. --no-metrics-isl drops ISL on both bases. - __main__ takes --token-count-source. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- .../services/metrics_aggregator/__main__.py | 30 +- .../services/metrics_aggregator/aggregator.py | 314 +++++--- .../metrics_aggregator/metrics_table.py | 232 +++++- .../services/metrics_aggregator/registry.py | 14 +- .../services/metrics_aggregator/snapshot.py | 5 +- .../metrics_aggregator/token_metrics.py | 3 +- .../metrics_aggregator/tokenization.py | 15 + .../metrics/token_count_basis.py | 295 +++++++ .../char_chat_bare_reasoning/tokenizer.json | 187 +++++ .../tokenizer_config.json | 10 + .../test_signal_handling.py | 4 + .../services/metrics_aggregator/conftest.py | 54 +- .../metrics_aggregator/test_aggregator.py | 754 ++++++++++++++++-- .../test_aggregator_error_handler.py | 7 +- .../test_early_stopping_integration.py | 10 +- .../test_main_signal_handler.py | 31 + .../metrics_aggregator/test_metrics_table.py | 10 +- .../test_output_tokenization_rule.py | 32 + .../test_response_boundary.py | 6 +- .../test_server_token_pairs.py | 334 ++++++++ .../test_steady_state_collection.py | 196 ++++- tests/unit/metrics/test_token_count_basis.py | 309 +++++++ 22 files changed, 2610 insertions(+), 242 deletions(-) create mode 100644 src/inference_endpoint/metrics/token_count_basis.py create mode 100644 tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer.json create mode 100644 tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer_config.json create mode 100644 tests/unit/async_utils/services/metrics_aggregator/test_server_token_pairs.py create mode 100644 tests/unit/metrics/test_token_count_basis.py diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py index e8ef68061..ad9da7b8c 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py @@ -27,12 +27,14 @@ from inference_endpoint.async_utils.loop_manager import LoopManager from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext from inference_endpoint.async_utils.transport.zmq.ready_check import send_ready_signal +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.metrics.early_stopping import EarlyStoppingSpec from inference_endpoint.metrics.steady_state_diagnostics import PROFILES +from inference_endpoint.metrics.token_count_basis import TOKEN_COUNT_SOURCE_METAVAR from inference_endpoint.utils.logging import setup_logging -from .aggregator import MetricCounterKey, MetricsAggregatorService -from .metrics_table import MetricsTable +from .aggregator import MetricsAggregatorService +from .metrics_table import MetricCounterKey, MetricsTable from .publisher import MetricsPublisher from .registry import MetricsRegistry from .snapshot import MetricsSnapshotCodec @@ -125,7 +127,7 @@ def _cap_live_lane_rayon_pool(tokenizer_workers: int) -> None: os.environ.setdefault("RAYON_NUM_THREADS", str(max(1, tokenizer_workers))) -async def main() -> None: +def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Metrics aggregator service - subscribes to EventRecords and computes real-time metrics" ) @@ -189,7 +191,7 @@ async def main() -> None: "--tokenizer", type=str, default=None, - help="HuggingFace tokenizer name for ISL/OSL/TPOT (e.g. 'gpt2'). If not set, token metrics are disabled.", + help="HuggingFace tokenizer for client-basis ISL/OSL/TPOT (e.g. 'gpt2'). If not set, the client basis counts only pre-tokenized prompts; the server basis needs no tokenizer.", ) parser.add_argument( "--tokenizer-workers", @@ -244,7 +246,24 @@ async def main() -> None: "standalone detector re-runs use the same profile." ), ) - args = parser.parse_args() + parser.add_argument( + "--token-count-source", + type=TokenCountSource, + choices=list(TokenCountSource), + metavar=TOKEN_COUNT_SOURCE_METAVAR, + required=True, + help=( + "Primary token-count basis: its isl/osl/tpot_ns series keep the base " + "name and feed early stopping and steady state; the other basis " + "records into the same series suffixed _. The parent passes " + "the run's resolved source." + ), + ) + return parser + + +async def main() -> None: + args = _build_parser().parse_args() setup_logging(level="INFO") if args.tokenizer_workers < 0: @@ -317,6 +336,7 @@ async def main() -> None: shutdown_event=shutdown_event, drain_timeout_s=args.drain_timeout, steady_state_profile=args.steady_state_profile, + token_count_source=args.token_count_source, ) aggregator.start() diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py index 6bab4609e..cbef59df0 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py @@ -19,7 +19,7 @@ import asyncio import logging -from enum import Enum +from dataclasses import dataclass from typing import Final from inference_endpoint.async_utils.transport.zmq.pubsub import ( @@ -32,7 +32,14 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import PhaseData, PhaseType +from inference_endpoint.core.types import ( + TOKEN_HDR_HIGH, + TOKEN_HDR_LOW, + PhaseData, + PhaseType, + TokenCountSource, + basis_key, +) from inference_endpoint.metrics.steady_state_diagnostics import ( PROFILES, SteadyState, @@ -41,58 +48,27 @@ ) from .metrics_table import ( + BasisSeriesNames, ChunkDeltaTrigger, IslTrigger, + MetricCounterKey, MetricSeriesKey, MetricsTable, OslTrigger, SampleField, SampleLatencyTrigger, + ServerUsageTrigger, TpotTrigger, TtftTrigger, ) from .publisher import MetricsPublisher -from .registry import TOKEN_HDR_HIGH, TOKEN_HDR_LOW, MetricsRegistry +from .registry import MetricsRegistry from .snapshot import SessionState from .token_metrics import BatchTokenizer, TokenBatchQueue logger = logging.getLogger(__name__) -class MetricCounterKey(str, Enum): - """Counter metric keys tracked by the aggregator. - - Total counters include all samples (warmup + tracked). - Tracked counters only include samples within performance tracking windows. - """ - - TOTAL_SAMPLES_ISSUED = "total_samples_issued" - TOTAL_SAMPLES_COMPLETED = "total_samples_completed" - TOTAL_SAMPLES_FAILED = "total_samples_failed" - TRACKED_SAMPLES_ISSUED = "tracked_samples_issued" - TRACKED_SAMPLES_COMPLETED = "tracked_samples_completed" - # Failed samples that were within a performance-tracking window. - # Counted at ERROR-event time; correctness depends on - # session.py:_handle_response emitting ERROR before COMPLETE so the - # tracked row still exists when the aggregator sees the ERROR. - TRACKED_SAMPLES_FAILED = "tracked_samples_failed" - TRACKED_FINISH_REASON_STOP = "tracked_finish_reason_stop" - TRACKED_FINISH_REASON_LENGTH = "tracked_finish_reason_length" - TRACKED_FINISH_REASON_TOOL_CALLS = "tracked_finish_reason_tool_calls" - TRACKED_FINISH_REASON_CONTENT_FILTER = "tracked_finish_reason_content_filter" - TRACKED_FINISH_REASON_FUNCTION_CALL = "tracked_finish_reason_function_call" - TRACKED_FINISH_REASON_OTHER = "tracked_finish_reason_other" - TRACKED_DURATION_NS = "tracked_duration_ns" - # Legacy MLPerf LoadGen Server "completed" window (final_query_all_samples_done_time). - LEGACY_LOADGEN_WINDOW_DURATION_NS = "legacy_loadgen_window_duration_ns" - # Total wall-clock duration since session start. Updated on every event as - # max(current, event_timestamp - session_start). Stored as a counter - # rather than computed from (now - start) at read time because - # time.monotonic_ns() has a process-local epoch — a reader in another - # process would get a meaningless value. - TOTAL_DURATION_NS = "total_duration_ns" - - _FINISH_REASON_COUNTERS: Final[dict[str, MetricCounterKey]] = { "stop": MetricCounterKey.TRACKED_FINISH_REASON_STOP, "length": MetricCounterKey.TRACKED_FINISH_REASON_LENGTH, @@ -119,6 +95,75 @@ class MetricCounterKey(str, Enum): _NS_HDR_HIGH: Final[int] = 3_600_000_000_000 # 1 hour in ns +@dataclass(frozen=True, slots=True) +class TokenSeriesSpec: + """One token metric, registered once per token-count basis.""" + + key: MetricSeriesKey + dtype: type + hdr_low: int + hdr_high: int + # Recorded only for streaming runs. + streaming_only: bool + # Dropped on both bases when ISL collection is off. + is_isl: bool + # Gets early-stopping estimates, on the primary basis only. + tail_latency: bool + + +TOKEN_SERIES: Final[tuple[TokenSeriesSpec, ...]] = ( + TokenSeriesSpec( + key=MetricSeriesKey.ISL, + dtype=int, + hdr_low=TOKEN_HDR_LOW, + hdr_high=TOKEN_HDR_HIGH, + streaming_only=False, + is_isl=True, + tail_latency=False, + ), + TokenSeriesSpec( + key=MetricSeriesKey.OSL, + dtype=int, + hdr_low=TOKEN_HDR_LOW, + hdr_high=TOKEN_HDR_HIGH, + streaming_only=False, + is_isl=False, + tail_latency=False, + ), + TokenSeriesSpec( + key=MetricSeriesKey.TPOT_NS, + dtype=float, + hdr_low=_NS_HDR_LOW, + hdr_high=_NS_HDR_HIGH, + streaming_only=True, + is_isl=False, + tail_latency=True, + ), +) + + +def series_names( + basis: TokenCountSource, + *, + primary: TokenCountSource, + enable_isl: bool, + streaming: bool, +) -> BasisSeriesNames: + """The ``TOKEN_SERIES`` names ``basis`` records into in this run. + + Names follow ``basis_key``: unsuffixed on the primary basis. A series the + run does not record is left out (``None``). + """ + return BasisSeriesNames( + **{ + spec.key.value: basis_key(spec.key.value, basis, primary=primary) + for spec in TOKEN_SERIES + if (enable_isl or not spec.is_isl) + and (streaming or not spec.streaming_only) + } + ) + + class MetricsAggregatorService(ZmqMessageSubscriber[EventRecord]): """Subscribes to EventRecords and computes per-sample metrics in real time. @@ -137,6 +182,7 @@ def __init__( publish_interval_s: float, sig_figs: int, n_histogram_buckets: int, + token_count_source: TokenCountSource, tokenizer: BatchTokenizer | None = None, live_flush_interval_s: float | None = None, streaming: bool = False, @@ -157,10 +203,11 @@ def __init__( self._registry = registry self._publisher = publisher self._publish_interval_s = publish_interval_s - # Token triggers enqueue onto this queue; it is flushed by the + # Tokenizer triggers enqueue onto this queue; it is flushed by the # queue's own live loop (start_live) and by the end-of-run drain. - # None when no tokenizer is set (token metrics disabled), in which - # case those triggers are no-ops. + # None when no tokenizer is set, in which case the client-basis + # triggers count only pre-tokenized prompts; the server-basis trigger + # needs no tokenizer. self._token_queue: TokenBatchQueue | None = ( TokenBatchQueue(tokenizer, self.loop) if tokenizer is not None else None ) @@ -175,6 +222,18 @@ def __init__( # below tags the snapshot state=interrupted, not COMPLETE. self._interrupted = False self._drain_timeout_s = drain_timeout_s + # Both bases record every token series; the primary keeps the base + # names and feeds early stopping and the steady-state collector. + self._token_count_source = token_count_source + self._series_names = { + basis: series_names( + basis, + primary=token_count_source, + enable_isl=enable_isl, + streaming=streaming, + ) + for basis in TokenCountSource + } self._session_start_ns: int | None = None self._total_duration_ns: int = 0 @@ -188,7 +247,7 @@ def __init__( # Pre-register all metrics on the registry. Tests can introspect via # registry.has_counter / has_series. - self._register_metrics(streaming, enable_isl, sig_figs, n_histogram_buckets) + self._register_metrics(streaming, sig_figs, n_histogram_buckets) # A profile name enables steady-state collection. The profile carries # the CoV bounds and warmup driver used for the verdict. None leaves @@ -217,14 +276,15 @@ def __init__( steady_state_profile, ) self._steady_state_profile = None - # PHASE_START supplies the bucket size. TpotTrigger needs the collector - # registered before any event can arrive, so the size is latched later. + # PHASE_START supplies the bucket size. The primary basis's TPOT + # trigger needs the collector registered before any event can arrive, + # so the size is latched later. self._collector = ( SuperPassCollector() if self._steady_state_profile is not None else None ) self._table = MetricsTable(self._registry) - self._register_triggers(streaming, enable_isl) + self._register_triggers(streaming) # ------------------------------------------------------------------ # Registration helpers @@ -233,88 +293,119 @@ def __init__( def _register_metrics( self, streaming: bool, - enable_isl: bool, sig_figs: int, n_histogram_buckets: int, ) -> None: - """Register all counters and series on the registry.""" + """Register all counters and series on the registry. + + Series are registered as ``sample_latency_ns``, primary ``isl`` and + ``osl``, ``ttft_ns``, ``chunk_delta_ns``, primary ``tpot_ns``, then the + other basis's ``TOKEN_SERIES``. Snapshot series follow registration + order, so a snapshot lists the primary basis's series before the other + basis's. + + Must stay in lockstep with ``_register_triggers``: recording into an + unregistered series raises inside ``process()``. Both take the token + series names from ``self._series_names``. + """ + registry = self._registry + primary = self._token_count_source for key in MetricCounterKey: - self._registry.register_counter(key.value) - - # Always-present series - self._registry.register_series( - MetricSeriesKey.SAMPLE_LATENCY_NS.value, - hdr_low=_NS_HDR_LOW, - hdr_high=_NS_HDR_HIGH, - tail_latency=True, - sig_figs=sig_figs, - n_histogram_buckets=n_histogram_buckets, - ) - if enable_isl: - self._registry.register_series( - MetricSeriesKey.ISL.value, - hdr_low=TOKEN_HDR_LOW, - hdr_high=TOKEN_HDR_HIGH, - sig_figs=sig_figs, - n_histogram_buckets=n_histogram_buckets, - ) - self._registry.register_series( - MetricSeriesKey.OSL.value, - hdr_low=TOKEN_HDR_LOW, - hdr_high=TOKEN_HDR_HIGH, - sig_figs=sig_figs, - n_histogram_buckets=n_histogram_buckets, - ) + registry.register_counter(key.value) - # Streaming-only series - if streaming: - self._registry.register_series( - MetricSeriesKey.TTFT_NS.value, + def register_timing(key: MetricSeriesKey, *, tail_latency: bool) -> None: + registry.register_series( + key.value, hdr_low=_NS_HDR_LOW, hdr_high=_NS_HDR_HIGH, - tail_latency=True, + tail_latency=tail_latency, sig_figs=sig_figs, n_histogram_buckets=n_histogram_buckets, ) - self._registry.register_series( - MetricSeriesKey.CHUNK_DELTA_NS.value, - hdr_low=_NS_HDR_LOW, - hdr_high=_NS_HDR_HIGH, - sig_figs=sig_figs, - n_histogram_buckets=n_histogram_buckets, - ) - self._registry.register_series( - MetricSeriesKey.TPOT_NS.value, - hdr_low=_NS_HDR_LOW, - hdr_high=_NS_HDR_HIGH, - tail_latency=True, - sig_figs=sig_figs, - n_histogram_buckets=n_histogram_buckets, - dtype=float, - ) - def _register_triggers(self, streaming: bool, enable_isl: bool) -> None: + def register_tokens( + basis: TokenCountSource, specs: tuple[TokenSeriesSpec, ...] + ) -> None: + names = self._series_names[basis] + for spec in specs: + name = getattr(names, spec.key.value) + if name is None: + continue + registry.register_series( + name, + hdr_low=spec.hdr_low, + hdr_high=spec.hdr_high, + tail_latency=spec.tail_latency and basis is primary, + sig_figs=sig_figs, + n_histogram_buckets=n_histogram_buckets, + dtype=spec.dtype, + ) + + register_timing(MetricSeriesKey.SAMPLE_LATENCY_NS, tail_latency=True) + register_tokens( + primary, tuple(spec for spec in TOKEN_SERIES if not spec.streaming_only) + ) + if streaming: + register_timing(MetricSeriesKey.TTFT_NS, tail_latency=True) + register_timing(MetricSeriesKey.CHUNK_DELTA_NS, tail_latency=False) + register_tokens( + primary, tuple(spec for spec in TOKEN_SERIES if spec.streaming_only) + ) + register_tokens(primary.other, TOKEN_SERIES) + + def _register_triggers(self, streaming: bool) -> None: """Register metric triggers on the table. - Streaming-only triggers (TTFT, chunk_delta, TPOT) are only registered - when ``streaming=True``; the ISL trigger only when ``enable_isl=True``. + Both bases always record: the tokenizer triggers into the client + names, and one ``ServerUsageTrigger`` into the server names. Only the + primary basis's TPOT trigger gets the steady-state collector. + Streaming-only triggers (TTFT, chunk_delta, client TPOT) are + registered only when ``streaming=True``; the ISL trigger only when ISL + is collected. """ table = self._table registry = self._registry queue = self._token_queue + primary = self._token_count_source + client = self._series_names[TokenCountSource.CLIENT] + server = self._series_names[TokenCountSource.SERVER] - # Latency and OSL are always registered; ISL is opt-out. - if enable_isl: - table.add_trigger(SampleField.ISSUED_NS, IslTrigger(registry, queue)) + if client.isl is not None: + table.add_trigger( + SampleField.ISSUED_NS, + IslTrigger(registry, queue, metric_name=client.isl), + ) table.add_trigger(SampleField.COMPLETE_NS, SampleLatencyTrigger(registry)) - table.add_trigger(SampleField.COMPLETE_NS, OslTrigger(registry, queue)) + table.add_trigger( + SampleField.COMPLETE_NS, + OslTrigger(registry, queue, metric_name=client.osl), + ) + table.add_trigger( + SampleField.COMPLETE_NS, + ServerUsageTrigger( + registry, + names=server, + collector=( + self._collector if primary is TokenCountSource.SERVER else None + ), + primary=primary, + ), + ) - # Streaming-only if streaming: table.add_trigger(SampleField.RECV_FIRST_NS, TtftTrigger(registry)) table.add_trigger(SampleField.LAST_RECV_NS, ChunkDeltaTrigger(registry)) + if client.tpot_ns is not None: table.add_trigger( - SampleField.COMPLETE_NS, TpotTrigger(registry, queue, self._collector) + SampleField.COMPLETE_NS, + TpotTrigger( + registry, + queue, + metric_name=client.tpot_ns, + collector=( + self._collector if primary is TokenCountSource.CLIENT else None + ), + ), ) @property @@ -599,8 +690,12 @@ async def _steady_state_verdict(self, n_pending: int) -> SteadyState | None: """Return the steady-state verdict for the terminal snapshot. Returns None when collection is disabled or the series does not describe - the completed run. An interrupted run is truncated. A pending drain means - token counts are missing, so TPOT and OSL would be wrong. + the completed run. An interrupted run is truncated. Otherwise the gate + follows the primary basis, which feeds the collector's TPOT and OSL. On + the server basis a tracked sample without usable usage leaves them + incomplete, while a pending tokenizer drain only affects the client + series and is ignored. On the client basis a pending drain leaves them + incomplete, while missing server usage is ignored. This runs before the snapshot is built. The single atomic final write then carries both the metrics and the verdict, so the analysis sits in @@ -621,7 +716,16 @@ async def _steady_state_verdict(self, n_pending: int) -> SteadyState | None: collector, profile = self._collector, self._steady_state_profile if collector is None or profile is None: return None - if self._interrupted or n_pending > 0 or not collector.superpass_size: + if self._token_count_source is TokenCountSource.SERVER: + incomplete_counts = ( + self._registry.counter_value( + MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE.value + ) + > 0 + ) + else: + incomplete_counts = n_pending > 0 + if self._interrupted or incomplete_counts or not collector.superpass_size: logger.info("Steady state not computed: the run is not described by it") return None try: diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index cf7046abe..9f3ffb4ab 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -32,9 +32,15 @@ TokenIdsInput, TokenizationInput, extract_tokenization_input, + has_countable_output, ) from inference_endpoint.core.record import SampleEventType, SessionEventType -from inference_endpoint.core.types import PromptData, TextModelOutput +from inference_endpoint.core.types import PromptData, TextModelOutput, TokenCountSource +from inference_endpoint.metrics.token_count_basis import ( + ServerUsage, + classify_server_usage, + missing_usage_message, +) if TYPE_CHECKING: from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( @@ -64,7 +70,11 @@ class SampleField(str, Enum): class MetricSeriesKey(str, Enum): - """Series metric keys written by triggers to the registry.""" + """Series metric keys written by triggers to the registry. + + ``ISL``, ``OSL`` and ``TPOT_NS`` are base names: each token-count basis + records them under ``basis_key`` (unsuffixed on the primary basis). + """ ISL = "isl" OSL = "osl" @@ -74,6 +84,66 @@ class MetricSeriesKey(str, Enum): TPOT_NS = "tpot_ns" +class MetricCounterKey(str, Enum): + """Counter metric keys tracked by the aggregator. + + Total counters include all samples (warmup + tracked). + Tracked counters only include samples within performance tracking windows. + """ + + TOTAL_SAMPLES_ISSUED = "total_samples_issued" + TOTAL_SAMPLES_COMPLETED = "total_samples_completed" + TOTAL_SAMPLES_FAILED = "total_samples_failed" + TRACKED_SAMPLES_ISSUED = "tracked_samples_issued" + TRACKED_SAMPLES_COMPLETED = "tracked_samples_completed" + # Failed samples that were within a performance-tracking window. + # Counted at ERROR-event time; correctness depends on + # session.py:_handle_response emitting ERROR before COMPLETE so the + # tracked row still exists when the aggregator sees the ERROR. + TRACKED_SAMPLES_FAILED = "tracked_samples_failed" + TRACKED_FINISH_REASON_STOP = "tracked_finish_reason_stop" + TRACKED_FINISH_REASON_LENGTH = "tracked_finish_reason_length" + TRACKED_FINISH_REASON_TOOL_CALLS = "tracked_finish_reason_tool_calls" + TRACKED_FINISH_REASON_CONTENT_FILTER = "tracked_finish_reason_content_filter" + TRACKED_FINISH_REASON_FUNCTION_CALL = "tracked_finish_reason_function_call" + TRACKED_FINISH_REASON_OTHER = "tracked_finish_reason_other" + TRACKED_DURATION_NS = "tracked_duration_ns" + # Legacy MLPerf LoadGen Server "completed" window (final_query_all_samples_done_time). + LEGACY_LOADGEN_WINDOW_DURATION_NS = "legacy_loadgen_window_duration_ns" + # Total wall-clock duration since session start. Updated on every event as + # max(current, event_timestamp - session_start). Stored as a counter + # rather than computed from (now - start) at read time because + # time.monotonic_ns() has a process-local epoch — a reader in another + # process would get a meaningless value. + TOTAL_DURATION_NS = "total_duration_ns" + # Tracked successful samples whose COMPLETE carried no usable server usage + # (``ServerUsage.MISSING``), whatever the primary basis, and the subset + # whose output was empty (in-band error frames, cut streams). + TRACKED_SAMPLES_MISSING_USAGE = "tracked_samples_missing_usage" + TRACKED_SAMPLES_MISSING_USAGE_EMPTY_OUTPUT = ( + "tracked_samples_missing_usage_empty_output" + ) + # Samples with both a tokenizer OSL and usable server usage, and the two + # token totals over exactly those samples. + OSL_PAIRED_SAMPLES = "osl_paired_samples" + OSL_PAIRED_CLIENT_TOKENS = "osl_paired_client_tokens" + OSL_PAIRED_SERVER_TOKENS = "osl_paired_server_tokens" + + +class BasisSeriesNames(msgspec.Struct, frozen=True, kw_only=True): # type: ignore[call-arg] + """The series one token-count basis records into, named at registration. + + ``None`` for a series the run does not record: ``isl`` without ISL + collection, ``tpot_ns`` without streaming. Field names must equal the + ``MetricSeriesKey`` values in ``TOKEN_SERIES``: the aggregator builds and + reads this struct by ``spec.key.value``. + """ + + osl: str + isl: str | None = None + tpot_ns: str | None = None + + # --------------------------------------------------------------------------- # SampleRow # --------------------------------------------------------------------------- @@ -316,8 +386,10 @@ def __init__( self, registry: MetricsRegistry, queue: TokenBatchQueue | None, + *, + metric_name: str, ): - super().__init__(MetricSeriesKey.ISL, registry, queue) + super().__init__(metric_name, registry, queue) def _extract_tokenization_input(self, ev_rec, row, pre_change): if not isinstance(ev_rec.data, PromptData): @@ -344,14 +416,39 @@ class OslTrigger(TokenTrigger): chat template renders all structured fields, including ordinary content. Otherwise, non-empty output content uses ``TextInput``. ``TextModelOutput`` does not carry server token IDs, so OSL has no token-ID/text precedence. + + A sample whose COMPLETE carries usable server usage + (``ServerUsage.USABLE``) also adds its tokenizer count and the server + ``completion_tokens`` to the ``osl_paired_*`` counters, so both totals + cover exactly the same samples. Usage is classified here only for outputs + sent to the tokenizer queue, so a run without a tokenizer never does it. """ def __init__( self, registry: MetricsRegistry, queue: TokenBatchQueue | None, + *, + metric_name: str, ): - super().__init__(MetricSeriesKey.OSL, registry, queue) + super().__init__(metric_name, registry, queue) + + def _make_recorder(self, ev_rec, row, pre_change): + server_count = ev_rec.completion_tokens + usage = classify_server_usage(ev_rec.prompt_tokens, server_count, ev_rec.data) + if usage is not ServerUsage.USABLE: + return super()._make_recorder(ev_rec, row, pre_change) + registry, name = self.registry, self.metric_name + + def record(count: int) -> None: + registry.record(name, count) + registry.increment(MetricCounterKey.OSL_PAIRED_SAMPLES.value) + registry.increment(MetricCounterKey.OSL_PAIRED_CLIENT_TOKENS.value, count) + registry.increment( + MetricCounterKey.OSL_PAIRED_SERVER_TOKENS.value, server_count + ) + + return record def _extract_tokenization_input(self, ev_rec, row, pre_change): if not isinstance(ev_rec.data, TextModelOutput): @@ -374,10 +471,12 @@ def __init__( self, registry: MetricsRegistry, queue: TokenBatchQueue | None, - collector: SuperPassCollector | None = None, + *, + metric_name: str, + collector: SuperPassCollector | None, ): super().__init__( - MetricSeriesKey.TPOT_NS, + metric_name, registry, queue, requires=(SampleField.RECV_FIRST_NS,), @@ -423,6 +522,127 @@ def _compute_value(self, token_count, ev_rec, pre_change): return (ev_rec.timestamp_ns - recv_first_ns) / token_count +# --------------------------------------------------------------------------- +# Server-usage trigger (sync) +# --------------------------------------------------------------------------- + + +class ServerUsageTrigger(EmitTrigger): + """Server-basis ISL, OSL and TPOT from the usage on the COMPLETE record. + + One ``classify_server_usage`` call per successful sample decides all + three, so the server series always cover the same samples: + + - ``USABLE``: records ``prompt_tokens`` as ISL, ``completion_tokens`` as + OSL and, when streaming, TPOT by LoadGen's rule + ``(complete_ns - recv_first_ns) / (completion_tokens - 1)``: the first + visible chunk is excluded from both the window and the count. TPOT is + skipped without a first chunk, for ``completion_tokens <= 1`` and for a + non-positive window. + - ``TRUE_ZERO``: records nothing; no server-basis series records an empty + completion, and client OSL and TPOT also skip empty output. + - ``MISSING``: counts the sample in ``tracked_samples_missing_usage`` (and + ``..._empty_output`` when its output is empty) and logs the first miss + of the run: a WARNING when the server basis is primary, else INFO. + + Errors and agentic timeouts complete with ``data=None`` and are skipped. + Usage exists only on COMPLETE, and no tokenizer is involved, so nothing + here adds to ``n_pending_tasks``. + + The TPOT window can be non-positive because the two ends are stamped in + different places: COMPLETE carries ``QueryResult.completed_at``, set when + the transport decodes the response, while RECV_FIRST carries the time the + session dequeues the chunk. A short stream decoded in one transport drain + can therefore "complete" before its first chunk. + + It fires inside ``set_field(COMPLETE)`` while the row is still in the + table, so it reads ``row.superpass_index`` directly and feeds ``collector`` + (given only when the server basis is primary) synchronously, with the + post-first-chunk count ``completion_tokens - 1``. + """ + + def __init__( + self, + registry: MetricsRegistry, + *, + names: BasisSeriesNames, + collector: SuperPassCollector | None, + primary: TokenCountSource, + ): + super().__init__(names.osl, registry, requires=(SampleField.RECV_FIRST_NS,)) + self._isl = names.isl + self._tpot = names.tpot_ns + self._collector = collector + self._primary = primary + self._miss_logged = False + + def fire(self, ev_rec, row, pre_change): + output = ev_rec.data + if not isinstance(output, TextModelOutput): + return + prompt, completion = ev_rec.prompt_tokens, ev_rec.completion_tokens + usage = classify_server_usage(prompt, completion, output) + if usage is not ServerUsage.USABLE: + if usage is ServerUsage.MISSING: + self._count_missing(ev_rec, output) + return + registry = self.registry + if self._isl is not None: + registry.record(self._isl, prompt) + registry.record(self.metric_name, completion) + if self._tpot is None or completion <= 1: + return + recv_first_ns = pre_change[SampleField.RECV_FIRST_NS] + if recv_first_ns is None: + return + delta = ev_rec.timestamp_ns - recv_first_ns + if delta <= 0: + return + value = delta / (completion - 1) + registry.record(self._tpot, value) + if self._collector is not None and row.superpass_index >= 0: + self._collector.add_tpot(row.superpass_index, value, completion - 1) + + def _count_missing(self, ev_rec: EventRecord, output: TextModelOutput) -> None: + """Count a sample without usable usage; log the run's first one. + + An empty output is counted on its own: in-band error frames and cut + streams look like that, and client counting would hide them rather + than fix them. + """ + registry = self.registry + registry.increment(MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE.value) + empty_output = not has_countable_output(output) + if empty_output: + registry.increment( + MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE_EMPTY_OUTPUT.value + ) + if self._miss_logged: + return + # Once per run, so a long run shows the problem early without logging + # per sample. The aggregator counts this sample's completion after its + # triggers fire, hence the + 1. + self._miss_logged = True + n_succeeded = ( + registry.counter_value(MetricCounterKey.TRACKED_SAMPLES_COMPLETED.value) + - registry.counter_value(MetricCounterKey.TRACKED_SAMPLES_FAILED.value) + + 1 + ) + logger.log( + ( + logging.WARNING + if self._primary is TokenCountSource.SERVER + else logging.INFO + ), + "Tracked sample %s completed without usable server token usage: %s. " + "Later misses are only counted (tracked_samples_missing_usage).", + ev_rec.sample_uuid, + missing_usage_message( + 1, int(empty_output), int(n_succeeded), primary=self._primary + ), + ) + + # --------------------------------------------------------------------------- # MetricsTable # --------------------------------------------------------------------------- diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py index 18f874ae9..b5294e5cd 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py @@ -41,6 +41,7 @@ import numpy as np from hdrh.histogram import HdrHistogram +from inference_endpoint.core.types import TOKEN_HDR_HIGH, TOKEN_HDR_LOW from inference_endpoint.metrics.early_stopping import ( EarlyStoppingSpec, es_percentile_estimate, @@ -536,6 +537,10 @@ def build_snapshot( def has_counter(self, name: str) -> bool: return name in self._counters + def counter_value(self, name: str) -> int | float: + """Current value of a counter, for cold-path readers.""" + return self._counters[name].value() + def has_series(self, name: str) -> bool: return name in self._series @@ -546,11 +551,10 @@ def has_series(self, name: str) -> bool: # Canonical parameters for a token-count series, shared by the aggregator's live # OSL/ISL series and the finalize-side accuracy OSL so both emit an identical -# block from one source. The HDR bounds size the live/HDR percentile view; the -# exact (raw-array) path used for COMPLETE snapshots and accuracy OSL ignores them -# and derives histogram edges from the observed range. -TOKEN_HDR_LOW: Final[int] = 1 -TOKEN_HDR_HIGH: Final[int] = 10_000_000 # 10M tokens +# block from one source. The HDR bounds (``TOKEN_HDR_LOW``/``TOKEN_HDR_HIGH`` in +# ``core.types``) size the live/HDR percentile view; the exact (raw-array) path +# used for COMPLETE snapshots and accuracy OSL ignores them and derives histogram +# edges from the observed range. TOKEN_SERIES_SIG_FIGS: Final[int] = 3 TOKEN_SERIES_HISTOGRAM_BUCKETS: Final[int] = 30 diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py index f32c59033..c7c01687f 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py @@ -162,8 +162,9 @@ class MetricsSnapshot( flushes) and during the drain phase. **An incomplete drain is detected as** ``state == COMPLETE and n_pending_tasks > 0``: the end-of-run - flush timed out or failed; the token-derived series - are missing those samples. + flush timed out or failed; the client-basis series + (unsuffixed when the client basis is primary, + ``*_client`` otherwise) are missing those samples. metrics: Tagged union of ``CounterStat`` and ``SeriesStat``, ordered counters-first then series, registration order within each. diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 36fcf16de..31c50210a 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -292,7 +292,8 @@ def load_reference_tokenizer(tokenizer_name: str) -> Any: ``trust_remote_code`` is required for tokenizers that ship custom code (e.g. DeepSeek-R1). The single construction point shared by the perf-window tokenizer, the sharded length-counting workers, and finalize-side accuracy - OSL, so every OSL number in a run comes from the same tokenizer. + OSL, so every client-basis OSL number in a run (unsuffixed when the client + basis is primary, ``*_client`` otherwise) comes from the same tokenizer. """ return AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py index 1132f1952..4e64238ce 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py @@ -76,3 +76,18 @@ def extract_tokenization_input( if text: return TextInput(text) return None + + +def has_countable_output(output: TextModelOutput) -> bool: + """Whether ``extract_tokenization_input(output)`` is not ``None``. + + Applies the same rule without building the input: it never joins the + output, and it stops at the first non-empty streamed chunk. Truthy tool + calls count without being parsed, so a malformed payload is countable. + """ + if output.reasoning or output.tool_calls: + return True + text = output.output + if isinstance(text, str): + return bool(text) + return isinstance(text, tuple) and any(text) diff --git a/src/inference_endpoint/metrics/token_count_basis.py b/src/inference_endpoint/metrics/token_count_basis.py new file mode 100644 index 000000000..1359f7435 --- /dev/null +++ b/src/inference_endpoint/metrics/token_count_basis.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rules that decide which token counts a run uses. + +- ``classify_server_usage``: the one rule for whether a request's server-reported + usage is usable, shared by the live aggregator and every post-run tool; + ``read_server_usage`` also reads the completion count it allows. +- ``effective_token_count_source``: the configured basis after the endpoint's + API type is taken into account. +- ``token_count_source_for_run``: the basis a finished run used, for standalone + tools that only see its report directory; ``replay_basis`` adds a + ``--token-count-source`` override (``TOKEN_COUNT_SOURCE_METAVAR``). +- ``missing_usage_message``: the one message for samples without usable usage. + +This module is imported by the metrics-aggregator subprocess, so it must not +import ``config.schema`` (pydantic and cyclopts startup cost). +""" + +from __future__ import annotations + +import json +import sys +from collections.abc import Callable +from enum import Enum, auto +from pathlib import Path +from typing import Any, Final, TypeVar, cast + +import yaml + +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + has_countable_output, +) +from inference_endpoint.core.types import ( + APIType, + TextModelOutput, + TokenCountSource, + well_formed_count, +) + +_E = TypeVar("_E", bound=Enum) + + +class ServerUsage(Enum): + """How one request's server-reported usage reads under the usage rule.""" + + USABLE = auto() # both counts well-formed, prompt >= 1, completion >= 1 + TRUE_ZERO = auto() # completion 0 beside empty output: a count, not a gap + MISSING = auto() # anything else, including 0 beside non-empty output + + +def classify_server_usage( + prompt_tokens: object, completion_tokens: object, output: TextModelOutput +) -> ServerUsage: + """Classify one request's server usage; the one rule for every consumer. + + ``USABLE`` iff both counts are well-formed (``well_formed_count``) and at + least 1. ``TRUE_ZERO`` iff the prompt count is well-formed and the server + reports 0 completion tokens for an output with no countable content + (``has_countable_output``). Everything else is ``MISSING``, including 0 + completion tokens beside non-empty output. ``output`` is inspected only + when the counts are not usable. + """ + prompt = well_formed_count(prompt_tokens) + completion = well_formed_count(completion_tokens) + if prompt is None or completion is None: + return ServerUsage.MISSING + if prompt >= 1 and completion >= 1: + return ServerUsage.USABLE + if completion == 0 and not has_countable_output(output): + return ServerUsage.TRUE_ZERO + return ServerUsage.MISSING + + +def read_server_usage( + prompt_tokens: object, completion_tokens: object, output: TextModelOutput +) -> tuple[ServerUsage, int | None]: + """``classify_server_usage``, with the completion count it allows. + + The count is ``completion_tokens`` when the usage is ``USABLE`` (a + well-formed int), else ``None``. + """ + usage = classify_server_usage(prompt_tokens, completion_tokens, output) + if usage is ServerUsage.USABLE: + return usage, cast(int, completion_tokens) + return usage, None + + +def effective_token_count_source( + configured: TokenCountSource, api_type: APIType +) -> TokenCountSource: + """The basis a run counts on: videogen endpoints report no usage, so client.""" + if api_type is APIType.VIDEOGEN: + return TokenCountSource.CLIENT + return configured + + +# The basis of a run that recorded none (written before the setting existed). +# Deliberately separate from the schema default so changing the default never +# reinterprets archived runs. +ARCHIVED_RUN_BASIS: Final = TokenCountSource.CLIENT + +# argparse metavar of every ``--token-count-source`` flag. +TOKEN_COUNT_SOURCE_METAVAR: Final = ( + "{" + ",".join(source.value for source in TokenCountSource) + "}" +) + + +class RunBasisError(Exception): + """A run's recorded token-count basis could not be read.""" + + +def _read_mapping(path: Path, parse: Callable[[str], Any]) -> dict[str, Any] | None: + """Parse ``path`` as a mapping; None when the file does not exist.""" + try: + parsed = parse(path.read_text()) + except FileNotFoundError: + return None + except (OSError, ValueError, yaml.YAMLError) as exc: + raise RunBasisError( + f"cannot read the token-count basis from {path} ({exc}); " + "pass --token-count-source to choose it" + ) from exc + if not isinstance(parsed, dict): + raise RunBasisError( + f"cannot read the token-count basis from {path} (not a mapping); " + "pass --token-count-source to choose it" + ) + return parsed + + +def _section(config: dict[str, Any], key: str, path: Path) -> dict[str, Any]: + section = config.get(key) + if section is None: + return {} + if not isinstance(section, dict): + raise RunBasisError( + f"cannot read the token-count basis from {path} ({key!r} is not a " + "mapping); pass --token-count-source to choose it" + ) + return section + + +def _parse_enum(enum_type: type[_E], value: object, field: str, path: Path) -> _E: + try: + return enum_type(value) + except ValueError: + raise RunBasisError( + f"{field} in {path} is {value!r}, not one of " + f"{[member.value for member in enum_type]}; " + "pass --token-count-source to choose the basis" + ) from None + + +def _recorded_run_basis(run_dir: Path, summary_path: Path) -> TokenCountSource | None: + """The basis ``run_dir`` records, or None when it records none.""" + summary = _read_mapping(summary_path, json.loads) + if summary is not None and summary.get("token_count_source") is not None: + return _parse_enum( + TokenCountSource, + summary["token_count_source"], + "token_count_source", + summary_path, + ) + + config_path = run_dir / "config.yaml" + config = _read_mapping(config_path, yaml.safe_load) + if config is None: + return None + configured = _section(config, "settings", config_path).get("token_count_source") + if configured is None: + return None + source = _parse_enum( + TokenCountSource, + configured, + "settings.token_count_source", + config_path, + ) + api_type = _section(config, "endpoint_config", config_path).get("api_type") + if api_type is None: + return source + return effective_token_count_source( + source, + _parse_enum(APIType, api_type, "endpoint_config.api_type", config_path), + ) + + +def _default_summary_path(run_dir: Path) -> Path: + return run_dir / "performance" / "result_summary.json" + + +def token_count_source_for_run( + run_dir: Path, *, summary_path: Path | None = None +) -> TokenCountSource: + """The token-count basis of the run that wrote ``run_dir``. + + Resolution order: + + 1. top-level ``token_count_source`` in ``summary_path``, by default + ``performance/result_summary.json`` under ``run_dir``; + 2. ``settings.token_count_source`` in ``config.yaml`` (written at setup, so + it survives a run that died before the summary), resolved with + ``effective_token_count_source`` against ``endpoint_config.api_type``; + 3. ``ARCHIVED_RUN_BASIS``, with a note on stderr. + + Raises ``RunBasisError`` for a file that exists but cannot be read or an + unknown basis value. + """ + if summary_path is None: + summary_path = _default_summary_path(run_dir) + recorded = _recorded_run_basis(run_dir, summary_path) + if recorded is not None: + return recorded + print( + f"no token-count basis recorded (neither token_count_source in " + f"{summary_path} nor settings.token_count_source in " + f"{run_dir / 'config.yaml'}); using the {ARCHIVED_RUN_BASIS.value} basis, " + "override with --token-count-source", + file=sys.stderr, + ) + return ARCHIVED_RUN_BASIS + + +def replay_basis( + run_dir: Path, + *, + override: TokenCountSource | None, + summary_path: Path | None = None, +) -> tuple[TokenCountSource | None, TokenCountSource]: + """``(primary, basis)`` for a post-run replay of the run in ``run_dir``. + + ``primary`` is the run's own basis (``token_count_source_for_run``) and + ``basis`` the one to replay on: ``override`` when given, else ``primary``. + Beside an override, a run that records no basis is ``ARCHIVED_RUN_BASIS`` + without the stderr note, and one whose basis cannot be read has a ``None`` + primary; each caller decides what that means for its output. Without an + override, ``RunBasisError`` propagates. + """ + if override is None: + primary = token_count_source_for_run(run_dir, summary_path=summary_path) + return primary, primary + if summary_path is None: + summary_path = _default_summary_path(run_dir) + try: + recorded = _recorded_run_basis(run_dir, summary_path) + except RunBasisError: + return None, override + return (ARCHIVED_RUN_BASIS if recorded is None else recorded), override + + +def missing_usage_message( + n_missing: int, + n_empty_output: int, + n_succeeded: int, + *, + primary: TokenCountSource, +) -> str: + """Why samples lack usable server usage, and what that means for the run. + + On the server basis the run is incomplete; the remedy is client counting, + which fixes samples that produced output but no usage and cannot fix + samples with empty output (server errors or cut streams). On the client + basis nothing failed: the ``*_server`` blocks just cover fewer samples. + """ + head = ( + f"{n_missing} of {n_succeeded} successful tracked samples returned no " + "server token usage" + ) + if primary is TokenCountSource.CLIENT: + return f"{head}; the *_server blocks are partial and the run is not affected" + parts = [head] + if n_missing > n_empty_output: + # `benchmark from-config` takes no --token-count-source; name the key too. + parts.append( + "set --token-count-source client if this server cannot report usage " + "(YAML: settings.token_count_source: client)" + ) + if n_empty_output > 0: + parts.append( + f"{n_empty_output} had empty output — likely server errors or cut " + "streams; client mode will not fix these" + ) + return "; ".join(parts) diff --git a/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer.json b/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer.json new file mode 100644 index 000000000..d9eeb653a --- /dev/null +++ b/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer.json @@ -0,0 +1,187 @@ +{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [ + { + "id": 0, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 1, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 2, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 3, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ], + "normalizer": null, + "pre_tokenizer": { + "type": "Split", + "pattern": { + "String": "" + }, + "behavior": "Isolated", + "invert": false + }, + "post_processor": { + "type": "TemplateProcessing", + "single": [ + { + "Sequence": { + "id": "A", + "type_id": 0 + } + } + ], + "pair": [ + { + "Sequence": { + "id": "A", + "type_id": 0 + } + }, + { + "Sequence": { + "id": "B", + "type_id": 1 + } + } + ], + "special_tokens": {} + }, + "decoder": null, + "model": { + "type": "WordLevel", + "vocab": { + "": 0, + "": 1, + "": 2, + "": 3, + "a": 4, + "b": 5, + "c": 6, + "d": 7, + "e": 8, + "f": 9, + "g": 10, + "h": 11, + "i": 12, + "j": 13, + "k": 14, + "l": 15, + "m": 16, + "n": 17, + "o": 18, + "p": 19, + "q": 20, + "r": 21, + "s": 22, + "t": 23, + "u": 24, + "v": 25, + "w": 26, + "x": 27, + "y": 28, + "z": 29, + "A": 30, + "B": 31, + "C": 32, + "D": 33, + "E": 34, + "F": 35, + "G": 36, + "H": 37, + "I": 38, + "J": 39, + "K": 40, + "L": 41, + "M": 42, + "N": 43, + "O": 44, + "P": 45, + "Q": 46, + "R": 47, + "S": 48, + "T": 49, + "U": 50, + "V": 51, + "W": 52, + "X": 53, + "Y": 54, + "Z": 55, + "0": 56, + "1": 57, + "2": 58, + "3": 59, + "4": 60, + "5": 61, + "6": 62, + "7": 63, + "8": 64, + "9": 65, + " ": 66, + "\t": 67, + "\n": 68, + "\r": 69, + "!": 70, + "\"": 71, + "#": 72, + "$": 73, + "%": 74, + "&": 75, + "'": 76, + "(": 77, + ")": 78, + "*": 79, + "+": 80, + ",": 81, + "-": 82, + ".": 83, + "/": 84, + ":": 85, + ";": 86, + "<": 87, + "=": 88, + ">": 89, + "?": 90, + "@": 91, + "[": 92, + "\\": 93, + "]": 94, + "^": 95, + "_": 96, + "`": 97, + "{": 98, + "|": 99, + "}": 100, + "~": 101 + }, + "unk_token": "" + } +} diff --git a/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer_config.json b/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer_config.json new file mode 100644 index 000000000..8a3ddac79 --- /dev/null +++ b/tests/assets/tokenizers/char_chat_bare_reasoning/tokenizer_config.json @@ -0,0 +1,10 @@ +{ + "backend": "tokenizers", + "bos_token": "", + "eos_token": "", + "model_max_length": 1000000000000000019884624838656, + "pad_token": "", + "tokenizer_class": "TokenizersBackend", + "unk_token": "", + "chat_template": "{% if tools %}<|tools|>{% for tool in tools %}{{ tool | tojson }}{% endfor %}<|/tools|>\n{% endif %}{% for message in messages %}<|im_start|>{{ message.role }}\n{% if message.reasoning_content %}{{ message.reasoning_content }}{% endif %}{{ message.content }}{% if message.tool_calls %}{% for tc in message.tool_calls %}\n{{ tc.function.name }}({{ tc.function.arguments | tojson }}){% endfor %}{% endif %}<|im_end|>\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" +} diff --git a/tests/integration/async_utils/services/metrics_aggregator/test_signal_handling.py b/tests/integration/async_utils/services/metrics_aggregator/test_signal_handling.py index 78ba91dcb..ea5a26950 100644 --- a/tests/integration/async_utils/services/metrics_aggregator/test_signal_handling.py +++ b/tests/integration/async_utils/services/metrics_aggregator/test_signal_handling.py @@ -36,6 +36,8 @@ import pytest +from tests.test_helpers import DEFAULT_SOURCE + def _spawn_aggregator( socket_dir: Path, @@ -73,6 +75,8 @@ def _spawn_aggregator( "5", "--tokenizer-workers", "0", + "--token-count-source", + DEFAULT_SOURCE.value, ] return subprocess.Popen( cmd, diff --git a/tests/unit/async_utils/services/metrics_aggregator/conftest.py b/tests/unit/async_utils/services/metrics_aggregator/conftest.py index 615c28269..d8f1a01e7 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/conftest.py +++ b/tests/unit/async_utils/services/metrics_aggregator/conftest.py @@ -27,8 +27,13 @@ from __future__ import annotations import asyncio +import contextlib +from collections.abc import Iterator +from pathlib import Path from unittest.mock import AsyncMock, MagicMock +import msgspec +import pytest from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( MetricsAggregatorService, ) @@ -52,7 +57,7 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import TextModelOutput +from inference_endpoint.core.types import TextModelOutput, TokenCountSource # --------------------------------------------------------------------------- # Mock BatchTokenizer — whitespace token counts; matches the BatchTokenizer @@ -79,8 +84,6 @@ def __init__(self, delay: float = 0.0) -> None: async def count_batch_async(self, inputs, _loop, live=False): if self._delay: await asyncio.sleep(self._delay) - import msgspec - outcomes = [] for item in inputs: if isinstance(item, TokenIdsInput): @@ -137,6 +140,25 @@ def sample_event( return EventRecord(event_type=ev_type, timestamp_ns=ts, sample_uuid=uuid, data=data) +def complete_event( + uuid: str, + ts: int, + data=None, + *, + prompt_tokens: int | msgspec.UnsetType = msgspec.UNSET, + completion_tokens: int | msgspec.UnsetType = msgspec.UNSET, +) -> EventRecord: + """A COMPLETE record carrying the server usage counts (UNSET = not reported).""" + return EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=ts, + sample_uuid=uuid, + data=data, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + + def text_output(s: str) -> TextModelOutput: return TextModelOutput(output=s) @@ -187,11 +209,18 @@ def snapshot_series_total(registry: MetricsRegistry, name: str) -> int | float: # --------------------------------------------------------------------------- +@pytest.fixture(params=list(TokenCountSource), ids=lambda source: source.value) +def primary(request: pytest.FixtureRequest) -> TokenCountSource: + """Each primary token-count basis in turn.""" + return request.param + + def make_aggregator( zmq_ctx: ManagedZMQContext, loop: asyncio.AbstractEventLoop, socket_name: str, *, + token_count_source: TokenCountSource, tokenizer=None, live_flush_interval_s: float | None = None, streaming: bool = True, @@ -199,6 +228,7 @@ def make_aggregator( shutdown_event: asyncio.Event | None = None, drain_timeout_s: float | None = None, steady_state_profile: str | None = None, + registry: MetricsRegistry | None = None, ) -> tuple[MetricsAggregatorService, MetricsRegistry, MagicMock]: """Construct an aggregator wired to a real SUB socket and a mocked publisher. @@ -210,7 +240,7 @@ def make_aggregator( Returns ``(agg, registry, publisher_mock)``. """ - registry = MetricsRegistry() + registry = registry or MetricsRegistry() # ``publish_final`` and ``aclose`` are awaited by the aggregator's # ENDED handler, so they must be AsyncMocks. The remaining surface # (``start``, ``close``) is synchronous and falls back to MagicMock's @@ -234,5 +264,21 @@ def make_aggregator( shutdown_event=shutdown_event, drain_timeout_s=drain_timeout_s, steady_state_profile=steady_state_profile, + token_count_source=token_count_source, ) return agg, registry, publisher + + +@contextlib.contextmanager +def aggregator_scope( + tmp_path: Path, socket_name: str, **kwargs +) -> Iterator[tuple[MetricsAggregatorService, MetricsRegistry, MagicMock]]: + """``make_aggregator`` inside a scoped ZMQ context; closes the aggregator.""" + with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: + agg, registry, publisher = make_aggregator( + ctx, asyncio.get_event_loop(), socket_name, **kwargs + ) + try: + yield agg, registry, publisher + finally: + agg.close() diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py index b2d5f1b64..2aabda4d3 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -26,14 +26,24 @@ from __future__ import annotations import asyncio +import logging +import msgspec import pytest from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( - MetricCounterKey, + MetricsAggregatorService, ) from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import ( + MetricCounterKey, MetricSeriesKey, ) +from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( + MetricsRegistry, +) +from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( + SeriesStat, + SessionState, +) from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext from inference_endpoint.core.record import ( ErrorEventType, @@ -41,10 +51,27 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import ErrorData, PromptData, TextModelOutput +from inference_endpoint.core.types import ( + ErrorData, + PhaseData, + PhaseType, + PromptData, + TextModelOutput, + TokenCountSource, + basis_key, +) +from inference_endpoint.metrics.early_stopping import EarlyStoppingSpec +from inference_endpoint.metrics.token_count_basis import ( + ServerUsage, + missing_usage_message, +) + +from tests.test_helpers import DEFAULT_SOURCE, SERVER_USAGE_CASES from .conftest import ( MockBatchTokenizer, + aggregator_scope, + complete_event, make_aggregator, sample_event, session_event, @@ -55,6 +82,16 @@ text_output, ) +CLIENT, SERVER = TokenCountSource.CLIENT, TokenCountSource.SERVER + + +def _name( + key: MetricSeriesKey, basis: TokenCountSource, *, primary: TokenCountSource +) -> str: + """``key``'s series name on ``basis`` in a run whose primary is ``primary``.""" + return basis_key(key.value, basis, primary=primary) + + # --------------------------------------------------------------------------- # Performance tracking window # --------------------------------------------------------------------------- @@ -66,7 +103,9 @@ class TestTrackingWindow: async def test_not_tracked_before_start(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_not_tracked_before") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_not_tracked_before", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -95,7 +134,9 @@ async def test_not_tracked_before_start(self, tmp_path): async def test_tracked_after_start(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_tracked_after_start") + agg, _, _ = make_aggregator( + ctx, loop, "agg_tracked_after_start", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -113,7 +154,12 @@ async def test_tracked_after_start(self, tmp_path): async def test_not_tracked_after_stop(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_not_tracked_after_stop") + agg, _, _ = make_aggregator( + ctx, + loop, + "agg_not_tracked_after_stop", + token_count_source=DEFAULT_SOURCE, + ) try: await agg.process( [ @@ -135,7 +181,9 @@ async def test_inflight_sample_continues_after_stop(self, tmp_path): """A sample issued during tracking completes normally after STOP.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_inflight_after_stop") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_inflight_after_stop", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -169,7 +217,9 @@ async def test_restart_tracking_window(self, tmp_path): """START -> STOP -> START creates a second tracking window.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_restart_tracking") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_restart_tracking", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -208,7 +258,9 @@ async def test_tracked_block_durations(self, tmp_path): """Tracked blocks extend to last sample completion.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_tracked_block_dur") + agg, _, _ = make_aggregator( + ctx, loop, "agg_tracked_block_dur", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -246,7 +298,9 @@ class TestTimingMetrics: async def test_ttft_and_sample_latency(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_ttft_latency") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_ttft_latency", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -277,7 +331,9 @@ async def test_ttft_and_sample_latency(self, tmp_path): async def test_chunk_deltas(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_chunk_deltas") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_chunk_deltas", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -309,13 +365,17 @@ async def test_chunk_deltas(self, tmp_path): agg.close() @pytest.mark.asyncio - async def test_non_streaming_latency_only(self, tmp_path): + async def test_non_streaming_latency_only(self, tmp_path, primary): """Non-streaming: emits sample_latency_ns + OSL, no TTFT/chunk_delta/TPOT.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.0) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_non_streaming", tokenizer=tokenizer + ctx, + loop, + "agg_non_streaming", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -341,7 +401,12 @@ async def test_non_streaming_latency_only(self, tmp_path): == 2000 ) # OSL = token_count("hello world") = 2 - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 2 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 2 + ) assert ( snapshot_series_count(registry, MetricSeriesKey.TTFT_NS.value) == 0 ) @@ -351,9 +416,9 @@ async def test_non_streaming_latency_only(self, tmp_path): ) == 0 ) - assert ( - snapshot_series_count(registry, MetricSeriesKey.TPOT_NS.value) == 0 - ) + for basis in TokenCountSource: + tpot = _name(MetricSeriesKey.TPOT_NS, basis, primary=primary) + assert snapshot_series_count(registry, tpot) == 0 finally: agg.close() @@ -362,7 +427,9 @@ async def test_chunk_delta_not_emitted_without_last_recv(self, tmp_path): """RECV_NON_FIRST without prior RECV_FIRST: no chunk_delta emitted.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_chunk_delta_no_recv") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_chunk_delta_no_recv", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -387,11 +454,13 @@ async def test_chunk_delta_not_emitted_without_last_recv(self, tmp_path): @pytest.mark.unit class TestIsl: @pytest.mark.asyncio - async def test_issued_with_token_ids_emits_isl_directly(self, tmp_path): + async def test_issued_with_token_ids_emits_isl_directly(self, tmp_path, primary): """SGLang path: PromptData with token_ids emits ISL = len(token_ids).""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_isl_token_ids") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_isl_token_ids", token_count_source=primary + ) try: await agg.process( [ @@ -406,15 +475,25 @@ async def test_issued_with_token_ids_emits_isl_directly(self, tmp_path): ), ] ) - assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) == 5 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + == 5 + ) + # Server ISL exists only on COMPLETE. + server_isl = _name(MetricSeriesKey.ISL, SERVER, primary=primary) + assert snapshot_series_count(registry, server_isl) == 0 finally: agg.close() @pytest.mark.asyncio - async def test_issued_without_data_no_isl(self, tmp_path): + async def test_issued_without_data_no_isl(self, tmp_path, primary): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_isl_no_data") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_isl_no_data", token_count_source=primary + ) try: await agg.process( [ @@ -424,13 +503,18 @@ async def test_issued_without_data_no_isl(self, tmp_path): sample_event(SampleEventType.ISSUED, "s1", ts=1000), ] ) - assert snapshot_series_count(registry, MetricSeriesKey.ISL.value) == 0 + for basis in TokenCountSource: + isl = _name(MetricSeriesKey.ISL, basis, primary=primary) + assert snapshot_series_count(registry, isl) == 0 finally: agg.close() @pytest.mark.asyncio - async def test_disabled_isl_skips_prompt_tokenization_keeps_osl(self, tmp_path): - """enable_isl=False: no ISL series, no prompt enqueued; OSL/TPOT intact.""" + async def test_disabled_isl_skips_prompt_tokenization_keeps_osl( + self, tmp_path, primary + ): + """enable_isl=False: no ISL series on either basis, no prompt enqueued; + OSL/TPOT intact.""" loop = asyncio.get_event_loop() messages = ( {"role": "system", "content": "follow instructions"}, @@ -441,6 +525,7 @@ async def test_disabled_isl_skips_prompt_tokenization_keeps_osl(self, tmp_path): ctx, loop, "agg_isl_disabled", + token_count_source=primary, tokenizer=MockBatchTokenizer(), enable_isl=False, ) @@ -458,7 +543,9 @@ async def test_disabled_isl_skips_prompt_tokenization_keeps_osl(self, tmp_path): ), ] ) - assert not registry.has_series(MetricSeriesKey.ISL.value) + for basis in TokenCountSource: + isl = _name(MetricSeriesKey.ISL, basis, primary=primary) + assert not registry.has_series(isl) assert agg.pending_tokens == 0 await agg.process( @@ -473,9 +560,14 @@ async def test_disabled_isl_skips_prompt_tokenization_keeps_osl(self, tmp_path): session_event(SessionEventType.ENDED, ts=5000), ] ) - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 3 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 3 + ) assert snapshot_series_total( - registry, MetricSeriesKey.TPOT_NS.value + registry, _name(MetricSeriesKey.TPOT_NS, CLIENT, primary=primary) ) == pytest.approx(1500.0) finally: agg.close() @@ -492,7 +584,9 @@ class TestEdgeCases: async def test_untracked_sample_events_ignored(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_untracked_ignored") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_untracked_ignored", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -529,7 +623,9 @@ async def test_duplicate_started_logs_error_and_preserves_state( """ loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_dup_started") + agg, _, _ = make_aggregator( + ctx, loop, "agg_dup_started", token_count_source=DEFAULT_SOURCE + ) try: with caplog.at_level("ERROR"): await agg.process( @@ -556,7 +652,9 @@ async def test_duplicate_started_logs_error_and_preserves_state( async def test_complete_removes_row(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_complete_removes") + agg, _, _ = make_aggregator( + ctx, loop, "agg_complete_removes", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -577,7 +675,9 @@ async def test_session_ended_calls_publish_final(self, tmp_path): """ENDED triggers ``publish_final`` and ``close`` on the publisher.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, publisher = make_aggregator(ctx, loop, "agg_ended_publish_final") + agg, _, publisher = make_aggregator( + ctx, loop, "agg_ended_publish_final", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -601,7 +701,9 @@ async def test_interrupted_marker_tags_final_snapshot(self, tmp_path): must never produce a COMPLETE final snapshot.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, publisher = make_aggregator(ctx, loop, "agg_interrupted_marker") + agg, _, publisher = make_aggregator( + ctx, loop, "agg_interrupted_marker", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -619,7 +721,9 @@ async def test_interrupted_marker_tags_final_snapshot(self, tmp_path): async def test_events_after_ended_are_dropped(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_events_after_ended") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_events_after_ended", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -643,7 +747,9 @@ async def test_events_after_ended_are_dropped(self, tmp_path): async def test_empty_sample_uuid_ignored(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_empty_uuid") + agg, _, _ = make_aggregator( + ctx, loop, "agg_empty_uuid", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -661,7 +767,9 @@ async def test_empty_sample_uuid_ignored(self, tmp_path): async def test_multiple_samples_independent(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_multi_independent") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_multi_independent", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -710,7 +818,9 @@ async def test_error_event_increments_total_failed(self, tmp_path): """ loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_error_total") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_error_total", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -741,7 +851,9 @@ async def test_error_event_increments_total_failed(self, tmp_path): async def test_session_started_stores_timestamp(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_started_ts") + agg, _, _ = make_aggregator( + ctx, loop, "agg_started_ts", token_count_source=DEFAULT_SOURCE + ) try: await agg.process([session_event(SessionEventType.STARTED, ts=42)]) assert agg._table.session_started_ns == 42 @@ -753,7 +865,9 @@ async def test_process_multiple_batches(self, tmp_path): """Two sequential process() calls maintain state correctly.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_multi_batch") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_multi_batch", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -790,7 +904,9 @@ async def test_ended_in_second_batch(self, tmp_path): """ENDED in a later batch still triggers finalize.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, publisher = make_aggregator(ctx, loop, "agg_ended_second_batch") + agg, _, publisher = make_aggregator( + ctx, loop, "agg_ended_second_batch", token_count_source=DEFAULT_SOURCE + ) try: await agg.process([session_event(SessionEventType.STARTED, ts=0)]) publisher.publish_final.assert_not_awaited() @@ -818,7 +934,9 @@ async def test_loadgen_window_duration_emitted(self, tmp_path): """ loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_loadgen_window") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_loadgen_window", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -858,7 +976,9 @@ async def test_total_vs_tracked_counters(self, tmp_path): """ loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_total_vs_tracked") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_total_vs_tracked", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -887,7 +1007,9 @@ async def test_total_vs_tracked_counters(self, tmp_path): async def test_finish_reason_counters_bucket_unknown_values(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, registry, _ = make_aggregator(ctx, loop, "agg_finish_reasons") + agg, registry, _ = make_aggregator( + ctx, loop, "agg_finish_reasons", token_count_source=DEFAULT_SOURCE + ) try: await agg.process( [ @@ -926,13 +1048,17 @@ async def test_finish_reason_counters_bucket_unknown_values(self, tmp_path): @pytest.mark.unit class TestAsyncTriggers: @pytest.mark.asyncio - async def test_isl_text_path_async(self, tmp_path): + async def test_isl_text_path_async(self, tmp_path, primary): """ISL with text prompt triggers async tokenization.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.01) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_isl_text_async", tokenizer=tokenizer + ctx, + loop, + "agg_isl_text_async", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -950,12 +1076,19 @@ async def test_isl_text_path_async(self, tmp_path): ) # ISL task is in-flight; drain it await agg._token_queue.drain_all() - assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) == 4 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + == 4 + ) finally: agg.close() @pytest.mark.asyncio - async def test_isl_structured_prompt_includes_reasoning_and_tools(self, tmp_path): + async def test_isl_structured_prompt_includes_reasoning_and_tools( + self, tmp_path, primary + ): loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer() messages = ( @@ -971,7 +1104,11 @@ async def test_isl_structured_prompt_includes_reasoning_and_tools(self, tmp_path tools = ({"function": {"name": "lookup"}},) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_isl_structured", tokenizer=tokenizer + ctx, + loop, + "agg_isl_structured", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -988,19 +1125,33 @@ async def test_isl_structured_prompt_includes_reasoning_and_tools(self, tmp_path ] ) await agg._token_queue.drain_all() - assert snapshot_series_count(registry, MetricSeriesKey.ISL.value) == 1 - assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) > 5 + assert ( + snapshot_series_count( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + == 1 + ) + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + > 5 + ) finally: agg.close() @pytest.mark.asyncio - async def test_osl_emitted_on_complete(self, tmp_path): + async def test_osl_emitted_on_complete(self, tmp_path, primary): """OSL is emitted via async tokenization when COMPLETE carries text.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.01) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_osl_complete", tokenizer=tokenizer + ctx, + loop, + "agg_osl_complete", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1026,18 +1177,27 @@ async def test_osl_emitted_on_complete(self, tmp_path): == 4000 ) # OSL = 4 tokens - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 4 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 4 + ) finally: agg.close() @pytest.mark.asyncio - async def test_tpot_emitted_for_streaming(self, tmp_path): + async def test_tpot_emitted_for_streaming(self, tmp_path, primary): """TPOT is emitted for streaming responses using text_after_first_chunk.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.0) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_tpot_streaming", tokenizer=tokenizer + ctx, + loop, + "agg_tpot_streaming", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1058,22 +1218,31 @@ async def test_tpot_emitted_for_streaming(self, tmp_path): ) await agg._token_queue.drain_all() # OSL = "hello world foo" = 3 tokens - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 3 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 3 + ) # tpot = (5000 - 2000) / token_count("world foo") = 3000 / 2 = 1500 assert snapshot_series_total( - registry, MetricSeriesKey.TPOT_NS.value + registry, _name(MetricSeriesKey.TPOT_NS, CLIENT, primary=primary) ) == pytest.approx(1500.0) finally: agg.close() @pytest.mark.asyncio - async def test_tpot_skipped_when_single_chunk(self, tmp_path): + async def test_tpot_skipped_when_single_chunk(self, tmp_path, primary): """TPOT is not emitted when there are no tokens after the first chunk.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.0) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_tpot_single_chunk", tokenizer=tokenizer + ctx, + loop, + "agg_tpot_single_chunk", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1093,15 +1262,24 @@ async def test_tpot_skipped_when_single_chunk(self, tmp_path): ] ) await agg._token_queue.drain_all() - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 1 assert ( - snapshot_series_count(registry, MetricSeriesKey.TPOT_NS.value) == 0 + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 1 + ) + assert ( + snapshot_series_count( + registry, + _name(MetricSeriesKey.TPOT_NS, CLIENT, primary=primary), + ) + == 0 ) finally: agg.close() @pytest.mark.asyncio - async def test_tpot_not_emitted_without_streaming_flag(self, tmp_path): + async def test_tpot_not_emitted_without_streaming_flag(self, tmp_path, primary): """When ``streaming=False``, TPOT/TTFT/chunk_delta series are NOT registered at all — the aggregator's snapshot has no entry for them. """ @@ -1114,6 +1292,7 @@ async def test_tpot_not_emitted_without_streaming_flag(self, tmp_path): "agg_tpot_no_streaming", tokenizer=tokenizer, streaming=False, + token_count_source=primary, ) try: await agg.process( @@ -1139,22 +1318,33 @@ async def test_tpot_not_emitted_without_streaming_flag(self, tmp_path): ) == 4000 ) - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 3 - # The streaming-only series are unregistered. - assert not registry.has_series(MetricSeriesKey.TPOT_NS.value) + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 3 + ) + # The streaming-only series are unregistered, on both bases. + for basis in TokenCountSource: + tpot = _name(MetricSeriesKey.TPOT_NS, basis, primary=primary) + assert not registry.has_series(tpot) assert not registry.has_series(MetricSeriesKey.TTFT_NS.value) assert not registry.has_series(MetricSeriesKey.CHUNK_DELTA_NS.value) finally: agg.close() @pytest.mark.asyncio - async def test_tpot_non_streaming_output_skipped(self, tmp_path): + async def test_tpot_non_streaming_output_skipped(self, tmp_path, primary): """TPOT is not emitted for non-streaming (str) TextModelOutput.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.0) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_tpot_str_output", tokenizer=tokenizer + ctx, + loop, + "agg_tpot_str_output", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1174,9 +1364,18 @@ async def test_tpot_non_streaming_output_skipped(self, tmp_path): ] ) await agg._token_queue.drain_all() - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 3 assert ( - snapshot_series_count(registry, MetricSeriesKey.TPOT_NS.value) == 0 + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 3 + ) + assert ( + snapshot_series_count( + registry, + _name(MetricSeriesKey.TPOT_NS, CLIENT, primary=primary), + ) + == 0 ) finally: agg.close() @@ -1190,6 +1389,7 @@ async def test_started_arms_the_live_flush_loop(self, tmp_path): ctx, loop, "agg_live_arm", + token_count_source=DEFAULT_SOURCE, tokenizer=MockBatchTokenizer(), live_flush_interval_s=0.01, ) @@ -1205,13 +1405,17 @@ async def test_started_arms_the_live_flush_loop(self, tmp_path): agg.close() @pytest.mark.asyncio - async def test_flush_records_buffered_tokenizations(self, tmp_path): + async def test_flush_records_buffered_tokenizations(self, tmp_path, primary): """fire() buffers tokenization; flush() tokenizes the batch and records.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_flush_records", tokenizer=tokenizer + ctx, + loop, + "agg_flush_records", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1233,18 +1437,27 @@ async def test_flush_records_buffered_tokenizations(self, tmp_path): await agg._token_queue.drain_all() assert agg._token_queue.pending == 0 - assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) == 5 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + == 5 + ) finally: agg.close() @pytest.mark.asyncio - async def test_shutdown_flushes_buffered_tokenizations(self, tmp_path): + async def test_shutdown_flushes_buffered_tokenizations(self, tmp_path, primary): """ENDED flushes buffered tokenizations before finalizing.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.02) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, publisher = make_aggregator( - ctx, loop, "agg_shutdown_drain", tokenizer=tokenizer + ctx, + loop, + "agg_shutdown_drain", + tokenizer=tokenizer, + token_count_source=primary, ) try: await agg.process( @@ -1262,7 +1475,12 @@ async def test_shutdown_flushes_buffered_tokenizations(self, tmp_path): ] ) # After ENDED, flush_remaining ran inside process() — ISL emitted. - assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) == 3 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.ISL, CLIENT, primary=primary) + ) + == 3 + ) publisher.publish_final.assert_awaited_once() finally: agg.close() @@ -1283,7 +1501,11 @@ async def count_batch_async(self, inputs, _loop, live=False): with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, _, publisher = make_aggregator( - ctx, loop, "agg_drain_failure", tokenizer=FailingBatchTokenizer() + ctx, + loop, + "agg_drain_failure", + token_count_source=DEFAULT_SOURCE, + tokenizer=FailingBatchTokenizer(), ) try: await agg.process( @@ -1331,6 +1553,7 @@ async def count_batch_async(self, inputs, _loop, live=False): ctx, loop, "agg_drain_timeout", + token_count_source=DEFAULT_SOURCE, tokenizer=BlockingBatchTokenizer(), ) agg._drain_timeout_s = 0.05 @@ -1364,13 +1587,17 @@ async def count_batch_async(self, inputs, _loop, live=False): agg.close() @pytest.mark.asyncio - async def test_tpot_osl_for_tool_call_complete(self, tmp_path): + async def test_tpot_osl_for_tool_call_complete(self, tmp_path, primary): """OSL and TPOT use message-path tokenization when COMPLETE carries tool_calls.""" loop = asyncio.get_event_loop() tokenizer = MockBatchTokenizer(delay=0.0) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, registry, _ = make_aggregator( - ctx, loop, "agg_tpot_osl_tool_call", tokenizer=tokenizer + ctx, + loop, + "agg_tpot_osl_tool_call", + tokenizer=tokenizer, + token_count_source=primary, ) try: tool_call = { @@ -1397,10 +1624,377 @@ async def test_tpot_osl_for_tool_call_complete(self, tmp_path): ) await agg._token_queue.drain_all() # OSL = token_count("ok" + tool_calls_json) = 2 - assert snapshot_series_total(registry, MetricSeriesKey.OSL.value) == 2 + assert ( + snapshot_series_total( + registry, _name(MetricSeriesKey.OSL, CLIENT, primary=primary) + ) + == 2 + ) # tpot = (5000 - 2000) / token_count(tool_calls_json) = 3000 / 1 = 3000 assert snapshot_series_total( - registry, MetricSeriesKey.TPOT_NS.value + registry, _name(MetricSeriesKey.TPOT_NS, CLIENT, primary=primary) ) == pytest.approx(3000.0) finally: agg.close() + + +# --------------------------------------------------------------------------- +# Token-count bases: primary series unsuffixed, the other basis suffixed +# --------------------------------------------------------------------------- + +_START = session_event(SessionEventType.START_PERFORMANCE_TRACKING, ts=0) +_TOKEN_KEYS = (MetricSeriesKey.ISL, MetricSeriesKey.OSL, MetricSeriesKey.TPOT_NS) +# One streamed sample: ISSUED 1000, first chunk 2000, COMPLETE 5000. Its output +# is 7 whitespace tokens, 6 after the first chunk, beside server usage 10 / 5. +_BOTH_BASES_SAMPLE = [ + sample_event( + SampleEventType.ISSUED, "s1", ts=1000, data=PromptData(text="w w w w") + ), + sample_event(SampleEventType.RECV_FIRST, "s1", ts=2000), + complete_event( + "s1", + 5000, + streaming_text("a", " b c d e f g"), + prompt_tokens=10, + completion_tokens=5, + ), +] +# TPOT of _BOTH_BASES_SAMPLE: 3000 ns over 6 tokenized tokens, and over +# completion_tokens - 1 = 4 server tokens. +_TPOT_BY_BASIS = {CLIENT: 500.0, SERVER: 750.0} +_CASES = [pytest.param(*case[1:], id=case[0]) for case in SERVER_USAGE_CASES] + + +@pytest.mark.unit +class TestTokenCountBases: + @pytest.mark.asyncio + async def test_series_registered_primary_unsuffixed_then_other( + self, tmp_path, primary + ): + with aggregator_scope(tmp_path, "agg_order", token_count_source=primary) as ( + _, + reg, + _, + ): + snap = reg.build_snapshot(state=SessionState.LIVE, n_pending_tasks=0) + names = [m.name for m in snap.metrics if isinstance(m, SeriesStat)] + assert names == [ + "sample_latency_ns", + "isl", + "osl", + "ttft_ns", + "chunk_delta_ns", + "tpot_ns", + *(_name(key, primary.other, primary=primary) for key in _TOKEN_KEYS), + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("streaming", [True, False]) + async def test_no_metrics_isl_drops_isl_on_both_bases( + self, tmp_path, primary, streaming + ): + with aggregator_scope( + tmp_path, + "agg_no_isl", + token_count_source=primary, + tokenizer=MockBatchTokenizer(), + enable_isl=False, + streaming=streaming, + ) as (agg, reg, _): + await agg.process([_START, *_BOTH_BASES_SAMPLE]) + # Only the output was enqueued; the prompt never was. + assert agg.pending_tokens == (2 if streaming else 1) + for basis in TokenCountSource: + assert not reg.has_series( + _name(MetricSeriesKey.ISL, basis, primary=primary) + ) + assert reg.has_series( + _name(MetricSeriesKey.OSL, basis, primary=primary) + ) + tpot = _name(MetricSeriesKey.TPOT_NS, basis, primary=primary) + assert reg.has_series(tpot) is streaming + + @pytest.mark.asyncio + async def test_each_basis_records_its_own_counts(self, tmp_path, primary): + with aggregator_scope( + tmp_path, + "agg_bases", + token_count_source=primary, + tokenizer=MockBatchTokenizer(), + ) as (agg, reg, _): + await agg.process([_START, *_BOTH_BASES_SAMPLE]) + await agg.token_queue.drain_all() + totals = { + (key, basis): snapshot_series_total( + reg, _name(key, basis, primary=primary) + ) + for key in _TOKEN_KEYS + for basis in TokenCountSource + } + assert totals == { + (MetricSeriesKey.ISL, CLIENT): 4, + (MetricSeriesKey.ISL, SERVER): 10, + (MetricSeriesKey.OSL, CLIENT): 7, + (MetricSeriesKey.OSL, SERVER): 5, + (MetricSeriesKey.TPOT_NS, CLIENT): _TPOT_BY_BASIS[CLIENT], + (MetricSeriesKey.TPOT_NS, SERVER): _TPOT_BY_BASIS[SERVER], + } + + @pytest.mark.asyncio + async def test_server_basis_needs_no_tokenizer(self, tmp_path, primary): + with aggregator_scope(tmp_path, "agg_srv", token_count_source=primary) as ( + agg, + reg, + _, + ): + await agg.process( + [ + sample_event(SampleEventType.ISSUED, "late", ts=1), # untracked + _START, + *( + sample_event(SampleEventType.ISSUED, u, ts=10) + for u in ("s1", "s2", "err") + ), + complete_event( + "s1", + 50, + text_output("x y"), + prompt_tokens=7, + completion_tokens=5, + ), + # A true zero: nothing is recorded on either basis. + complete_event( + "s2", 50, text_output(""), prompt_tokens=3, completion_tokens=0 + ), + complete_event( + "err", 50, None, prompt_tokens=4, completion_tokens=4 + ), + complete_event( + "late", + 50, + text_output("q"), + prompt_tokens=50, + completion_tokens=50, + ), + ] + ) + isl = _name(MetricSeriesKey.ISL, SERVER, primary=primary) + osl = _name(MetricSeriesKey.OSL, SERVER, primary=primary) + assert ( + snapshot_series_count(reg, isl), + snapshot_series_total(reg, isl), + ) == ( + 1, + 7, + ) + assert ( + snapshot_series_count(reg, osl), + snapshot_series_total(reg, osl), + ) == ( + 1, + 5, + ) + assert agg.pending_tokens == 0 + assert snapshot_counters(reg)["tracked_samples_missing_usage"] == 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize(("prompt", "completion", "output", "expected"), _CASES) + async def test_missing_usage_counted_for_both_primaries( + self, tmp_path, primary, prompt, completion, output, expected + ): + with aggregator_scope(tmp_path, "agg_miss", token_count_source=primary) as ( + agg, + reg, + _, + ): + await agg.process( + [ + _START, + sample_event(SampleEventType.ISSUED, "s1", ts=10), + sample_event(SampleEventType.ISSUED, "err", ts=10), + complete_event( + "s1", + 20, + output, + prompt_tokens=prompt, + completion_tokens=completion, + ), + # An error completes without data or usage: never missing. + complete_event("err", 20, None), + ] + ) + c = snapshot_counters(reg) + recorded = { + key: snapshot_series_count(reg, _name(key, SERVER, primary=primary)) + for key in (MetricSeriesKey.ISL, MetricSeriesKey.OSL) + } + assert c["tracked_samples_missing_usage"] == (expected is ServerUsage.MISSING) + assert c["tracked_samples_missing_usage_empty_output"] == 0 + # Server ISL and OSL cover exactly the usable samples. + usable = int(expected is ServerUsage.USABLE) + assert recorded == {MetricSeriesKey.ISL: usable, MetricSeriesKey.OSL: usable} + + @pytest.mark.asyncio + async def test_first_miss_logged_once_at_the_primary_level( + self, tmp_path, primary, caplog + ): + events = [sample_event(SampleEventType.ISSUED, "untracked", ts=1), _START] + events += [ + sample_event(SampleEventType.ISSUED, u, ts=10) + for u in ("ok", "a", "b", "c", "err") + ] + events += [ + complete_event("untracked", 20, text_output("x")), + complete_event( + "ok", 20, text_output("w"), prompt_tokens=1, completion_tokens=1 + ), + complete_event("a", 20, text_output("x y")), + complete_event("b", 20, text_output("")), # empty output + complete_event("c", 20, text_output("z"), prompt_tokens=3), # one UNSET + complete_event("err", 20, None), # error: excluded + ] + with ( + caplog.at_level(logging.INFO), + aggregator_scope(tmp_path, "agg_miss_log", token_count_source=primary) as ( + agg, + reg, + _, + ), + ): + await agg.process(events) + c = snapshot_counters(reg) + assert c["tracked_samples_missing_usage"] == 3 + assert c["tracked_samples_missing_usage_empty_output"] == 1 + # "a" is the first miss and the second successful tracked sample. + message = missing_usage_message(1, 0, 2, primary=primary) + level = logging.WARNING if primary is SERVER else logging.INFO + logged = [ + r.levelno + for r in caplog.records + if "returned no server token usage" in r.getMessage() + ] + assert logged == [level] + assert any(message in r.getMessage() for r in caplog.records) + + @pytest.mark.asyncio + async def test_paired_osl_counts_only_usable_samples(self, tmp_path, primary): + events = [_START] + for case_id, prompt, completion, output, _ in SERVER_USAGE_CASES: + events += [ + sample_event(SampleEventType.ISSUED, case_id, ts=10), + complete_event( + case_id, + 20, + output, + prompt_tokens=prompt, + completion_tokens=completion, + ), + ] + usable = [case for case in SERVER_USAGE_CASES if case[4] is ServerUsage.USABLE] + with aggregator_scope( + tmp_path, + "agg_pair", + token_count_source=primary, + tokenizer=MockBatchTokenizer(), + ) as (agg, reg, _): + await agg.process(events) + # The client OSL recorder runs at the token-queue flush. + assert snapshot_counters(reg)["osl_paired_samples"] == 0 + await agg.token_queue.drain_all() + c = snapshot_counters(reg) + assert ( + c["osl_paired_samples"], + c["osl_paired_client_tokens"], + c["osl_paired_server_tokens"], + ) == ( + len(usable), + sum(len(str(case[3]).split()) for case in usable), + sum(case[2] for case in usable), + ) + + @pytest.mark.asyncio + async def test_early_stopping_and_steady_state_attach_to_primary_tpot_only( + self, tmp_path, primary + ): + reg = MetricsRegistry(early_stopping=EarlyStoppingSpec()) + phase = EventRecord( + event_type=SessionEventType.PHASE_START, + timestamp_ns=0, + data=PhaseData( + phase_type=PhaseType.PERFORMANCE, drain_after=True, num_turns=1 + ), + ) + with aggregator_scope( + tmp_path, + "agg_primary_tpot", + token_count_source=primary, + tokenizer=MockBatchTokenizer(), + steady_state_profile="concurrency", + registry=reg, + ) as (agg, _, _): + await agg.process([phase, _START, *_BOTH_BASES_SAMPLE]) + await agg.token_queue.drain_all() + collected = list(agg._collector.series()[0].tpot_ns) + snap = reg.build_snapshot(state=SessionState.COMPLETE, n_pending_tasks=0) + with_es = { + m.name + for m in snap.metrics + if isinstance(m, SeriesStat) and m.early_stopping_percentiles is not None + } + assert with_es == {"sample_latency_ns", "ttft_ns", "tpot_ns"} + assert collected == [_TPOT_BY_BASIS[primary]] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "completion_tokens, recv_first_ts, expected", + [ + (4, 2000, 1000.0), # (5000 - 2000) / (4 - 1) + (1, 2000, None), # <= 1: skipped (LoadGen would log a negative) + (msgspec.UNSET, 2000, None), + (4, None, None), # streaming only + (4, 5000, None), # delta == 0: skipped + (4, 6000, None), # RECV_FIRST stamped after COMPLETE: skipped + ], + ) + async def test_server_tpot_uses_loadgen_rule( + self, tmp_path, primary, completion_tokens, recv_first_ts, expected + ): + events = [_START, sample_event(SampleEventType.ISSUED, "s1", ts=1000)] + if recv_first_ts is not None: + events.append( + sample_event(SampleEventType.RECV_FIRST, "s1", ts=recv_first_ts) + ) + events.append( + complete_event( + "s1", + 5000, + streaming_text("a", " b"), + prompt_tokens=2, + completion_tokens=completion_tokens, + ) + ) + tpot = _name(MetricSeriesKey.TPOT_NS, SERVER, primary=primary) + with aggregator_scope(tmp_path, "agg_tpot", token_count_source=primary) as ( + agg, + reg, + _, + ): + await agg.process(events) + assert snapshot_series_count(reg, tpot) == (expected is not None) + if expected is not None: + assert snapshot_series_total(reg, tpot) == pytest.approx(expected) + + +@pytest.mark.unit +def test_aggregator_requires_token_count_source(): + with pytest.raises(TypeError, match="token_count_source"): + MetricsAggregatorService( + "unused", + None, + None, + registry=MetricsRegistry(), + publisher=None, + publish_interval_s=0.25, + sig_figs=3, + n_histogram_buckets=10, + ) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator_error_handler.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator_error_handler.py index 40e6eb91b..ba9c3814a 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator_error_handler.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator_error_handler.py @@ -27,9 +27,11 @@ import pytest from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( - MetricCounterKey, MetricsAggregatorService, ) +from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import ( + MetricCounterKey, +) from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( MetricsRegistry, ) @@ -46,6 +48,8 @@ SessionEventType, ) +from tests.test_helpers import DEFAULT_SOURCE + def _counters(registry: MetricsRegistry) -> dict[str, int | float]: """Read all counter values via a snapshot. State/n_pending don't matter @@ -81,6 +85,7 @@ def _make_aggregator( publish_interval_s=0.25, sig_figs=3, n_histogram_buckets=10, + token_count_source=DEFAULT_SOURCE, streaming=streaming, drain_timeout_s=None, ) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_early_stopping_integration.py b/tests/unit/async_utils/services/metrics_aggregator/test_early_stopping_integration.py index 1ac153c4b..9aa98a6e3 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_early_stopping_integration.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_early_stopping_integration.py @@ -19,7 +19,9 @@ snapshot_to_dict, ) from inference_endpoint.metrics.early_stopping import EarlyStoppingSpec -from inference_endpoint.metrics.report import _series_to_metric_dict +from inference_endpoint.metrics.report import Report, _series_to_metric_dict + +from tests.test_helpers import DEFAULT_SOURCE pytestmark = pytest.mark.unit @@ -146,9 +148,9 @@ def test_report_from_snapshot_preserves_empty_series_map(): # (Report.from_snapshot): a count==0 short-circuit there drops the map even # when _series_to_metric_dict preserves it, so exercising the helper alone # would prove nothing. - from inference_endpoint.metrics.report import Report - - report = Report.from_snapshot(_snap(_registry(EarlyStoppingSpec(), n=0))) + report = Report.from_snapshot( + _snap(_registry(EarlyStoppingSpec(), n=0)), token_count_source=DEFAULT_SOURCE + ) esp = report.ttft["early_stopping_percentiles"] assert list(esp) == _GRID_ES_KEYS assert all(v is None for v in esp.values()) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_main_signal_handler.py b/tests/unit/async_utils/services/metrics_aggregator/test_main_signal_handler.py index 230c3ef2a..6e7285515 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_main_signal_handler.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_main_signal_handler.py @@ -35,6 +35,7 @@ from inference_endpoint.async_utils.services.metrics_aggregator import ( __main__ as agg_main, ) +from inference_endpoint.core.types import TokenCountSource @pytest.mark.unit @@ -168,3 +169,33 @@ def test_drain_timeout_must_be_positive(value): @pytest.mark.unit def test_positive_drain_timeout_parses(): assert agg_main._positive_float("0.25") == 0.25 + + +_ARGV = [ + "--socket-dir", + "d", + "--socket-name", + "s", + "--metrics-socket", + "m", + "--metrics-output-dir", + "o", +] + + +@pytest.mark.unit +@pytest.mark.parametrize("source", list(TokenCountSource), ids=lambda s: s.value) +def test_token_count_source_argv(source): + args = agg_main._build_parser().parse_args( + [*_ARGV, "--token-count-source", source.value] + ) + assert args.token_count_source is source + + +@pytest.mark.unit +@pytest.mark.parametrize( + "extra", [[], ["--token-count-source", "both"]], ids=["absent", "unknown"] +) +def test_main_requires_token_count_source_flag(extra): + with pytest.raises(SystemExit): + agg_main._build_parser().parse_args([*_ARGV, *extra]) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py index b869ab3f8..8587cc90f 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py @@ -306,7 +306,7 @@ async def test_osl_with_tool_calls_uses_message_path(self): registry.register_series("osl", hdr_low=1, hdr_high=100_000) loop = asyncio.get_running_loop() queue = TokenBatchQueue(MockBatchTokenizer(), loop) - trigger = OslTrigger(registry, queue) + trigger = OslTrigger(registry, queue, metric_name="osl") tool_calls = ( { @@ -342,7 +342,7 @@ async def test_osl_without_tool_calls_uses_text_path(self): registry.register_series("osl", hdr_low=1, hdr_high=100_000) loop = asyncio.get_running_loop() queue = TokenBatchQueue(MockBatchTokenizer(), loop) - trigger = OslTrigger(registry, queue) + trigger = OslTrigger(registry, queue, metric_name="osl") tmo = TextModelOutput(output="hello world") ev = EventRecord( @@ -379,7 +379,7 @@ async def count_batch_async(self, inputs, _loop, live=False): registry.register_series("osl", hdr_low=1, hdr_high=100_000) tokenizer = CapturingTokenizer() queue = TokenBatchQueue(tokenizer, asyncio.get_running_loop()) - trigger = OslTrigger(registry, queue) + trigger = OslTrigger(registry, queue, metric_name="osl") output = TextModelOutput(output="answer", reasoning="private reasoning") trigger.fire( @@ -419,7 +419,7 @@ async def test_tpot_tools_after_first_content_delta(self): ) loop = asyncio.get_running_loop() queue = TokenBatchQueue(MockBatchTokenizer(), loop) - trigger = TpotTrigger(registry, queue) + trigger = TpotTrigger(registry, queue, metric_name="tpot_ns", collector=None) tool_calls = ( { @@ -459,7 +459,7 @@ async def test_tpot_uses_tool_call_deltas_after_first_chunk(self): ) loop = asyncio.get_running_loop() queue = TokenBatchQueue(MockBatchTokenizer(), loop) - trigger = TpotTrigger(registry, queue) + trigger = TpotTrigger(registry, queue, metric_name="tpot_ns", collector=None) tool_call_chunks = ( ( diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_output_tokenization_rule.py b/tests/unit/async_utils/services/metrics_aggregator/test_output_tokenization_rule.py index 297d249b8..3aca91a48 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_output_tokenization_rule.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_output_tokenization_rule.py @@ -15,6 +15,7 @@ MessageInput, TextInput, extract_tokenization_input, + has_countable_output, ) from inference_endpoint.core.types import TextModelOutput @@ -71,3 +72,34 @@ def test_empty_output_is_excluded() -> None: it as 0 would drag min/avg down on both the window and the full run. """ assert extract_tokenization_input(TextModelOutput(output="")) is None + + +_TOOL_CALL = { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": "{}"}, +} + + +@pytest.mark.parametrize( + "output", + [ + pytest.param(TextModelOutput(output=""), id="empty-str"), + pytest.param(TextModelOutput(output="a"), id="str"), + pytest.param(TextModelOutput(output=("", "")), id="empty-chunks"), + pytest.param(TextModelOutput(output=("", "x")), id="late-chunk"), + pytest.param( + TextModelOutput(output="", reasoning=("",)), id="empty-reasoning-chunk" + ), + pytest.param( + TextModelOutput(output="", tool_calls=(_TOOL_CALL,)), id="tool-calls" + ), + pytest.param( + TextModelOutput(output="", reasoning="thinking"), id="reasoning-only" + ), + ], +) +def test_has_countable_output_matches_the_osl_rule(output: TextModelOutput) -> None: + assert has_countable_output(output) == ( + extract_tokenization_input(output) is not None + ) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_response_boundary.py b/tests/unit/async_utils/services/metrics_aggregator/test_response_boundary.py index 8e41e5822..fa7697d87 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_response_boundary.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_response_boundary.py @@ -132,8 +132,10 @@ async def test_response_metrics_use_full_output_and_actual_tail( queue = TokenBatchQueue(response_tokenizer, loop) row = SampleRow(sample_uuid="qid") pre_change = {SampleField.RECV_FIRST_NS: 1000} - OslTrigger(registry, queue).fire(event, row, pre_change) - TpotTrigger(registry, queue).fire(event, row, pre_change) + OslTrigger(registry, queue, metric_name="osl").fire(event, row, pre_change) + TpotTrigger(registry, queue, metric_name="tpot_ns", collector=None).fire( + event, row, pre_change + ) await queue.drain_all() assert snapshot_series_total(registry, "osl") == full_count if any(tail_parts): diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_server_token_pairs.py b/tests/unit/async_utils/services/metrics_aggregator/test_server_token_pairs.py new file mode 100644 index 000000000..889a59d81 --- /dev/null +++ b/tests/unit/async_utils/services/metrics_aggregator/test_server_token_pairs.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Server token counts are invariant across inline and structured responses. + +The same completion arrives either with reasoning and tool calls inline in +``content`` (``...``, ``...``) or split into +structured ``reasoning_content`` / ``tool_calls`` fields. The server reports +one usage count for both, so the server-basis ISL, OSL and TPOT must match +across each pair while the tokenizer-basis OSL follows the representation it +was given. Both hold whichever basis is primary. +""" + +import asyncio +import json +import tempfile +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import pytest +from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( + MetricsRegistry, +) +from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( + BatchTokenizer, +) +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, + TextInput, + extract_tokenization_input, +) +from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext +from inference_endpoint.core.record import ( + EventRecord, + EventRecordCodec, + SampleEventType, + SessionEventType, +) +from inference_endpoint.core.types import TextModelOutput, TokenCountSource, basis_key +from inference_endpoint.openai.accumulator import OpenAISSEAccumulator +from inference_endpoint.openai.openai_msgspec_adapter import OpenAIMsgspecAdapter +from inference_endpoint.openai.types import SSEChoice, SSEDelta + +from .conftest import ( + make_aggregator, + sample_event, + session_event, + snapshot_counters, + snapshot_series_count, + snapshot_series_total, +) + +_TOKENIZERS = Path(__file__).resolve().parents[4] / "assets/tokenizers" +_PROMPT, _COMPLETION = 5, 6 # identical server usage for both members of a pair +_TOOL = { + "index": 0, + "id": "call-0", + "type": "function", + "function": {"name": "grep", "arguments": '{"q": "x"}'}, +} +_SPLIT = [{"reasoning_content": "c"}, {"reasoning_content": "d"}, {"content": "ab"}] +# inline deltas, structured deltas, client OSL (inline, structured), +# client TPOT tail tokens (inline, structured) +_PAIRS = [ + pytest.param( + [{"content": "ab"}, {"content": "cd"}], + [{"content": "ab"}, {"content": "cd"}], + (4, 4), + (2, 2), + id="content", + ), + pytest.param( + [{"content": "cd"}, {"content": "ef"}], + [{"reasoning_content": "cd"}, {"reasoning_content": "ef"}], + (19, 4), + (10, 2), + id="reasoning", + ), + pytest.param( + [{"content": "cd"}, {"content": "ab"}], + [{"reasoning_content": "cd"}, {"content": "ab"}], + (19, 4), + (10, 2), + id="both", + ), + pytest.param( + [ + { + "content": '{"name": "grep", "arguments": {"q": "x"}}' + } + ], + [{"tool_calls": [_TOOL]}], + (64, 39), + (0, 0), + id="tool", + ), + pytest.param( + [ + {"content": ""}, + {"content": "c"}, + {"content": "dab"}, + ], + _SPLIT, + (19, 4), + (12, 3), + id="split", + ), + pytest.param( + [{"content": "cdab"}], + _SPLIT, + (19, 4), + (16, 3), + id="first", + ), +] + + +@pytest.fixture(scope="module") +def bare_tokenizer() -> Iterator[BatchTokenizer]: + with BatchTokenizer( + str(_TOKENIZERS / "char_chat_bare_reasoning"), n_workers=0, live_workers=1 + ) as tok: + yield tok + + +@pytest.fixture +def socket_dir() -> Iterator[str]: + # macOS tmp_path breaks the 103-char IPC socket path limit. + with tempfile.TemporaryDirectory(prefix="pair") as d: + yield d + + +def _output(deltas: list[dict[str, Any]]) -> TextModelOutput: + acc = OpenAISSEAccumulator("s1", stream_all_chunks=False) + for delta in deltas: + acc.add_chunk(SSEChoice(delta=SSEDelta(**delta))) + out = acc.get_final_output().response_output + assert isinstance(out, TextModelOutput) + return out + + +async def _aggregate( + tokenizer: BatchTokenizer, + socket_dir: str, + name: str, + output: TextModelOutput, + *, + primary: TokenCountSource, + prompt_tokens: int = _PROMPT, + completion_tokens: int = _COMPLETION, +) -> MetricsRegistry: + """Run one tracked sample (ISSUED 500, first chunk 1000, COMPLETE 11000).""" + with ManagedZMQContext.scoped(socket_dir=socket_dir) as ctx: + agg, registry, _ = make_aggregator( + ctx, + asyncio.get_running_loop(), + name, + tokenizer=tokenizer, + token_count_source=primary, + ) + try: + await agg.process( + [ + session_event(SessionEventType.START_PERFORMANCE_TRACKING, ts=0), + sample_event(SampleEventType.ISSUED, "s1", ts=500), + sample_event(SampleEventType.RECV_FIRST, "s1", ts=1_000), + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=11_000, + sample_uuid="s1", + data=output, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ), + ] + ) + assert agg.token_queue is not None # a tokenizer was supplied + await agg.token_queue.drain_all() + finally: + agg.close() + return registry + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("asset", "structured"), [("char_chat", 19), ("char_chat_bare_reasoning", 4)] +) +def test_asset_frames_reasoning_without_inline_markers(asset, structured): + with BatchTokenizer(str(_TOKENIZERS / asset), n_workers=0, live_workers=1) as tok: + assert tok.count_batch_sync( + [TextInput("cdab"), MessageInput("ab", "cd", None)] + ) == [19, structured] + + +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize(("inline", "structured", "osl_client", "tail"), _PAIRS) +async def test_server_series_match_across_representations( + bare_tokenizer, socket_dir, primary, inline, structured, osl_client, tail +): + server = { + base: basis_key(base, TokenCountSource.SERVER, primary=primary) + for base in ("isl", "osl", "tpot_ns") + } + client = { + base: basis_key(base, TokenCountSource.CLIENT, primary=primary) + for base in ("osl", "tpot_ns") + } + for i, (deltas, expected, tail_tokens) in enumerate( + zip((inline, structured), osl_client, tail, strict=True) + ): + output = _output(deltas) + item = extract_tokenization_input(output) + assert item is not None and bare_tokenizer.count_batch_sync([item]) == [ + expected + ] + registry = await _aggregate( + bare_tokenizer, socket_dir, f"pair{i}", output, primary=primary + ) + assert snapshot_series_total(registry, server["osl"]) == _COMPLETION + assert snapshot_series_total(registry, server["isl"]) == _PROMPT + assert snapshot_series_total(registry, server["tpot_ns"]) == pytest.approx( + 10_000 / (_COMPLETION - 1) + ) + assert snapshot_series_total(registry, client["osl"]) == expected + if tail_tokens: + assert snapshot_series_total(registry, client["tpot_ns"]) == pytest.approx( + 10_000 / tail_tokens + ) + else: + assert snapshot_series_count(registry, client["tpot_ns"]) == 0 + c = snapshot_counters(registry) + assert ( + c["osl_paired_samples"], + c["osl_paired_server_tokens"], + c["osl_paired_client_tokens"], + ) == (1, _COMPLETION, expected) + assert c["tracked_samples_missing_usage"] == 0 + + +_FRAMES = [ + { + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "say literally"}, + } + ] + }, + { + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call-0", + "type": "function", + "function": { + "name": "grep", + "arguments": '{"pattern": ""}', + }, + } + ] + }, + } + ] + }, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]}, + { + "choices": [], + "usage": {"prompt_tokens": 3, "completion_tokens": 9, "total_tokens": 12}, + }, +] + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_literal_markers_survive_capture_verbatim( + bare_tokenizer, socket_dir, primary +): + buffer = ( + b"".join(b"data: " + json.dumps(f).encode() + b"\n\n" for f in _FRAMES) + + b"data: [DONE]\n\n" + ) + acc = OpenAISSEAccumulator("s1", stream_all_chunks=False) + for chunk in OpenAIMsgspecAdapter.parse_sse_chunk(buffer, len(buffer)): + acc.add_chunk(chunk) + result = acc.get_final_output() + codec = EventRecordCodec() + _, payload = codec.encode( + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=1, + sample_uuid="s1", + data=result.response_output, + ) + ) + output = codec.decode(payload).data + assert isinstance(output, TextModelOutput) + content, reasoning, tool_calls = output.as_message_parts() + assert (content, reasoning) == ("say literally", None) + assert tool_calls is not None + assert tool_calls[0]["function"]["arguments"] == '{"pattern": ""}' + md = result.metadata + registry = await _aggregate( + bare_tokenizer, + socket_dir, + "verbatim", + output, + primary=primary, + prompt_tokens=md["prompt_tokens"], + completion_tokens=md["completion_tokens"], + ) + server_osl = basis_key("osl", TokenCountSource.SERVER, primary=primary) + client_osl = basis_key("osl", TokenCountSource.CLIENT, primary=primary) + assert snapshot_series_total(registry, server_osl) == 9 + # Markers are counted as ordinary characters. + assert snapshot_series_total(registry, client_osl) == 73 diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_steady_state_collection.py b/tests/unit/async_utils/services/metrics_aggregator/test_steady_state_collection.py index 13db16fec..21380058b 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_steady_state_collection.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_steady_state_collection.py @@ -26,6 +26,7 @@ import asyncio import dataclasses +import msgspec import pytest from inference_endpoint.async_utils.services.event_logger.file_writer import JSONLWriter from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext @@ -34,12 +35,19 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import PhaseData, PhaseType, TextModelOutput +from inference_endpoint.core.types import ( + PhaseData, + PhaseType, + TextModelOutput, + TokenCountSource, +) from inference_endpoint.metrics.steady_state_diagnostics import ( build_super_pass_series, compute_steady_state_metrics, ) +from tests.test_helpers import DEFAULT_SOURCE + from .conftest import ( MockBatchTokenizer, count_text_tokens, @@ -62,7 +70,13 @@ def phase_start( def _sample(uuid: str, issue: int, ttft: int, decode: int, chunks, turn=None): - """ISSUED -> RECV_FIRST -> COMPLETE for one streamed sample.""" + """ISSUED -> RECV_FIRST -> COMPLETE for one streamed sample. + + The COMPLETE carries server usage with ``completion_tokens`` equal to the + whitespace word count of the whole output. Every fixture's first chunk is + one word, so ``completion_tokens - 1`` equals the post-first-chunk count the + tokenizing replay derives, and server-mode TPOT/OSL match it exactly. + """ return [ EventRecord( event_type=SampleEventType.ISSUED, @@ -83,11 +97,13 @@ def _sample(uuid: str, issue: int, ttft: int, decode: int, chunks, turn=None): sample_uuid=uuid, turn=turn, data=TextModelOutput(output=tuple(chunks)), + prompt_tokens=1, + completion_tokens=len("".join(chunks).split()), ), ] -def _event_stream() -> list[EventRecord]: +def _event_stream(*, ns_usage: bool = True) -> list[EventRecord]: """A run exercising every collector rule that equivalence can break on.""" records = [ session_event(SessionEventType.STARTED, ts=0), @@ -112,6 +128,7 @@ def _event_stream() -> list[EventRecord]: # A retry: re-ISSUED inside the window, so only its issue timestamp moves. records += _sample("s3", issue, 12, 900, ["r ", "z z"]) # Non-streaming output: no post-first-chunk text, so no TPOT contribution. + # Without usage it is the run's one tracked sample missing server counts. records += [ EventRecord( event_type=SampleEventType.ISSUED, @@ -123,6 +140,8 @@ def _event_stream() -> list[EventRecord]: timestamp_ns=issue + 900, sample_uuid="ns", data=TextModelOutput(output="whole answer"), + prompt_tokens=1 if ns_usage else msgspec.UNSET, + completion_tokens=2 if ns_usage else msgspec.UNSET, ), ] # Still in flight when the window closes. @@ -144,6 +163,16 @@ def _event_stream() -> list[EventRecord]: return records +def _tokenized_replay(events_path, superpass_size): + """The event-log replay on the client basis with the aggregator's test tokenizer.""" + return build_super_pass_series( + events_path, + superpass_size, + count_text_tokens, + basis=TokenCountSource.CLIENT, + ).super_passes + + def _write_log(tmp_path, records): writer = JSONLWriter(tmp_path / "events") for record in records: @@ -152,7 +181,9 @@ def _write_log(tmp_path, records): return str(writer.file_path) -async def _collect(tmp_path, records, socket_name): +async def _collect( + tmp_path, records, socket_name, *, token_count_source: TokenCountSource +): """Drive ``records`` through a steady-state-enabled aggregator.""" loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: @@ -163,6 +194,7 @@ async def _collect(tmp_path, records, socket_name): tokenizer=MockBatchTokenizer(), steady_state_profile="concurrency", streaming=True, + token_count_source=token_count_source, ) try: await agg.process(records) @@ -210,12 +242,12 @@ def _steady_run_records(n_super_passes: int = 14) -> list[EventRecord]: @pytest.mark.unit class TestProducerEquivalence: @pytest.mark.asyncio - async def test_live_series_equals_the_event_log_parse(self, tmp_path): + async def test_live_series_equals_the_event_log_parse(self, tmp_path, primary): records = _event_stream() - collector = await _collect(tmp_path, records, "ss_equiv") - replayed = build_super_pass_series( - _write_log(tmp_path, records), SUPERPASS, count_text_tokens + collector = await _collect( + tmp_path, records, "ss_equiv", token_count_source=primary ) + replayed = _tokenized_replay(_write_log(tmp_path, records), SUPERPASS) live = collector.series() assert live, "the collector produced no super-passes" assert [dataclasses.asdict(sp) for sp in live] == [ @@ -223,25 +255,27 @@ async def test_live_series_equals_the_event_log_parse(self, tmp_path): ] @pytest.mark.asyncio - async def test_both_producers_reach_the_same_verdict(self, tmp_path): + async def test_both_producers_reach_the_same_verdict(self, tmp_path, primary): records = _event_stream() - collector = await _collect(tmp_path, records, "ss_equiv_verdict") - replayed = build_super_pass_series( - _write_log(tmp_path, records), SUPERPASS, count_text_tokens + collector = await _collect( + tmp_path, records, "ss_equiv_verdict", token_count_source=primary ) + replayed = _tokenized_replay(_write_log(tmp_path, records), SUPERPASS) kwargs = {"superpass_size": SUPERPASS, "warmup": 0} assert compute_steady_state_metrics( collector.series(), **kwargs ) == compute_steady_state_metrics(replayed, **kwargs) @pytest.mark.asyncio - async def test_a_certified_window_is_identical_on_both_producers(self, tmp_path): + async def test_a_certified_window_is_identical_on_both_producers( + self, tmp_path, primary + ): """Equivalence where it counts: a run the detector actually certifies.""" records = _steady_run_records() - collector = await _collect(tmp_path, records, "ss_equiv_window") - replayed = build_super_pass_series( - _write_log(tmp_path, records), STEADY_SUPERPASS, count_text_tokens + collector = await _collect( + tmp_path, records, "ss_equiv_window", token_count_source=primary ) + replayed = _tokenized_replay(_write_log(tmp_path, records), STEADY_SUPERPASS) live_verdict = compute_steady_state_metrics( collector.series(), superpass_size=STEADY_SUPERPASS ) @@ -253,11 +287,49 @@ async def test_a_certified_window_is_identical_on_both_producers(self, tmp_path) ) +def _with_usage(records): + """Server count = whitespace words + 1 (EOS): differs from the tokenizer count, so a + replay that fell back to the tokenizer cannot match.""" + return [ + ( + msgspec.structs.replace(r, completion_tokens=len(str(r.data).split()) + 1) + if r.event_type is SampleEventType.COMPLETE + and isinstance(r.data, TextModelOutput) + else r + ) + for r in records + ] + + +@pytest.mark.unit +class TestServerProducerEquivalence: + @pytest.mark.asyncio + async def test_live_server_series_equals_the_event_log_parse(self, tmp_path): + records = _with_usage(_event_stream()) + collector = await _collect( + tmp_path, + records, + "ss_equiv_srv", + token_count_source=TokenCountSource.SERVER, + ) + replayed = build_super_pass_series( + _write_log(tmp_path, records), + SUPERPASS, + None, + basis=TokenCountSource.SERVER, + ) + live = collector.series() + assert any(sp.tpot_ns for sp in live) + assert [dataclasses.asdict(sp) for sp in live] == [ + dataclasses.asdict(sp) for sp in replayed.super_passes + ] + + @pytest.mark.unit class TestOrderingInvariants: @pytest.mark.asyncio async def test_tpot_lands_in_the_bucket_the_sample_was_issued_into(self, tmp_path): - """Pins both ordering rules in the aggregator's TPOT path. + """Pins both ordering rules in the aggregator's tokenizer (client) TPOT path. ``on_complete`` must run AFTER ``set_field`` (``TpotTrigger`` fires from inside it and reads the sample's super-pass), and ``TpotTrigger`` must @@ -273,7 +345,9 @@ async def test_tpot_lands_in_the_bucket_the_sample_was_issued_into(self, tmp_pat *_sample("b", 200, 10, 900, ["c ", "x y z"]), *_sample("c", 900, 10, 900, ["c ", "x y z"]), ] - collector = await _collect(tmp_path, records, "ss_order") + collector = await _collect( + tmp_path, records, "ss_order", token_count_source=TokenCountSource.CLIENT + ) series = collector.series() assert [sp.n_issued for sp in series] == [2, 1] # 3 words after the first chunk -> tpot = 900 / 3. @@ -304,6 +378,7 @@ async def test_the_row_is_gone_by_the_time_the_token_count_arrives(self, tmp_pat tokenizer=MockBatchTokenizer(), steady_state_profile="concurrency", streaming=True, + token_count_source=TokenCountSource.CLIENT, ) try: await agg.process(records) @@ -314,6 +389,37 @@ async def test_the_row_is_gone_by_the_time_the_token_count_arrives(self, tmp_pat finally: agg.close() + @pytest.mark.asyncio + async def test_server_tpot_lands_at_complete_without_a_flush(self, tmp_path): + records = [ + session_event(SessionEventType.STARTED, ts=0), + phase_start(), + session_event(SessionEventType.START_PERFORMANCE_TRACKING, ts=10), + *_sample("a", 100, 10, 900, ["c ", "x y z"]), + ] + loop = asyncio.get_event_loop() + with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: + # No tokenizer and no flush: only the synchronous server path can feed TPOT. + agg, _, _ = make_aggregator( + ctx, + loop, + "ss_srv_sync", + steady_state_profile="concurrency", + streaming=True, + token_count_source=TokenCountSource.SERVER, + ) + try: + await agg.process(records) + assert agg.token_queue is None and agg.table.get_row("a") is None + sp = agg._collector.series()[0] + assert (list(sp.tpot_ns), list(sp.osl), sp.out_tokens) == ( + [300.0], + [3.0], + 3, + ) + finally: + agg.close() + @pytest.mark.unit class TestCollectionGate: @@ -321,7 +427,9 @@ class TestCollectionGate: async def test_disabled_by_default(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "ss_off") + agg, _, _ = make_aggregator( + ctx, loop, "ss_off", token_count_source=DEFAULT_SOURCE + ) try: await agg.process(_event_stream()) assert agg._collector is None @@ -345,6 +453,7 @@ async def test_an_unsupported_profile_collects_nothing(self, tmp_path, profile): f"ss_unsupported_{profile}", tokenizer=MockBatchTokenizer(), steady_state_profile=profile, + token_count_source=DEFAULT_SOURCE, ) try: await agg.process(_event_stream()) @@ -365,7 +474,9 @@ async def test_a_warmup_announcement_never_sizes_the_series(self, tmp_path): session_event(SessionEventType.STARTED, ts=0), phase_start(num_turns=99, phase_type=PhaseType.WARMUP), ] - collector = await _collect(tmp_path, records, "ss_warmup_size") + collector = await _collect( + tmp_path, records, "ss_warmup_size", token_count_source=DEFAULT_SOURCE + ) assert collector.superpass_size == 0 @pytest.mark.asyncio @@ -378,7 +489,9 @@ async def test_a_later_accuracy_phase_is_refused_by_phase_type(self, tmp_path): *_sample("a", 100, 10, 900, ["c ", "x"]), phase_start(num_turns=77, phase_type=PhaseType.ACCURACY), ] - collector = await _collect(tmp_path, records, "ss_size") + collector = await _collect( + tmp_path, records, "ss_size", token_count_source=DEFAULT_SOURCE + ) assert collector.superpass_size == 3 @pytest.mark.asyncio @@ -400,7 +513,9 @@ async def test_a_second_performance_phase_cannot_resize_a_started_series( phase_start(num_turns=77), *_sample("b", 2000, 10, 900, ["d ", "y"]), ] - collector = await _collect(tmp_path, records, "ss_resize") + collector = await _collect( + tmp_path, records, "ss_resize", token_count_source=DEFAULT_SOURCE + ) assert collector.superpass_size == 3 @@ -418,6 +533,7 @@ async def _finalize(tmp_path, socket_name, *, extra=()): socket_name, tokenizer=MockBatchTokenizer(), steady_state_profile="concurrency", + token_count_source=DEFAULT_SOURCE, ) try: await agg.process( @@ -458,6 +574,7 @@ async def test_no_verdict_when_tokenizations_did_not_drain(self, tmp_path): "ss_final_pending", tokenizer=MockBatchTokenizer(), steady_state_profile="concurrency", + token_count_source=TokenCountSource.CLIENT, ) try: await agg.process(_event_stream()) @@ -466,6 +583,40 @@ async def test_no_verdict_when_tokenizations_did_not_drain(self, tmp_path): finally: agg.close() + @pytest.mark.asyncio + @pytest.mark.parametrize( + "source, ns_usage, n_pending, has_verdict", + [ + # A pending tokenizer drain only touches the client basis. + (TokenCountSource.SERVER, True, 3, True), + # A tracked sample without usage withholds it. + (TokenCountSource.SERVER, False, 0, False), + # Missing server usage only touches the server basis. + (TokenCountSource.CLIENT, False, 0, True), + # A pending tokenizer drain withholds it. + (TokenCountSource.CLIENT, True, 3, False), + ], + ) + async def test_steady_state_verdict_gate_follows_primary( + self, tmp_path, source, ns_usage, n_pending, has_verdict + ): + loop = asyncio.get_event_loop() + with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: + agg, _, _ = make_aggregator( + ctx, + loop, + "ss_gate", + tokenizer=MockBatchTokenizer(), + steady_state_profile="concurrency", + token_count_source=source, + ) + try: + await agg.process(_event_stream(ns_usage=ns_usage)) + verdict = await agg._steady_state_verdict(n_pending=n_pending) + assert (verdict is not None) is has_verdict + finally: + agg.close() + @pytest.mark.asyncio async def test_a_detector_failure_never_costs_the_run_its_snapshot( self, tmp_path, monkeypatch @@ -478,6 +629,7 @@ async def test_a_detector_failure_never_costs_the_run_its_snapshot( "ss_final_boom", tokenizer=MockBatchTokenizer(), steady_state_profile="concurrency", + token_count_source=DEFAULT_SOURCE, ) try: await agg.process(_event_stream()) diff --git a/tests/unit/metrics/test_token_count_basis.py b/tests/unit/metrics/test_token_count_basis.py new file mode 100644 index 000000000..545575931 --- /dev/null +++ b/tests/unit/metrics/test_token_count_basis.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ``inference_endpoint.metrics.token_count_basis``.""" + +import json + +import pytest +import yaml +from inference_endpoint.config.schema import BenchmarkConfig +from inference_endpoint.config.schema import TestType as BenchmarkTestType +from inference_endpoint.core.types import APIType, TextModelOutput, TokenCountSource +from inference_endpoint.metrics.token_count_basis import ( + ARCHIVED_RUN_BASIS, + RunBasisError, + ServerUsage, + classify_server_usage, + effective_token_count_source, + missing_usage_message, + read_server_usage, + replay_basis, + token_count_source_for_run, +) + +from tests.test_helpers import SERVER_USAGE_CASES + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), + [case[1:] for case in SERVER_USAGE_CASES], + ids=[case[0] for case in SERVER_USAGE_CASES], +) +def test_classify_server_usage(prompt, completion, output, expected): + assert classify_server_usage(prompt, completion, output) is expected + + +@pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), + [case[1:] for case in SERVER_USAGE_CASES], + ids=[case[0] for case in SERVER_USAGE_CASES], +) +def test_read_server_usage(prompt, completion, output, expected): + assert read_server_usage(prompt, completion, output) == ( + expected, + completion if expected is ServerUsage.USABLE else None, + ) + + +@pytest.mark.parametrize( + "output", + [ + TextModelOutput(output="", reasoning="thinking"), + TextModelOutput(output="", tool_calls=({"type": "function"},)), + # A malformed tool-call payload is non-empty output, not a true zero. + TextModelOutput(output="", tool_calls=(("not-a-dict",),)), + TextModelOutput(output="", tool_calls=(({"index": ["unhashable"]},),)), + ], + ids=["reasoning", "tool_calls", "malformed_not_a_dict", "malformed_index"], +) +def test_zero_completion_beside_non_text_output_is_missing(output): + assert classify_server_usage(10, 0, output) is ServerUsage.MISSING + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +def test_effective_basis_forces_client_for_videogen(configured): + assert ( + effective_token_count_source(configured, APIType.VIDEOGEN) + is TokenCountSource.CLIENT + ) + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +@pytest.mark.parametrize( + "api_type", [t for t in APIType if t is not APIType.VIDEOGEN], ids=lambda t: t.value +) +def test_effective_basis_keeps_configured_otherwise(configured, api_type): + assert effective_token_count_source(configured, api_type) is configured + + +_USAGE = "3 of 10 successful tracked samples returned no server token usage" +# The CLI flag, then the YAML key: `benchmark from-config` takes no such flag. +_FIX = ( + "set --token-count-source client if this server cannot report usage " + "(YAML: settings.token_count_source: client)" +) +_EMPTY = ( + "had empty output — likely server errors or cut streams; " + "client mode will not fix these" +) + + +@pytest.mark.parametrize( + ("n_missing", "n_empty", "expected"), + [ + (3, 0, f"{_USAGE}; {_FIX}"), + (3, 3, f"{_USAGE}; 3 {_EMPTY}"), + (3, 1, f"{_USAGE}; {_FIX}; 1 {_EMPTY}"), + ], +) +def test_missing_usage_message_server_primary_is_the_fatal_text( + n_missing, n_empty, expected +): + assert ( + missing_usage_message(n_missing, n_empty, 10, primary=TokenCountSource.SERVER) + == expected + ) + + +def test_missing_usage_message_by_primary(): + server = missing_usage_message(3, 0, 10, primary=TokenCountSource.SERVER) + client = missing_usage_message(3, 0, 10, primary=TokenCountSource.CLIENT) + assert "settings.token_count_source" in server + assert "_server" in client + assert "exit" not in client + assert "3 of 10" in client + + +def _write_summary(run_dir, payload): + (run_dir / "performance").mkdir(parents=True, exist_ok=True) + text = payload if isinstance(payload, str) else json.dumps(payload) + (run_dir / "performance" / "result_summary.json").write_text(text) + + +def _write_config(run_dir, payload): + text = payload if isinstance(payload, str) else yaml.safe_dump(payload) + (run_dir / "config.yaml").write_text(text) + + +def _config_yaml(source, api_type=APIType.OPENAI): + return { + "settings": {"token_count_source": source.value}, + "endpoint_config": {"api_type": api_type.value}, + } + + +@pytest.mark.parametrize("recorded", list(TokenCountSource)) +def test_run_basis_reads_summary_first(tmp_path, recorded): + _write_summary(tmp_path, {"token_count_source": recorded.value}) + _write_config(tmp_path, _config_yaml(recorded.other)) + assert token_count_source_for_run(tmp_path) is recorded + + +@pytest.mark.parametrize("recorded", list(TokenCountSource)) +def test_run_basis_reads_the_given_summary(tmp_path, recorded): + summary_path = tmp_path / "elsewhere" / "result_summary.json" + summary_path.parent.mkdir() + summary_path.write_text(json.dumps({"token_count_source": recorded.value})) + _write_summary(tmp_path, {"token_count_source": recorded.other.value}) + assert token_count_source_for_run(tmp_path, summary_path=summary_path) is recorded + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +def test_run_basis_given_summary_without_key_falls_back_to_config_yaml( + tmp_path, configured +): + summary_path = tmp_path / "elsewhere.json" + summary_path.write_text(json.dumps({"tpot": {}})) + _write_summary(tmp_path, {"token_count_source": configured.other.value}) + _write_config(tmp_path, _config_yaml(configured)) + assert token_count_source_for_run(tmp_path, summary_path=summary_path) is configured + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +@pytest.mark.parametrize( + "summary", [None, {"run_config": {}}], ids=["absent", "no_key"] +) +def test_run_basis_falls_back_to_config_yaml(tmp_path, configured, summary): + if summary is not None: + _write_summary(tmp_path, summary) + _write_config(tmp_path, _config_yaml(configured)) + assert token_count_source_for_run(tmp_path) is configured + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +def test_run_basis_config_yaml_videogen_is_client(tmp_path, configured): + _write_config(tmp_path, _config_yaml(configured, APIType.VIDEOGEN)) + assert token_count_source_for_run(tmp_path) is TokenCountSource.CLIENT + + +@pytest.mark.parametrize("configured", list(TokenCountSource)) +@pytest.mark.parametrize("api_type", [APIType.OPENAI, APIType.VIDEOGEN]) +def test_run_basis_reads_what_a_run_writes(tmp_path, configured, api_type): + config = BenchmarkConfig.create_default_config(BenchmarkTestType.OFFLINE) + config = config.with_updates( + settings=config.settings.model_copy(update={"token_count_source": configured}), + endpoint_config=config.endpoint_config.model_copy( + update={"api_type": api_type} + ), + ) + config.to_yaml_file(tmp_path / "config.yaml") + assert token_count_source_for_run(tmp_path) is effective_token_count_source( + configured, api_type + ) + + +def test_run_basis_archived_is_client(tmp_path, capsys): + assert token_count_source_for_run(tmp_path) is ARCHIVED_RUN_BASIS + assert "--token-count-source" in capsys.readouterr().err + + +def test_run_basis_config_without_recorded_basis_is_archived(tmp_path, capsys): + _write_config(tmp_path, {"settings": {}, "endpoint_config": {}}) + assert token_count_source_for_run(tmp_path) is ARCHIVED_RUN_BASIS + assert "--token-count-source" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "summary", + ["{not json", "[1, 2]", json.dumps({"token_count_source": 3})], + ids=["bad_json", "not_an_object", "not_a_string"], +) +def test_run_basis_unreadable_summary_raises(tmp_path, summary): + _write_summary(tmp_path, summary) + with pytest.raises(RunBasisError, match="--token-count-source"): + token_count_source_for_run(tmp_path) + + +@pytest.mark.parametrize( + "config", + ["settings: [unclosed", "- a\n- b\n", {"settings": "x"}], + ids=["bad_yaml", "not_a_mapping", "settings_not_a_mapping"], +) +def test_run_basis_unreadable_config_raises(tmp_path, config): + _write_config(tmp_path, config) + with pytest.raises(RunBasisError, match="--token-count-source"): + token_count_source_for_run(tmp_path) + + +def test_run_basis_unknown_value_raises(tmp_path): + _write_summary(tmp_path, {"token_count_source": "bogus"}) + with pytest.raises(RunBasisError, match="--token-count-source"): + token_count_source_for_run(tmp_path) + + +@pytest.mark.parametrize( + "config", + [ + {"settings": {"token_count_source": "bogus"}}, + { + "settings": {"token_count_source": TokenCountSource.SERVER.value}, + "endpoint_config": {"api_type": "bogus"}, + }, + ], + ids=["bogus_source", "bogus_api_type"], +) +def test_run_basis_unknown_config_value_raises(tmp_path, config): + _write_config(tmp_path, config) + with pytest.raises(RunBasisError, match="--token-count-source"): + token_count_source_for_run(tmp_path) + + +def _override_id(override): + return "none" if override is None else override.value + + +@pytest.mark.parametrize("recorded", list(TokenCountSource)) +@pytest.mark.parametrize("override", [None, *TokenCountSource], ids=_override_id) +def test_replay_basis_is_the_override_else_the_run_basis(tmp_path, recorded, override): + _write_summary(tmp_path, {"token_count_source": recorded.value}) + basis = recorded if override is None else override + assert replay_basis(tmp_path, override=override) == (recorded, basis) + + +@pytest.mark.parametrize("recorded", list(TokenCountSource)) +def test_replay_basis_reads_the_given_summary(tmp_path, recorded): + summary_path = tmp_path / "elsewhere.json" + summary_path.write_text(json.dumps({"token_count_source": recorded.value})) + assert replay_basis(tmp_path, override=None, summary_path=summary_path) == ( + recorded, + recorded, + ) + + +def test_replay_basis_notes_an_archived_run_only_without_override(tmp_path, capsys): + assert replay_basis(tmp_path, override=None) == ( + ARCHIVED_RUN_BASIS, + ARCHIVED_RUN_BASIS, + ) + assert "--token-count-source" in capsys.readouterr().err + for override in TokenCountSource: + assert replay_basis(tmp_path, override=override) == ( + ARCHIVED_RUN_BASIS, + override, + ) + assert capsys.readouterr().err == "" + + +@pytest.mark.parametrize("override", list(TokenCountSource)) +def test_replay_basis_unreadable_run_basis_is_none_beside_an_override( + tmp_path, override +): + _write_summary(tmp_path, "{not json") + assert replay_basis(tmp_path, override=override) == (None, override) + with pytest.raises(RunBasisError, match="--token-count-source"): + replay_basis(tmp_path, override=None) From bb4e2b2a45fddd8132b5c27fbb6b2a6c3d54d6c9 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:02 -0500 Subject: [PATCH 07/26] feat(report): report the other token basis and gate completeness on the primary Report.from_snapshot(snapshot, token_count_source=) takes the basis as a required argument; the snapshot wire format is unchanged. Unsuffixed fields come from the primary basis, so under the client default result_summary.json and report.txt keep every main key, value and line. The other basis is reported beside it: - result_summary.json: _ blocks (input/output_sequence_lengths, tpot, tps), with "partial": true when that basis covered fewer samples, plus osl_paired and the missing-usage counts; - report.txt: main's rows unchanged, then an "Other token basis ():" block, with a " (partial)" suffix on short sections. Completeness follows the primary. A client primary keeps main's rule (n_pending_tasks == 0) and only marks server blocks partial. A server primary requires every tracked successful sample to carry usable usage; otherwise run_benchmark exits 4 with the cause-aware message from Report.incomplete_reason(). pipeline.resolve_token_count_source is the only reader of the setting; videogen resolves to client. Under a server primary, setup fails before any load on a streaming OpenAI run that opts out of usage, and execute.py fails on rows with n > 1 (usage sums all choices). Under a client primary, n > 1 rows only warn. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- src/inference_endpoint/commands/audit.py | 2 +- .../commands/benchmark/execute.py | 105 ++- .../commands/benchmark/pipeline.py | 99 ++- src/inference_endpoint/metrics/report.py | 397 +++++++-- tests/integration/commands/conftest.py | 65 ++ .../integration/commands/test_run_timeout.py | 81 +- .../commands/test_server_token_counts.py | 395 +++++++++ tests/integration/commands/test_sigint.py | 43 +- tests/integration/commands/test_warmup.py | 5 + .../videogen/test_videogen_benchmark.py | 89 ++ tests/unit/commands/test_benchmark.py | 639 +++++++++++++- .../commands/test_benchmark_final_snapshot.py | 204 ++++- tests/unit/compliance/test_output_caching.py | 57 +- tests/unit/metrics/test_report_builder.py | 820 ++++++++++++++++-- .../validation/test_shipped_cohort_golden.py | 8 +- 15 files changed, 2778 insertions(+), 231 deletions(-) create mode 100644 tests/integration/commands/conftest.py create mode 100644 tests/integration/commands/test_server_token_counts.py create mode 100644 tests/integration/videogen/test_videogen_benchmark.py diff --git a/src/inference_endpoint/commands/audit.py b/src/inference_endpoint/commands/audit.py index c1fba46b6..aaed66b6c 100644 --- a/src/inference_endpoint/commands/audit.py +++ b/src/inference_endpoint/commands/audit.py @@ -202,7 +202,7 @@ def _run_phases( if not report.complete: raise ExecutionError( f"Audit phase '{spec.label}' did not complete cleanly " - "(metrics drain timed out); " + f"({report.incomplete_reason()}); " "refusing to certify a result from partial data" ) # Without a fixed spec count, record what this phase actually issued. diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index ff7f95ade..6fbab5fb1 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -57,7 +57,10 @@ write_accuracy_results, ) from inference_endpoint.commands.benchmark.full_run_osl import full_run_osl_for_report -from inference_endpoint.commands.benchmark.pipeline import MetricsPipeline +from inference_endpoint.commands.benchmark.pipeline import ( + MetricsPipeline, + resolve_token_count_source, +) from inference_endpoint.commands.benchmark.profiling import ( ProfileController, write_profiling_section, @@ -80,7 +83,7 @@ TestMode, TestType, ) -from inference_endpoint.core.types import PhaseType, QueryResult +from inference_endpoint.core.types import PhaseType, QueryResult, TokenCountSource from inference_endpoint.dataset_manager.agentic_inference_dataset import ( AgenticInferenceDataset, ) @@ -108,6 +111,7 @@ SessionResult, ) from inference_endpoint.metrics.report import Report +from inference_endpoint.openai.types import requests_multiple_choices if TYPE_CHECKING: from inference_endpoint.async_utils.event_publisher import EventPublisherService @@ -191,6 +195,7 @@ class BenchmarkContext: dataloader: Dataset | None rt_settings: RuntimeSettings | None total_samples: int + token_count_source: TokenCountSource eval_configs: list[AccuracyConfiguration] = field(default_factory=list) affinity_plan: AffinityPlan | None = None @@ -253,9 +258,6 @@ def _check_tokenizer_exists(model_name: str) -> bool: return True except Exception as e: logger.warning(f"Could not verify tokenizer for {model_name}: {e}") - logger.warning( - "Continuing without tokenizer (ISL/OSL/TPOT metrics will be unavailable)" - ) return False @@ -302,6 +304,58 @@ def _validate_accuracy_config_for_scorer( ) +def _issued_datasets( + config: BenchmarkConfig, + dataloader: Dataset | None, + eval_configs: list[AccuracyConfiguration], +) -> list[tuple[str, Dataset]]: + """Each distinct dataset whose rows are sent to the endpoint, with its name.""" + issued: list[tuple[str, Dataset]] = [] + if dataloader is not None: + perf_name = next( + (d.name for d in config.datasets if d.type == DatasetType.PERFORMANCE), + "performance", + ) + issued.append((perf_name, dataloader)) + for ec in eval_configs: + # The performance dataset's inline scorer reuses the loader above, and + # an external scorer never sends its rows to the endpoint. + if ec.dataset is not dataloader and not ec.scorer.SKIP_ENDPOINT_PHASE: + issued.append((ec.dataset_name, ec.dataset)) + return issued + + +def _check_multi_choice_rows( + datasets: list[tuple[str, Dataset]], *, token_count_source: TokenCountSource +) -> None: + """Reject (server basis) or count and warn about (client basis) ``n > 1`` rows. + + Server usage sums the tokens of all ``n`` choices, but only ``choices[0]`` + is measured, so those rows have no usable server counts. + """ + n_rows = 0 + for name, dataset in datasets: + for index, row in enumerate(dataset.data or ()): + n = row.get("n") + if not requests_multiple_choices(n): + continue + if token_count_source is TokenCountSource.SERVER: + raise SetupError( + f"Dataset '{name}': sample {index} sets n={n}. " + "Server token usage sums all n choices but only choices[0] is " + "measured; remove 'n' or set --token-count-source client " + "(YAML: settings.token_count_source: client)." + ) + n_rows += 1 + if n_rows: + logger.warning( + "%d sample(s) set n > 1: server token usage sums all n choices but " + "only choices[0] is measured, so they get no server-reported " + "(*_server) token counts", + n_rows, + ) + + def _load_datasets( config: BenchmarkConfig, report_dir: Path, @@ -525,6 +579,21 @@ def setup_benchmark( tokenizer_name = tokenizer_override else: tokenizer_name = model_name if _check_tokenizer_exists(model_name) else None + token_count_source = resolve_token_count_source(config) + if tokenizer_name is None: + if token_count_source is TokenCountSource.CLIENT: + logger.warning( + "No tokenizer: ISL/OSL/TPOT are unavailable while the client " + "basis is primary; the *_server ISL/OSL/TPOT blocks are " + "reported when the server returns usage. Set --tokenizer " + "(model_params.tokenizer_name) or --token-count-source server " + "(settings.token_count_source)" + ) + else: + logger.warning( + "No tokenizer: the *_client ISL/OSL/TPOT blocks are unavailable; " + "ISL/OSL/TPOT use server-reported usage" + ) # Streaming logger.info( @@ -534,6 +603,10 @@ def setup_benchmark( # Datasets dataloader, eval_configs = _load_datasets(config, report_dir, test_mode) + _check_multi_choice_rows( + _issued_datasets(config, dataloader, eval_configs), + token_count_source=token_count_source, + ) rt_settings: RuntimeSettings | None = None total_samples = 0 @@ -602,6 +675,7 @@ def setup_benchmark( total_samples=total_samples, eval_configs=eval_configs, affinity_plan=affinity_plan, + token_count_source=token_count_source, ) @@ -853,6 +927,7 @@ async def _run_benchmark_async( event_log_dir=event_log_dir, metrics_output_dir=metrics_output_dir, loop=loop, + token_count_source=ctx.token_count_source, ) report: Report | None = None profiler: ProfileController @@ -1229,7 +1304,11 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: if aborted: logger.warning("Run aborted — skipping accuracy scoring on partial data") else: - full_run_osl = full_run_osl_for_report(ctx.report_dir, ctx.tokenizer_name) + full_run_osl = full_run_osl_for_report( + ctx.report_dir, + ctx.tokenizer_name, + token_count_source=ctx.token_count_source, + ) accuracy_scores = score_accuracy(ctx, result) except KeyboardInterrupt: if report is not None: @@ -1266,7 +1345,9 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: # Emit the accuracy results as a focused artifact under accuracy/. Written # after the report artifacts so a write failure here can't discard them. - write_accuracy_results(ctx.report_dir, accuracy_scores) + write_accuracy_results( + ctx.report_dir, accuracy_scores, token_count_source=ctx.token_count_source + ) def run_benchmark( @@ -1319,14 +1400,8 @@ def run_benchmark( ) if bench.report is None: raise ExecutionError("Benchmark produced no usable metrics report") - if not (report := bench.report).complete: - if report.state == "complete" and report.n_samples_dropped == 0: - raise ExecutionError("Metrics tokenization did not finish") - raise ExecutionError( - "Benchmark report is incomplete " - f"(state={report.state}, " - f"n_samples_dropped={report.n_samples_dropped})" - ) + if not bench.report.complete: + raise ExecutionError(bench.report.incomplete_reason()) except KeyboardInterrupt: # Salvage results (finally), then propagate to main.py -> exit 130. logger.warning("Benchmark interrupted by user") diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index 78146def2..31ca8e958 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -59,9 +59,12 @@ MetricsSnapshotSubscriber, ) from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext -from inference_endpoint.config.schema import LoadPatternType +from inference_endpoint.config.schema import LoadPatternType, StreamingMode +from inference_endpoint.core.types import APIType, TokenCountSource +from inference_endpoint.exceptions import SetupError from inference_endpoint.metrics.report import Report from inference_endpoint.metrics.steady_state_diagnostics import profile_for_load_pattern +from inference_endpoint.metrics.token_count_basis import effective_token_count_source if TYPE_CHECKING: from inference_endpoint.config.schema import BenchmarkConfig @@ -99,6 +102,7 @@ def steady_state_profile( accuracy_only: bool, enable_streaming: bool, tokenizer_name: str | None, + token_count_source: TokenCountSource, ) -> str | None: """Workload profile for live steady-state collection; None when ineligible. @@ -106,10 +110,11 @@ def steady_state_profile( parent gate and the aggregator's profile from drifting apart. The decision must happen before the run starts because collection happens during the run. - Accuracy-only runs have no performance phase to collect. Without a tokenizer - there is no TPOT to bucket. Without streaming there is no ``TpotTrigger``; - every plateau gate would see an empty TPOT series, and every verdict would - be ``found: false``. + Accuracy-only runs have no performance phase to collect. TPOT is counted + from server usage in server mode, so only a client-mode run needs a + tokenizer; without one there is no TPOT to bucket. Without streaming there + is no ``TpotTrigger``; every plateau gate would see an empty TPOT series, + and every verdict would be ``found: false``. Agentic is refused by name as well as by profile. Its steady-state metric is per-trajectory NATL, which this collector does not produce. @@ -123,8 +128,11 @@ def steady_state_profile( reason = "the run is accuracy-only, so it has no performance window" elif not enable_streaming: reason = "streaming is off, so there is no TPOT series to gate on" - elif not tokenizer_name: - reason = "the run resolved no tokenizer, so TPOT cannot be counted" + elif not tokenizer_name and token_count_source is TokenCountSource.CLIENT: + reason = ( + "the run resolved no tokenizer and the primary basis is client, " + "so TPOT cannot be counted" + ) elif load_pattern is LoadPatternType.AGENTIC_INFERENCE: reason = "agentic runs are measured by per-trajectory NATL, not a steady window" else: @@ -138,6 +146,51 @@ def steady_state_profile( return None +# API types whose streamed responses carry usage only when stream_options asks. +_STREAM_OPTIONS_API_TYPES = frozenset({APIType.OPENAI, APIType.OPENAI_COMPLETIONS}) + + +def resolve_token_count_source(config: BenchmarkConfig) -> TokenCountSource: + """Effective token-count source for the run; logs it once. + + This is the only place that reads the configured token-count setting. + Videogen endpoints report no token usage, so they always count with the + client tokenizer, whatever the setting says. + + Raises: + SetupError: The server basis is primary on a streaming OpenAI run that + does not ask for usage, so every streamed sample would miss it. + """ + configured = config.settings.token_count_source + api_type = config.endpoint_config.api_type + effective = effective_token_count_source(configured, api_type) + if ( + effective is TokenCountSource.SERVER + and not config.settings.client.stream_include_usage + and api_type in _STREAM_OPTIONS_API_TYPES + and config.model_params.streaming == StreamingMode.ON + ): + raise SetupError( + "settings.token_count_source: server needs usage on every streamed " + "response, but settings.client.stream_include_usage is false " + "(--no-stream-include-usage), so every streamed sample would miss " + "it. Set --token-count-source client (YAML: " + "settings.token_count_source: client) or remove " + "--no-stream-include-usage." + ) + if effective is configured: + logger.info("Token counts: %s", configured.value) + else: + logger.info( + "Token counts: %s (api_type %s reports no token usage; " + "configured %s ignored)", + effective.value, + api_type.value, + configured.value, + ) + return effective + + def _build_aggregator_args( *, socket_dir: str, @@ -151,6 +204,7 @@ def _build_aggregator_args( enable_isl: bool, early_stopping: bool, steady_state_profile: str | None, + token_count_source: TokenCountSource, ) -> list[str]: """CLI args for the metrics_aggregator subprocess.""" args: list[str] = [ @@ -175,6 +229,7 @@ def _build_aggregator_args( args.extend(["--steady-state-profile", steady_state_profile]) args.extend(["--tokenizer-workers", str(tokenizer_workers)]) args.append("--metrics-isl" if enable_isl else "--no-metrics-isl") + args.extend(["--token-count-source", token_count_source.value]) return args @@ -195,7 +250,10 @@ def _build_event_logger_args( def _build_report_from_snapshot( - snap_dict: dict[str, Any], config: BenchmarkConfig + snap_dict: dict[str, Any], + config: BenchmarkConfig, + *, + token_count_source: TokenCountSource, ) -> Report | None: """Build the Report from a metrics snapshot dict; best-effort (None on failure). @@ -223,12 +281,21 @@ def _build_report_from_snapshot( load_pattern.type == LoadPatternType.POISSON and load_pattern.use_legacy_loadgen_qps_metrics ), + token_count_source=token_count_source, ) if not report.complete: - logger.warning( - "Report is incomplete (state=%s, n_pending_tasks=%d)", - report.state, - snap_dict.get("n_pending_tasks", 0), + logger.warning("Report is incomplete: %s", report.incomplete_reason()) + if (partial_reason := report.other_basis_partial_reason()) is not None: + # Same role rule as the aggregator's first-miss log: a shortfall + # on the server basis is informational when client is primary. + logger.log( + ( + logging.WARNING + if token_count_source is TokenCountSource.SERVER + else logging.INFO + ), + "%s", + partial_reason, ) if report.legacy_loadgen_window_duration_ns is not None: logger.warning( @@ -262,6 +329,7 @@ def __init__( event_log_dir: Path, metrics_output_dir: Path, loop: asyncio.AbstractEventLoop, + token_count_source: TokenCountSource, ) -> None: self._config = config self._tokenizer_name = tokenizer_name @@ -270,6 +338,7 @@ def __init__( self._event_log_dir = event_log_dir self._metrics_output_dir = metrics_output_dir self._loop = loop + self._token_count_source = token_count_source self._stack: contextlib.ExitStack | None = None self._launcher: ServiceLauncher | None = None @@ -354,7 +423,9 @@ async def start(self) -> None: accuracy_only=self._accuracy_only, enable_streaming=self._enable_streaming, tokenizer_name=self._tokenizer_name, + token_count_source=self._token_count_source, ), + token_count_source=self._token_count_source, ) event_logger_args = _build_event_logger_args( event_log_dir=self._event_log_dir, @@ -445,7 +516,9 @@ async def drain_and_build_report( logger.error("No metrics snapshot available; cannot build report") report = ( - _build_report_from_snapshot(snap_dict, self._config) + _build_report_from_snapshot( + snap_dict, self._config, token_count_source=self._token_count_source + ) if snap_dict is not None else None ) diff --git a/src/inference_endpoint/metrics/report.py b/src/inference_endpoint/metrics/report.py index 98834de34..3582722cc 100644 --- a/src/inference_endpoint/metrics/report.py +++ b/src/inference_endpoint/metrics/report.py @@ -22,21 +22,24 @@ import os from collections.abc import Callable, Iterable from pathlib import Path -from typing import Any, Final +from typing import Any, Final, NamedTuple import msgspec.json import msgspec.structs -from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( +from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import ( MetricCounterKey, + MetricSeriesKey, ) from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( build_token_series_dict, ) +from inference_endpoint.core.types import TokenCountSource, basis_key from inference_endpoint.evaluation.accuracy_results import ( samples_weighted_average_accuracy, ) from inference_endpoint.metrics.steady_state_diagnostics import SteadyState +from inference_endpoint.metrics.token_count_basis import missing_usage_message from inference_endpoint.utils.version import get_version_info from ..utils import monotime_to_datetime @@ -51,6 +54,38 @@ } +class TokenMetric(NamedTuple): + """A token-count metric reported on both bases, by its base names.""" + + series: str # aggregator series + field: str # result_summary.json key + label: str # report.txt section + unit: str # report.txt unit + scale: float # series value -> unit + + +# The token-count metrics. The primary basis reports them under these names; +# the other basis under the names ``basis_key`` gives it (``tpot_server``). +# Existing tooling parses report.txt by section label, so labels stay fixed. +TOKEN_METRICS: Final = ( + TokenMetric( + MetricSeriesKey.ISL.value, + "input_sequence_lengths", + "Input sequence lengths", + "tokens", + 1.0, + ), + TokenMetric( + MetricSeriesKey.OSL.value, + "output_sequence_lengths", + "Output sequence lengths", + "tokens", + 1.0, + ), + TokenMetric(MetricSeriesKey.TPOT_NS.value, "tpot", "TPOT", "ms", 1e-6), +) + + logger = logging.getLogger(__name__) @@ -206,7 +241,50 @@ def _steady_state_of(snap: dict) -> SteadyState | None: return None -class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] +class OtherBasisMetrics(msgspec.Struct, frozen=True, kw_only=True): # type: ignore[call-arg] + """The token-count metrics of a run's non-primary basis. + + A block is None when the basis recorded no sample. A recorded block + carries ``"partial": True`` when the basis missed some samples: the server + basis any sample without usable usage, the client basis any tokenization + still pending. ``tps`` is None for a partial basis, since a throughput over + an incomplete token total misleads. ``Report.to_json`` writes each field + under its ``basis_key`` name. + """ + + input_sequence_lengths: dict[str, Any] | None + output_sequence_lengths: dict[str, Any] | None + tpot: dict[str, Any] | None + tps: float | None + + @property + def partial(self) -> bool: + """True when a recorded block covers only some samples.""" + return any( + block is not None and block.get("partial", False) + for block in ( + self.input_sequence_lengths, + self.output_sequence_lengths, + self.tpot, + ) + ) + + +class OslPaired(msgspec.Struct, frozen=True, kw_only=True): # type: ignore[call-arg] + """OSL over the samples counted on both bases. + + ``mean_diff`` is the mean per-sample (client - server) difference; + ``ratio`` is client over server tokens, None when the server total is 0. + """ + + n_samples: int + client_tokens: int + server_tokens: int + mean_diff: float + ratio: float | None + + +class Report(msgspec.Struct, frozen=True, kw_only=True): # type: ignore[call-arg] """Summarized benchmark report.""" version: str @@ -222,13 +300,17 @@ class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] # without re-parsing the source dict, and so JSON round-trips don't # depend on Report importing the SessionState enum. state: str - # True iff the aggregator completed, no metrics work remains, and every - # issued sample reached a terminal outcome with valid counter ordering. - # False signals dropped samples, partial async metrics, an interrupted run, - # or a live-snapshot fallback. + # True iff the aggregator completed, every issued sample reached a terminal + # outcome with valid counter ordering, and the primary basis counted every + # successful sample: on the client basis no tokenization work remains, on + # the server basis every successful tracked sample reported usable usage. + # The other basis falling short leaves this True and marks its blocks + # partial. False signals dropped samples, a primary-basis shortfall, an + # interrupted run, or a live-snapshot fallback. complete: bool - # Per-metric rollup dicts (output of _series_to_metric_dict) + # Per-metric rollup dicts (output of _series_to_metric_dict). tpot and the + # sequence-length blocks count tokens on the token_count_source basis. ttft: dict[str, Any] tpot: dict[str, Any] latency: dict[str, Any] @@ -245,8 +327,9 @@ class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] # Derived throughput, computed once in from_snapshot so the serialized # report (result_summary.json) is self-complete. qps is None without a - # duration; tps is also None when no OSL was recorded (non-streaming or - # tokenizer unavailable). + # duration; tps is also None when no OSL was recorded. tps counts tokens + # on the token_count_source basis; other_basis.tps is the other basis + # over the same window. qps: float | None = None tps: float | None = None e2e_avg_interactivity: float | None = None @@ -270,16 +353,23 @@ class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] # OSL over EVERY completed turn, including turns issued after the # performance window closed. ``output_sequence_lengths`` above covers only # the window, whose length depends on concurrency — so it is not - # comparable across Pareto points, while this block is. Same token-counting - # rule (``extract_tokenization_input``) on both. Attached at finalize; None - # only when the block was skipped (no tokenizer, or an unreadable/missing - # performance sample map). A run that scanned turns but counted none returns - # a block with ``n_turns_counted == 0``. ``partial`` is True when any turn was - # errored, undecodable, or missing (a perf UUID with no COMPLETE record) — the - # mean is then over a subset. A run is invalid for the OSL accuracy gate when - # this is None OR ``n_turns_counted == 0`` OR ``partial`` OR not complete. - # Carries ``output_sequence_lengths``/``n_turns_counted``/``n_empty``/ - # ``n_errors``/``n_undecodable``/``n_missing``/``partial``. + # comparable across Pareto points, while this block is. Counted on both + # bases like the window: the client basis with the tokenizer through + # ``extract_tokenization_input``, the server basis from each turn's usable + # ``completion_tokens``. The primary basis takes ``output_sequence_lengths``/ + # ``n_turns_counted``/``n_empty`` (and ``n_errors`` when it is the client + # basis); the other basis sits beside it under ``basis_key`` names, its + # series None when it counted nothing and carrying ``"partial": True`` when + # it fell short. ``n_undecodable``/``n_missing``/``n_missing_usage`` are + # basis-independent. Attached at finalize; None only when the block was + # skipped (a client-primary run without a tokenizer, or an + # unreadable/missing performance sample map). A run that scanned turns but + # counted none returns a block with ``n_turns_counted == 0``. ``partial`` + # follows the primary basis only: True when any turn was errored, + # undecodable, missing (a perf UUID with no COMPLETE record), or, on the + # server basis, lacked usable usage — the mean is then over a subset. A run + # is invalid for the OSL accuracy gate when this is None OR + # ``n_turns_counted == 0`` OR ``partial`` OR not complete. # See mlcommons/endpoints#500. output_sequence_lengths_full_run: dict[str, Any] | None = None @@ -288,6 +378,23 @@ class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] # ``to_json`` (unlike ``accuracy``), so it reaches result_summary.json. steady_state: SteadyState | None = None + # The primary basis: every unsuffixed token count above (tpot, the + # sequence-length blocks, tps, e2e_avg_interactivity, the early-stopping + # maps, the steady block) is on it, and only its shortfall can clear + # ``complete``. + token_count_source: TokenCountSource + # Successful tracked samples whose response reported no usable usage, and + # the subset with empty output: the server basis's shortfall. + n_samples_missing_usage: int = 0 + n_samples_missing_usage_empty_output: int = 0 + # Tokenization tasks still pending at the final snapshot: the client + # basis's shortfall. + client_token_tasks_pending: int = 0 + # The token-count metrics on the other basis. + other_basis: OtherBasisMetrics + # None when no sample was counted on both bases. + osl_paired: OslPaired | None = None + @property def n_samples_succeeded(self) -> int: return max(0, self.n_samples_completed - self.n_samples_failed) @@ -296,11 +403,58 @@ def n_samples_succeeded(self) -> int: def n_samples_dropped(self) -> int: return max(0, self.n_samples_issued - self.n_samples_completed) + def incomplete_reason(self) -> str | None: + """The cause of an incomplete report, or None when it is complete. + + The run state and dropped samples come first, then the primary basis's + shortfall: an interrupted run or a live-snapshot fallback holds partial + counters, so its state is the cause. + """ + if self.complete: + return None + if self.state == "complete" and self.n_samples_dropped == 0: + if self.token_count_source is TokenCountSource.CLIENT: + return "Metrics tokenization did not finish" + if self.n_samples_missing_usage > 0: + return missing_usage_message( + self.n_samples_missing_usage, + self.n_samples_missing_usage_empty_output, + self.n_samples_succeeded, + primary=TokenCountSource.SERVER, + ) + return ( + "Benchmark report is incomplete " + f"(state={self.state}, n_samples_dropped={self.n_samples_dropped})" + ) + + def other_basis_partial_reason(self) -> str | None: + """Why the other basis's blocks are partial, or None when none is. + + Information, not a failure: only the primary basis decides + ``complete``. + """ + if not self.other_basis.partial: + return None + primary = self.token_count_source + if primary.other is TokenCountSource.SERVER: + return missing_usage_message( + self.n_samples_missing_usage, + self.n_samples_missing_usage_empty_output, + self.n_samples_succeeded, + primary=primary, + ) + return ( + f"{self.client_token_tasks_pending} client tokenization tasks were " + "still pending; the *_client blocks are partial and the run is not " + "affected" + ) + @classmethod def from_snapshot( cls, snap: dict[str, Any], *, + token_count_source: TokenCountSource, run_config: dict[str, Any] | None = None, use_legacy_loadgen_qps_metrics: bool = True, ) -> Report: @@ -311,6 +465,11 @@ def from_snapshot( report so result_summary.json is self-describing; it is config, not part of the metrics snapshot. + ``token_count_source`` is the run's primary token-count basis, which + the snapshot does not record. Its series keep their base names and its + shortfall alone decides ``complete``. The other basis's series, named + by ``basis_key``, fill ``other_basis``. + Input is the dict form produced by ``inference_endpoint.async_utils.services.metrics_aggregator.snapshot .snapshot_to_dict``, which is also the shape persisted to @@ -362,8 +521,34 @@ def _series_dict(key: str) -> dict[str, Any]: n_issued = _counter("tracked_samples_issued") n_failed = _counter("tracked_samples_failed") n_completed = _counter("tracked_samples_completed") - isl = _series_dict("isl") - osl = _series_dict("osl") + n_missing = _counter("tracked_samples_missing_usage") + n_missing_empty = _counter("tracked_samples_missing_usage_empty_output") + n_pending_tasks = snap.get("n_pending_tasks", 0) + # The primary basis keeps the base series names. + isl = _series_dict(MetricSeriesKey.ISL.value) + osl = _series_dict(MetricSeriesKey.OSL.value) + + primary = token_count_source + other = primary.other + # Whether each basis missed some successful samples: the client basis + # while tokenization is pending, the server basis for each sample + # without usable usage. + fell_short = { + TokenCountSource.CLIENT: n_pending_tasks > 0, + TokenCountSource.SERVER: n_missing > 0, + } + + def _other_block(metric: TokenMetric) -> dict[str, Any] | None: + stat = series.get(basis_key(metric.series, other, primary=primary)) + if stat is None or stat.get("count", 0) == 0: + return None + block = _series_to_metric_dict(stat) + if fell_short[other]: + block["partial"] = True + return block + + other_blocks = {metric.field: _other_block(metric) for metric in TOKEN_METRICS} + other_osl = other_blocks["output_sequence_lengths"] # Legacy MLPerf LoadGen Server "completed" window (poisson only): first # issued request -> completion of the last-issued request @@ -376,9 +561,9 @@ def _series_dict(key: str) -> dict[str, Any]: # self-complete. The legacy LoadGen window drives the headline QPS/TPS # only when it is enabled, available, AND there are >=2 completions # (QPS = (completed-1)/window is undefined below 2). If any of those - # fail, BOTH QPS and TPS fall back to the native window so they always - # share one window, and legacy_loadgen_window_duration_ns stays None so - # the serialized report honestly records which view it holds. + # fail, QPS and both bases' TPS fall back to the native window so they + # always share one window, and legacy_loadgen_window_duration_ns stays + # None so the serialized report honestly records which view it holds. use_legacy_window = ( use_legacy_loadgen_qps_metrics and raw_loadgen_window_ns > 0 @@ -387,17 +572,35 @@ def _series_dict(key: str) -> dict[str, Any]: legacy_loadgen_window_duration_ns = ( raw_loadgen_window_ns if use_legacy_window else None ) + window_s: float | None + qps: float | None if use_legacy_window: window_s = raw_loadgen_window_ns / 1e9 qps = (n_completed - 1) / window_s - tps = (osl.get("total", 0) / window_s) if osl else None elif duration_ns is not None: - duration_s = duration_ns / 1e9 - qps = n_completed / duration_s - tps = (osl.get("total", 0) / duration_s) if osl else None + window_s = duration_ns / 1e9 + qps = n_completed / window_s else: - qps = None - tps = None + window_s = qps = None + tps = osl.get("total", 0) / window_s if osl and window_s is not None else None + other_tps = ( + other_osl["total"] / window_s + if other_osl is not None and not fell_short[other] and window_s is not None + else None + ) + + n_paired = _counter("osl_paired_samples") + osl_paired: OslPaired | None = None + if n_paired > 0: + paired_client = _counter("osl_paired_client_tokens") + paired_server = _counter("osl_paired_server_tokens") + osl_paired = OslPaired( + n_samples=n_paired, + client_tokens=paired_client, + server_tokens=paired_server, + mean_diff=(paired_client - paired_server) / n_paired, + ratio=paired_client / paired_server if paired_server else None, + ) total_sample_latency_ns = series.get("sample_latency_ns", {}).get("total", 0) # OSL is recorded only for successful responses, while sample latency @@ -416,7 +619,6 @@ def _series_dict(key: str) -> dict[str, Any]: # clean completion). Drives complete=False and the interrupted # indicator in display(). state = snap.get("state", "interrupted") - n_pending_tasks = snap.get("n_pending_tasks", 0) finish_reason_prefix = "tracked_finish_reason_" finish_reason_counts = { reason.value.removeprefix(finish_reason_prefix): int( @@ -437,9 +639,9 @@ def _series_dict(key: str) -> dict[str, Any]: state=state, complete=( state == "complete" - and n_pending_tasks == 0 and accounting_valid and n_completed == n_issued + and not fell_short[primary] ), **{ field: _series_dict(series) @@ -454,6 +656,12 @@ def _series_dict(key: str) -> dict[str, Any]: finish_reason_counts=finish_reason_counts, run_config=run_config, steady_state=_steady_state_of(snap), + token_count_source=primary, + n_samples_missing_usage=n_missing, + n_samples_missing_usage_empty_output=n_missing_empty, + client_token_tasks_pending=n_pending_tasks, + other_basis=OtherBasisMetrics(**other_blocks, tps=other_tps), + osl_paired=osl_paired, ) def to_json(self, save_to: os.PathLike | None = None) -> bytes: @@ -462,6 +670,10 @@ def to_json(self, save_to: os.PathLike | None = None) -> bytes: # accuracy field stays on the struct so report.txt / the console summary # still render it; it is just dropped from this serialized perf summary. payload = msgspec.structs.asdict(self) + # The other basis's fields go flat under their ``basis_key`` names. + primary = self.token_count_source + for name, value in msgspec.structs.asdict(payload.pop("other_basis")).items(): + payload[basis_key(name, primary.other, primary=primary)] = value payload["n_samples_succeeded"] = self.n_samples_succeeded payload["n_samples_dropped"] = self.n_samples_dropped payload.pop("accuracy") @@ -523,10 +735,16 @@ def display( f"metrics below are best-effort partial data.{newline}" ) elif not self.complete: - fn( - "WARNING: report is incomplete (drain timed out or no " - f"final snapshot received) — some async metrics may be missing.{newline}" - ) + # A client-primary report keeps one fixed wording for report.txt + # readers; a server-primary one names the cause. + if self.token_count_source is TokenCountSource.CLIENT: + warning = ( + "report is incomplete (drain timed out or no final snapshot " + "received) — some async metrics may be missing." + ) + else: + warning = f"report is incomplete: {self.incomplete_reason()}" + fn(f"WARNING: {warning}{newline}") fn(f"Version: {self.version}{newline}") if self.git_sha: fn(f"Git SHA: {self.git_sha}{newline}") @@ -622,14 +840,24 @@ def display( fr = self.output_sequence_lengths_full_run n_counted = (fr or {}).get("n_turns_counted", 0) + fr_counts = "" + if fr: + fr_counts = ( + f"({fr.get('n_empty', 0)} empty, {fr.get('n_errors', 0)} errored, " + f"{fr.get('n_undecodable', 0)} undecodable, " + f"{fr.get('n_missing', 0)} missing" + ) + # Missing usage decides the block only on a server primary; a + # client-primary line reads as a tokenizer-only run's. + if self.token_count_source is TokenCountSource.SERVER: + fr_counts += f", {fr.get('n_missing_usage', 0)} missing usage" + fr_counts += ")" if fr and n_counted > 0: osl = fr.get("output_sequence_lengths", {}) fn( f" OSL per-turn mean (accuracy, all turns): " f"{osl.get('avg', 0):.1f} tokens over {n_counted} turns " - f"({fr.get('n_empty', 0)} empty, {fr.get('n_errors', 0)} errored, " - f"{fr.get('n_undecodable', 0)} undecodable, " - f"{fr.get('n_missing', 0)} missing){newline}" + f"{fr_counts}{newline}" ) if fr.get("partial"): fn( @@ -644,37 +872,64 @@ def display( # Block present but no countable turns — do not print a fake 0.0 mean. fn( f" OSL per-turn mean (accuracy): no countable turns " - f"({fr.get('n_empty', 0)} empty, {fr.get('n_errors', 0)} errored, " - f"{fr.get('n_undecodable', 0)} undecodable, " - f"{fr.get('n_missing', 0)} missing){newline}" + f"{fr_counts}{newline}" ) elif self.output_sequence_lengths and self.complete: # A complete performance run (windowed OSL present) but no full-run # block: the accuracy OSL field is missing, so the run is not usable # for that gate. - fn( - f" OSL per-turn mean (accuracy): not computed " - f"(no tokenizer or unreadable sample map){newline}" + # Only a client-primary run needs a tokenizer for the block; the + # block is also skipped without an event log or performance map. + cause = ( + "no tokenizer or unreadable sample map" + if self.token_count_source is TokenCountSource.CLIENT + else "no event log or unreadable sample map" ) + fn(f" OSL per-turn mean (accuracy): not computed ({cause}){newline}") elif not self.complete: fn( f" OSL per-turn mean (accuracy): not computed — run incomplete{newline}" ) + # The other basis follows the primary's summary rows, which stay + # contiguous for readers that expect them in a fixed order. + primary = self.token_count_source + other = primary.other + fn(f"Other token basis ({other.value}):{newline}") + fn(f" Token count source: {primary.value}{newline}") + if (other_tps := self.other_basis.tps) is not None: + fn(f" TPS ({other.value}): {other_tps:.2f}{newline}") + if (paired := self.osl_paired) is not None: + ratio_str = f"{paired.ratio:.4f}" if paired.ratio is not None else "N/A" + fn( + f" OSL paired (client - server): mean {paired.mean_diff:.2f} " + f"tokens/sample, ratio {ratio_str} over {paired.n_samples} " + f"samples{newline}" + ) + if (partial_reason := self.other_basis_partial_reason()) is not None: + fn(f" {partial_reason}{newline}") + if summary_only: fn(f"----------------- End of Summary -----------------{newline}") return fn(f"\n------------------- Latency Breakdowns -------------------{newline}") - for section_name, metric_dict, unit, scale_factor in [ + token_sections = { + metric.field: ( + metric.label, + getattr(self, metric.field), + metric.unit, + metric.scale, + ) + for metric in TOKEN_METRICS + } + sections: list[tuple[str, dict[str, Any], str, float]] = [ ("TTFT", self.ttft, "ms", 1e-6), - ("TPOT", self.tpot, "ms", 1e-6), + token_sections["tpot"], ("Latency", self.latency, "ms", 1e-6), - ("Input sequence lengths", self.input_sequence_lengths, "tokens", 1.0), - # Label deliberately unchanged: existing tooling parses report.txt - # by section name. The full-run block is added alongside it. - ("Output sequence lengths", self.output_sequence_lengths, "tokens", 1.0), + token_sections["input_sequence_lengths"], + token_sections["output_sequence_lengths"], ( "Output sequence lengths (full run, all turns)", (self.output_sequence_lengths_full_run or {}).get( @@ -683,7 +938,35 @@ def display( "tokens", 1.0, ), - ]: + ] + for metric in TOKEN_METRICS: + block = getattr(self.other_basis, metric.field) + if block is not None: + sections.append( + ( + _other_basis_title(metric.label, other, block), + block, + metric.unit, + metric.scale, + ) + ) + full_run_other = (self.output_sequence_lengths_full_run or {}).get( + basis_key("output_sequence_lengths", other, primary=primary) + ) or {} + sections.append( + ( + _other_basis_title( + "Output sequence lengths (full run, all turns)", + other, + full_run_other, + ), + full_run_other, + "tokens", + 1.0, + ) + ) + + for section_name, metric_dict, unit, scale_factor in sections: if not metric_dict: continue fn(f"{section_name}:{newline}") @@ -704,6 +987,14 @@ def display( # --------------------------------------------------------------------------- +def _other_basis_title( + label: str, other: TokenCountSource, block: dict[str, Any] +) -> str: + """An other-basis section title, marked when its block misses some samples.""" + title = f"{label} ({other.value})" + return f"{title} (partial)" if block.get("partial") else title + + def _display_metric( metric_dict: dict[str, Any], fn: Callable[[str], None], diff --git a/tests/integration/commands/conftest.py b/tests/integration/commands/conftest.py new file mode 100644 index 000000000..7c2766a37 --- /dev/null +++ b/tests/integration/commands/conftest.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fixtures shared by the CLI command integration tests.""" + +import json +from collections.abc import Callable, Iterator +from typing import Any + +import pytest +from aiohttp import web +from inference_endpoint.testing.echo_server import EchoServer + + +class RecordingEchoServer(EchoServer): + """Echo server that keeps every chat and text-completion request body.""" + + def __init__(self, *, report_usage: bool = True) -> None: + super().__init__(port=0, report_usage=report_usage) + self.bodies: list[dict[str, Any]] = [] + + async def _handle_echo_chat_completions_request( + self, request: web.Request + ) -> web.Response: + # aiohttp caches the body, so the parent handler can read it again. + self.bodies.append(json.loads(await request.read())) + return await super()._handle_echo_chat_completions_request(request) + + async def _handle_echo_completions_request( + self, request: web.Request + ) -> web.StreamResponse: + self.bodies.append(json.loads(await request.read())) + return await super()._handle_echo_completions_request(request) + + +@pytest.fixture +def recording_echo_server_factory() -> Iterator[Callable[..., RecordingEchoServer]]: + """Yields ``factory(*, report_usage=True) -> RecordingEchoServer``. + + Each server is started on a free port and stopped at test teardown. + """ + servers: list[RecordingEchoServer] = [] + + def factory(*, report_usage: bool = True) -> RecordingEchoServer: + server = RecordingEchoServer(report_usage=report_usage) + server.start() + servers.append(server) + return server + + yield factory + + for server in servers: + server.stop() diff --git a/tests/integration/commands/test_run_timeout.py b/tests/integration/commands/test_run_timeout.py index bc9a48c41..3c503bbee 100644 --- a/tests/integration/commands/test_run_timeout.py +++ b/tests/integration/commands/test_run_timeout.py @@ -56,11 +56,12 @@ Timeouts, WarmupConfig, ) +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.endpoint_client.config import HTTPClientConfig from inference_endpoint.exceptions import ExecutionError from inference_endpoint.metrics.report import Report -from tests.test_helpers import write_large_prompt_dataset +from tests.test_helpers import DEFAULT_SOURCE, write_large_prompt_dataset # Local character-level tokenizer: lets the metrics aggregator tokenize # ISL/OSL without a HuggingFace Hub download (same trick as @@ -154,11 +155,13 @@ def _make_config( timeouts: Timeouts | None = None, streaming: StreamingMode = StreamingMode.OFF, metrics_tokenizer_workers: int | None = None, + token_count_source: TokenCountSource = DEFAULT_SOURCE, ) -> BenchmarkConfig: settings_kwargs: dict[str, Any] = { "load_pattern": load_pattern, "client": _FAST_CLIENT, "warmup": WarmupConfig(enabled=False), + "token_count_source": token_count_source, } if runtime is not None: settings_kwargs["runtime"] = runtime @@ -318,6 +321,9 @@ def test_metrics_drain_timeout_fails_run(mock_http_echo_server, tmp_path): # watchdog: the drain deadline itself must fail the run. metrics_tokenizer_workers=0, timeouts=Timeouts(metrics_drain_timeout_s=1.0), + # In server mode a pending client drain is non-fatal; only the + # client-mode rule fails the run on it. + token_count_source=TokenCountSource.CLIENT, ) report_dir = tmp_path / "report" @@ -332,6 +338,54 @@ def test_metrics_drain_timeout_fails_run(mock_http_echo_server, tmp_path): assert summary["complete"] is False +@pytest.mark.integration +@pytest.mark.parametrize("streaming", [StreamingMode.OFF, StreamingMode.ON]) +def test_metrics_drain_timeout_is_partial_not_fatal_in_server_mode( + mock_http_echo_server, tmp_path, streaming +): + """In server mode an expired client drain leaves the run complete. + + Headline ISL/OSL/TPOT come from server usage, so the pending tokenizer + backlog only marks the *_client blocks partial and the run exits 0. + """ + # The echo server streams every word (~42k frames per sample), so the + # streaming case uses fewer samples; their backlog still outlasts the drain. + n_samples = 30 if streaming is StreamingMode.ON else 100 + dataset_path = write_large_prompt_dataset(tmp_path, n_samples) + config = _make_config( + mock_http_echo_server.url, + dataset_path, + tmp_path / "report", + model_name=str(_CHAR_TOKENIZER_DIR), + metrics_tokenizer_workers=0, + timeouts=Timeouts(metrics_drain_timeout_s=1.0), + token_count_source=TokenCountSource.SERVER, + streaming=streaming, + ) + report_dir = tmp_path / "report" + + run_benchmark(config, TestMode.PERF) + + snapshot = _read_final_snapshot(report_dir) + summary = _read_result_summary(report_dir) + assert snapshot["state"] == "complete" and snapshot["n_pending_tasks"] > 0 + assert summary["complete"] is True and summary["token_count_source"] == "server" + assert "token_count_source" not in summary["run_config"] + assert summary["n_samples_missing_usage"] == 0 + assert summary["client_token_tasks_pending"] == snapshot["n_pending_tasks"] + assert summary["output_sequence_lengths"]["total"] > 0 + if streaming is StreamingMode.ON: + assert summary["tpot"]["total"] > 0 + # A client block is absent when the expired drain recorded nothing, and + # partial otherwise; its throughput is withheld either way. + keys = ["input_sequence_lengths_client", "output_sequence_lengths_client"] + if streaming is StreamingMode.ON: + keys.append("tpot_client") + for key in keys: + assert summary[key] is None or summary[key]["partial"] is True, key + assert summary["tps_client"] is None + + @pytest.mark.integration def test_run_timeout_during_service_launch_aborts_promptly( mock_http_echo_server, ds_dataset_path, tmp_path @@ -363,3 +417,28 @@ def test_run_timeout_during_service_launch_aborts_promptly( # Prompt unwind: nowhere near the 30 s service_ready_timeout_s the # pre-fix behavior would have waited out. assert elapsed < 15.0, f"launch abort took {elapsed:.1f}s" + + +@pytest.mark.integration +def test_run_timeout_outranks_missing_usage( + recording_echo_server_factory, ds_dataset_path, tmp_path +): + """A fired run watchdog reports the timeout, not the missing usage.""" + server = recording_echo_server_factory(report_usage=False) + config = _make_config( + server.url, + ds_dataset_path, + tmp_path, + test_type=TestType.ONLINE, + load_pattern=LoadPattern(type=LoadPatternType.POISSON, target_qps=5), + runtime=RuntimeConfig(n_samples_to_issue=600), + timeouts=Timeouts(run_timeout_s=6.0), + streaming=StreamingMode.ON, + token_count_source=TokenCountSource.SERVER, + ) + + with pytest.raises(ExecutionError, match="Run timeout"): + run_benchmark(config, TestMode.PERF) + + summary = _read_result_summary(tmp_path) + assert summary["complete"] is False and summary["n_samples_missing_usage"] > 0 diff --git a/tests/integration/commands/test_server_token_counts.py b/tests/integration/commands/test_server_token_counts.py new file mode 100644 index 000000000..45d09952d --- /dev/null +++ b/tests/integration/commands/test_server_token_counts.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Server-reported token counts end to end, against the echo server. + +The echo server reports ``prompt_tokens = len(user content)`` and +``completion_tokens = len(emitted text)``, which equal the hermetic char +tokenizer's counts for plain text. Covers the missing-usage failure paths, +the ``client`` opt-out, and which requests carry ``stream_options``. +""" + +import json +import shutil +import subprocess +from pathlib import Path +from typing import Any + +import pytest +from inference_endpoint.commands.benchmark.execute import run_benchmark +from inference_endpoint.config.schema import ( + AccuracyConfig, + AgenticInferenceConfig, + BenchmarkConfig, + Dataset, + DatasetType, + EndpointConfig, + LoadPattern, + LoadPatternType, + ModelParams, + RuntimeConfig, + Settings, + StreamingMode, + TestMode, + TestType, + WarmupConfig, +) +from inference_endpoint.core.types import APIType, TokenCountSource +from inference_endpoint.endpoint_client.config import HTTPClientConfig +from inference_endpoint.exceptions import ExecutionError + +_CHAR = Path(__file__).resolve().parents[2] / "assets/tokenizers/char" +_FAST_CLIENT = HTTPClientConfig(num_workers=1, warmup_connections=0, max_connections=10) +# 12 characters each: char-tokenizer count == echo usage == 96 in total. +_PROMPTS = [f"alpha beta {i}" for i in range(8)] + +_USAGE = "8 of 8 successful tracked samples returned no server token usage" +_FIX = ( + "set --token-count-source client if this server cannot report usage " + "(YAML: settings.token_count_source: client)" +) +_EMPTY = ( + "8 had empty output — likely server errors or cut streams; " + "client mode will not fix these" +) + +# The echo answers its prompt, so the first three rows score and the last two +# do not. +_ACC = [ + ("alpha", "alpha"), + ("beta", "beta"), + ("gamma", "gamma"), + ("What is the answer?", "42"), + ("Tell me a joke", "knock knock"), +] + + +def _rows(path: Path, rows: list[dict[str, Any]]) -> Path: + path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + return path + + +def _summary(report_dir: Path) -> dict[str, Any]: + return json.loads((report_dir / "performance/result_summary.json").read_text()) + + +def _counters(report_dir: Path) -> dict[str, int]: + snapshot = json.loads((report_dir / "metrics/final_snapshot.json").read_text()) + return { + m["name"]: m["value"] for m in snapshot["metrics"] if m["type"] == "counter" + } + + +def _config( + url: str, + dataset: Path, + report_dir: Path, + *, + source: TokenCountSource = TokenCountSource.SERVER, + streaming: StreamingMode = StreamingMode.ON, + api_type: APIType = APIType.OPENAI, + model: str = str(_CHAR), + n_samples: int | None = 8, + agentic: bool = False, + stream_include_usage: bool = True, +) -> BenchmarkConfig: + if agentic: + test_type = TestType.ONLINE + load_pattern = LoadPattern( + type=LoadPatternType.AGENTIC_INFERENCE, target_concurrency=2 + ) + dataset_config = Dataset( + path=str(dataset), agentic_inference=AgenticInferenceConfig() + ) + else: + test_type = TestType.OFFLINE + load_pattern = LoadPattern() + dataset_config = Dataset(path=str(dataset)) + return BenchmarkConfig( + type=test_type, + endpoint_config=EndpointConfig(endpoints=[url], api_type=api_type), + model_params=ModelParams(name=model, streaming=streaming), + datasets=[dataset_config], + report_dir=report_dir, + settings=Settings( + load_pattern=load_pattern, + client=_FAST_CLIENT.with_updates(stream_include_usage=stream_include_usage), + warmup=WarmupConfig(enabled=False), + runtime=RuntimeConfig(n_samples_to_issue=n_samples), + token_count_source=source, + ), + ) + + +@pytest.mark.integration +def test_missing_usage_exits_4_and_writes_artifacts( + recording_echo_server_factory, tmp_path +): + server = recording_echo_server_factory(report_usage=False) + report_dir, config_path = tmp_path / "report", tmp_path / "bench.yaml" + dataset = _rows(tmp_path / "p.jsonl", [{"prompt": p} for p in _PROMPTS]) + _config(server.url, dataset, report_dir).to_yaml_file(config_path) + cli = shutil.which("inference-endpoint") + assert cli is not None + + proc = subprocess.run( + [cli, "benchmark", "from-config", "-c", str(config_path)], + capture_output=True, + text=True, + timeout=240, + ) + + log = proc.stdout + proc.stderr + assert proc.returncode == 4, log[-3000:] + assert f"{_USAGE}; {_FIX}" in log + assert "tokenization did not finish" not in log and "drain timed out" not in log + s, c = _summary(report_dir), _counters(report_dir) + assert s["complete"] is False + assert ( + s["n_samples_missing_usage"], + s["n_samples_missing_usage_empty_output"], + ) == (8, 0) + assert ( + c["tracked_samples_missing_usage"], + c["tracked_samples_missing_usage_empty_output"], + ) == (8, 0) + assert _USAGE in (report_dir / "report.txt").read_text() + assert (report_dir / "events.jsonl").exists() + + +@pytest.mark.integration +def test_missing_usage_with_empty_output_names_server_errors( + recording_echo_server_factory, tmp_path +): + server = recording_echo_server_factory(report_usage=False) + dataset = _rows(tmp_path / "p.jsonl", [{"prompt": ""}] * 8) + + with pytest.raises(ExecutionError) as exc: + run_benchmark(_config(server.url, dataset, tmp_path / "r"), TestMode.PERF) + + assert str(exc.value) == f"{_USAGE}; {_EMPTY}" + s = _summary(tmp_path / "r") + assert ( + s["n_samples_missing_usage"], + s["n_samples_missing_usage_empty_output"], + ) == (8, 8) + + +@pytest.mark.integration +@pytest.mark.parametrize("stream_include_usage", [True, False]) +def test_client_mode_passes_without_usage( + recording_echo_server_factory, tmp_path, stream_include_usage +): + server = recording_echo_server_factory(report_usage=False) + dataset = _rows(tmp_path / "p.jsonl", [{"prompt": p} for p in _PROMPTS]) + + run_benchmark( + _config( + server.url, + dataset, + tmp_path / "r", + source=TokenCountSource.CLIENT, + stream_include_usage=stream_include_usage, + ), + TestMode.PERF, + ) + + s = _summary(tmp_path / "r") + # Missing usage is counted, never fatal, and leaves no *_server block. + assert s["complete"] is True and s["n_samples_missing_usage"] == 8 + assert s["token_count_source"] == "client" + assert "token_count_source" not in s["run_config"] + assert s["output_sequence_lengths"]["total"] == 96 + assert s["output_sequence_lengths_server"] is None and s["tps_server"] is None + # The run setting alone decides whether streaming requests ask for usage. + expected = {"include_usage": True} if stream_include_usage else None + assert len(server.bodies) == 8 + assert all(b["stream"] is True for b in server.bodies) + assert [b.get("stream_options") for b in server.bodies] == [expected] * 8 + + +@pytest.mark.integration +def test_server_mode_matches_char_tokenizer_on_plain_text( + recording_echo_server_factory, tmp_path +): + server = recording_echo_server_factory() + dataset = _rows(tmp_path / "p.jsonl", [{"prompt": p} for p in _PROMPTS]) + + run_benchmark(_config(server.url, dataset, tmp_path / "r"), TestMode.PERF) + + s = _summary(tmp_path / "r") + assert s["complete"] is True and s["token_count_source"] == "server" + assert s["n_samples_missing_usage"] == 0 + for key in ( + "input_sequence_lengths", + "input_sequence_lengths_client", + "output_sequence_lengths", + "output_sequence_lengths_client", + ): + assert s[key]["total"] == 96 + assert s["osl_paired"] == { + "n_samples": 8, + "client_tokens": 96, + "server_tokens": 96, + "mean_diff": 0.0, + "ratio": 1.0, + } + assert [b.get("stream_options") for b in server.bodies] == [ + {"include_usage": True} + ] * 8 + + +@pytest.mark.integration +def test_agentic_rows_request_usage_only_when_streaming( + recording_echo_server_factory, tmp_path +): + server = recording_echo_server_factory() + # Every row sets stream: a stream column with gaps loads as float64 and is + # sent as "stream": 1.0 / 0.0. + rows = [ + { + "conversation_id": "c1", + "turn": 1, + "role": "user", + "content": "alpha", + "stream": True, + }, + { + "conversation_id": "c1", + "turn": 2, + "role": "assistant", + "content": "beta", + "stream": False, + }, + { + "conversation_id": "c1", + "turn": 3, + "role": "user", + "content": "gamma", + "stream": True, + }, + { + "conversation_id": "c2", + "turn": 1, + "role": "user", + "content": "delta", + "stream": False, + }, + ] + + run_benchmark( + _config( + server.url, + _rows(tmp_path / "a.jsonl", rows), + tmp_path / "r", + model="echo-server", + n_samples=None, + agentic=True, + ), + TestMode.PERF, + ) + + streamed = [b for b in server.bodies if b["stream"] is True] + unstreamed = [b for b in server.bodies if b["stream"] is False] + assert (len(streamed), len(unstreamed)) == (2, 1) + assert all(b["stream_options"] == {"include_usage": True} for b in streamed) + assert "stream_options" not in unstreamed[0] + s = _summary(tmp_path / "r") + assert s["complete"] is True and s["n_samples_failed"] == 0 + assert s["n_samples_missing_usage"] == 0 + + +@pytest.mark.integration +@pytest.mark.parametrize("streaming", [StreamingMode.ON, StreamingMode.OFF]) +def test_completions_run_requests_usage_only_when_streaming( + recording_echo_server_factory, tmp_path, streaming +): + # The echo answers stream_options without stream with HTTP 400. + server = recording_echo_server_factory() + + run_benchmark( + _config( + server.url, + _rows(tmp_path / "t.jsonl", [{"input_tokens": [1, 2, 3]}] * 4), + tmp_path / "r", + api_type=APIType.OPENAI_COMPLETIONS, + model="echo-server", + streaming=streaming, + n_samples=4, + ), + TestMode.PERF, + ) + + expected = {"include_usage": True} if streaming is StreamingMode.ON else None + assert [b.get("stream_options") for b in server.bodies] == [expected] * 4 + s = _summary(tmp_path / "r") + assert s["complete"] is True and s["n_samples_missing_usage"] == 0 + # Echo usage per sample: P = len([1, 2, 3]) = 3, C = len("1 2 3") = 5. + assert ( + s["input_sequence_lengths"]["total"], + s["output_sequence_lengths"]["total"], + ) == (12, 20) + # Token-ID prompts give a client ISL without a tokenizer. + assert s["input_sequence_lengths_client"]["total"] == 12 + + +@pytest.mark.integration +def test_accuracy_only_run_without_usage_succeeds( + recording_echo_server_factory, tmp_path +): + # An accuracy-only run never opens performance tracking, so missing usage + # cannot fail it. + server = recording_echo_server_factory(report_usage=False) + acc = Dataset( + name="echo_accuracy", + type=DatasetType.ACCURACY, + path=str( + _rows( + tmp_path / "acc.jsonl", + [{"prompt": p, "ground_truth": g} for p, g in _ACC], + ) + ), + accuracy_config=AccuracyConfig( + eval_method="string_match", + ground_truth="ground_truth", + extractor="identity_extractor", + ), + ) + config = _config( + server.url, + _rows(tmp_path / "p.jsonl", [{"prompt": p} for p in _PROMPTS]), + tmp_path / "r", + n_samples=None, + ) + + run_benchmark(config.with_updates(datasets=[*config.datasets, acc]), TestMode.ACC) + + s = _summary(tmp_path / "r") + assert s["complete"] is True and s["n_samples_missing_usage"] == 0 + entry = json.loads((tmp_path / "r/accuracy/accuracy_results.json").read_text())[ + "accuracy_scores" + ][0] + assert entry["score"] == pytest.approx(0.6) + assert entry["response_counts"]["scored"] == 5 + # The char tokenizer supplies the client counts the server did not. + assert ( + entry["n_osl_samples"], + entry["n_osl_samples_client"], + entry["n_missing_usage"], + ) == (0, 5, 5) + # Requests still ask for usage. + assert [b.get("stream_options") for b in server.bodies] == [ + {"include_usage": True} + ] * 5 diff --git a/tests/integration/commands/test_sigint.py b/tests/integration/commands/test_sigint.py index 24caab057..3ddbdda08 100644 --- a/tests/integration/commands/test_sigint.py +++ b/tests/integration/commands/test_sigint.py @@ -63,13 +63,22 @@ def _write_config( endpoint_url: str, config_path: Path, interrupted_teardown_grace_s: float | None = None, + token_count_source: str | None = None, ) -> None: - """~120 s workload (600 samples @ 5 QPS): only the ^C can end the run.""" + """~120 s workload (600 samples @ 5 QPS): only the ^C can end the run. + + ``token_count_source`` pins the primary basis; unset keeps the default. + """ grace = ( f"\n timeouts:\n interrupted_teardown_grace_s: {interrupted_teardown_grace_s}" if interrupted_teardown_grace_s is not None else "" ) + basis = ( + f"\n token_count_source: {token_count_source}" + if token_count_source is not None + else "" + ) config_path.write_text( f""" type: online @@ -93,7 +102,7 @@ def _write_config( runtime: n_samples_to_issue: 600 warmup: - enabled: false{grace} + enabled: false{grace}{basis} """ ) @@ -255,6 +264,36 @@ def test_group_sigint_exits_gracefully(mock_http_echo_server, tmp_path): assert "KeyboardInterrupt" not in log_text +@pytest.mark.integration +def test_group_sigint_outranks_missing_usage(recording_echo_server_factory, tmp_path): + """A group ^C against a server without usage still exits 130.""" + server = recording_echo_server_factory(report_usage=False) + report_dir = tmp_path / "report" + config_path = tmp_path / "bench.yaml" + _write_config(report_dir, server.url, config_path, token_count_source="server") + + with _benchmark_proc( + [_cli(), "benchmark", "from-config", "-c", str(config_path)] + ) as proc: + _wait_services_ready(proc, report_dir) + # ^C only once samples without usage have completed, so the run is + # one that would otherwise fail on missing usage. + deadline = time.monotonic() + 60.0 + while len(server.bodies) < 3: + assert proc.poll() is None, "benchmark died before issuing" + assert time.monotonic() < deadline, "no requests reached the server" + time.sleep(0.1) + os.killpg(proc.pid, signal.SIGINT) + rc = proc.wait(timeout=60.0) + + assert rc == 130, f"user abort must exit 130, got {rc}" + _assert_interrupted_artifacts(report_dir) + summary = json.loads( + (report_dir / "performance" / "result_summary.json").read_text() + ) + assert summary["n_samples_missing_usage"] > 0 + + @pytest.mark.integration def test_second_sigint_force_quits_without_artifacts(mock_http_echo_server, tmp_path): report_dir = tmp_path / "report" diff --git a/tests/integration/commands/test_warmup.py b/tests/integration/commands/test_warmup.py index 3c26302a0..5a2fc3c59 100644 --- a/tests/integration/commands/test_warmup.py +++ b/tests/integration/commands/test_warmup.py @@ -72,6 +72,11 @@ def _echo_response(prompt: str) -> dict: result = QueryResult(id=req_id, response_output=TextModelOutput(output=prompt)) body = OpenAIAdapter.to_endpoint_response(result).model_dump(mode="json") body["id"] = req_id + body["usage"] = { + "prompt_tokens": len(prompt), + "completion_tokens": len(prompt), + "total_tokens": 2 * len(prompt), + } return body diff --git a/tests/integration/videogen/test_videogen_benchmark.py b/tests/integration/videogen/test_videogen_benchmark.py new file mode 100644 index 000000000..e072cc0d8 --- /dev/null +++ b/tests/integration/videogen/test_videogen_benchmark.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test: a videogen benchmark and its TEST04 audit end to end.""" + +import json + +import pytest +from inference_endpoint.commands.audit import run_audit +from inference_endpoint.commands.benchmark.execute import run_benchmark +from inference_endpoint.config.schema import ( + AuditConfig, + AuditTestId, + BenchmarkConfig, + Dataset, + DatasetType, + EndpointConfig, + ModelParams, + Settings, + StreamingMode, + TestMode, + TestType, + WarmupConfig, +) +from inference_endpoint.core.types import APIType, TokenCountSource +from inference_endpoint.endpoint_client.config import HTTPClientConfig + + +@pytest.mark.integration +def test_videogen_run_and_audit_resolve_to_client_counts(mock_trtllm_serve, tmp_path): + """Videogen endpoints report no usage, so every phase counts with the client.""" + dataset = tmp_path / "prompts.jsonl" + dataset.write_text( + "".join(json.dumps({"prompt": f"a cat {i}"}) + "\n" for i in range(5)) + ) + config = BenchmarkConfig( + type=TestType.OFFLINE, + audit=AuditConfig( + test=AuditTestId.OUTPUT_CACHING_TEST, + samples=5, + audit_samples=5, + sample_index=0, + threshold=0.9, + ), + endpoint_config=EndpointConfig( + endpoints=[mock_trtllm_serve.url], api_type=APIType.VIDEOGEN + ), + model_params=ModelParams(name="wan22", streaming=StreamingMode.OFF), + datasets=[Dataset(path=str(dataset), type=DatasetType.PERFORMANCE)], + settings=Settings( + client=HTTPClientConfig( + num_workers=1, warmup_connections=0, max_connections=10 + ), + warmup=WarmupConfig(enabled=False), + token_count_source=TokenCountSource.SERVER, + ), + report_dir=str(tmp_path / "r"), + ) + # The setting asks for the server basis; the resolver overrides it. + assert config.settings.token_count_source is TokenCountSource.SERVER + + report_dir = run_benchmark(config, TestMode.PERF) + audit_dir = report_dir / "audit" + result = run_audit(config, audit_dir) + + # verify() ran on complete phases and wrote its verdict. + assert result.test_id == AuditTestId.OUTPUT_CACHING_TEST.value + verdict = json.loads((audit_dir / "audit_output_caching_test.json").read_text()) + assert (verdict["test"], verdict["passed"]) == (result.test_id, result.passed) + assert (audit_dir / "verify_OUTPUT_CACHING_TEST.txt").is_file() + for phase in ("", "audit/reference/", "audit/output_caching/"): + s = json.loads( + (report_dir / f"{phase}performance/result_summary.json").read_text() + ) + assert s["complete"] is True + assert s["token_count_source"] == "client" + assert "token_count_source" not in s["run_config"] diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 9b4729a32..3f75f684c 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -24,9 +24,11 @@ import tempfile from pathlib import Path from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch from urllib import error as urllib_error +import cyclopts import inference_endpoint.commands.benchmark.execute as execute_mod import pandas as pd import pytest @@ -48,6 +50,7 @@ finalize_benchmark, setup_benchmark, ) +from inference_endpoint.commands.benchmark.pipeline import resolve_token_count_source from inference_endpoint.commands.benchmark.profiling import ( ProfileController, _derive_profile_urls, @@ -84,7 +87,12 @@ OnlineBenchmarkConfig as OnlineConfig, ) from inference_endpoint.config.utils import cli_error_formatter as _error_formatter -from inference_endpoint.core.types import APIType, PhaseType, QueryResult +from inference_endpoint.core.types import ( + APIType, + PhaseType, + QueryResult, + TokenCountSource, +) from inference_endpoint.dataset_manager.agentic_inference_dataset import ( AgenticInferenceDataset, ) @@ -93,6 +101,7 @@ from inference_endpoint.endpoint_client.config import HTTPClientConfig from inference_endpoint.evaluation.scoring import ( AgenticInferenceInlineScorer, + PassAt1Scorer, Scorer, SWEBenchScorer, ) @@ -113,6 +122,8 @@ from inference_endpoint.metrics.report import Report from pydantic import ValidationError +from tests.test_helpers import DEFAULT_SOURCE, make_report + TEMPLATE_DIR = ( Path(__file__).resolve().parents[3] / "src" @@ -219,6 +230,7 @@ def _make_benchmark_context( dataloader=dataloader, rt_settings=rt_settings, total_samples=dataloader.num_samples(), + token_count_source=DEFAULT_SOURCE, eval_configs=eval_configs or [], ) @@ -592,6 +604,237 @@ def test_accuracy_only_skips_salt_validation(self, tmp_path): assert perf_loader is None +@pytest.mark.unit +class TestMultiChoiceRows: + """Rows with ``n > 1`` fail setup on the server basis and warn on the client basis.""" + + _ACC_CONFIG = { + "eval_method": "pass_at_1", + "ground_truth": "ground_truth", + "extractor": "boxed_math_extractor", + } + + def _config( + self, + tmp_path: Path, + source: TokenCountSource, + perf_rows: list[dict], + acc_rows: list[dict], + *, + perf_accuracy: bool = False, + ) -> BenchmarkConfig: + for name, rows in (("perf", perf_rows), ("acc", acc_rows)): + (tmp_path / f"{name}.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in rows) + ) + perf: dict[str, Any] = { + "name": "perf", + "type": "performance", + "path": str(tmp_path / "perf.jsonl"), + } + if perf_accuracy: + perf["accuracy_config"] = self._ACC_CONFIG + return BenchmarkConfig( + type=TestType.OFFLINE, + model_params={"name": "test-model"}, + endpoint_config={"endpoints": ["http://localhost:8000"]}, + report_dir=str(tmp_path / "report"), + settings={"token_count_source": source}, + datasets=[ + perf, + { + "name": "acc", + "type": "accuracy", + "path": str(tmp_path / "acc.jsonl"), + "accuracy_config": self._ACC_CONFIG, + }, + ], + ) + + @staticmethod + def _setup(config: BenchmarkConfig, mode: TestMode) -> BenchmarkContext: + with ( + patch.object(execute_mod, "pin_loadgen", return_value=None), + patch.object(execute_mod, "_check_tokenizer_exists", return_value=False), + ): + return setup_benchmark(config, mode) + + @staticmethod + def _n_warnings(caplog) -> list[str]: + return [ + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "set n > 1" in r.getMessage() + ] + + @pytest.mark.parametrize("source", list(TokenCountSource)) + @pytest.mark.parametrize( + "perf_row, acc_row, name, mode", + [ + ( + {"prompt": "p", "n": 2}, + {"prompt": "q", "ground_truth": "42"}, + "perf", + TestMode.BOTH, + ), + ( + {"prompt": "p"}, + {"prompt": "q", "ground_truth": "42", "n": 2}, + "acc", + TestMode.BOTH, + ), + ( # PERF still issues accuracy phases + {"prompt": "p"}, + {"prompt": "q", "ground_truth": "42", "n": 2}, + "acc", + TestMode.PERF, + ), + ], + ) + def test_n_gt_1_rows_rejected_only_for_server_primary( + self, tmp_path, caplog, source, perf_row, acc_row, name, mode + ): + config = self._config(tmp_path, source, [perf_row], [acc_row]) + if source is TokenCountSource.SERVER: + with pytest.raises( + SetupError, + match=rf"Dataset '{name}': sample 0 sets n=2\. " + r".*remove 'n' or set --token-count-source client " + r"\(YAML: settings\.token_count_source: client\)\.$", + ): + self._setup(config, mode) + return + with caplog.at_level(logging.WARNING): + self._setup(config, mode) + [warning] = self._n_warnings(caplog) + assert warning.startswith("1 sample(s) set n > 1") + assert "*_server" in warning + + def test_client_warning_counts_each_issued_row_once(self, tmp_path, caplog): + # The performance dataset's inline scorer reuses the same loaded rows. + config = self._config( + tmp_path, + TokenCountSource.CLIENT, + [{"prompt": "p", "n": 2, "ground_truth": "42"}] * 2, + [{"prompt": "q", "ground_truth": "42", "n": 3}], + perf_accuracy=True, + ) + with caplog.at_level(logging.WARNING): + self._setup(config, TestMode.BOTH) + [warning] = self._n_warnings(caplog) + assert warning.startswith("3 sample(s) set n > 1") + + @pytest.mark.parametrize("mode", [TestMode.ACC, TestMode.BOTH]) + def test_n_allowed_on_external_scorer_dataset(self, tmp_path, mode): + config = self._config( + tmp_path, + TokenCountSource.SERVER, + [{"prompt": "p"}], + [{"prompt": "q", "ground_truth": "42", "n": 2}], + ) + with patch.object(PassAt1Scorer, "SKIP_ENDPOINT_PHASE", True): + ctx = self._setup(config, mode) + assert ctx.eval_configs[0].dataset.load_sample(0)["n"] == 2 + + def test_n_of_one_allowed_on_server_basis(self, tmp_path, caplog): + config = self._config( + tmp_path, + TokenCountSource.SERVER, + [{"prompt": "p", "n": 1}], + [{"prompt": "q", "ground_truth": "42"}], + ) + with caplog.at_level(logging.WARNING): + ctx = self._setup(config, TestMode.BOTH) + assert ctx.dataloader.load_sample(0)["n"] == 1 + assert self._n_warnings(caplog) == [] + + +def _usage_opt_out_config( + *, + api_type: APIType, + streaming: StreamingMode, + source: TokenCountSource, + stream_include_usage: bool, +) -> OfflineConfig: + return OfflineConfig( + **( + _OFFLINE_KWARGS + | { + "endpoint_config": { + "endpoints": ["http://test:8000"], + "api_type": api_type, + }, + "model_params": {"name": "test-model", "streaming": streaming}, + } + ), + settings=OfflineSettings( + token_count_source=source, + client=HTTPClientConfig(stream_include_usage=stream_include_usage), + ), + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("api_type", [APIType.OPENAI, APIType.OPENAI_COMPLETIONS]) +def test_server_primary_with_usage_opt_out_rejected(api_type): + config = _usage_opt_out_config( + api_type=api_type, + streaming=StreamingMode.ON, + source=TokenCountSource.SERVER, + stream_include_usage=False, + ) + with pytest.raises(SetupError) as exc: + resolve_token_count_source(config) + assert "settings.token_count_source" in str(exc.value) + assert "settings.client.stream_include_usage" in str(exc.value) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "api_type, streaming, source, stream_include_usage", + [ + pytest.param( + APIType.OPENAI, + StreamingMode.OFF, + TokenCountSource.SERVER, + False, + id="non-streaming-body-carries-usage", + ), + pytest.param( + APIType.SGLANG, + StreamingMode.ON, + TokenCountSource.SERVER, + False, + id="sglang-always-reports-usage", + ), + pytest.param( + APIType.OPENAI, + StreamingMode.ON, + TokenCountSource.CLIENT, + False, + id="client-primary", + ), + pytest.param( + APIType.OPENAI, + StreamingMode.ON, + TokenCountSource.SERVER, + True, + id="usage-requested", + ), + ], +) +def test_usage_opt_out_allowed_when_primary_counts_need_no_stream_usage( + api_type, streaming, source, stream_include_usage +): + config = _usage_opt_out_config( + api_type=api_type, + streaming=streaming, + source=source, + stream_include_usage=stream_include_usage, + ) + assert resolve_token_count_source(config) is source + + class TestCommandHandlers: """Test offline/online/from_config handlers (mock run_benchmark).""" @@ -701,6 +944,54 @@ def test_metrics_isl_cli_default_and_disable(self): ) assert bound.arguments["config"].settings.metrics_isl is False + @pytest.mark.unit + def test_token_count_source_cli_default_and_override(self): + base = [ + "offline", + "--endpoints", + "http://h:80", + "--model", + "m", + "--dataset", + "d.jsonl", + ] + _, bound, _ = benchmark_app.parse_args(base, exit_on_error=False) + assert bound.arguments["config"].settings.token_count_source is DEFAULT_SOURCE + + _, bound, _ = benchmark_app.parse_args( + [*base, "--token-count-source", DEFAULT_SOURCE.other.value], + exit_on_error=False, + ) + assert ( + bound.arguments["config"].settings.token_count_source + is DEFAULT_SOURCE.other + ) + + with pytest.raises(cyclopts.ValidationError, match="token_count_source"): + benchmark_app.parse_args( + [*base, "--token-count-source", "tokenizer"], exit_on_error=False + ) + + @pytest.mark.unit + def test_cli_opt_out_flag(self): + base = [ + "offline", + "--endpoints", + "http://h:80", + "--model", + "m", + "--dataset", + "d.jsonl", + ] + for extra, expected in ( + ([], True), + (["--no-stream-include-usage"], False), + (["--stream-include-usage"], True), + ): + _, bound, _ = benchmark_app.parse_args([*base, *extra], exit_on_error=False) + client = bound.arguments["config"].settings.client + assert client.stream_include_usage is expected, extra + @pytest.mark.unit @pytest.mark.parametrize( ("flag", "value"), @@ -1363,8 +1654,48 @@ def _make_ctx(self, config, tmp_path): dataloader=ds, rt_settings=rt, total_samples=1, + token_count_source=DEFAULT_SOURCE, ) + async def _aggregator_args(self, ctx, tmp_path) -> list[str]: + """Run the benchmark up to service launch; return the aggregator argv.""" + captured: list = [] + + async def _capture_launch(service_configs, *, timeout): + captured.extend(service_configs) + raise KeyboardInterrupt("stop after launch") + + mock_zmq = MagicMock() + mock_zmq.socket_dir = str(tmp_path / "sockets") + + with ( + patch( + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" + ) as MockZMQ, + patch( + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" + ) as MockPub, + patch( + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" + ) as MockSub, + patch( + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" + ) as MockLauncher, + patch("inference_endpoint.commands.benchmark.execute.tqdm"), + ): + MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) + MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) + MockPub.return_value.socket_name = "test_pub" + MockSub.return_value.start = MagicMock() + MockLauncher.return_value.launch = _capture_launch + + loop = asyncio.get_event_loop() + with pytest.raises(KeyboardInterrupt): + await _run_benchmark_async(ctx, loop, sigint=SigintGovernor()) + + aggregator_cfg = next(c for c in captured if "metrics_aggregator" in c.module) + return aggregator_cfg.args + @pytest.mark.unit @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1491,48 +1822,24 @@ async def test_metrics_isl_forwarded_to_aggregator_args( **_OFFLINE_KWARGS, settings=OfflineSettings(metrics_isl=metrics_isl), ) - ctx = self._make_ctx(config, tmp_path) - - captured: list = [] - - async def _capture_launch(service_configs, *, timeout): - captured.extend(service_configs) - raise KeyboardInterrupt("stop after launch") - - mock_zmq = MagicMock() - mock_zmq.socket_dir = str(tmp_path / "sockets") - - with ( - patch( - "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" - ) as MockZMQ, - patch( - "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" - ) as MockPub, - patch( - "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" - ) as MockSub, - patch( - "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" - ) as MockLauncher, - patch("inference_endpoint.commands.benchmark.execute.tqdm"), - ): - MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) - MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) - MockPub.return_value.socket_name = "test_pub" - MockSub.return_value.start = MagicMock() - MockLauncher.return_value.launch = _capture_launch - - loop = asyncio.get_event_loop() - with pytest.raises(KeyboardInterrupt): - await _run_benchmark_async(ctx, loop, sigint=SigintGovernor()) - - aggregator_cfg = next(c for c in captured if "metrics_aggregator" in c.module) - args = aggregator_cfg.args + args = await self._aggregator_args(self._make_ctx(config, tmp_path), tmp_path) assert expected_flag in args # Exactly one spelling of the flag is emitted. assert sum(a.endswith("metrics-isl") for a in args) == 1 + @pytest.mark.unit + @pytest.mark.asyncio + @pytest.mark.parametrize("source", list(TokenCountSource)) + async def test_token_count_source_forwarded_to_aggregator_args( + self, tmp_path, source + ): + ctx = dataclasses.replace( + self._make_ctx(OfflineConfig(**_OFFLINE_KWARGS), tmp_path), + token_count_source=source, + ) + args = await self._aggregator_args(ctx, tmp_path) + assert args[args.index("--token-count-source") + 1] == source.value + @pytest.mark.unit @pytest.mark.asyncio async def test_tmpfs_dir_cleaned_up_on_mid_run_crash(self, tmp_path, monkeypatch): @@ -2012,6 +2319,7 @@ def _make_ctx(self, config, rt_settings, dataloader): dataloader=dataloader, rt_settings=rt_settings, total_samples=dataloader.num_samples(), + token_count_source=DEFAULT_SOURCE, eval_configs=[], ) @@ -2648,6 +2956,28 @@ def test_skip_endpoint_phase_scorer_gets_empty_sample_idx_map(self, tmp_path): assert results["accuracy_scores"][0]["dataset_name"] == "external_accuracy" assert results["accuracy_scores"][0]["score"] == 1.0 + @pytest.mark.unit + @pytest.mark.parametrize("source", list(TokenCountSource)) + def test_accuracy_results_carry_the_run_token_basis(self, tmp_path, source): + dataset = _make_loaded_dataset() + ctx = _make_benchmark_context( + config=OfflineConfig(**_OFFLINE_KWARGS), + report_dir=tmp_path, + test_mode=TestMode.ACC, + dataloader=dataset, + eval_configs=[ + self._make_eval_config(_SelfContainedScorer, dataset, tmp_path) + ], + ) + ctx.token_count_source = source + + finalize_benchmark(ctx, _make_benchmark_result(tmp_path)) + + results = json.loads( + (tmp_path / "accuracy" / "accuracy_results.json").read_text() + ) + assert results["token_count_source"] == source.value + @pytest.mark.unit @pytest.mark.parametrize(("dataset_size", "expected"), [(3, 2), (1, 1)]) def test_skip_endpoint_phase_scorer_reports_external_sample_count( @@ -2744,6 +3074,24 @@ def test_sigint_during_scoring_invalidates_run(self, tmp_path, monkeypatch): assert summary["complete"] is False assert summary["state"] == "interrupted" + @pytest.mark.unit + @pytest.mark.parametrize("source", list(TokenCountSource)) + def test_full_run_osl_gets_the_effective_token_basis( + self, tmp_path, monkeypatch, source + ): + ctx = dataclasses.replace( + _make_benchmark_context( + config=OfflineConfig(**_OFFLINE_KWARGS), report_dir=tmp_path + ), + token_count_source=source, + ) + full_run = MagicMock(return_value=None) + monkeypatch.setattr(execute_mod, "full_run_osl_for_report", full_run) + + finalize_benchmark(ctx, _make_benchmark_result(tmp_path)) + + full_run.assert_called_once_with(tmp_path, None, token_count_source=source) + @staticmethod def _make_report(state: str) -> Report: return Report.from_snapshot( @@ -2766,7 +3114,8 @@ def _make_report(state: str) -> Report: }, {"type": "counter", "name": "tracked_samples_failed", "value": 0}, ], - } + }, + token_count_source=DEFAULT_SOURCE, ) @@ -2796,6 +3145,98 @@ def test_stall_finalizes_then_raises(self, tmp_path, monkeypatch): finalize.assert_called_once_with(ctx, bench) +_M = { + "complete": False, + "token_count_source": TokenCountSource.SERVER, + "n_samples_missing_usage": 3, +} + + +class TestRunBenchmarkFailurePrecedence: + """Which failure ``run_benchmark`` reports when several apply. + + Order: interrupt > run timeout/stall/no report > the incomplete report's + ``incomplete_reason``. + """ + + @staticmethod + def _run(tmp_path, monkeypatch, **bench_over): + config = OfflineConfig(**_OFFLINE_KWARGS) + ctx = _make_benchmark_context(config, tmp_path) + bench = _make_benchmark_result(tmp_path) + for name, value in bench_over.items(): + setattr(bench, name, value) + monkeypatch.setattr(execute_mod, "setup_benchmark", MagicMock(return_value=ctx)) + monkeypatch.setattr( + execute_mod, "run_benchmark_async", MagicMock(return_value=bench) + ) + monkeypatch.setattr(execute_mod, "finalize_benchmark", MagicMock()) + return execute_mod.run_benchmark(config, TestMode.PERF) + + @pytest.mark.unit + def test_interrupt_outranks_missing_usage(self, tmp_path, monkeypatch): + with pytest.raises(KeyboardInterrupt): + self._run( + tmp_path, monkeypatch, user_interrupted=True, report=make_report(**_M) + ) + + @pytest.mark.unit + @pytest.mark.parametrize( + ("over", "match"), + [ + ({"run_timed_out": True}, "Run timeout"), + ({"stall_error": "endpoint stalled"}, "endpoint stalled"), + ({"report": None}, "no usable metrics report"), + ], + ) + def test_timeout_stall_no_report_outrank_missing_usage( + self, tmp_path, monkeypatch, over, match + ): + with pytest.raises(ExecutionError, match=match): + self._run(tmp_path, monkeypatch, **({"report": make_report(**_M)} | over)) + + @pytest.mark.unit + @pytest.mark.parametrize( + ("over", "match"), + [ + ( + {}, + r"^3 of 10 successful tracked samples returned no server token " + r"usage; set --token-count-source client if this server cannot " + r"report usage \(YAML: settings\.token_count_source: client\)$", + ), + ( + {"n_samples_missing_usage_empty_output": 3}, + "; 3 had empty output — likely server errors or cut streams; " + "client mode will not fix these$", + ), + ( + { + "token_count_source": TokenCountSource.CLIENT, + "n_samples_missing_usage": 0, + "client_token_tasks_pending": 4, + }, + "^Metrics tokenization did not finish$", + ), + ( + {"state": "interrupted"}, + r"^Benchmark report is incomplete \(state=interrupted, " + r"n_samples_dropped=0\)$", + ), + ], + ) + def test_incomplete_report_messages(self, tmp_path, monkeypatch, over, match): + with pytest.raises(ExecutionError, match=match): + self._run(tmp_path, monkeypatch, report=make_report(**(_M | over))) + + @pytest.mark.unit + def test_server_mode_pending_client_drain_exits_zero(self, tmp_path, monkeypatch): + report = make_report( + token_count_source=TokenCountSource.SERVER, client_token_tasks_pending=4 + ) + assert self._run(tmp_path, monkeypatch, report=report) == tmp_path + + class TestScorerMethodSync: """Ensure ScorerMethod enum stays in sync with the scorer registry.""" @@ -3029,6 +3470,59 @@ def test_no_override_yields_none_when_model_has_no_tokenizer( assert ctx.tokenizer_name is None + @pytest.mark.unit + @pytest.mark.parametrize( + ("source", "expected"), + [ + ( + TokenCountSource.SERVER, + "No tokenizer: the *_client ISL/OSL/TPOT blocks are unavailable; " + "ISL/OSL/TPOT use server-reported usage", + ), + ( + TokenCountSource.CLIENT, + "No tokenizer: ISL/OSL/TPOT are unavailable while the client " + "basis is primary; the *_server ISL/OSL/TPOT blocks are " + "reported when the server returns usage. Set --tokenizer " + "(model_params.tokenizer_name) or --token-count-source server " + "(settings.token_count_source)", + ), + ], + ) + def test_no_tokenizer_warning_names_the_token_basis( + self, + tmp_path, + caplog, + _base_patches, + _simple_dataset, + _rt_settings, + source, + expected, + ): + config = OfflineConfig( + **_OFFLINE_KWARGS, + report_dir=str(tmp_path), + settings=OfflineSettings(token_count_source=source), + ) + with ( + patch( + "inference_endpoint.commands.benchmark.execute._check_tokenizer_exists", + return_value=False, + ), + patch( + "inference_endpoint.commands.benchmark.execute._load_datasets", + return_value=(_simple_dataset, []), + ), + patch( + "inference_endpoint.commands.benchmark.execute.RuntimeSettings.from_config", + return_value=_rt_settings, + ), + caplog.at_level(logging.WARNING), + ): + setup_benchmark(config, TestMode.PERF) + + assert expected in caplog.text + class TestSetupBenchmark: """Tests for setup-time config normalization and sample accounting.""" @@ -3300,6 +3794,71 @@ def test_agentic_trajectory_repeats_are_included_in_total_samples( # Issuing five trajectories repeats c1 and c2, adding 2 + 1 more. assert ctx.total_samples == 9 + @pytest.mark.unit + @pytest.mark.parametrize( + ("api_type", "configured", "expected", "log_line"), + [ + ( + APIType.OPENAI, + TokenCountSource.SERVER, + TokenCountSource.SERVER, + "Token counts: server", + ), + ( + APIType.OPENAI, + TokenCountSource.CLIENT, + TokenCountSource.CLIENT, + "Token counts: client", + ), + ( + APIType.VIDEOGEN, + TokenCountSource.SERVER, + TokenCountSource.CLIENT, + "Token counts: client (api_type videogen reports no token usage; " + "configured server ignored)", + ), + ], + ) + def test_setup_sets_effective_token_count_source( + self, + tmp_path, + caplog, + _base_patches, + _simple_dataset, + _rt_settings, + api_type, + configured, + expected, + log_line, + ): + config = OfflineConfig( + **( + _OFFLINE_KWARGS + | { + "endpoint_config": { + "endpoints": ["http://test:8000"], + "api_type": api_type, + } + } + ), + settings=OfflineSettings(token_count_source=configured), + report_dir=str(tmp_path), + ) + with caplog.at_level(logging.INFO): + ctx = self._setup( + config, TestMode.PERF, (_simple_dataset, []), _rt_settings + ) + + assert ctx.token_count_source is expected + assert [m for m in caplog.messages if m.startswith("Token counts:")] == [ + log_line + ] + + @pytest.mark.unit + def test_context_takes_the_default_source(self, tmp_path): + ctx = _make_benchmark_context(OfflineConfig(**_OFFLINE_KWARGS), tmp_path) + assert ctx.token_count_source is DEFAULT_SOURCE + class TestReportConfigSecretRedaction: @pytest.mark.unit diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index b89999b93..fc0810933 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -43,11 +43,18 @@ _build_aggregator_args, _build_report_from_snapshot, _load_final_snapshot_from_disk, + resolve_token_count_source, steady_state_profile, ) -from inference_endpoint.config.schema import LoadPatternType +from inference_endpoint.config.schema import ( + LoadPatternType, + OfflineBenchmarkConfig, +) +from inference_endpoint.core.types import APIType, TokenCountSource, basis_key from inference_endpoint.metrics.report import Report +from tests.test_helpers import DEFAULT_SOURCE + # Dotted path for monkeypatch string targets — keeps a single import style (the # module is imported once via ``from ... import`` and patched by path here). _PIPE = "inference_endpoint.commands.benchmark.pipeline" @@ -127,18 +134,25 @@ class TestReportFromLoadedSnapshot: relies on across the three terminal states.""" @pytest.mark.parametrize( - "state, n_pending, expected_complete", + "state, n_pending, primary, expected_complete", [ - (SessionState.COMPLETE.value, 0, True), - # Drain-timeout: COMPLETE state but tasks still pending. - (SessionState.COMPLETE.value, 3, False), + (SessionState.COMPLETE.value, 0, DEFAULT_SOURCE, True), + # Drain-timeout: COMPLETE state but tokenization still pending, + # which leaves only the client basis short. + (SessionState.COMPLETE.value, 3, TokenCountSource.CLIENT, False), + (SessionState.COMPLETE.value, 3, TokenCountSource.SERVER, True), # Interrupted: signal-handler-written snapshot. - (SessionState.INTERRUPTED.value, 0, False), - (SessionState.INTERRUPTED.value, 7, False), + (SessionState.INTERRUPTED.value, 0, DEFAULT_SOURCE, False), + (SessionState.INTERRUPTED.value, 7, DEFAULT_SOURCE, False), ], ) def test_report_complete_flag_matches_state_and_pending( - self, tmp_path: Path, state: str, n_pending: int, expected_complete: bool + self, + tmp_path: Path, + state: str, + n_pending: int, + primary: TokenCountSource, + expected_complete: bool, ): target = tmp_path / "final_snapshot.json" target.write_text( @@ -146,7 +160,7 @@ def test_report_complete_flag_matches_state_and_pending( ) loaded = _load_final_snapshot_from_disk(target) assert loaded is not None - report = Report.from_snapshot(loaded) + report = Report.from_snapshot(loaded, token_count_source=primary) assert report.state == state assert report.complete is expected_complete @@ -159,7 +173,10 @@ def test_interrupted_display_surfaces_signal_warning(self, tmp_path: Path): target.write_text( json.dumps(_snapshot_dict(state=SessionState.INTERRUPTED.value)) ) - report = Report.from_snapshot(_load_final_snapshot_from_disk(target) or {}) + report = Report.from_snapshot( + _load_final_snapshot_from_disk(target) or {}, + token_count_source=DEFAULT_SOURCE, + ) lines: list[str] = [] report.display(fn=lines.append, summary_only=True) output = "\n".join(lines) @@ -203,7 +220,11 @@ class TestBuildReportFromSnapshot: def test_incomplete_report_logs_warning(self, monkeypatch, caplog): fake_report = SimpleNamespace( - complete=False, state="complete", legacy_loadgen_window_duration_ns=None + complete=False, + state="complete", + legacy_loadgen_window_duration_ns=None, + incomplete_reason=lambda: "Metrics tokenization did not finish", + other_basis_partial_reason=lambda: None, ) monkeypatch.setattr( f"{_PIPE}.Report", @@ -211,16 +232,52 @@ def test_incomplete_report_logs_warning(self, monkeypatch, caplog): ) with caplog.at_level("WARNING"): out = _build_report_from_snapshot( - _snapshot_dict(n_pending_tasks=3), _fake_config() + _snapshot_dict(n_pending_tasks=3), + _fake_config(), + token_count_source=TokenCountSource.CLIENT, ) assert out is fake_report - assert any("incomplete" in r.message.lower() for r in caplog.records) + assert ( + "Report is incomplete: Metrics tokenization did not finish" in caplog.text + ) + + @pytest.mark.parametrize("source", list(TokenCountSource)) + def test_report_gets_the_effective_token_count_source(self, monkeypatch, source): + seen: dict = {} + + def _from_snapshot(*a, **k): + seen.update(k) + return SimpleNamespace( + complete=True, + legacy_loadgen_window_duration_ns=None, + other_basis_partial_reason=lambda: None, + ) + + monkeypatch.setattr( + f"{_PIPE}.Report", SimpleNamespace(from_snapshot=_from_snapshot) + ) + _build_report_from_snapshot( + _snapshot_dict(), _fake_config(), token_count_source=source + ) + assert seen["token_count_source"] is source + + @pytest.mark.parametrize("source", list(TokenCountSource)) + def test_result_summary_has_single_basis_key(self, source): + """The basis is recorded once, as the summary's top-level key.""" + report = _build_report_from_snapshot( + _snapshot_dict(), _fake_config(), token_count_source=source + ) + assert report is not None + summary = json.loads(report.to_json()) + assert summary["token_count_source"] == source.value + assert "token_count_source" not in summary["run_config"] def test_legacy_loadgen_qps_warning(self, monkeypatch, caplog): fake_report = SimpleNamespace( complete=True, state="complete", legacy_loadgen_window_duration_ns=5_000_000_000, + other_basis_partial_reason=lambda: None, ) monkeypatch.setattr( f"{_PIPE}.Report", @@ -228,7 +285,9 @@ def test_legacy_loadgen_qps_warning(self, monkeypatch, caplog): ) with caplog.at_level("WARNING"): out = _build_report_from_snapshot( - _snapshot_dict(), _fake_config(poisson=True, use_legacy=True) + _snapshot_dict(), + _fake_config(poisson=True, use_legacy=True), + token_count_source=TokenCountSource.CLIENT, ) assert out is fake_report assert any("legacy" in r.message.lower() for r in caplog.records) @@ -239,12 +298,62 @@ def _boom(*a, **k): monkeypatch.setattr(f"{_PIPE}.Report", SimpleNamespace(from_snapshot=_boom)) with caplog.at_level("WARNING"): - out = _build_report_from_snapshot(_snapshot_dict(), _fake_config()) + out = _build_report_from_snapshot( + _snapshot_dict(), + _fake_config(), + token_count_source=TokenCountSource.CLIENT, + ) assert out is None assert any( "Failed to build report from snapshot" in r.message for r in caplog.records ) + @pytest.mark.parametrize( + ("primary", "level"), + [ + # Missing server usage beside a client primary is informational. + (TokenCountSource.CLIENT, logging.INFO), + (TokenCountSource.SERVER, logging.WARNING), + ], + ids=["client-info", "server-warning"], + ) + def test_other_basis_shortfall_is_complete_and_logged(self, primary, level, caplog): + other = primary.other + # The client basis falls short through pending tokenization, the + # server basis through samples without usage. + n_pending, n_missing = (3, 0) if other is TokenCountSource.CLIENT else (0, 2) + snap = _snapshot_dict(n_pending_tasks=n_pending) + snap["metrics"] += [ + { + "type": "counter", + "name": "tracked_samples_missing_usage", + "value": n_missing, + }, + { + "type": "series", + "name": basis_key("osl", other, primary=primary), + "count": 2, + "total": 10, + "min": 4, + "max": 6, + "sum_sq": 52, + "percentiles": {"50.0": 5}, + "histogram": [[[4.0, 6.0], 2]], + }, + ] + with caplog.at_level(logging.INFO, logger=_PIPE): + out = _build_report_from_snapshot( + snap, _fake_config(), token_count_source=primary + ) + assert out.complete is True and out.token_count_source is primary + assert out.other_basis.output_sequence_lengths["partial"] is True + [record] = [ + r + for r in caplog.records + if f"the *_{other.value} blocks are partial" in r.getMessage() + ] + assert record.levelno == level + def _make_pipe(tmp_path: Path) -> MetricsPipeline: return MetricsPipeline( @@ -255,6 +364,7 @@ def _make_pipe(tmp_path: Path) -> MetricsPipeline: event_log_dir=tmp_path / "events", metrics_output_dir=tmp_path / "metrics", loop=asyncio.get_event_loop(), + token_count_source=TokenCountSource.SERVER, ) @@ -273,7 +383,7 @@ async def test_falls_back_to_subscriber_when_no_disk_snapshot( pipe.subscriber = MagicMock(latest=object()) monkeypatch.setattr(f"{_PIPE}.snapshot_to_dict", lambda s: {"k": "v"}) monkeypatch.setattr( - f"{_PIPE}._build_report_from_snapshot", lambda d, c: "REPORT" + f"{_PIPE}._build_report_from_snapshot", lambda d, c, **kw: "REPORT" ) with caplog.at_level("WARNING"): report = await pipe.drain_and_build_report() @@ -389,12 +499,14 @@ def _profile( acc=False, streaming=True, tokenizer="tok", + source=TokenCountSource.CLIENT, ): return steady_state_profile( self._config(load_pattern, enabled), accuracy_only=acc, enable_streaming=streaming, tokenizer_name=tokenizer, + token_count_source=source, ) def test_off_by_default_and_when_opted_out(self): @@ -454,6 +566,23 @@ def test_needs_a_performance_phase_a_tokenizer_and_streaming(self): # plateau gate runs on would be empty and every verdict found: false. assert self._profile(LoadPatternType.CONCURRENCY, streaming=False) is None + def test_server_mode_needs_no_tokenizer(self): + server = TokenCountSource.SERVER + assert ( + self._profile(LoadPatternType.CONCURRENCY, tokenizer=None, source=server) + == "concurrency" + ) + assert ( + self._profile(LoadPatternType.CONCURRENCY, streaming=False, source=server) + is None + ) + assert ( + self._profile( + LoadPatternType.AGENTIC_INFERENCE, tokenizer=None, source=server + ) + is None + ) + def test_flag_reaches_the_aggregator_only_when_eligible(self): args = _build_aggregator_args( socket_dir="/tmp", @@ -467,6 +596,7 @@ def test_flag_reaches_the_aggregator_only_when_eligible(self): enable_isl=True, early_stopping=False, steady_state_profile="poisson", + token_count_source=TokenCountSource.SERVER, ) assert args[args.index("--steady-state-profile") + 1] == "poisson" assert "--steady-state-profile" not in _build_aggregator_args( @@ -481,7 +611,49 @@ def test_flag_reaches_the_aggregator_only_when_eligible(self): enable_isl=True, early_stopping=False, steady_state_profile=None, + token_count_source=TokenCountSource.SERVER, + ) + + +@pytest.mark.unit +class TestResolveTokenCountSource: + @pytest.mark.parametrize( + ("api_type", "configured", "expected"), + [ + (APIType.OPENAI, "server", TokenCountSource.SERVER), + (APIType.OPENAI, "client", TokenCountSource.CLIENT), + (APIType.OPENAI_COMPLETIONS, "server", TokenCountSource.SERVER), + (APIType.SGLANG, "server", TokenCountSource.SERVER), + (APIType.VIDEOGEN, "server", TokenCountSource.CLIENT), + (APIType.VIDEOGEN, "client", TokenCountSource.CLIENT), + ], + ) + def test_effective_source(self, api_type, configured, expected): + config = OfflineBenchmarkConfig( + endpoint_config={"endpoints": ["http://h:80"], "api_type": api_type}, + model_params={"name": "m"}, + datasets=[{"path": "d.jsonl"}], + settings={"token_count_source": configured}, + ) + assert resolve_token_count_source(config) is expected + + @pytest.mark.parametrize("source", list(TokenCountSource)) + def test_token_count_source_reaches_the_aggregator(self, source): + args = _build_aggregator_args( + socket_dir="/tmp", + pub_socket_name="p", + metrics_socket_name="m", + metrics_output_dir=Path("/tmp"), + enable_streaming=True, + tokenizer_name="tok", + drain_timeout_s=None, + tokenizer_workers=1, + enable_isl=True, + early_stopping=False, + steady_state_profile="poisson", + token_count_source=source, ) + assert args[args.index("--token-count-source") + 1] == source.value @pytest.mark.unit diff --git a/tests/unit/compliance/test_output_caching.py b/tests/unit/compliance/test_output_caching.py index 2eb633804..866bcdb6d 100644 --- a/tests/unit/compliance/test_output_caching.py +++ b/tests/unit/compliance/test_output_caching.py @@ -46,6 +46,8 @@ ) from inference_endpoint.exceptions import ExecutionError, SetupError +from tests.test_helpers import make_report + # --------------------------------------------------------------------------- # verify_output_caching — pure function, no I/O # --------------------------------------------------------------------------- @@ -482,13 +484,66 @@ def test_refuses_result_on_incomplete_phase(self, tmp_path, monkeypatch): config = self._audit_config() incomplete = MagicMock() incomplete.complete = False + incomplete.state = "complete" + incomplete.incomplete_reason.return_value = ( + "Metrics tokenization did not finish" + ) bench = MagicMock() bench.run_timed_out = False bench.report = incomplete self._patch_phase(monkeypatch, num_samples=100, bench=bench) - with pytest.raises(ExecutionError, match="did not complete cleanly"): + with pytest.raises( + ExecutionError, + match=r"did not complete cleanly \(Metrics tokenization did not finish\)", + ): + run_audit(config, tmp_path) + + @staticmethod + def _missing_usage_bench(**over) -> MagicMock: + bench = MagicMock() + bench.run_timed_out = False + bench.report = make_report( + complete=False, + n_samples_issued=4, + n_samples_completed=4, + token_count_source="server", + n_samples_missing_usage=4, + n_samples_missing_usage_empty_output=4, + **over, + ) + return bench + + @pytest.mark.unit + def test_refuses_missing_usage_phase_with_the_usage_message( + self, tmp_path, monkeypatch + ): + config = self._audit_config() + bench = self._missing_usage_bench() + self._patch_phase(monkeypatch, num_samples=100, bench=bench) + + with pytest.raises( + ExecutionError, + match=r"did not complete cleanly \(4 of 4 successful tracked samples " + r"returned no server token usage; 4 had empty output", + ): + run_audit(config, tmp_path) + + @pytest.mark.unit + def test_interrupted_missing_usage_phase_raises_keyboard_interrupt( + self, tmp_path, monkeypatch + ): + """An interrupt outranks missing usage: the phase exits as interrupted + and the usage message is never surfaced.""" + config = self._audit_config() + bench = self._missing_usage_bench(state="interrupted") + self._patch_phase(monkeypatch, num_samples=100, bench=bench) + + with pytest.raises( + KeyboardInterrupt, match=r"Audit phase '.*' interrupted" + ) as e: run_audit(config, tmp_path) + assert "token usage" not in str(e.value) @pytest.mark.unit def test_interrupted_phase_raises_keyboard_interrupt(self, tmp_path, monkeypatch): diff --git a/tests/unit/metrics/test_report_builder.py b/tests/unit/metrics/test_report_builder.py index b284f983d..fef0bf6e7 100644 --- a/tests/unit/metrics/test_report_builder.py +++ b/tests/unit/metrics/test_report_builder.py @@ -24,13 +24,12 @@ import json import math from pathlib import Path +from typing import Any import msgspec.structs import pytest -from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( - MetricCounterKey, -) from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import ( + MetricCounterKey, MetricSeriesKey, ) from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( @@ -42,7 +41,12 @@ SessionState, snapshot_to_dict, ) -from inference_endpoint.metrics.report import Report, series_metric_dict +from inference_endpoint.core.types import TokenCountSource, basis_key +from inference_endpoint.metrics.report import ( + OtherBasisMetrics, + Report, + series_metric_dict, +) from inference_endpoint.metrics.steady_state_diagnostics import ( Anomaly, SteadyState, @@ -50,6 +54,8 @@ TpsBlock, ) +from tests.test_helpers import DEFAULT_SOURCE, make_report + # 1 hour in ns — same as the aggregator's default bound for time-series. _NS_HIGH = 3_600_000_000_000 @@ -183,6 +189,7 @@ def _build_report( state: SessionState = SessionState.COMPLETE, n_pending_tasks: int = 0, use_legacy_loadgen_qps_metrics: bool = True, + token_count_source: TokenCountSource = DEFAULT_SOURCE, ) -> Report: """Build a Report from a snapshot dict (matches the consumer contract). @@ -196,6 +203,7 @@ def _build_report( return Report.from_snapshot( snapshot_to_dict(snap), use_legacy_loadgen_qps_metrics=use_legacy_loadgen_qps_metrics, + token_count_source=token_count_source, ) @@ -305,7 +313,10 @@ def test_run_config_keyword_only_passthrough(self): snap = snapshot_to_dict( registry.build_snapshot(state=SessionState.COMPLETE, n_pending_tasks=0) ) - assert Report.from_snapshot(snap).run_config is None + assert ( + Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE).run_config + is None + ) run_config = { "load_pattern": {"type": "poisson", "target_qps": 14.75}, "warmup": {"enabled": False, "warmup_random_seed": 42}, @@ -313,7 +324,10 @@ def test_run_config_keyword_only_passthrough(self): "dataloader_random_seed": 42, } assert ( - Report.from_snapshot(snap, run_config=run_config).run_config == run_config + Report.from_snapshot( + snap, run_config=run_config, token_count_source=DEFAULT_SOURCE + ).run_config + == run_config ) def test_additive_accounting_uses_tracked_failures_and_rejects_drops(self): @@ -361,13 +375,20 @@ def test_complete_flag_true_when_state_complete_and_no_pending(self): report = _build_report(registry, state=SessionState.COMPLETE, n_pending_tasks=0) assert report.complete is True - def test_complete_flag_false_when_drain_timeout(self): - """COMPLETE state but n_pending_tasks > 0 → drain timed out, report - is partial. + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_complete_flag_on_drain_timeout_follows_primary(self, primary): + """COMPLETE state but n_pending_tasks > 0 → the tokenizer drain timed + out. Only a client-primary report is then partial: on the server basis + the pending work feeds the *_client blocks alone. """ registry = _make_registry(n_samples=5) - report = _build_report(registry, state=SessionState.COMPLETE, n_pending_tasks=2) - assert report.complete is False + report = _build_report( + registry, + state=SessionState.COMPLETE, + n_pending_tasks=2, + token_count_source=primary, + ) + assert report.complete is (primary is TokenCountSource.SERVER) def test_complete_flag_false_when_state_live(self): """LIVE/DRAINING snapshots produce reports with ``complete=False``.""" @@ -470,6 +491,38 @@ def test_display_full(self): assert "Histogram" in output assert "Percentiles" in output + def test_client_primary_sections_render_pinned_lines(self): + """The rendering of the primary sections is pinned: order, label, unit + and scale (TPOT is recorded in ns and shown in ms).""" + registry = _make_registry(n_samples=3) + for value in (2_000_000.0, 3_000_000.0, 4_000_000.0): + registry.record(MetricSeriesKey.TPOT_NS.value, value) + report = _build_report(registry, token_count_source=TokenCountSource.CLIENT) + + lines: list[str] = [] + report.display(fn=lines.append) + breakdowns = lines[ + lines.index( + "\n------------------- Latency Breakdowns -------------------" + ) : + ] + assert [ + line + for line in breakdowns + if (line[:1].isalpha() and line.endswith(":")) or line.startswith(" Avg.:") + ] == [ + "TTFT:", + " Avg.: 1.01 ms", + "TPOT:", + " Avg.: 3.00 ms", + "Latency:", + " Avg.: 5.05 ms", + "Input sequence lengths:", + " Avg.: 51.00 tokens", + "Output sequence lengths:", + " Avg.: 101.00 tokens", + ] + def test_to_json(self): registry = _make_registry(n_samples=5) report = _build_report(registry) @@ -523,7 +576,9 @@ def test_to_json_and_display_carry_run_config(self): "scheduler_random_seed": 42, "dataloader_random_seed": 42, } - report = Report.from_snapshot(snap, run_config=run_config) + report = Report.from_snapshot( + snap, run_config=run_config, token_count_source=DEFAULT_SOURCE + ) assert json.loads(report.to_json())["run_config"] == run_config lines: list[str] = [] @@ -532,7 +587,8 @@ def test_to_json_and_display_carry_run_config(self): assert any("load_pattern:" in ln and "poisson" in ln for ln in lines) # Absent run_config -> null, not omitted. - assert json.loads(Report.from_snapshot(snap).to_json())["run_config"] is None + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) + assert json.loads(report.to_json())["run_config"] is None def test_to_json_save(self, tmp_path: Path): registry = _make_registry(n_samples=5) @@ -546,21 +602,8 @@ def test_to_json_save(self, tmp_path: Path): def test_display_no_started_at(self): """test_started_at=0 should not display a timestamp.""" - report = Report( - version="test", - git_sha=None, - test_started_at=0, - n_samples_issued=0, - n_samples_completed=0, - n_samples_failed=0, - duration_ns=None, - state="complete", - complete=True, - ttft={}, - tpot={}, - latency={}, - input_sequence_lengths={}, - output_sequence_lengths={}, + report = make_report( + n_samples_issued=0, n_samples_completed=0, duration_ns=None ) lines: list[str] = [] report.display(fn=lines.append, summary_only=True) @@ -569,44 +612,23 @@ def test_display_no_started_at(self): def test_display_warns_when_incomplete(self): """Reports with ``complete=False`` surface a WARNING in display().""" - report = Report( - version="test", - git_sha=None, - test_started_at=0, - n_samples_issued=10, - n_samples_completed=10, - n_samples_failed=0, - duration_ns=1_000_000_000, - state="complete", # drain-timeout case: complete state, n_pending>0 - complete=False, - ttft={}, - tpot={}, - latency={}, - input_sequence_lengths={}, - output_sequence_lengths={}, - ) + # Drain-timeout case: complete state, n_pending>0. + report = make_report(duration_ns=1_000_000_000, complete=False) lines: list[str] = [] report.display(fn=lines.append, summary_only=True) output = "\n".join(lines) assert "WARNING" in output or "incomplete" in output.lower() - def _report_with_full_run_osl(self, fr, *, osl_windowed=None, complete=True): - return Report( - version="test", - git_sha=None, - test_started_at=0, - n_samples_issued=10, - n_samples_completed=10, - n_samples_failed=0, + def _report_with_full_run_osl( + self, fr, *, osl_windowed=None, complete=True, primary=DEFAULT_SOURCE + ): + return make_report( duration_ns=1_000_000_000, state="complete" if complete else "interrupted", complete=complete, - ttft={}, - tpot={}, - latency={}, - input_sequence_lengths={}, output_sequence_lengths=osl_windowed or {}, output_sequence_lengths_full_run=fr, + token_count_source=primary, ) def test_display_full_run_osl_mean(self): @@ -663,45 +685,136 @@ def test_display_full_run_osl_partial_is_flagged(self): assert "PARTIAL" in output assert "NOT valid for the OSL accuracy gate" in output - def test_display_full_run_osl_missing_on_complete_performance_run(self): + @pytest.mark.parametrize( + ("primary", "expected"), + [ + # Fixed wording that report.txt readers match on. A missing + # tokenizer can only cost a client-primary run its block. + ( + TokenCountSource.CLIENT, + " OSL per-turn mean (accuracy): not computed " + "(no tokenizer or unreadable sample map)", + ), + ( + TokenCountSource.SERVER, + " OSL per-turn mean (accuracy): not computed " + "(no event log or unreadable sample map)", + ), + ], + ids=["client", "server"], + ) + def test_display_full_run_osl_missing_on_complete_performance_run( + self, primary, expected + ): """A COMPLETE perf run (windowed OSL present) with no full-run block - blames the tokenizer/sample-map, not incompleteness.""" - report = self._report_with_full_run_osl(None, osl_windowed={"avg": 100.0}) - lines: list[str] = [] - report.display(fn=lines.append, summary_only=True) - output = "".join(lines) - assert "no tokenizer or unreadable sample map" in output - assert "run incomplete" not in output + blames the sample map (and, client-primary, the tokenizer), not + incompleteness.""" + report = self._report_with_full_run_osl( + None, osl_windowed={"avg": 100.0}, primary=primary + ) + lines = _display_lines(report, summary_only=True) + assert expected in lines + assert not any("run incomplete" in line for line in lines) - def test_display_full_run_osl_incomplete_run_with_windowed_osl(self): + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_display_full_run_osl_incomplete_run_with_windowed_osl(self, primary): """An INTERRUPTED perf run (windowed OSL present, no block) must report incompleteness — not a false tokenizer/sample-map diagnosis.""" report = self._report_with_full_run_osl( - None, osl_windowed={"avg": 100.0}, complete=False + None, osl_windowed={"avg": 100.0}, complete=False, primary=primary ) - lines: list[str] = [] - report.display(fn=lines.append, summary_only=True) - output = "".join(lines) + output = "".join(_display_lines(report, summary_only=True)) assert "run incomplete" in output - assert "no tokenizer or unreadable sample map" not in output + assert "unreadable sample map" not in output + + @pytest.mark.parametrize( + ("primary", "counts"), + [ + # Missing usage does not decide a client-primary block, so its + # counts line omits it. + ( + TokenCountSource.CLIENT, + "(1 empty, 0 errored, 0 undecodable, 0 missing)", + ), + ( + TokenCountSource.SERVER, + "(1 empty, 0 errored, 0 undecodable, 0 missing, 1 missing usage)", + ), + ], + ) + def test_display_full_run_osl_counts_by_primary(self, primary, counts): + """The summary lines name missing usage only on a server primary.""" + fr = { + "output_sequence_lengths": {"avg": 5.0}, + "n_turns_counted": 2, + "n_empty": 1, + "n_errors": 0, + "n_undecodable": 0, + "n_missing": 0, + "n_missing_usage": 1, + "partial": False, + } + uncounted = {**fr, "output_sequence_lengths": {}, "n_turns_counted": 0} + lines: list[str] = [] + for block in (fr, uncounted): + self._report_with_full_run_osl(block, primary=primary).display( + fn=lines.append, summary_only=True + ) + assert [line for line in lines if "OSL per-turn mean" in line] == [ + f" OSL per-turn mean (accuracy, all turns): 5.0 tokens over 2 turns " + f"{counts}", + f" OSL per-turn mean (accuracy): no countable turns {counts}", + ] + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + @pytest.mark.parametrize("other_block", [None, {"avg": 3.0}]) + def test_display_full_run_osl_by_basis(self, primary, other_block): + """The other basis's full-run breakdown is shown from its suffixed key.""" + fr = { + "output_sequence_lengths": {"avg": 5.0}, + "n_turns_counted": 2, + "n_empty": 1, + "n_errors": 0, + "n_undecodable": 0, + "n_missing": 0, + "n_missing_usage": 1, + basis_key( + "output_sequence_lengths", primary.other, primary=primary + ): other_block, + } + lines: list[str] = [] + self._report_with_full_run_osl(fr, primary=primary).display(fn=lines.append) + out = "".join(lines) + label = ( + f"Output sequence lengths (full run, all turns) ({primary.other.value}):" + ) + assert (label in out) is (other_block is not None) + assert other_block is None or " Avg.: 3.00 tokens" in out + + def test_display_full_run_osl_no_usage_at_all(self): + """A server block whose turns all lacked usage shows the count, not a mean.""" + lines: list[str] = [] + self._report_with_full_run_osl( + { + "output_sequence_lengths": {}, + "n_turns_counted": 0, + "n_missing_usage": 4, + "partial": True, + }, + primary=TokenCountSource.SERVER, + ).display(fn=lines.append, summary_only=True) + assert ( + "no countable turns " + "(0 empty, 0 errored, 0 undecodable, 0 missing, 4 missing usage)" + ) in "".join(lines) def test_display_warns_when_interrupted(self): """Reports with ``state == "interrupted"`` surface a distinct WARNING.""" - report = Report( - version="test", - git_sha=None, - test_started_at=0, - n_samples_issued=10, + report = make_report( n_samples_completed=5, - n_samples_failed=0, duration_ns=1_000_000_000, state="interrupted", complete=False, - ttft={}, - tpot={}, - latency={}, - input_sequence_lengths={}, - output_sequence_lengths={}, ) lines: list[str] = [] report.display(fn=lines.append, summary_only=True) @@ -730,7 +843,7 @@ def test_minimal_dict_yields_empty_report(self): "n_pending_tasks": 0, "metrics": [], } - report = Report.from_snapshot(snap) + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) assert report.n_samples_issued == 0 assert report.n_samples_completed == 0 assert report.n_samples_failed == 0 @@ -744,7 +857,7 @@ def test_empty_dict_defaults_to_interrupted_incomplete(self): produces a non-crashing Report tagged interrupted and incomplete. Defaults: state→interrupted, counters→0, series→empty. """ - report = Report.from_snapshot({}) + report = Report.from_snapshot({}, token_count_source=DEFAULT_SOURCE) assert report.state == "interrupted" assert report.complete is False assert report.n_samples_issued == 0 @@ -762,7 +875,7 @@ def test_interrupted_state_round_trips_to_report(self): {"type": "counter", "name": "tracked_samples_completed", "value": 80}, ], } - report = Report.from_snapshot(snap) + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) assert report.state == "interrupted" assert report.complete is False # Partial counters still surface through. @@ -781,7 +894,7 @@ def test_missing_metric_type_is_skipped_not_crashed(self): {"type": "counter", "name": "tracked_samples_issued", "value": 5}, ], } - report = Report.from_snapshot(snap) + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) assert report.n_samples_issued == 5 def test_display_handles_scrubbed_nan_percentiles(self): @@ -828,7 +941,7 @@ def test_display_handles_scrubbed_nan_percentiles(self): }, ], } - report = Report.from_snapshot(snap) + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) lines: list[str] = [] # Currently crashes with TypeError on val * scale_factor. @@ -995,7 +1108,7 @@ def _report(self, verdict): state=SessionState.COMPLETE, n_pending_tasks=0, steady_state=verdict ) ) - return Report.from_snapshot(snap) + return Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) def test_absent_by_default(self): assert self._report(None).steady_state is None @@ -1023,7 +1136,7 @@ def test_a_scrubbed_verdict_is_dropped_rather_than_rendered_wrong(self): registry.build_snapshot(state=SessionState.COMPLETE, n_pending_tasks=0) ) snap["steady_state"] = raw - report = Report.from_snapshot(snap) + report = Report.from_snapshot(snap, token_count_source=DEFAULT_SOURCE) assert report.steady_state is None lines: list[str] = [] report.display(lines.append) # must not raise @@ -1048,3 +1161,534 @@ def test_display_reports_a_missing_window(self): assert "Steady state: not found (no admissible steady plateau)" in "".join( lines ) + + +# --------------------------------------------------------------------------- +# Token-count bases: the primary keeps the base names, the other is suffixed +# --------------------------------------------------------------------------- + +# Snapshot inputs that leave one basis short of some samples: the server basis +# through samples without usable usage, the client basis through tokenization +# still pending at the final snapshot. +_SHORTFALL: dict[TokenCountSource, dict[str, int]] = { + TokenCountSource.SERVER: {"n_missing": 2, "n_missing_empty": 1}, + TokenCountSource.CLIENT: {"n_pending": 3}, +} +_TOKEN_SERIES = ("isl", "osl", "tpot_ns") +_SERVER_USAGE_COUNTERS = frozenset( + key.value + for key in ( + MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE, + MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE_EMPTY_OUTPUT, + MetricCounterKey.OSL_PAIRED_SAMPLES, + MetricCounterKey.OSL_PAIRED_CLIENT_TOKENS, + MetricCounterKey.OSL_PAIRED_SERVER_TOKENS, + ) +) +# Report fields that describe server usage, which a snapshot without server +# series or usage counters cannot fill. +_SERVER_USAGE_FIELDS = frozenset( + { + "other_basis", + "osl_paired", + "n_samples_missing_usage", + "n_samples_missing_usage_empty_output", + } +) +_NO_OTHER_BASIS = OtherBasisMetrics( + input_sequence_lengths=None, output_sequence_lengths=None, tpot=None, tps=None +) + + +def _basis_snapshot( + primary: TokenCountSource, + *, + other_series: bool = True, + other_records: bool = True, + n_pending: int = 0, + n_missing: int = 0, + n_missing_empty: int = 0, + paired: tuple[int, int, int] = (0, 0, 0), + loadgen_window_ns: int = 0, + state: SessionState = SessionState.COMPLETE, +) -> dict[str, Any]: + """A 5-sample snapshot dict with the other basis's series beside the primary's. + + The primary basis records ISL 50..54 and OSL 100..104 (total 510) over a + 10 s tracked window. The other basis records ISL 40..44 (total 210), OSL + 90..94 (total 460) and a 2 ms TPOT per sample. + """ + registry = _make_registry(n_samples=5) + if other_series: + isl, osl, tpot = ( + basis_key(base, primary.other, primary=primary) for base in _TOKEN_SERIES + ) + for name, dtype in ((isl, int), (osl, int), (tpot, float)): + registry.register_series( + name, + hdr_low=1, + hdr_high=_NS_HIGH, + n_histogram_buckets=10, + percentiles=(50.0, 90.0, 99.0), + dtype=dtype, + ) + for i in range(5 if other_records else 0): + registry.record(isl, 40 + i) + registry.record(osl, 90 + i) + registry.record(tpot, 2_000_000.0) + for key, value in ( + (MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE, n_missing), + (MetricCounterKey.TRACKED_SAMPLES_MISSING_USAGE_EMPTY_OUTPUT, n_missing_empty), + (MetricCounterKey.OSL_PAIRED_SAMPLES, paired[0]), + (MetricCounterKey.OSL_PAIRED_CLIENT_TOKENS, paired[1]), + (MetricCounterKey.OSL_PAIRED_SERVER_TOKENS, paired[2]), + ): + registry.set_counter(key.value, value) + if loadgen_window_ns: + _set_loadgen_window(registry, duration_ns=loadgen_window_ns) + return snapshot_to_dict( + registry.build_snapshot(state=state, n_pending_tasks=n_pending) + ) + + +def _basis_report(primary: TokenCountSource, **snapshot_over: Any) -> Report: + return Report.from_snapshot( + _basis_snapshot(primary, **snapshot_over), token_count_source=primary + ) + + +def _display_lines(report: Report, **display_over: Any) -> list[str]: + lines: list[str] = [] + report.display(fn=lines.append, **display_over) + return lines + + +def _with_full_run(report: Report, *, other_partial: bool) -> Report: + """``report`` with a full-run OSL block on both bases. + + The primary basis counts 2 turns (mean 6.0) and one empty turn; the other + basis counts 1 turn, marked partial when ``other_partial``. + """ + primary = report.token_count_source + other_series = series_metric_dict([5]) + if other_partial: + other_series["partial"] = True + block = { + "output_sequence_lengths": series_metric_dict([3, 9]), + "n_turns_counted": 2, + "n_empty": 1, + "n_errors": 0, + "n_undecodable": 0, + "n_missing": 0, + "partial": False, + "n_missing_usage": 0, + basis_key("output_sequence_lengths", primary.other, primary=primary): ( + other_series + ), + } + return msgspec.structs.replace(report, output_sequence_lengths_full_run=block) + + +@pytest.mark.unit +class TestTokenCountSourceRoles: + @pytest.mark.parametrize( + ("state", "n_pending"), + [ + (SessionState.COMPLETE, 0), + (SessionState.COMPLETE, 3), + (SessionState.INTERRUPTED, 0), + ], + ) + def test_client_primary_unsuffixed_fields_ignore_server_series( + self, state, n_pending + ): + """S+ adds the *_server series and usage counters to S; under the client + basis that changes only the fields describing server usage, and the + display only gains lines.""" + client = TokenCountSource.CLIENT + s_plus = _basis_snapshot( + client, + n_missing=2, + n_missing_empty=1, + paired=(4, 380, 400), + n_pending=n_pending, + state=state, + ) + server_only = _SERVER_USAGE_COUNTERS | { + basis_key(base, TokenCountSource.SERVER, primary=client) + for base in _TOKEN_SERIES + } + s = { + **s_plus, + "metrics": [m for m in s_plus["metrics"] if m["name"] not in server_only], + } + base = Report.from_snapshot(s, token_count_source=client) + plus = Report.from_snapshot(s_plus, token_count_source=client) + + fields = [f for f in Report.__struct_fields__ if f not in _SERVER_USAGE_FIELDS] + assert {f: getattr(plus, f) for f in fields} == { + f: getattr(base, f) for f in fields + } + assert plus.incomplete_reason() == base.incomplete_reason() + assert base.other_basis == _NO_OTHER_BASIS != plus.other_basis + # Every line of base's display appears in plus's, in the same order. + plus_lines = iter(_display_lines(plus)) + assert all(line in plus_lines for line in _display_lines(base)) + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_other_basis_json_keys_follow_naming_rule(self, primary): + data = json.loads(_basis_report(primary).to_json()) + other = primary.other + for field in ( + "input_sequence_lengths", + "output_sequence_lengths", + "tpot", + "tps", + ): + assert basis_key(field, other, primary=primary) in data, field + assert [key for key in data if key.endswith(f"_{primary.value}")] == [] + assert "other_basis" not in data + assert ( + data[basis_key("output_sequence_lengths", other, primary=primary)]["total"] + == 460 + ) + assert data[basis_key("tps", other, primary=primary)] == pytest.approx(46.0) + assert data["token_count_source"] == primary.value + assert { + "n_samples_missing_usage", + "n_samples_missing_usage_empty_output", + "client_token_tasks_pending", + "osl_paired", + } <= data.keys() + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + @pytest.mark.parametrize( + ("window_ns", "tps", "other_tps"), + [(0, 51.0, 46.0), (5_000_000_000, 102.0, 92.0)], + ) + def test_other_basis_tps_shares_the_tps_window( + self, primary, window_ns, tps, other_tps + ): + r = _basis_report(primary, loadgen_window_ns=window_ns) + other = r.other_basis + assert r.tps == pytest.approx(tps) + assert other.tps == pytest.approx(other_tps) + assert other.input_sequence_lengths["total"] == 210 + assert other.output_sequence_lengths["total"] == 460 + assert other.tpot["min"] == pytest.approx(2_000_000.0) + assert "partial" not in other.output_sequence_lengths + # From the primary OSL. + assert r.e2e_avg_interactivity == pytest.approx(20_000.0) + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_other_block_partial_inside_dict_and_tps_none(self, primary): + r = _basis_report(primary, **_SHORTFALL[primary.other]) + other = r.other_basis + for block in ( + other.input_sequence_lengths, + other.output_sequence_lengths, + other.tpot, + ): + assert block["partial"] is True + assert other.tps is None + assert "partial" not in r.output_sequence_lengths + assert r.tps == pytest.approx(51.0) + data = json.loads(r.to_json()) + assert ( + data[basis_key("output_sequence_lengths", primary.other, primary=primary)][ + "partial" + ] + is True + ) + assert data[basis_key("tps", primary.other, primary=primary)] is None + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + @pytest.mark.parametrize("other_series", [False, True], ids=["absent", "empty"]) + def test_other_block_none_when_nothing_recorded(self, primary, other_series): + """No recorded sample means no block, even when the basis fell short.""" + r = _basis_report( + primary, + other_series=other_series, + other_records=False, + **_SHORTFALL[primary.other], + ) + assert r.other_basis == _NO_OTHER_BASIS + assert r.other_basis_partial_reason() is None + data = json.loads(r.to_json()) + for field in ( + "input_sequence_lengths", + "output_sequence_lengths", + "tpot", + "tps", + ): + assert data[basis_key(field, primary.other, primary=primary)] is None + label = f"({primary.other.value})" + assert [line for line in _display_lines(r) if label in line] == [ + f"Other token basis {label}:" + ] + + @pytest.mark.parametrize( + ("primary", "short", "complete"), + [ + (TokenCountSource.CLIENT, TokenCountSource.SERVER, True), + (TokenCountSource.SERVER, TokenCountSource.CLIENT, True), + (TokenCountSource.CLIENT, TokenCountSource.CLIENT, False), + (TokenCountSource.SERVER, TokenCountSource.SERVER, False), + ], + ) + def test_complete_follows_primary_only(self, primary, short, complete): + assert _basis_report(primary, **_SHORTFALL[short]).complete is complete + + @pytest.mark.parametrize( + ("primary", "reason"), + [ + ( + TokenCountSource.CLIENT, + "2 of 5 successful tracked samples returned no server token " + "usage; the *_server blocks are partial and the run is not " + "affected", + ), + ( + TokenCountSource.SERVER, + "3 client tokenization tasks were still pending; the *_client " + "blocks are partial and the run is not affected", + ), + ], + ) + def test_partial_other_basis_is_shown_as_information(self, primary, reason): + r = _basis_report(primary, **_SHORTFALL[primary.other]) + assert r.other_basis_partial_reason() == reason + out = "\n".join(_display_lines(r, summary_only=True)) + assert reason in out + assert "WARNING" not in out + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_other_basis_sections_are_labelled_with_the_basis(self, primary): + other = primary.other.value + out = "\n".join(_display_lines(_basis_report(primary, paired=(4, 380, 400)))) + for text in ( + "Input sequence lengths:", + "Output sequence lengths:", + f"Input sequence lengths ({other}):", + f"Output sequence lengths ({other}):", + f"TPOT ({other}):", + f"Token count source: {primary.value}", + f"TPS ({other}): 46.00", + "OSL paired (client - server): mean -5.00 tokens/sample, " + "ratio 0.9500 over 4 samples", + ): + assert text in out + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_other_basis_lines_follow_the_primary_summary(self, primary): + other = primary.other.value + report = _basis_report( + primary, paired=(4, 380, 400), **_SHORTFALL[primary.other] + ) + lines = _display_lines( + _with_full_run(report, other_partial=False), summary_only=True + ) + heading = lines.index(f"Other token basis ({other}):") + assert lines[heading + 1 : -1] == [ + f" Token count source: {primary.value}", + " OSL paired (client - server): mean -5.00 tokens/sample, " + "ratio 0.9500 over 4 samples", + f" {report.other_basis_partial_reason()}", + ] + assert lines[heading - 1].startswith( + " OSL per-turn mean (accuracy, all turns)" + ) + assert lines[-1] == "----------------- End of Summary -----------------" + assert f" TPS ({other}): 46.00" in _display_lines( + _basis_report(primary), summary_only=True + ) + + def test_client_primary_summary_is_mains_up_to_the_other_basis_heading(self): + report = msgspec.structs.replace( + _with_full_run( + _basis_report( + TokenCountSource.CLIENT, + paired=(4, 380, 400), + **_SHORTFALL[TokenCountSource.SERVER], + ), + other_partial=True, + ), + version="9.9.9", + git_sha="abc1234", + steady_state=TestSteadyStateOnTheReport._verdict(), + accuracy=[ + { + "dataset_name": "gpqa", + "score": 0.5, + "unit_samples": 5, + "num_repeats": 1, + "total_samples": 5, + } + ], + ) + lines = _display_lines(report, summary_only=True) + heading = lines.index("Other token basis (server):") + # Byte-identical to main's display() for this snapshot: 5 samples over + # a 10 s window, primary OSL 100..104, latency 5.00..5.20 ms. + assert lines[:heading] == [ + "----------------- Summary -----------------", + "Version: 9.9.9", + "Git SHA: abc1234", + "Total samples issued: 5", + "Total samples succeeded: 5", + "Total samples failed: 0", + "Total samples dropped: 0", + "Duration: 10.00 seconds", + "QPS: 0.50", + "TPS: 51.00", + "E2E average interactivity: 20000.00 tokens/s", + "Steady state: super-passes 1...8 (post-warmup), 4000 samples over 700s", + " Steady TPS: 1234.50 system, 42.00 per user", + " WARNING: ttft_p50 drifting up over the rest of the run", + " ANOMALY: level shift at super-pass 7, TPOT +12.5% toward end of run", + "Accuracy:", + " gpqa: 0.5 (unit=5, repeats=1, total=5)", + " Average: 0.5", + " OSL per-turn mean (accuracy, all turns): 6.0 tokens over 2 turns " + "(1 empty, 0 errored, 0 undecodable, 0 missing)", + ] + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + @pytest.mark.parametrize("short", [False, True], ids=["whole", "partial"]) + def test_partial_other_basis_sections_are_marked(self, primary, short): + other = primary.other.value + report = _with_full_run( + _basis_report(primary, **(_SHORTFALL[primary.other] if short else {})), + other_partial=short, + ) + titles = [line for line in _display_lines(report) if line.endswith(":")] + mark = " (partial)" if short else "" + for label in ( + "Input sequence lengths", + "Output sequence lengths", + "TPOT", + "Output sequence lengths (full run, all turns)", + ): + assert f"{label} ({other}){mark}:" in titles + for label in ( + "Input sequence lengths", + "Output sequence lengths", + "Output sequence lengths (full run, all turns)", + ): + assert f"{label}:" in titles + assert sum("(partial)" in title for title in titles) == (4 if short else 0) + + def test_osl_paired_json_shape(self): + golden = { + "n_samples": 4, + "client_tokens": 380, + "server_tokens": 400, + "mean_diff": -5.0, + "ratio": 0.95, + } + paired = json.loads( + _basis_report(DEFAULT_SOURCE, paired=(4, 380, 400)).to_json() + )["osl_paired"] + assert paired == golden + assert list(paired) == list(golden) + assert _basis_report(DEFAULT_SOURCE).osl_paired is None + no_server_tokens = _basis_report(DEFAULT_SOURCE, paired=(2, 10, 0)) + assert json.loads(no_server_tokens.to_json())["osl_paired"]["ratio"] is None + + def test_from_snapshot_requires_basis(self): + with pytest.raises(TypeError, match="token_count_source"): + Report.from_snapshot(_basis_snapshot(DEFAULT_SOURCE)) + + +_USAGE = "3 of 10 successful tracked samples returned no server token usage" +# The CLI flag, then the YAML key: `benchmark from-config` takes no such flag. +_FIX = ( + "set --token-count-source client if this server cannot report usage " + "(YAML: settings.token_count_source: client)" +) +_MISSING = { + "complete": False, + "token_count_source": TokenCountSource.SERVER, + "n_samples_missing_usage": 3, +} +# Report fields that leave each basis short when it is primary, and the cause +# ``incomplete_reason`` then names. +_PRIMARY_SHORTFALL: dict[TokenCountSource, tuple[dict[str, int], str]] = { + TokenCountSource.CLIENT: ( + {"client_token_tasks_pending": 4}, + "Metrics tokenization did not finish", + ), + TokenCountSource.SERVER: ({"n_samples_missing_usage": 3}, f"{_USAGE}; {_FIX}"), +} + + +@pytest.mark.unit +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_incomplete_reason_names_state_then_drops_then_primary(primary): + """State, then dropped samples, then the primary basis's shortfall.""" + shortfall, reason = _PRIMARY_SHORTFALL[primary] + incomplete = {"complete": False, "token_count_source": primary, **shortfall} + assert make_report(**incomplete).incomplete_reason() == reason + assert make_report(**incomplete, n_samples_completed=8).incomplete_reason() == ( + "Benchmark report is incomplete (state=complete, n_samples_dropped=2)" + ) + assert make_report(**incomplete, state="interrupted").incomplete_reason() == ( + "Benchmark report is incomplete (state=interrupted, n_samples_dropped=0)" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("over", "expected"), + [ + ({}, None), + ( + {"complete": False, "state": "live"}, + "Benchmark report is incomplete (state=live, n_samples_dropped=0)", + ), + # Every server count is whole, so there is no shortfall to name. + ( + {"complete": False, "token_count_source": TokenCountSource.SERVER}, + "Benchmark report is incomplete (state=complete, n_samples_dropped=0)", + ), + ], +) +def test_incomplete_reason_without_a_primary_shortfall(over, expected): + assert make_report(**over).incomplete_reason() == expected + + +@pytest.mark.unit +def test_display_names_the_cause_not_a_drain_timeout(): + lines: list[str] = [] + make_report(**_MISSING).display(fn=lines.append, summary_only=True) + out = "\n".join(lines) + assert f"WARNING: report is incomplete: {_USAGE}" in out + assert "drain timed out" not in out and "tokenization did not finish" not in out + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("primary", "expected"), + [ + # Fixed wording that report.txt readers match on. + ( + TokenCountSource.CLIENT, + "WARNING: report is incomplete (drain timed out or no final " + "snapshot received) — some async metrics may be missing.", + ), + ( + TokenCountSource.SERVER, + f"WARNING: report is incomplete: {_USAGE}; {_FIX}", + ), + ], + ids=["client", "server"], +) +def test_display_incomplete_warning_wording_by_primary(primary, expected): + shortfall, _ = _PRIMARY_SHORTFALL[primary] + report = make_report( + duration_ns=1_000_000_000, + complete=False, + token_count_source=primary, + **shortfall, + ) + assert expected in _display_lines(report, summary_only=True) diff --git a/tests/unit/validation/test_shipped_cohort_golden.py b/tests/unit/validation/test_shipped_cohort_golden.py index 115c45858..e599674b6 100644 --- a/tests/unit/validation/test_shipped_cohort_golden.py +++ b/tests/unit/validation/test_shipped_cohort_golden.py @@ -38,6 +38,8 @@ from inference_endpoint.metrics.report import Report from inference_endpoint.validation.__main__ import main +from tests.test_helpers import DEFAULT_SOURCE + COHORT = Path(inference_endpoint.__file__).parent / "config" / "cohorts" / "2026-10-C1" SYSTEM = "sys" MODEL = "gpt-oss-120b" @@ -113,7 +115,11 @@ def _summary(run_dir: Path, *, failed: int, with_tpot: bool) -> None: } ) snapshot = {"state": "complete", "n_pending_tasks": 0, "metrics": metrics} - report = Report.from_snapshot(snapshot, run_config={"load_pattern": "concurrency"}) + report = Report.from_snapshot( + snapshot, + token_count_source=DEFAULT_SOURCE, + run_config={"load_pattern": "concurrency"}, + ) report.to_json(save_to=run_dir / "result_summary.json") From 78d7b8b9aeadb6897421ca3186307c7af5cfb4ef Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:03 -0500 Subject: [PATCH 08/26] feat(benchmark): post-run tools count tokens on the run's basis Full-run OSL (the agentic accuracy-gate input), the accuracy OSL rollup, the scorer, the steady-state replay CLI and scripts/early_stopping_estimate_from_events.py now share the live usage rule (read_server_usage). Each writes the primary basis unsuffixed with the other basis's _ keys beside it. A block that covers fewer turns than were completed carries "partial": true. Under a server primary, report.txt also says ", N missing usage"; under the client default, the full-run line keeps main's wording. A replay takes its basis from replay_basis(run_dir, *, override): the summary's top-level token_count_source, else config.yaml, else client for archived runs. --token-count-source overrides it, and the early-stopping script's --summary names the summary to read. Readers of events.jsonl stop silently dropping data: the steady-state replays classify lines with leading_event_topic, as the full-run OSL and scorer readers do, so a line outside the compact layout is decoded whole rather than skipped; the replays also count and warn on undecodable COMPLETE lines. Full-run OSL keeps batched client counting. The archived pre-usage run test moves to the new call signatures; every pinned number is unchanged. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- .../early_stopping_estimate_from_events.py | 139 ++++- scripts/synth_events.py | 19 +- .../commands/benchmark/accuracy.py | 210 +++++-- .../commands/benchmark/full_run_osl.py | 323 +++++++--- src/inference_endpoint/core/record.py | 25 - .../evaluation/bfcl_v4_scorer.py | 5 +- src/inference_endpoint/evaluation/scoring.py | 59 +- .../metrics/steady_state_diagnostics.py | 510 +++++++++++---- tests/unit/commands/test_full_run_osl.py | 406 ++++++++++-- tests/unit/commands/test_score_accuracy.py | 388 ++++++++++-- tests/unit/evaluation/test_scoring.py | 87 +++ .../unit/metrics/test_archived_run_compat.py | 17 +- .../metrics/test_steady_state_diagnostics.py | 587 ++++++++++++++++-- ...est_early_stopping_estimate_from_events.py | 309 ++++++++- tests/unit/scripts/test_synth_events.py | 47 +- .../validation/test_shipped_cohort_golden.py | 1 + 16 files changed, 2648 insertions(+), 484 deletions(-) diff --git a/scripts/early_stopping_estimate_from_events.py b/scripts/early_stopping_estimate_from_events.py index 009b8d36d..52fc0618c 100644 --- a/scripts/early_stopping_estimate_from_events.py +++ b/scripts/early_stopping_estimate_from_events.py @@ -17,10 +17,19 @@ completions after stop still count for rows that exist. - ttft_ns = recv_first.ts - issued.ts - sample_latency_ns = complete.ts - issued.ts - - tpot_ns = (complete.ts - recv_first.ts) / tokens(text_after_first_chunk) - (streaming only; token counts via the raw ``tokenizers`` backend - with add_special_tokens=False, identical to - ``token_metrics.encode_lengths``). Tool-call samples are skipped + - tpot_ns = (complete.ts - recv_first.ts) / tokens after the first chunk + (streaming only). The denominator follows the run's token-count + basis, read by ``token_count_source_for_run`` from --summary + when given, else from the event log's directory + (--token-count-source overrides it): + server: ``completion_tokens - 1`` from the COMPLETE record's usage + (``classify_server_usage``, then LoadGen's rule, as the + aggregator). Samples without usable usage, with + completion_tokens <= 1 or a non-positive window are skipped. + Tool-call samples are included. No tokenizer is needed. + client: tokens(text_after_first_chunk) via the raw ``tokenizers`` + backend with add_special_tokens=False, identical to + ``token_metrics.encode_lengths``. Tool-call samples are skipped (their chat-template tokenization path is not replicated). Events are decoded with the product's own typed ``EventRecord`` (the exact wire types @@ -35,7 +44,10 @@ [--json ] # write the summary AUGMENTED with # early_stopping_percentiles maps — exactly the # shape an ES-enabled run would have produced - [--tokenizer ] # enables TPOT + [--token-count-source server|client] # overrides the run's recorded basis; + # --json refuses one that differs from it + # (it would mislabel the TPOT maps) + [--tokenizer ] # enables client-basis TPOT [--percentiles 50,90,95,99.9] [--confidence 0.99] # stdout-analysis overrides """ @@ -44,6 +56,7 @@ import argparse import json import sys +from pathlib import Path import msgspec import msgspec.json @@ -60,7 +73,7 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import TextModelOutput +from inference_endpoint.core.types import TextModelOutput, TokenCountSource from inference_endpoint.metrics.early_stopping import ( CONFIDENCE, es_percentile_estimate, @@ -70,6 +83,13 @@ SERIES_TO_SUMMARY_FIELD, place_early_stopping_percentiles, ) +from inference_endpoint.metrics.token_count_basis import ( + TOKEN_COUNT_SOURCE_METAVAR, + RunBasisError, + ServerUsage, + read_server_usage, + replay_basis, +) # mirror of the JSONL writer: msgspec.json.Encoder(enc_hook=EventType.encode_hook) _DECODER = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) @@ -99,17 +119,31 @@ def extract_tpot_text(data) -> str | None: return data.text_after_first_chunk() -def compute_series(path, count_tokens=None, progress=False) -> dict[str, list[float]]: +def compute_series( + events_path, + *, + count_tokens, + basis: TokenCountSource, + progress=False, +) -> dict[str, list[float]]: """Stream the log and rebuild the {ttft_ns, tpot_ns, sample_latency_ns} series. - ``count_tokens``: callable(list[str]) -> list[int], or None to skip TPOT. + ``basis`` is the TPOT basis replayed on. SERVER: TPOT from the COMPLETE + record's usable ``completion_tokens`` (``read_server_usage``, as the + aggregator); ``count_tokens`` is then unused. Tracked ``TextModelOutput`` + completions without usable usage are counted as missing usage (warned on + stderr), never tokenized. + CLIENT: ``count_tokens`` is callable(list[str]) -> list[int], or None to + skip TPOT. """ + server_basis = basis is TokenCountSource.SERVER ttft: list[float] = [] latency: list[float] = [] tpot: list[float] = [] rows: dict[str, list] = {} # uuid -> [issued_ns, recv_first_ns] tracking = False tool_call_skipped = 0 + missing_usage = 0 malformed = 0 batch_deltas: list[float] = [] batch_texts: list[str] = [] @@ -122,7 +156,7 @@ def flush() -> None: batch_deltas.clear() batch_texts.clear() - with open(path, "rb") as f: + with open(events_path, "rb") as f: for i, line in enumerate(f): if progress and i and i % 200_000 == 0: print(f" ...{i} events", file=sys.stderr) @@ -174,6 +208,21 @@ def flush() -> None: if row is None: continue latency.append(ts - row[0]) + if server_basis: + if isinstance(rec.data, TextModelOutput): + usage, n = read_server_usage( + rec.prompt_tokens, rec.completion_tokens, rec.data + ) + if usage is ServerUsage.MISSING: + missing_usage += 1 + if ( + n is not None + and n > 1 + and row[1] is not None + and ts - row[1] > 0 + ): + tpot.append((ts - row[1]) / (n - 1)) + continue if count_tokens is None or row[1] is None: continue text = extract_tpot_text(rec.data) @@ -196,6 +245,12 @@ def flush() -> None: f"WARN: {tool_call_skipped} tool-call samples skipped for TPOT", file=sys.stderr, ) + if missing_usage: + print( + f"WARN: {missing_usage} tracked samples returned no server token usage; " + "they are counted as missing (no TPOT), never tokenized", + file=sys.stderr, + ) return {"ttft_ns": ttft, "tpot_ns": tpot, "sample_latency_ns": latency} @@ -252,7 +307,18 @@ def main(argv=None): ), ) ap.add_argument( - "--tokenizer", help="HF model dir (or tokenizer.json) — enables TPOT" + "--token-count-source", + type=TokenCountSource, + metavar=TOKEN_COUNT_SOURCE_METAVAR, + help=( + "token-count basis for TPOT (default: the run's own, from --summary or " + "the performance/result_summary.json beside the event log, else the " + "config.yaml beside it)" + ), + ) + ap.add_argument( + "--tokenizer", + help="HF model dir (or tokenizer.json) — enables client-basis TPOT", ) ap.add_argument( "--percentiles", @@ -274,6 +340,8 @@ def main(argv=None): raise SystemExit( "FATAL: --json requires --summary (the augmented output IS the summary)" ) + if args.summary and not Path(args.summary).is_file(): + raise SystemExit(f"FATAL: --summary {args.summary} not found") values = {float(tok) for tok in args.percentiles.split(",")} if not all(0.0 < v < 100.0 for v in values): raise SystemExit( @@ -293,17 +361,52 @@ def main(argv=None): if not 0.0 < args.confidence < 1.0: raise SystemExit(f"FATAL: confidence must be in (0, 1), got {args.confidence}") - counter = _make_counter(args.tokenizer) if args.tokenizer else None - if counter is None: - print("NOTE: no --tokenizer given -> TPOT skipped", file=sys.stderr) - - series = compute_series(args.events, count_tokens=counter, progress=True) + try: + primary, basis = replay_basis( + Path(args.events).parent, + override=args.token_count_source, + summary_path=None if args.summary is None else Path(args.summary), + ) + except RunBasisError as exc: + ap.error(str(exc)) + if primary is None: + if args.json_out: + # The summary's own basis is unknown, so the chosen one cannot label it. + raise SystemExit( + "FATAL: --json augments the run's summary, whose token-count basis " + "could not be read; drop --json to replay on the chosen basis" + ) + # The run's basis is unreadable; the one the user chose stands in. + primary = basis + # The summary is labelled with the run's basis: TPOT replayed on the other + # one is neither cross-checked against it nor written into it. + off_primary = basis is not primary + if off_primary and args.json_out: + raise SystemExit( + f"FATAL: --json augments the run's summary, which is on the " + f"{primary.value} token basis; it cannot carry {basis.value}-basis TPOT " + "estimates (drop --token-count-source or --json)" + ) + print(f"NOTE: token basis: {basis.value}", file=sys.stderr) + counter = None + if basis is TokenCountSource.CLIENT: + counter = _make_counter(args.tokenizer) if args.tokenizer else None + if counter is None: + print("NOTE: no --tokenizer given -> TPOT skipped", file=sys.stderr) + + series = compute_series( + args.events, + count_tokens=counter, + basis=basis, + progress=True, + ) summary = None if args.summary: with open(args.summary) as f: summary = json.load(f) augmented = json.loads(json.dumps(summary)) if summary is not None else {} + print(f"token basis: {basis.value}") rows = [] for name, values in series.items(): if not values: @@ -354,6 +457,12 @@ def main(argv=None): f" XCHECK count: ours={len(values)} " f"n_samples_completed={completed}{note}" ) + if off_primary and name == "tpot_ns": + print( + f" XCHECK in-band: skipped (TPOT replayed on the {basis.value} " + f"basis; the run reported {primary.value})" + ) + continue _cross_check(summary, name, dict(esp)) if metric: # augment only metrics the run actually reported augmented[field] = place_early_stopping_percentiles( diff --git a/scripts/synth_events.py b/scripts/synth_events.py index 6e6dfaab9..22e62b5a0 100644 --- a/scripts/synth_events.py +++ b/scripts/synth_events.py @@ -36,6 +36,10 @@ HF tokenizer; for synthetic validation feed it the whitespace word counter so 1 chunk == 1 word == ~1 token, making TTFT/TPOT recover the planted base values. +COMPLETE lines also carry the server usage a live run records: ``prompt_tokens`` (the ISL +draw) and ``completion_tokens`` (== OSL). The server-basis replay divides by +``completion_tokens - 1``, so it recovers the planted TPOT without a tokenizer. + The steady-state ground truth is the *base* metric level in the flat region: after the ramp has decayed, with no drift and before any staircase step. It is written to a sidecar ``.groundtruth.json``. @@ -172,7 +176,13 @@ def _emit_turn(self, issue_ns: int, conv_id: str | None, turn: int | None) -> in self._add(EV_ISSUED, issue_ns, base) self._add(EV_RECV_FIRST, recv_ns, base) chunks = ["w "] * osl # OSL single-word chunks -> whitespace count == OSL - self._add(EV_COMPLETE, complete_ns, base, data=["TextModelOutput", chunks]) + self._add( + EV_COMPLETE, + complete_ns, + base, + data=["TextModelOutput", chunks], + extra={"prompt_tokens": isl, "completion_tokens": osl}, + ) self.n_samples += 1 self.last_complete_ns = max(self.last_complete_ns, complete_ns) return complete_ns @@ -183,12 +193,15 @@ def _add( ts: int, base: dict[str, Any] | None = None, data: Any | None = None, + extra: dict[str, Any] | None = None, ) -> None: rec: dict[str, Any] = {"event_type": et, "timestamp_ns": ts} if base: rec.update(base) if data is not None: rec["data"] = data + if extra: + rec.update(extra) self.events.append(rec) def _run_conversation(self, start_ns: int) -> int: @@ -288,7 +301,9 @@ def groundtruth(planted: Planted, gen: _Gen) -> dict[str, Any]: def write_run(out_path: str, gen: _Gen, gt: dict[str, Any]) -> str: with open(out_path, "w") as fh: for rec in gen.events: - fh.write(json.dumps(rec)) + # Compact, like the JSONL writer: the replays find COMPLETE lines by + # their leading bytes (``core.record.leading_event_topic``). + fh.write(json.dumps(rec, separators=(",", ":"))) fh.write("\n") sidecar = out_path + ".groundtruth.json" with open(sidecar, "w") as fh: diff --git a/src/inference_endpoint/commands/benchmark/accuracy.py b/src/inference_endpoint/commands/benchmark/accuracy.py index e709e8a73..a426d89cd 100644 --- a/src/inference_endpoint/commands/benchmark/accuracy.py +++ b/src/inference_endpoint/commands/benchmark/accuracy.py @@ -27,18 +27,24 @@ import logging import numbers import time -from collections.abc import Iterable +from collections.abc import Collection, Iterable from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any +import msgspec import msgspec.json from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( encode_lengths, load_reference_backend, ) -from inference_endpoint.config.schema import DatasetType, ScorerMethod, TestMode +from inference_endpoint.config.schema import ( + DatasetType, + ScorerMethod, + TestMode, +) +from inference_endpoint.core.types import TokenCountSource, basis_key from inference_endpoint.dataset_manager.dataset import Dataset from inference_endpoint.evaluation import Extractor from inference_endpoint.evaluation.accuracy_results import ( @@ -91,14 +97,14 @@ def _phase_osl_stats( backend: Any, batch_size: int = 256, ) -> dict[str, Any] | None: - """Output-token-length rollup over one accuracy phase's completions. + """Client-basis output-token-length rollup over one accuracy phase. - Counts tokens on each sample's response text via the shared reference - tokenizer backend — the server's ``completion_tokens`` is not persisted, only - the text is (in ``events.jsonl``) — then shapes the lengths via + Counts tokens on each sample's response text (from ``events.jsonl``) via the + shared reference tokenizer backend, then shapes the lengths via ``series_metric_dict`` so the block matches the perf report's - ``output_sequence_lengths`` exactly. Returns ``None`` when the phase has no - completed outputs. + ``output_sequence_lengths`` exactly; :func:`_phase_server_osl_stats` is the + server-basis counterpart. Returns ``None`` when the phase has no completed + outputs. ``batch_size`` bounds each ``encode_batch`` pass: accuracy outputs can be tens of thousands of tokens each (e.g. gpt-oss lcb at 32768), so counting the whole @@ -120,6 +126,129 @@ def _phase_osl_stats( return series_metric_dict(lengths) or None +class PhaseServerOsl(msgspec.Struct, frozen=True): # type: ignore[call-arg] + """Server-basis output-token lengths of one accuracy phase. + + ``block`` has the shape of :func:`_phase_osl_stats` and is ``None`` when no + scored sample carried usable usage; ``n_counted`` is the number of samples + in it, and ``n_missing_usage`` the scored samples without usable usage. + """ + + block: dict[str, Any] | None + n_counted: int + n_missing_usage: int + + +def _phase_server_osl_stats( + sample_uuids: Iterable[str], + uuid_to_text: dict[str, str], + uuid_to_server_tokens: dict[str, int], +) -> PhaseServerOsl: + """Server-basis output-token-length rollup over one accuracy phase. + + The population is the phase's scored samples: non-empty output, the + ``scored`` class of :func:`_phase_response_counts`. Each is counted at its + usable ``completion_tokens`` (``uuid_to_server_tokens`` holds only those). + A scored sample without usable usage is counted as missing usage, which is + not an error: an accuracy-only run may target a server that reports none, + and its score stands. + """ + values: list[int] = [] + n_missing_usage = 0 + for u in sample_uuids: + if not uuid_to_text.get(u): + continue + tokens = uuid_to_server_tokens.get(u) + if tokens is None: + n_missing_usage += 1 + else: + values.append(tokens) + return PhaseServerOsl( + block=series_metric_dict(values) or None, + n_counted=len(values), + n_missing_usage=n_missing_usage, + ) + + +def _phase_osl_entry( + sample_uuids: Collection[str], + uuid_to_text: dict[str, str], + uuid_to_server_tokens: dict[str, int], + backend: Any | None, + *, + primary: TokenCountSource, + dataset_name: str, +) -> dict[str, Any]: + """Response accounting and output-token lengths of one accuracy phase. + + Each basis's block is named by ``basis_key("output_sequence_lengths", …)`` + and is present only when that basis counted a sample: the client basis + needs ``backend``, the server basis usable usage. A server block that is + not primary carries ``"partial": True`` when some scored samples lacked + usable usage, as other-basis blocks do. ``n_osl_samples`` counts the + samples behind each block and ``n_missing_usage`` the scored samples + without usable usage, which are never filled in from the tokenizer. A + tokenizer failure drops only the client block. + """ + response_counts = _phase_response_counts(sample_uuids, uuid_to_text) + server = _phase_server_osl_stats(sample_uuids, uuid_to_text, uuid_to_server_tokens) + blocks: dict[TokenCountSource, dict[str, Any] | None] = { + TokenCountSource.CLIENT: None, + TokenCountSource.SERVER: server.block, + } + tokenize_s = None + if backend is not None: + t0 = time.perf_counter() + try: + blocks[TokenCountSource.CLIENT] = _phase_osl_stats( + sample_uuids, uuid_to_text, backend + ) + except Exception as e: # noqa: BLE001 - optional block; keep the server one + logger.warning( + "Accuracy client-basis OSL skipped for %s: %s", dataset_name, e + ) + # Wall-clock of just this phase's tokenization (seconds); summed across + # datasets for the accuracy report's total. + tokenize_s = round(time.perf_counter() - t0, 3) + server_is_other = primary is not TokenCountSource.SERVER + if server.block is not None and server.n_missing_usage and server_is_other: + blocks[TokenCountSource.SERVER] = {**server.block, "partial": True} + n_osl_samples = { + TokenCountSource.CLIENT: ( + response_counts["scored"] + if blocks[TokenCountSource.CLIENT] is not None + else 0 + ), + TokenCountSource.SERVER: server.n_counted, + } + + entry: dict[str, Any] = {"response_counts": response_counts} + for basis in (primary, primary.other): + block = blocks[basis] + if block is not None: + # Same shape as the perf report's output_sequence_lengths. + entry[basis_key("output_sequence_lengths", basis, primary=primary)] = block + if basis is TokenCountSource.CLIENT: + entry["osl_tokenize_s"] = tokenize_s + entry[basis_key("n_osl_samples", basis, primary=primary)] = n_osl_samples[basis] + entry["n_missing_usage"] = server.n_missing_usage + return entry + + +def _read_accuracy_outputs( + scorer: Scorer, uuids: set[str] | None +) -> tuple[dict[str, str], dict[str, int]]: + """One read of ``events.jsonl``: each sample's response text, and the + ``completion_tokens`` of each sample whose server usage is usable.""" + uuid_to_text: dict[str, str] = {} + uuid_to_server_tokens: dict[str, int] = {} + for uuid, text, completion_tokens in scorer.iter_complete_outputs(uuids): + uuid_to_text[uuid] = text + if completion_tokens is not None: + uuid_to_server_tokens[uuid] = completion_tokens + return uuid_to_text, uuid_to_server_tokens + + def _phase_response_counts( sample_uuids: Iterable[str], uuid_to_text: dict[str, str], @@ -246,10 +375,14 @@ def score_accuracy( ) # Accuracy-phase output-token lengths (finalize-side, off the hot path): the - # aggregator only tokenizes perf-window samples, so count the accuracy - # responses (already in events.jsonl) here, using the same reference tokenizer - # as the perf side. (Counts still differ from perf for tool-call responses — - # client-side OSL is approximate for structured output.) + # aggregator only counts perf-window samples, so the accuracy responses + # (already in events.jsonl) are counted here on both token bases. The client + # basis uses the same reference tokenizer as the perf side; the server basis + # each sample's usable completion_tokens and needs no tokenizer. The run's + # primary basis keeps the base key names. (Tokenizer counts still differ from + # perf for tool-call responses — client-side OSL is approximate for + # structured output.) + primary = ctx.token_count_source has_accuracy = any(ec.dataset_type == DatasetType.ACCURACY for ec in eval_configs) osl_backend = _load_osl_backend(has_accuracy, ctx.tokenizer_name) # Bound the raw-output read to the accuracy population so finalize never holds @@ -257,7 +390,7 @@ def score_accuracy( accuracy_uuids = ( _accuracy_uuid_bound(ctx.report_dir, eval_configs) if has_accuracy else set() ) - uuid_to_text: dict[str, str] | None = None + outputs: tuple[dict[str, str], dict[str, int]] | None = None for eval_cfg in eval_configs: try: @@ -337,38 +470,29 @@ def score_accuracy( # perf entry (its OSL / failure counts live in result_summary.json). The # counts are computed independent of the tokenizer and of OSL returning a # block — an all-failed phase must still publish scored=0 rather than - # silently omitting everything. OSL stays tokenizer-gated. A read/tokenize - # failure only drops these blocks — it never fails scoring. + # silently omitting everything. A read failure only drops these blocks — + # it never fails scoring. if eval_cfg.dataset_type == DatasetType.ACCURACY: try: - if uuid_to_text is None: + if outputs is None: # Built once from the first scorer and reused for every - # dataset. get_raw_outputs() returns the model's actual - # completion text (not the scorer's scoring-normalized form) + # dataset: the model's actual completion text (not the + # scorer's scoring-normalized form) and usable server counts # for *all* phases' COMPLETE events, bounded to the accuracy - # population; intersecting it with each dataset's + # population; intersecting them with each dataset's # sample_index_map yields correct per-dataset counts. - out_df = scorer_instance.get_raw_outputs(accuracy_uuids) - uuid_to_text = dict( - zip(out_df["sample_uuid"], out_df["output"], strict=False) + outputs = _read_accuracy_outputs(scorer_instance, accuracy_uuids) + uuid_to_text, uuid_to_server_tokens = outputs + entry.update( + _phase_osl_entry( + scorer_instance.sample_index_map, + uuid_to_text, + uuid_to_server_tokens, + osl_backend, + primary=primary, + dataset_name=eval_cfg.dataset_name, ) - # Drop the DataFrame so finalize doesn't hold both it and the - # dict (each carrying the response-text corpus). - del out_df - entry["response_counts"] = _phase_response_counts( - scorer_instance.sample_index_map, uuid_to_text ) - if osl_backend is not None: - t0 = time.perf_counter() - osl = _phase_osl_stats( - scorer_instance.sample_index_map, uuid_to_text, osl_backend - ) - if osl is not None: - # Same shape/key as the perf report output_sequence_lengths. - entry["output_sequence_lengths"] = osl - # Wall-clock of just this phase's tokenization (seconds); - # summed across datasets for the accuracy report's total. - entry["osl_tokenize_s"] = round(time.perf_counter() - t0, 3) except Exception as e: # noqa: BLE001 - optional blocks; never fail scoring logger.warning( "Accuracy response counts/OSL skipped for %s: %s", @@ -386,7 +510,10 @@ def score_accuracy( def write_accuracy_results( - report_dir: Path, accuracy_scores: list[dict[str, Any]] + report_dir: Path, + accuracy_scores: list[dict[str, Any]], + *, + token_count_source: TokenCountSource, ) -> None: """Emit ``accuracy/accuracy_results.json`` from the per-dataset score entries. @@ -394,6 +521,10 @@ def write_accuracy_results( ``performance/result_summary.json`` and response/error text lives in ``events.jsonl``, so neither is duplicated here. A no-op when ``accuracy_scores`` is empty — a perf-only run leaves no ``accuracy/`` folder. + + ``token_count_source`` is the run's effective token basis, written top level + as the label of every entry's ``output_sequence_lengths`` (the file is read + on its own, without the sibling ``result_summary.json``). """ if not accuracy_scores: return @@ -420,6 +551,7 @@ def write_accuracy_results( accuracy_payload["average_accuracy"] = avg_accuracy if osl_computed: accuracy_payload["osl_tokenization_s"] = osl_tokenization_s + accuracy_payload["token_count_source"] = token_count_source.value accuracy_payload["accuracy_scores"] = accuracy_scores # Atomic write so a crash mid-write can't leave truncated JSON the # compliance checker would read as corrupt. Not swallowed: if scoring diff --git a/src/inference_endpoint/commands/benchmark/full_run_osl.py b/src/inference_endpoint/commands/benchmark/full_run_osl.py index 5a081ad14..b4668ec94 100644 --- a/src/inference_endpoint/commands/benchmark/full_run_osl.py +++ b/src/inference_endpoint/commands/benchmark/full_run_osl.py @@ -11,10 +11,20 @@ is measured, even when the model produces identical output for every turn. This module recomputes OSL over every completed turn from ``events.jsonl``, -which retains the full ``TextModelOutput`` (reasoning and tool calls included) -regardless of tracking state. Counting goes through the same -``extract_tokenization_input`` rule the perf-side ``OslTrigger`` uses, so the -two statistics are comparable rather than merely adjacent. +which retains each turn's COMPLETE record regardless of tracking state. One +pass counts each turn on both token bases, as the windowed OSL does: + +- Client: the full ``TextModelOutput`` (reasoning and tool calls included) + through the same ``extract_tokenization_input`` rule the perf-side + ``OslTrigger`` uses, so the two statistics are comparable rather than merely + adjacent. Needs a tokenizer. +- Server: each turn's reported ``completion_tokens``, read by + ``read_server_usage``; a turn without usable usage is counted under + ``n_missing_usage``. + +The run's primary basis keeps the base key names and alone sets the block's +``partial``; the other basis sits beside it under suffixed names (see +``basis_key``). See mlcommons/endpoints#500. """ @@ -23,7 +33,10 @@ import logging import os +from array import array from collections.abc import Callable, Collection, Sequence +from contextlib import nullcontext +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -43,12 +56,21 @@ SampleEventType, leading_event_topic, ) -from inference_endpoint.core.types import PhaseType, TextModelOutput +from inference_endpoint.core.types import ( + PhaseType, + TextModelOutput, + TokenCountSource, + basis_key, +) from inference_endpoint.endpoint_client.cpu_affinity import ( UnsupportedPlatformError, get_all_online_cpus, ) from inference_endpoint.metrics.report import series_metric_dict +from inference_endpoint.metrics.token_count_basis import ( + ServerUsage, + read_server_usage, +) logger = logging.getLogger(__name__) @@ -57,8 +79,6 @@ # to find the COMPLETE records. Any other layout is decoded and classified by its # decoded event_type (see compute_full_run_osl). _COMPLETE_TOPIC: bytes = SampleEventType.COMPLETE.topic_bytes # type: ignore[attr-defined] -# AgenticInferenceInlineScorer's dataset_name; the key its turns live under in -# sample_idx_map.json. # Counts a batch of turns, returning a count or the exception for each turn. BatchTokenCounter = Callable[ @@ -73,6 +93,66 @@ _MAX_COUNT_THREADS = 32 +@dataclass(slots=True) +class _Tally: + """Turn lengths and exclusion counters for one token basis. + + ``n_errors`` (turns the tokenizer could not count, and unexpected + payloads) and ``n_missing_usage`` (turns without usable server usage) are + the turns this basis should have counted but could not. + """ + + lengths: array[int] = field(default_factory=lambda: array("q")) + n_empty: int = 0 + n_errors: int = 0 + n_missing_usage: int = 0 + first_error: str | None = None + + def error(self, reason: str) -> None: + self.n_errors += 1 + if self.first_error is None: + self.first_error = reason + + def queue_tokenization( + self, data: TextModelOutput, pending: list[MessageInput | TextInput] + ) -> None: + """Queue one turn's output on ``pending`` for the next counted batch.""" + # A malformed tool-call payload can raise in extract_tokenization_input, + # and one bad turn must not void the whole statistic. + try: + tokenization_input = extract_tokenization_input(data) + except Exception as e: # noqa: BLE001 - one bad turn must not void the stat + self.error(str(e)) + return + if tokenization_input is None: + self.n_empty += 1 + return + pending.append(tokenization_input) + + def add_counts(self, outcomes: Sequence[int | Exception]) -> None: + """Record one counted batch: a length, or the error, per turn.""" + for outcome in outcomes: + if isinstance(outcome, Exception): + self.error(str(outcome)) + else: + self.lengths.append(outcome) + + def count_usage(self, record: EventRecord, data: TextModelOutput) -> None: + usage, completion_tokens = read_server_usage( + record.prompt_tokens, record.completion_tokens, data + ) + if completion_tokens is not None: + self.lengths.append(completion_tokens) + elif usage is ServerUsage.TRUE_ZERO: + self.n_empty += 1 + else: + self.n_missing_usage += 1 + + def missed_any(self) -> bool: + """Whether this basis could not count some turn it saw.""" + return bool(self.n_errors or self.n_missing_usage) + + def _in_population( record: EventRecord, performance_uuids: Collection[str] | None, @@ -95,50 +175,85 @@ def _in_population( - Salted turns (``agentic_inference.enable_salt``). Salt perturbs the prompt to defeat prefix caching; it does not create extra turns, so there is nothing to exclude. - - Empty completions. Those are excluded by ``extract_tokenization_input`` - returning ``None`` — the same rule the perf side applies — and counted - under ``n_empty`` rather than dropped silently. + - Empty completions. Those are excluded — by ``extract_tokenization_input`` + returning ``None`` on the client basis (the rule the perf side applies), + by a reported ``completion_tokens`` of 0 for an empty output on the + server basis — and counted under ``n_empty`` rather than dropped + silently. """ if performance_uuids is None: return True return record.sample_uuid in performance_uuids +def _basis_keys( + tally: _Tally, basis: TokenCountSource, *, primary: TokenCountSource +) -> dict[str, Any]: + """One basis's keys, named by ``basis_key``. + + The primary series is ``{}`` when nothing was counted, the other series + ``None``. ``n_errors`` exists only on the client basis. + """ + series: dict[str, Any] | None = series_metric_dict(tally.lengths) + if basis is not primary: + series = series or None + keys = { + basis_key("output_sequence_lengths", basis, primary=primary): series, + basis_key("n_turns_counted", basis, primary=primary): len(tally.lengths), + basis_key("n_empty", basis, primary=primary): tally.n_empty, + } + if basis is TokenCountSource.CLIENT: + keys[basis_key("n_errors", basis, primary=primary)] = tally.n_errors + return keys + + def compute_full_run_osl( events_path: Path, - count_batch: BatchTokenCounter, + count_batch: BatchTokenCounter | None, *, + token_count_source: TokenCountSource, performance_uuids: Collection[str] | None = None, batch_size: int = _DEFAULT_BATCH_SIZE, ) -> dict[str, Any] | None: """Roll up OSL over every completed turn in ``events_path``. + Counts each turn on both bases in one pass: the client basis with + ``count_batch``, ``batch_size`` turns at a time, the server basis from the + turn's usage. The primary basis, ``token_count_source``, takes the base key + names; the client basis needs ``count_batch`` when it is primary. Without + ``count_batch`` a client basis that is the other one counts nothing. + Returns ``None`` only when the file held no relevant COMPLETE lines at all and no expected turn was missing. A file with records but nothing countable - (all empty, timed out, errored, or undecodable) — or one missing an expected - performance turn — returns a block with an empty ``output_sequence_lengths`` - and non-zero counters so the run stays visible rather than looking uncomputed. - - ``partial`` is True when any turn was errored, undecodable, or missing - (reconciled against ``performance_uuids``): the mean is then over a subset of - the population and must not feed the OSL accuracy gate. Empty turns do not - make it partial — they are the shared rule's legitimate zero-output exclusion. + (all empty, timed out, errored, undecodable, or without usage) — or one + missing an expected performance turn — returns a block with an empty + ``output_sequence_lengths`` and non-zero counters so the run stays visible + rather than looking uncomputed. + + ``partial`` is True when the primary basis fell short: a turn errored, + was undecodable or missing (reconciled against ``performance_uuids``), or, + on the server basis, lacked usable usage. The mean is then over a subset + of the population and must not feed the OSL accuracy gate. Empty turns do + not make it partial — they are a legitimate zero-output exclusion. The + other basis's shortfall marks only its own series. """ + client_primary = token_count_source is TokenCountSource.CLIENT + if count_batch is None and client_primary: + raise ValueError("count_batch is required for client-basis full-run OSL") decoder = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) - lengths: list[int] = [] - n_empty = 0 - n_errors = 0 + tallies = {basis: _Tally() for basis in TokenCountSource} + client, server = tallies[TokenCountSource.CLIENT], tallies[TokenCountSource.SERVER] + # A client basis without a tokenizer counts nothing at all. + counting = (client, server) if count_batch is not None else (server,) n_undecodable = 0 - first_error: str | None = None seen_uuids: set[str] = set() batch_size = max(1, batch_size) pending: list[MessageInput | TextInput] = [] - def count_pending() -> None: - nonlocal n_errors, first_error + def count_pending(count: BatchTokenCounter) -> None: n_pending = len(pending) try: - outcomes = count_batch(pending) + outcomes = count(pending) except Exception as e: # noqa: BLE001 - charge the batch, keep the stat outcomes = [e] * n_pending if len(outcomes) != n_pending: @@ -146,17 +261,10 @@ def count_pending() -> None: # cannot be matched back to its turns, so charge the whole batch # (the block turns partial) rather than average an unknown subset. mismatch = RuntimeError( - f"token counter returned {len(outcomes)} outcomes for " - f"{n_pending} turns" + f"token counter returned {len(outcomes)} outcomes for {n_pending} turns" ) outcomes = [mismatch] * n_pending - for outcome in outcomes: - if isinstance(outcome, Exception): - n_errors += 1 - if first_error is None: - first_error = str(outcome) - else: - lengths.append(outcome) + client.add_counts(outcomes) pending.clear() with events_path.open("rb") as events_file: @@ -182,32 +290,25 @@ def count_pending() -> None: if not _in_population(record, performance_uuids): continue seen_uuids.add(record.sample_uuid) - if record.data is None: - n_empty += 1 # e.g. a timed-out turn logs COMPLETE with data=None + data = record.data + if data is None: + # e.g. a timed-out turn logs COMPLETE with data=None + for tally in counting: + tally.n_empty += 1 continue - if not isinstance(record.data, TextModelOutput): + if not isinstance(data, TextModelOutput): # Unexpected payload (ErrorData/PromptData) on a COMPLETE event. - n_errors += 1 - if first_error is None: - first_error = f"unexpected {type(record.data).__name__} payload" + reason = f"unexpected {type(data).__name__} payload" + for tally in counting: + tally.error(reason) continue - # Output preparation stays inside the per-turn try: a malformed - # tool-call payload can raise in extract_tokenization_input, and one - # bad turn must not void the whole statistic. - try: - tokenization_input = extract_tokenization_input(record.data) - if tokenization_input is None: - n_empty += 1 - continue - pending.append(tokenization_input) + server.count_usage(record, data) + if count_batch is not None: + client.queue_tokenization(data, pending) if len(pending) >= batch_size: - count_pending() - except Exception as e: # noqa: BLE001 - one bad turn must not void the stat - n_errors += 1 - if first_error is None: - first_error = str(e) - if pending: - count_pending() + count_pending(count_batch) + if count_batch is not None and pending: + count_pending(count_batch) # A performance turn whose COMPLETE record never reached the log is invisible # to the counters above; reconcile observed UUIDs against the expected @@ -223,39 +324,50 @@ def count_pending() -> None: logger.warning( "Full-run OSL: skipped %d undecodable COMPLETE line(s)", n_undecodable ) - if n_errors: - logger.warning( - "Full-run OSL: %d turn(s) could not be counted (first: %s)", - n_errors, - first_error, - ) + for basis, tally in tallies.items(): + if tally.n_errors: + logger.warning( + "Full-run OSL: %d turn(s) could not be counted on the %s basis " + "(first: %s)", + tally.n_errors, + basis.value, + tally.first_error, + ) if n_missing: logger.warning( "Full-run OSL: %d performance turn(s) had no COMPLETE record in the log", n_missing, ) + if server.n_missing_usage: + # A shortfall decides the block only on the primary basis. + logger.log( + logging.INFO if client_primary else logging.WARNING, + "Full-run OSL: %d turn(s) carried no usable server token usage", + server.n_missing_usage, + ) - if ( - not lengths - and not n_empty - and not n_errors - and not n_undecodable - and not n_missing - ): + if not seen_uuids and not n_undecodable and not n_missing: return None - # ``partial`` is the machine-readable gate signal: errored / undecodable / - # missing turns mean the mean is over a subset and must not feed the OSL - # accuracy gate. Empty turns are a legitimate shared-rule exclusion. - partial = bool(n_errors or n_undecodable or n_missing) - return { - "output_sequence_lengths": series_metric_dict(lengths), - "n_turns_counted": len(lengths), - "n_empty": n_empty, - "n_errors": n_errors, + + # Every basis falls short of the population when a line was undecodable or + # an expected turn is missing. + population_short = bool(n_undecodable or n_missing) + primary, other = token_count_source, token_count_source.other + block = { + **_basis_keys(tallies[primary], primary, primary=primary), "n_undecodable": n_undecodable, "n_missing": n_missing, - "partial": partial, + # The machine-readable gate signal: the primary mean is over a subset + # and must not feed the OSL accuracy gate. + "partial": population_short or tallies[primary].missed_any(), + "n_missing_usage": server.n_missing_usage, + **_basis_keys(tallies[other], other, primary=primary), } + # The other basis's shortfall marks only its own series. + other_series = block[basis_key("output_sequence_lengths", other, primary=primary)] + if other_series is not None and (population_short or tallies[other].missed_any()): + other_series["partial"] = True + return block def load_performance_uuids(report_dir: Path) -> set[str] | None: @@ -289,21 +401,35 @@ def load_performance_uuids(report_dir: Path) -> set[str] | None: return set(per_phase) +def _count_threads() -> int: + """Tokenizer threads for the finalize pass: the online CPUs, capped.""" + try: + threads = len(get_all_online_cpus()) or 1 + except UnsupportedPlatformError: + threads = os.cpu_count() or 1 + return min(threads, _MAX_COUNT_THREADS) + + def full_run_osl_for_report( report_dir: Path, tokenizer_name: str | None, + *, + token_count_source: TokenCountSource, ) -> dict[str, Any] | None: """Compute the full-run OSL block for a finished run's report directory. - Best-effort by design: a tokenizer that cannot be loaded drops the block - rather than failing finalize. Counting runs in-process (``n_workers=0``), - which works with or without a fast (Rust) backend — the same fallback the - live ``OslTrigger`` uses — so the block is not gated on backend - availability. Returns ``None`` when there is no performance population to - bound to (accuracy-only / perf-less runs). + Best-effort by design: it never fails finalize. A client-primary run + without a tokenizer, or whose tokenizer cannot be loaded, drops the block; + a server-primary run needs no tokenizer, and without one counts nothing on + the client basis. Counting runs in-process (``n_workers=0``), which works + with or without a fast (Rust) backend — the same fallback the live + ``OslTrigger`` uses — so the block is not gated on backend availability. + Returns ``None`` when there is no performance population to bound to + (accuracy-only / perf-less runs). """ + client_primary = token_count_source is TokenCountSource.CLIENT events_path = report_dir / "events.jsonl" - if tokenizer_name is None or not events_path.is_file(): + if (tokenizer_name is None and client_primary) or not events_path.is_file(): return None performance_uuids = load_performance_uuids(report_dir) if performance_uuids is None: @@ -313,18 +439,23 @@ def full_run_osl_for_report( # runs or an unreadable map. logger.info("Full-run OSL skipped: no performance-phase samples to bound to") return None - try: + tokenizer_cm: BatchTokenizer | nullcontext[None] = nullcontext() + if tokenizer_name is not None: try: - threads = len(get_all_online_cpus()) or 1 - except UnsupportedPlatformError: - threads = os.cpu_count() or 1 - threads = min(threads, _MAX_COUNT_THREADS) - with BatchTokenizer( - tokenizer_name, live_workers=threads, n_workers=0 - ) as tokenizer: + tokenizer_cm = BatchTokenizer( + tokenizer_name, live_workers=_count_threads(), n_workers=0 + ) + except Exception as e: # noqa: BLE001 - optional block; never fail finalize + if client_primary: + logger.warning("Full-run OSL skipped: %s", e) + return None + logger.warning("Full-run OSL: client counts skipped: %s", e) + try: + with tokenizer_cm as tokenizer: return compute_full_run_osl( events_path, - tokenizer.count_batch_sync, + tokenizer.count_batch_sync if tokenizer is not None else None, + token_count_source=token_count_source, performance_uuids=performance_uuids, ) except Exception as e: # noqa: BLE001 - optional block; never fail finalize diff --git a/src/inference_endpoint/core/record.py b/src/inference_endpoint/core/record.py index 70b3a0934..5dc064389 100644 --- a/src/inference_endpoint/core/record.py +++ b/src/inference_endpoint/core/record.py @@ -172,31 +172,6 @@ class EventRecord(msgspec.Struct, kw_only=True, frozen=True, gc=False): # type: completion_tokens: int | msgspec.UnsetType = msgspec.UNSET -# ``event_type`` is EventRecord's first field, so in an ``events.jsonl`` line -# (msgspec JSON) the marker sits in the leading bytes. A JSON string value -# escapes its quotes, so the marker cannot occur inside one. -_COMPLETE_LINE_MARKER: Final = b'"event_type":"sample.complete"' -_MARKER_WINDOW: Final = 256 - - -def is_complete_line(line: bytes) -> bool: - """Whether one ``events.jsonl`` line holds a COMPLETE record. - - Searches only the line's leading bytes, so a reader skips every other - record without decoding it or scanning its prompt and response payload. - - Precondition: the line is compact msgspec JSON with ``event_type`` first, - as ``JSONLWriter`` writes it. Any other encoding (e.g. ``json.dumps`` - with its default ``": "`` separator) never matches. - """ - return line.find(_COMPLETE_LINE_MARKER, 0, _MARKER_WINDOW) != -1 - - -class EventLogFormatError(ValueError): - """An event log a reader cannot read without losing samples, such as one - not in the compact form ``is_complete_line`` requires.""" - - EVENT_LINE_LEAD: Final[bytes] = b'{"event_type":"' """bytes: How an events.jsonl line begins when it is written compactly. diff --git a/src/inference_endpoint/evaluation/bfcl_v4_scorer.py b/src/inference_endpoint/evaluation/bfcl_v4_scorer.py index 5bb6690db..927720a46 100644 --- a/src/inference_endpoint/evaluation/bfcl_v4_scorer.py +++ b/src/inference_endpoint/evaluation/bfcl_v4_scorer.py @@ -125,8 +125,9 @@ def get_scoring_outputs(self) -> pd.DataFrame: the prose is chatter. Plain-text responses (no tool calls, e.g. hallucination refusals) fall back to the full string. - OSL / response accounting read the base ``get_raw_outputs()`` instead, so - they count the full generated text rather than this normalized form. + OSL / response accounting read the base ``iter_complete_outputs()`` + instead, so they count the full generated text rather than this + normalized form. """ outputs: list[dict[str, str]] = [] for sample_uuid, data in self._iter_complete(): diff --git a/src/inference_endpoint/evaluation/scoring.py b/src/inference_endpoint/evaluation/scoring.py index 61bec6575..7c1d60346 100644 --- a/src/inference_endpoint/evaluation/scoring.py +++ b/src/inference_endpoint/evaluation/scoring.py @@ -54,6 +54,7 @@ from ..dataset_manager.agentic_inference_dataset import AgenticInferenceDataset from ..dataset_manager.dataset import Dataset from ..dataset_manager.predefined.shopify_product_catalogue import ProductMetadata +from ..metrics.token_count_basis import read_server_usage from .accuracy_results import build_breakdown from .extractor import ( Extractor, @@ -132,6 +133,15 @@ def guess_and_decode(line: bytes) -> EventRecord | _IssuedKey | None: return None +# The output a COMPLETE event without a text payload is classified against. +_NO_OUTPUT = TextModelOutput(output="") + + +def _generated_text(data: Any) -> str: + """The text the model generated: ``str(TextModelOutput)``, "" without output.""" + return str(data) if data is not None else "" + + class Scorer(ABC): """Scorers will read in a dataset and outputs from a log and compute an accuracy score. An optional extractor can be provided to post-process the output to extract values that @@ -232,12 +242,15 @@ def _load_sample_index_map(self): d = msgspec.json.decode(f.read()) return d[self.dataset_name] # Implicitly raises KeyError - def _iter_complete(self) -> Iterator[tuple[str, Any]]: - """Yield ``(sample_uuid, data)`` for each COMPLETE event in events.jsonl. + def _iter_complete_records(self) -> Iterator[EventRecord]: + """Yield each COMPLETE ``EventRecord`` in events.jsonl. - The single events reader shared by ``get_raw_outputs`` and any scorer - that needs the structured event data (e.g. BFCL's ``tool_calls``). The - EventLoggerService writes EventRecord objects serialized via msgspec. + The single events reader behind ``_iter_complete``, + ``get_raw_outputs`` and ``iter_complete_outputs``. The + EventLoggerService writes EventRecord objects serialized via msgspec; + in its compact layout only COMPLETE lines are decoded + (``leading_event_topic``), and a line in any other layout is decoded + whole and classified by its ``event_type``. """ events_log_path = self.report_dir / "events.jsonl" if not events_log_path.exists(): @@ -268,7 +281,39 @@ def _iter_complete(self) -> Iterator[tuple[str, Any]]: # skipped undecoded. continue if record.event_type == SampleEventType.COMPLETE: - yield record.sample_uuid, record.data + yield record + + def _iter_complete(self) -> Iterator[tuple[str, Any]]: + """Yield ``(sample_uuid, data)`` for each COMPLETE event in events.jsonl. + + For ``get_raw_outputs`` and any scorer that needs the structured event + data (e.g. BFCL's ``tool_calls``). + """ + for record in self._iter_complete_records(): + yield record.sample_uuid, record.data + + def iter_complete_outputs( + self, uuids: set[str] | None + ) -> Iterator[tuple[str, str, int | None]]: + """Yield ``(uuid, text, completion_tokens)`` per COMPLETE event. + + One read serves accuracy OSL and response accounting on both token + bases. ``text`` is the model's generated text, as in + ``get_raw_outputs``. ``completion_tokens`` is the event's + server-reported count when its usage is usable (``read_server_usage``), + else ``None``. ``uuids`` bounds the read to a target population; + ``None`` reads every event. + """ + for record in self._iter_complete_records(): + if uuids is not None and record.sample_uuid not in uuids: + continue + data = record.data + _, completion_tokens = read_server_usage( + record.prompt_tokens, + record.completion_tokens, + data if isinstance(data, TextModelOutput) else _NO_OUTPUT, + ) + yield record.sample_uuid, _generated_text(data), completion_tokens def get_raw_outputs(self, wanted_uuids: set[str] | None = None) -> pd.DataFrame: """uuid -> the text the model generated (``str(TextModelOutput)``). @@ -280,7 +325,7 @@ def get_raw_outputs(self, wanted_uuids: set[str] | None = None) -> pd.DataFrame: response-text corpus. """ rows = [ - {"sample_uuid": u, "output": str(d) if d is not None else ""} + {"sample_uuid": u, "output": _generated_text(d)} for u, d in self._iter_complete() if wanted_uuids is None or u in wanted_uuids ] diff --git a/src/inference_endpoint/metrics/steady_state_diagnostics.py b/src/inference_endpoint/metrics/steady_state_diagnostics.py index db300bc82..56c247d97 100644 --- a/src/inference_endpoint/metrics/steady_state_diagnostics.py +++ b/src/inference_endpoint/metrics/steady_state_diagnostics.py @@ -9,13 +9,24 @@ For each performance-tracked sample it computes: - ttft_ns = recv_first.ts - issued.ts - - tpot_ns = (complete.ts - recv_first.ts) / tokens(text_after_first_chunk) + - tpot_ns = (complete.ts - recv_first.ts) / tokens after the first chunk -Token counts use plain tokenization of ``text_after_first_chunk``. The live -metrics aggregator uses the chat-template path for reasoning and tool-call -outputs. For reasoning models, absolute TPOT can differ from the run report. -CoV and trend tests are scale-invariant, so the steady/drift verdict is the -same. +The TPOT denominator follows a token-count basis. By default it is the run's +primary basis, read by ``token_count_source_for_run`` from the event log's +directory; ``--token-count-source`` replays on either basis: + + - server: ``completion_tokens - 1`` from the COMPLETE line's usage, by the + live aggregator's rule (``classify_server_usage``, then LoadGen's + formula). Samples without usable usage, with ``completion_tokens <= 1`` + or a non-positive window get no TPOT. No tokenizer is needed. + - client: plain tokenization of ``text_after_first_chunk``. The live + metrics aggregator uses the chat-template path for reasoning and + tool-call outputs, so for reasoning models absolute TPOT can differ from + the run report. CoV and trend tests are scale-invariant, so the + steady/drift verdict is the same. + +The JSON verdict key follows the basis naming rule: ``steady_state`` on the +run's primary basis, ``steady_state_`` on the other. Samples are grouped by issue order into super-passes. ``--superpass-size`` sets the samples per bucket. By default it is the dataset size. Each metric and @@ -42,13 +53,16 @@ p99 and end-to-end latency are also diagnostic. Latency variation tracks the OSL mix. -usage (auto-detects the tokenizer and workload profile from config.yaml, and -the super-pass size from the phase_start event; see MODEL_REGISTRY + PROFILES): +usage (auto-detects the token-count basis from result_summary.json or +config.yaml, the tokenizer for client-basis runs and the workload profile from +config.yaml, and the super-pass size from the phase_start event; see +MODEL_REGISTRY + PROFILES): python -m inference_endpoint.metrics.steady_state_diagnostics # 0 flags python -m inference_endpoint.metrics.steady_state_diagnostics \ --model kimi-k3 # 1 flag -Every derived setting has an explicit override (--tokenizer, --dataset-size, ---superpass-size, --profile, --cov-bounds, --warmup, --json, ...). +Every derived setting has an explicit override (--token-count-source, +--tokenizer, --dataset-size, --superpass-size, --profile, --cov-bounds, +--warmup, --json, ...). """ from __future__ import annotations @@ -61,14 +75,34 @@ from array import array from collections.abc import Callable, Sequence from dataclasses import dataclass, field +from pathlib import Path from statistics import NormalDist, median, pstdev -from typing import Final, Literal, NamedTuple, TypedDict +from typing import Final, Literal, NamedTuple import msgspec import yaml from transformers import AutoTokenizer -from inference_endpoint.core.types import PhaseData, PhaseType +from inference_endpoint.core.record import ( + EventRecord, + EventType, + SampleEventType, + leading_event_topic, +) +from inference_endpoint.core.types import ( + PhaseData, + PhaseType, + TextModelOutput, + TokenCountSource, + basis_key, +) +from inference_endpoint.metrics.token_count_basis import ( + TOKEN_COUNT_SOURCE_METAVAR, + RunBasisError, + ServerUsage, + read_server_usage, + replay_basis, +) # --------------------------------------------------------------------------- # # Event wire constants (mirror core/record.py category.value topics) @@ -456,36 +490,99 @@ def superpass_size_from_events(events_path: str) -> int | None: return None +# --------------------------------------------------------------------------- # +# COMPLETE records +# --------------------------------------------------------------------------- # +# A COMPLETE line carries the output and server usage a replay counts tokens +# from, so it is decoded as the typed record the JSONL writer encoded. The +# replays find these lines with ``_is_complete`` and read every other event as +# an untyped JSON object. +_EVENT_DECODER = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) +_JSON_DECODER = msgspec.json.Decoder() +_COMPLETE_TOPIC: bytes = SampleEventType.COMPLETE.topic_bytes # type: ignore[attr-defined] + + +def _is_complete(line: bytes) -> bool: + """Whether ``line`` is a COMPLETE event. + + A line in the writer's compact layout is classified from its leading bytes + (``leading_event_topic``). A line in any other layout is decoded to read its + ``event_type``, so a change in how the log is written cannot silently drop + COMPLETE events. + """ + topic = leading_event_topic(line) + if topic is not None: + return topic == _COMPLETE_TOPIC + try: + return _JSON_DECODER.decode(line).get("event_type") == EV_COMPLETE + except msgspec.DecodeError: + return False + + +def _decode_complete(line: bytes) -> EventRecord | None: + """The COMPLETE record on ``line``; None when the line does not decode.""" + try: + return _EVENT_DECODER.decode(line) + except (msgspec.DecodeError, NotImplementedError): + # A truncated or corrupt line (e.g. the last line of a killed run). + # EventType.decode_hook raises NotImplementedError for a non-string + # event_type, which msgspec does not wrap. + return None + + # --------------------------------------------------------------------------- # # NATL (agentic per-trajectory throughput) -- EXPERIMENTAL, see the CLI warning # --------------------------------------------------------------------------- # -def full_output_text(data: object) -> str: - """Return all generated text for a turn, for OSL counting.""" - if not isinstance(data, list): - return "" +def full_output_text(output: TextModelOutput) -> str: + """Return all generated text for a turn, for OSL counting. + + Reasoning, then output, then the tool calls as logged (JSON). + """ parts: list[str] = [] - reasoning = data[2] if len(data) > 2 else None - output = data[1] if len(data) > 1 else "" - if reasoning: - parts.extend(reasoning if isinstance(reasoning, list) else [reasoning]) - if output: - parts.extend(output if isinstance(output, list) else [output]) - tool_calls = data[3] if len(data) > 3 else None - if tool_calls: - parts.append(json.dumps(tool_calls)) - return "".join(str(p) for p in parts) + for text in (output.reasoning, output.output): + if isinstance(text, tuple): + parts.extend(text) + elif text: + parts.append(text) + if output.tool_calls: + parts.append(json.dumps(output.tool_calls)) + return "".join(parts) + + +class ReplayTrajectory(msgspec.Struct, frozen=True): # type: ignore[call-arg] + """:func:`build_trajectory_natl` result. + + ``pairs`` are ``(last_complete_ns, natl)`` sorted by completion time. + ``n_missing_usage`` and ``n_undecodable`` are counted as in + :class:`ReplaySeries`; a turn without usable usage adds no tokens. + """ + + pairs: list[tuple[int, float]] + n_missing_usage: int | None + n_undecodable: int def build_trajectory_natl( events_path: str, - count_tokens: Callable[[list[str]], list[int]], + count_tokens: Callable[[list[str]], list[int]] | None, flush_size: int = 512, -) -> list[tuple[int, float]]: + *, + basis: TokenCountSource, +) -> ReplayTrajectory: """Compute per-trajectory NATL values sorted by completion time. NATL is sum(output tokens) / sum(e2e latency seconds), grouped by - ``conversation_id``. Each returned pair is ``(last_complete_ns, natl)``. + ``conversation_id``. + + Output tokens per turn follow ``basis``, the basis replayed on: the usable + server ``completion_tokens`` (``read_server_usage``; no tokenizer), or + ``count_tokens`` over the turn's full output text. """ + server_basis = basis is TokenCountSource.SERVER + if not server_basis and count_tokens is None: + raise ValueError("count_tokens is required for a client-basis replay") + n_missing_usage = 0 + n_undecodable = 0 issue: dict[str, tuple[str, int]] = {} conv_lat_s: dict[str, float] = {} conv_tokens: dict[str, int] = {} @@ -495,20 +592,50 @@ def build_trajectory_natl( tracking = False def flush() -> None: - if b_text: + if b_text and count_tokens is not None: for conv, cnt in zip(b_conv, count_tokens(b_text), strict=True): conv_tokens[conv] = conv_tokens.get(conv, 0) + cnt b_conv.clear() b_text.clear() - with open(events_path) as f: + with open(events_path, "rb") as f: for line in f: - line = line.strip() - if not line: + if _is_complete(line): + record = _decode_complete(line) + if record is None: + n_undecodable += 1 + continue + iv = issue.pop(record.sample_uuid, None) + if iv is None: + continue + ts = record.timestamp_ns + conv, its = iv + conv_lat_s[conv] = conv_lat_s.get(conv, 0.0) + (ts - its) / 1e9 + conv_end_ns[conv] = max(conv_end_ns.get(conv, 0), ts) + output = record.data + if not isinstance(output, TextModelOutput): + continue # errors and agentic timeouts complete without output + if server_basis: + usage, n = read_server_usage( + record.prompt_tokens, record.completion_tokens, output + ) + if usage is ServerUsage.MISSING: + n_missing_usage += 1 + if n is not None: + conv_tokens[conv] = conv_tokens.get(conv, 0) + n + continue + text = full_output_text(output) + if text: + b_conv.append(conv) + b_text.append(text) + if len(b_text) >= flush_size: + flush() continue try: - rec = json.loads(line) - except json.JSONDecodeError: + rec = _JSON_DECODER.decode(line) + except msgspec.DecodeError: + # Skip a truncated or partial line (e.g. the last line of a + # killed run); an undecodable COMPLETE line is counted above. continue et = rec.get("event_type") ts = rec.get("timestamp_ns") @@ -520,19 +647,6 @@ def flush() -> None: uuid = rec.get("sample_uuid") if tracking and uuid and ts is not None: issue[uuid] = (rec.get("conversation_id") or "", ts) - elif et == EV_COMPLETE: - iv = issue.pop(rec.get("sample_uuid"), None) - if iv is None or ts is None: - continue - conv, its = iv - conv_lat_s[conv] = conv_lat_s.get(conv, 0.0) + (ts - its) / 1e9 - conv_end_ns[conv] = max(conv_end_ns.get(conv, 0), ts) - text = full_output_text(rec.get("data")) - if text: - b_conv.append(conv) - b_text.append(text) - if len(b_text) >= flush_size: - flush() flush() pairs = [ (conv_end_ns[c], conv_tokens[c] / conv_lat_s[c]) @@ -540,7 +654,11 @@ def flush() -> None: if conv_lat_s[c] > 0 and conv_tokens.get(c, 0) > 0 ] pairs.sort(key=lambda p: p[0]) - return pairs + return ReplayTrajectory( + pairs=pairs, + n_missing_usage=n_missing_usage if server_basis else None, + n_undecodable=n_undecodable, + ) def build_natl_result( @@ -581,29 +699,23 @@ def build_natl_result( # --------------------------------------------------------------------------- # -# TextModelOutput.text_after_first_chunk for the parsed JSON array +# TextModelOutput.text_after_first_chunk without the tool calls # --------------------------------------------------------------------------- # -def text_after_first_chunk(data: object) -> str: +def text_after_first_chunk(output: TextModelOutput) -> str: """Return output text after the first streamed chunk. - This text is tokenized for the TPOT denominator. ``data`` is the COMPLETE - event payload: ``[tag, output, reasoning?, tool_calls?]``. - - ``output`` and ``reasoning`` are each either a string for non-streaming - output or a list of streamed chunks. The logic mirrors + This text is tokenized for the TPOT denominator. ``output`` and + ``reasoning`` are each either a string for non-streaming output or a tuple + of streamed chunks. The logic mirrors ``TextModelOutput.text_after_first_chunk`` in core/types.py. """ - if not isinstance(data, list) or not data: - return "" - output = data[1] if len(data) > 1 else "" - reasoning = data[2] if len(data) > 2 else None parts: list[str] = [] - if isinstance(reasoning, list): - parts.extend(reasoning[1:]) - if isinstance(output, list): - parts.extend(output[1:]) - elif isinstance(output, str) and isinstance(reasoning, list) and reasoning: - parts.append(output) + if isinstance(output.reasoning, tuple): + parts.extend(output.reasoning[1:]) + if isinstance(output.output, tuple): + parts.extend(output.output[1:]) + elif isinstance(output.reasoning, tuple) and output.reasoning: + parts.append(output.output) # Tool-call reconstruction is intentionally omitted: tool-call samples use a # chat-template tokenization path this diagnostic does not replicate. return "".join(parts) @@ -764,16 +876,41 @@ def add_tpot(self, superpass_index: int, tpot_ns: float, token_count: int) -> No sp.out_tokens += token_count +class ReplaySeries(msgspec.Struct, frozen=True): # type: ignore[call-arg] + """:func:`build_super_pass_series` result. + + ``n_missing_usage`` counts tracked ``TextModelOutput`` completions without + usable server usage on a server-basis replay (no TPOT, never tokenized). It + is None on a client-basis replay, which never reads usage. + ``n_undecodable`` counts COMPLETE lines that do not decode as an + ``EventRecord``; their samples are missing from the series. + """ + + super_passes: list[SuperPassRollup] + n_missing_usage: int | None + n_undecodable: int + + def build_super_pass_series( events_path: str, superpass_size: int, - count_tokens: Callable[[list[str]], list[int]], + count_tokens: Callable[[list[str]], list[int]] | None, flush_size: int = TOKENIZE_BATCH_SIZE, -) -> list[SuperPassRollup]: + *, + basis: TokenCountSource, +) -> ReplaySeries: """Replay an ``events.jsonl`` through a :class:`SuperPassCollector`. - ``count_tokens`` maps a text batch to token counts. It is injected so tests - can avoid a real tokenizer and callers can swap tokenizers. + ``basis`` is the basis replayed on, that of the TPOT token counts: + + - SERVER: TPOT is ``(complete - recv_first) / (completion_tokens - 1)`` + from the COMPLETE line's usage, with the live aggregator's rule and + guards (usable usage by ``read_server_usage``, streamed, + ``completion_tokens > 1``, positive window). ``count_tokens`` may be + None. + - CLIENT: ``count_tokens`` maps a text batch to token counts. It is + injected so tests can avoid a real tokenizer and callers can swap + tokenizers. ``flush_size`` caps buffered output texts before tokenizer flush. Lower values bound peak memory for long reasoning outputs at the cost of smaller @@ -781,6 +918,11 @@ def build_super_pass_series( """ if superpass_size <= 0: raise ValueError("superpass_size must be positive") + server_basis = basis is TokenCountSource.SERVER + if not server_basis and count_tokens is None: + raise ValueError("count_tokens is required for a client-basis replay") + n_missing_usage = 0 + n_undecodable = 0 collector = SuperPassCollector(superpass_size) # The parse keeps its own in-flight rows: it has no MetricsTable to hang the # super-pass index off, which is where the live producer stores it. @@ -792,7 +934,7 @@ def build_super_pass_series( batch_texts: list[str] = [] def flush_tpot() -> None: - if batch_texts: + if batch_texts and count_tokens is not None: counts = count_tokens(batch_texts) for (superpass_index, delta), cnt in zip(pending, counts, strict=True): if cnt > 0: @@ -800,14 +942,46 @@ def flush_tpot() -> None: pending.clear() batch_texts.clear() - with open(events_path) as f: + with open(events_path, "rb") as f: for line in f: - line = line.strip() - if not line: + if _is_complete(line): + record = _decode_complete(line) + if record is None: + n_undecodable += 1 + continue + row = rows.pop(record.sample_uuid, None) + if row is None: + continue + delta = collector.on_complete( + row.superpass_index, + row.issue_ns, + row.recv_first_ns, + record.timestamp_ns, + ) + output = record.data + if not isinstance(output, TextModelOutput): + continue # errors and agentic timeouts complete without output + if server_basis: + usage, n = read_server_usage( + record.prompt_tokens, record.completion_tokens, output + ) + if usage is ServerUsage.MISSING: + n_missing_usage += 1 + if n is not None and n > 1 and delta is not None and delta > 0: + collector.add_tpot(row.superpass_index, delta / (n - 1), n - 1) + continue + if delta is None: + continue + text = text_after_first_chunk(output) + if text: + pending.append((row.superpass_index, delta)) + batch_texts.append(text) + if len(batch_texts) >= flush_size: + flush_tpot() continue try: - rec = json.loads(line) - except json.JSONDecodeError: + rec = _JSON_DECODER.decode(line) + except msgspec.DecodeError: continue # skip a truncated/partial line (e.g. last line of a killed run) et = rec.get("event_type") if et == EV_START_TRACKING: @@ -845,23 +1019,12 @@ def flush_tpot() -> None: ) if row.recv_first_ns is None: row.recv_first_ns = ts - elif et == EV_COMPLETE: - row = rows.pop(uuid, None) - if row is None: - continue - delta = collector.on_complete( - row.superpass_index, row.issue_ns, row.recv_first_ns, ts - ) - if delta is None: - continue - text = text_after_first_chunk(rec.get("data")) - if text: - pending.append((row.superpass_index, delta)) - batch_texts.append(text) - if len(batch_texts) >= flush_size: - flush_tpot() flush_tpot() - return collector.series() + return ReplaySeries( + super_passes=collector.series(), + n_missing_usage=n_missing_usage if server_basis else None, + n_undecodable=n_undecodable, + ) # --------------------------------------------------------------------------- # @@ -1843,16 +2006,43 @@ def compute_steady_state_metrics( # --------------------------------------------------------------------------- # # Top-level orchestration # --------------------------------------------------------------------------- # -class DiagnosticsResult(TypedDict): +class DiagnosticsResult(msgspec.Struct, frozen=True): # type: ignore[call-arg] """Standalone CLI result: the verdict plus debug tables. Benchmark runs produce only ``steady_state``. The CLI also includes ``drift``, a whole-run trend scan used to explain ``found: false`` verdicts. + The token-count labels sit beside the verdict, not inside it, so the block + stays comparable with a run's ``result_summary.json``. """ + basis: TokenCountSource # basis the TPOT series was replayed on + primary: TokenCountSource # the run's basis, which keeps the verdict's name + # Server basis: tracked TextModelOutput completions without usable usage + # (no TPOT, never tokenized). None on the client basis, which never reads it. + n_missing_usage: int | None + # COMPLETE lines that did not decode; their samples are not in the series. + n_undecodable: int steady_state: SteadyState drift: dict[str, dict[str, Verdict]] # metric key -> algorithm -> verdict + def to_json(self) -> bytes: + """The CLI's JSON blob; the verdict key follows ``basis_key``. + + The replayed basis is written as ``token_count_source``, and + ``n_undecodable`` only when a line was skipped. + """ + verdict_key = basis_key("steady_state", self.basis, primary=self.primary) + # msgspec encodes the Struct tree and validates its shape. + blob: dict[str, object] = { + "token_count_source": self.basis, + "n_missing_usage": self.n_missing_usage, + } + if self.n_undecodable: + blob["n_undecodable"] = self.n_undecodable + blob[verdict_key] = self.steady_state + blob["drift"] = self.drift + return msgspec.json.format(msgspec.json.encode(blob)) + def _drift_verdicts(trajectory: Sequence[float]) -> dict[str, Verdict]: return {name: fn(trajectory).verdict for name, fn in ALGORITHMS.items()} @@ -1861,7 +2051,7 @@ def _drift_verdicts(trajectory: Sequence[float]) -> dict[str, Verdict]: def run( events_path: str, superpass_size: int, - count_tokens: Callable[[list[str]], list[int]], + count_tokens: Callable[[list[str]], list[int]] | None, warmup: int | str = "auto", cov_bounds: Sequence[float] = DEFAULT_COV_BOUNDS, trend_gate: str = "mk_hamed_rao", @@ -1869,14 +2059,24 @@ def run( warmup_band: float = 0.05, warmup_driver: str = "tpot_p50", enforce_min_duration: bool = True, + *, + basis: TokenCountSource, + primary: TokenCountSource, ) -> DiagnosticsResult: """Reconstruct the series from an event log, then analyse and diagnose it. - ``superpass_size`` is a count of samples. + ``superpass_size`` is a count of samples. ``basis`` is the basis TPOT is + replayed on (see :func:`build_super_pass_series`; on the server basis + ``count_tokens`` may be None) and ``primary`` is the run's own basis. """ - series = build_super_pass_series( - events_path, superpass_size, count_tokens, tokenize_batch_size + replay = build_super_pass_series( + events_path, + superpass_size, + count_tokens, + tokenize_batch_size, + basis=basis, ) + series = replay.super_passes ss = compute_steady_state_metrics( series, superpass_size=superpass_size, @@ -1888,15 +2088,19 @@ def run( enforce_min_duration=enforce_min_duration, ) post = series[ss.warmup :] - return { - "steady_state": ss, - "drift": { + return DiagnosticsResult( + basis=basis, + primary=primary, + n_missing_usage=replay.n_missing_usage, + n_undecodable=replay.n_undecodable, + steady_state=ss, + drift={ m.key: _drift_verdicts( super_pass_percentile_series(post, m.source_attr, m.percentile) ) for m in TRACKED_METRICS }, - } + ) # --------------------------------------------------------------------------- # @@ -1993,8 +2197,9 @@ def _render_cov_table( def render_text(result: DiagnosticsResult, cov_bounds: Sequence[float]) -> str: - ss = result["steady_state"] + ss = result.steady_state lines = [ + f"token basis: {result.basis.value}", f"super-passes: {ss.n_super_passes} " f"(size {ss.superpass_size}, warmup {ss.warmup})", "", @@ -2019,7 +2224,7 @@ def render_text(result: DiagnosticsResult, cov_bounds: Sequence[float]) -> str: algos = list(ALGORITHMS) lines.append(f" {'metric':<12} " + " ".join(f"{a:>16}" for a in algos)) for m in TRACKED_METRICS: - whole = result["drift"][m.key] + whole = result.drift[m.key] cells = " ".join(f"{_VERDICT_GLYPH[whole[a]]:>16}" for a in algos) lines.append(f" {m.key:<12} {cells}") return "\n".join(lines) @@ -2071,9 +2276,11 @@ def _agentic_warning() -> str: def render_natl(r: dict) -> str: + """Render a :func:`build_natl_result` dict labelled with ``token_count_source``.""" d = r["distribution"] return "\n".join( [ + f"token basis: {r['token_count_source']}", "=== NATL (agentic per-trajectory throughput) ===", f" trajectories: {r['n_trajectories']} (super-pass = " f"{r['superpass_trajectories']} trajectories -> {r['n_super_passes']} " @@ -2087,6 +2294,26 @@ def render_natl(r: dict) -> str: ) +def _warn_missing_usage(n_missing_usage: int | None) -> None: + if n_missing_usage: + print( + f"WARNING: {n_missing_usage} tracked samples returned no server token " + "usage; they add latency but no tokens, so NATL and system TPS are " + "understated for the affected trajectories and super-passes", + file=sys.stderr, + ) + + +def _warn_undecodable(n_undecodable: int) -> None: + if n_undecodable: + print( + f"WARNING: {n_undecodable} sample.complete lines could not be decoded " + "and were skipped; their samples are missing from the replay " + "(truncated or corrupt log?)", + file=sys.stderr, + ) + + def main(argv: Sequence[str] | None = None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) ap.add_argument("target", help="run directory OR path to events.jsonl") @@ -2096,7 +2323,18 @@ def main(argv: Sequence[str] | None = None) -> int: help="model name (auto-detected from config if omitted)", ) ap.add_argument( - "--tokenizer", default=None, help="HF tokenizer id/dir (overrides the registry)" + "--token-count-source", + type=TokenCountSource, + default=None, + metavar=TOKEN_COUNT_SOURCE_METAVAR, + help="token-count basis for TPOT/NATL (default: the run's own, from " + "performance/result_summary.json or config.yaml beside the event log). " + "server replays completion_tokens and needs no tokenizer", + ) + ap.add_argument( + "--tokenizer", + default=None, + help="HF tokenizer id/dir for client-basis runs (overrides the registry)", ) ap.add_argument( "--trust-remote-code", @@ -2147,20 +2385,35 @@ def main(argv: Sequence[str] | None = None) -> int: events, cfg_path = find_run_files(args.target) cfg = read_run_config(cfg_path) - - # tokenizer: explicit flag > model registry > config tokenizer path - model = args.model or cfg["model"] - reg = resolve_tokenizer(model) if model else None - if args.tokenizer: - tokenizer, trust = args.tokenizer, args.trust_remote_code - elif reg is not None: - tokenizer, trust = reg - elif model and "/" in model: - tokenizer, trust = model, args.trust_remote_code - else: - ap.error( - "could not resolve a tokenizer; pass --model (a known model) or --tokenizer" + try: + primary, basis = replay_basis( + Path(events).parent, override=args.token_count_source ) + except RunBasisError as exc: + ap.error(str(exc)) + if primary is None: + # The run's basis is unreadable; the one the user chose stands in. + primary = basis + print(f"[token basis: {basis.value}]", file=sys.stderr) + + # Server-basis replays read completion_tokens, so they resolve no tokenizer. + tokenizer: str | None = None + trust = False + if basis is TokenCountSource.CLIENT: + # tokenizer: explicit flag > model registry > config tokenizer path + model = args.model or cfg["model"] + reg = resolve_tokenizer(model) if model else None + if args.tokenizer: + tokenizer, trust = args.tokenizer, args.trust_remote_code + elif reg is not None: + tokenizer, trust = reg + elif model and "/" in model: + tokenizer, trust = model, args.trust_remote_code + else: + ap.error( + "could not resolve a tokenizer; pass --model (a known model) or " + "--tokenizer" + ) profile = ( PROFILES[args.profile] @@ -2178,12 +2431,22 @@ def main(argv: Sequence[str] | None = None) -> int: cov_bounds = args.cov_bounds or list(profile.cov_bounds) warmup_driver = args.warmup_driver or profile.warmup_driver flush = args.tokenize_batch_size or profile.tokenize_batch_size - count_tokens = _make_token_counter(tokenizer, trust) + count_tokens = ( + _make_token_counter(tokenizer, trust) if tokenizer is not None else None + ) if profile.metric == "natl": sp_traj = args.superpass_size or profile.superpass_size or 32 - pairs = build_trajectory_natl(events, count_tokens, flush) - natl = build_natl_result(pairs, sp_traj, tuple(cov_bounds)) + trajectory = build_trajectory_natl(events, count_tokens, flush, basis=basis) + natl: dict = { + "token_count_source": basis.value, + "n_missing_usage": trajectory.n_missing_usage, + } + if trajectory.n_undecodable: + natl["n_undecodable"] = trajectory.n_undecodable + natl.update(build_natl_result(trajectory.pairs, sp_traj, tuple(cov_bounds))) + _warn_missing_usage(trajectory.n_missing_usage) + _warn_undecodable(trajectory.n_undecodable) print(render_natl(natl)) if args.json_out: with open(args.json_out, "w") as fh: @@ -2212,12 +2475,15 @@ def main(argv: Sequence[str] | None = None) -> int: warmup_band=args.warmup_band, warmup_driver=warmup_driver, enforce_min_duration=args.enforce_min_duration, + basis=basis, + primary=primary, ) + _warn_missing_usage(result.n_missing_usage) + _warn_undecodable(result.n_undecodable) print(render_text(result, cov_bounds)) if args.json_out: - # msgspec encodes the Struct tree and validates its shape. with open(args.json_out, "wb") as fh: - fh.write(msgspec.json.format(msgspec.json.encode(result))) + fh.write(result.to_json()) print(f"\nwrote {args.json_out}", file=sys.stderr) return 0 diff --git a/tests/unit/commands/test_full_run_osl.py b/tests/unit/commands/test_full_run_osl.py index c826cabb2..a5adea0e6 100644 --- a/tests/unit/commands/test_full_run_osl.py +++ b/tests/unit/commands/test_full_run_osl.py @@ -11,17 +11,19 @@ from __future__ import annotations +import functools import json import logging from pathlib import Path +from typing import Any +import inference_endpoint.commands.benchmark.full_run_osl as full_run_osl_mod import msgspec import pytest from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( MessageInput, TextInput, ) -from inference_endpoint.commands.benchmark import full_run_osl as full_run_osl_module from inference_endpoint.commands.benchmark.full_run_osl import ( compute_full_run_osl, full_run_osl_for_report, @@ -33,12 +35,33 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import PhaseType, TextModelOutput +from inference_endpoint.core.types import TextModelOutput, TokenCountSource, basis_key from inference_endpoint.endpoint_client.cpu_affinity import UnsupportedPlatformError +from inference_endpoint.metrics.token_count_basis import ServerUsage + +from tests.test_helpers import LOGGED_SERVER_USAGE_CASES pytestmark = pytest.mark.unit _ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) +CLIENT, SERVER = TokenCountSource.CLIENT, TokenCountSource.SERVER + +# The statistic on the client basis, whose behavior these tests share with a +# tokenizer-only run. +_client_full_run_osl = functools.partial( + compute_full_run_osl, token_count_source=CLIENT +) + +# The block keys a client-primary run shares with a tokenizer-only run. +_CLIENT_PRIMARY_KEYS = ( + "output_sequence_lengths", + "n_turns_counted", + "n_empty", + "n_errors", + "n_undecodable", + "n_missing", + "partial", +) def _word_count(tok_input: MessageInput | TextInput) -> int: @@ -74,13 +97,22 @@ def _write_events(path: Path, records: list[EventRecord]) -> Path: return path -def _complete(turn: int, text: str, uuid: str | None = None) -> EventRecord: +def _complete( + turn: int, + text: str, + uuid: str | None = None, + *, + prompt_tokens: int | msgspec.UnsetType = msgspec.UNSET, + completion_tokens: int | msgspec.UnsetType = msgspec.UNSET, +) -> EventRecord: return EventRecord( event_type=SampleEventType.COMPLETE, sample_uuid=uuid or f"u{turn}", conversation_id="c0", turn=turn, data=TextModelOutput(output=text), + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, ) @@ -109,7 +141,7 @@ def test_counts_turns_completed_after_the_performance_window(tmp_path: Path) -> ], ) - result = compute_full_run_osl(events, _word_counts) + result = _client_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 4 @@ -141,8 +173,8 @@ def build(name: str, stop_after: int) -> Path: ], ) - early = compute_full_run_osl(build("early.jsonl", 2), _word_counts) - late = compute_full_run_osl(build("late.jsonl", 5), _word_counts) + early = _client_full_run_osl(build("early.jsonl", 2), _word_counts) + late = _client_full_run_osl(build("late.jsonl", 5), _word_counts) assert early is not None and late is not None assert early["n_turns_counted"] == late["n_turns_counted"] == 6 @@ -169,7 +201,7 @@ def test_excludes_completions_from_other_phases(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) + result = _client_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 2 @@ -192,7 +224,7 @@ def test_counts_empty_completions_separately(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts) + result = _client_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -253,10 +285,10 @@ def windowed_mean(stop_after: int) -> float: counted = [_word_count(TextInput(str(r.data))) for r in turns[:stop_after]] return sum(counted) / len(counted) - early = compute_full_run_osl( + early = _client_full_run_osl( build("e.jsonl", 2), _word_counts, performance_uuids=uuids ) - late = compute_full_run_osl( + late = _client_full_run_osl( build("l.jsonl", 5), _word_counts, performance_uuids=uuids ) @@ -282,7 +314,7 @@ def test_data_none_completion_counts_as_empty(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts) + result = _client_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -307,7 +339,7 @@ def flaky(tok_input: MessageInput | TextInput) -> int: raise ValueError("tokenizer blew up") return _word_count(tok_input) - result = compute_full_run_osl(events, _batched(flaky)) + result = _client_full_run_osl(events, _batched(flaky)) assert result is not None assert result["n_turns_counted"] == 1 @@ -336,8 +368,8 @@ def recording(items): batch_sizes.append(len(items)) return _word_counts(items) - batched = compute_full_run_osl(events, recording, batch_size=2) - single = compute_full_run_osl(events, _word_counts, batch_size=1024) + batched = _client_full_run_osl(events, recording, batch_size=2) + single = _client_full_run_osl(events, _word_counts, batch_size=1024) assert batch_sizes == [2, 2, 1] assert batched is not None and single is not None @@ -357,7 +389,7 @@ def first_batch_fails(items): raise RuntimeError("encoder crashed") return _word_counts(items) - result = compute_full_run_osl(events, first_batch_fails, batch_size=2) + result = _client_full_run_osl(events, first_batch_fails, batch_size=2) assert batch_sizes == [2, 2, 1] # the failed batch is not retried assert result is not None @@ -384,7 +416,7 @@ def wrong_count_on_first_batch(items): return _word_counts(items) with caplog.at_level(logging.WARNING): - result = compute_full_run_osl(events, wrong_count_on_first_batch, batch_size=2) + result = _client_full_run_osl(events, wrong_count_on_first_batch, batch_size=2) assert batch_sizes == [2, 2, 1] assert result is not None @@ -413,7 +445,7 @@ def test_all_empty_population_returns_a_visible_block(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts) + result = _client_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 0 @@ -433,7 +465,7 @@ def test_undecodable_complete_line_is_skipped( path.write_bytes(good + b"\n" + corrupt + b"\n") with caplog.at_level(logging.WARNING): - result = compute_full_run_osl(path, _word_counts) + result = _client_full_run_osl(path, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -469,7 +501,7 @@ def test_complete_lines_in_other_valid_layouts_are_counted(tmp_path: Path) -> No ] path.write_bytes(b"\n".join(lines) + b"\n") - result = compute_full_run_osl( + result = _client_full_run_osl( path, _word_counts, performance_uuids={"u1", "u3", "u5"} ) @@ -490,7 +522,7 @@ def test_unreadable_line_that_does_not_decode_is_charged_as_undecodable( path = tmp_path / "events.jsonl" path.write_bytes(_ENCODER.encode(_complete(1, "one")) + b"\nnot json\n") - result = compute_full_run_osl(path, _word_counts) + result = _client_full_run_osl(path, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -505,7 +537,7 @@ def test_all_undecodable_returns_visible_block(tmp_path: Path) -> None: path = tmp_path / "events.jsonl" path.write_bytes(corrupt + b"\n" + corrupt + b"\n") - result = compute_full_run_osl(path, _word_counts) + result = _client_full_run_osl(path, _word_counts) assert result is not None assert result["n_turns_counted"] == 0 @@ -513,13 +545,19 @@ def test_all_undecodable_returns_visible_block(tmp_path: Path) -> None: assert result["output_sequence_lengths"] == {} -def test_for_report_fails_closed_without_performance_map(tmp_path: Path) -> None: +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_for_report_fails_closed_without_performance_map( + tmp_path: Path, primary: TokenCountSource +) -> None: """full_run_osl_for_report skips (returns None) when no performance uuid set can be loaded — accuracy-phase turns must never count as performance OSL.""" (tmp_path / "events.jsonl").write_bytes(b"") # exists so the events gate passes # No sample_idx_map.json → load_performance_uuids None → fail closed, before # any tokenizer is constructed. - assert full_run_osl_for_report(tmp_path, "any-tokenizer") is None + assert ( + full_run_osl_for_report(tmp_path, "any-tokenizer", token_count_source=primary) + is None + ) class _SpyTokenizer: @@ -548,13 +586,7 @@ def count_batch_sync(self, items): def _perf_report(report_dir: Path) -> None: """A two-turn performance report that full_run_osl_for_report will accept.""" - _write_events( - report_dir / "events.jsonl", - [_complete(1, "a"), _complete(3, "b c")], - ) - (report_dir / "sample_idx_map.json").write_bytes( - msgspec.json.encode({PhaseType.PERFORMANCE.value: {"u1": 0, "u3": 1}}) - ) + _perf_run(report_dir, [_complete(1, "a"), _complete(3, "b c")]) @pytest.mark.parametrize("n_cpus", [4, 32, 352]) @@ -565,18 +597,18 @@ def test_for_report_caps_counting_threads( threads stops helping, and without tokenizer shards (n_workers=0).""" _perf_report(tmp_path) _SpyTokenizer.built = [] - monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _SpyTokenizer) monkeypatch.setattr( - full_run_osl_module, "get_all_online_cpus", lambda: list(range(n_cpus)) + full_run_osl_mod, "get_all_online_cpus", lambda: list(range(n_cpus)) ) - result = full_run_osl_for_report(tmp_path, "tok") + result = full_run_osl_for_report(tmp_path, "tok", token_count_source=CLIENT) assert result is not None and result["n_turns_counted"] == 2 assert _SpyTokenizer.built == [ { "tokenizer_name": "tok", - "live_workers": min(n_cpus, full_run_osl_module._MAX_COUNT_THREADS), + "live_workers": min(n_cpus, full_run_osl_mod._MAX_COUNT_THREADS), "n_workers": 0, } ] @@ -584,7 +616,7 @@ def test_for_report_caps_counting_threads( @pytest.mark.parametrize( ("cpu_count", "expected"), - [(64, full_run_osl_module._MAX_COUNT_THREADS), (8, 8), (None, 1)], + [(64, full_run_osl_mod._MAX_COUNT_THREADS), (8, 8), (None, 1)], ) def test_for_report_threads_without_cpu_affinity( tmp_path: Path, @@ -595,15 +627,17 @@ def test_for_report_threads_without_cpu_affinity( """Where the affinity API is unsupported, fall back to os.cpu_count().""" _perf_report(tmp_path) _SpyTokenizer.built = [] - monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _SpyTokenizer) def _unsupported() -> list[int]: raise UnsupportedPlatformError("no affinity here") - monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", _unsupported) - monkeypatch.setattr(full_run_osl_module.os, "cpu_count", lambda: cpu_count) + monkeypatch.setattr(full_run_osl_mod, "get_all_online_cpus", _unsupported) + monkeypatch.setattr(full_run_osl_mod.os, "cpu_count", lambda: cpu_count) - assert full_run_osl_for_report(tmp_path, "tok") is not None + assert ( + full_run_osl_for_report(tmp_path, "tok", token_count_source=CLIENT) is not None + ) assert _SpyTokenizer.built[0]["live_workers"] == expected @@ -613,20 +647,15 @@ def test_for_report_empty_cpu_list_counts_on_one_thread( """A probe that finds no CPUs must not ask for zero counting threads.""" _perf_report(tmp_path) _SpyTokenizer.built = [] - monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) - monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", lambda: []) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_mod, "get_all_online_cpus", lambda: []) - assert full_run_osl_for_report(tmp_path, "tok") is not None + assert ( + full_run_osl_for_report(tmp_path, "tok", token_count_source=CLIENT) is not None + ) assert _SpyTokenizer.built[0]["live_workers"] == 1 -class _FailingTokenizer(_SpyTokenizer): - """A tokenizer that cannot be constructed.""" - - def __init__(self, *args: object, **kwargs: object) -> None: - raise RuntimeError("tokenizer could not be loaded") - - @pytest.mark.parametrize("failure", ["cpu_probe", "tokenizer"]) def test_for_report_failure_drops_the_block_with_a_warning( tmp_path: Path, @@ -642,22 +671,29 @@ def test_for_report_failure_drops_the_block_with_a_warning( def _broken_probe() -> list[int]: raise RuntimeError("probe broke") - monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) - monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", _broken_probe) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_mod, "get_all_online_cpus", _broken_probe) else: - monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _FailingTokenizer) - monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", lambda: [0, 1]) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _BrokenTokenizer) + monkeypatch.setattr(full_run_osl_mod, "get_all_online_cpus", lambda: [0, 1]) with caplog.at_level(logging.WARNING): - assert full_run_osl_for_report(tmp_path, "tok") is None + assert ( + full_run_osl_for_report(tmp_path, "tok", token_count_source=CLIENT) is None + ) assert any("Full-run OSL skipped" in r.getMessage() for r in caplog.records) -def test_for_report_skips_without_tokenizer_or_events(tmp_path: Path) -> None: - """No tokenizer, or a missing events.jsonl, short-circuits to None.""" - assert full_run_osl_for_report(tmp_path, None) is None # no tokenizer - assert full_run_osl_for_report(tmp_path, "any-tokenizer") is None # no events.jsonl +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_for_report_skips_without_events( + tmp_path: Path, primary: TokenCountSource +) -> None: + """A missing events.jsonl short-circuits to None on either basis.""" + assert ( + full_run_osl_for_report(tmp_path, "any-tokenizer", token_count_source=primary) + is None + ) def test_missing_turn_is_detected_via_uuid_reconciliation(tmp_path: Path) -> None: @@ -672,7 +708,7 @@ def test_missing_turn_is_detected_via_uuid_reconciliation(tmp_path: Path) -> Non ], ) - result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) + result = _client_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 1 @@ -698,7 +734,7 @@ def flaky(tok_input: MessageInput | TextInput) -> int: raise ValueError("tokenizer blew up") return _word_count(tok_input) - result = compute_full_run_osl( + result = _client_full_run_osl( events, _batched(flaky), performance_uuids={"u1", "u3"} ) @@ -720,7 +756,7 @@ def test_clean_run_is_not_partial(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) + result = _client_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 2 @@ -741,9 +777,255 @@ def test_empty_turns_do_not_make_block_partial(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) + result = _client_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_empty"] == 1 assert result["n_missing"] == 0 assert result["partial"] is False + + +def _perf_run(tmp_path: Path, records: list[EventRecord]) -> None: + _write_events(tmp_path / "events.jsonl", records) + perf = {r.sample_uuid: i for i, r in enumerate(records)} + (tmp_path / "sample_idx_map.json").write_bytes( + msgspec.json.encode({"performance": perf}) + ) + + +class _BrokenTokenizer: + """BatchTokenizer stand-in whose tokenizer cannot be loaded.""" + + def __init__(self, *a: object, **k: object) -> None: + raise OSError("no such tokenizer") + + +def _with_usage(turn: int, text: str, completion_tokens: int) -> EventRecord: + """A turn reporting ``completion_tokens`` beside a prompt count of 1.""" + return _complete(turn, text, prompt_tokens=1, completion_tokens=completion_tokens) + + +# One turn of each kind. Client counts 1/2/3 tokens and one empty turn; the +# server counts 2/3 tokens, one true zero, and one turn without usage. +_MIXED_TURNS = [ + _with_usage(1, "one", 2), + _with_usage(3, "two words", 3), + _complete(5, "a b c"), # no usage + _with_usage(7, "", 0), # true zero on both bases +] +_MIXED_UUIDS = {"u1", "u3", "u5", "u7"} +# Per basis: (total tokens, turns counted, empty turns) over _MIXED_TURNS. +_MIXED_COUNTS = {CLIENT: (6, 3, 1), SERVER: (5, 2, 1)} + + +def _strip_usage(record: EventRecord) -> EventRecord: + return msgspec.structs.replace( + record, prompt_tokens=msgspec.UNSET, completion_tokens=msgspec.UNSET + ) + + +def test_full_run_unsuffixed_block_identical_with_and_without_usage_keys( + tmp_path: Path, +) -> None: + """On the client basis, server usage only adds keys: the shared keys hold + the same values whether or not the log carries usage.""" + turns = [ + *_MIXED_TURNS, + _with_usage(9, "x y", 0), # 0 beside output: missing usage + EventRecord(event_type=SampleEventType.COMPLETE, sample_uuid="u11"), + ] + uuids = {r.sample_uuid for r in turns} | {"u13"} # u13 never completed + with_usage = _write_events(tmp_path / "with.jsonl", turns) + without_usage = _write_events( + tmp_path / "without.jsonl", [_strip_usage(r) for r in turns] + ) + + with_block = _client_full_run_osl(with_usage, _word_counts, performance_uuids=uuids) + without_block = _client_full_run_osl( + without_usage, _word_counts, performance_uuids=uuids + ) + + assert with_block is not None and without_block is not None + assert {k: with_block[k] for k in _CLIENT_PRIMARY_KEYS} == { + k: without_block[k] for k in _CLIENT_PRIMARY_KEYS + } + assert with_block["n_turns_counted"] == 4 and with_block["partial"] is True + + +def test_full_run_missing_usage_never_sets_top_level_partial_for_client_primary( + tmp_path: Path, +) -> None: + """Turns without usable usage leave the client-primary block whole: they + are counted, and only the server series is marked partial.""" + events = _write_events(tmp_path / "events.jsonl", _MIXED_TURNS) + without_usage = _write_events( + tmp_path / "bare.jsonl", [_strip_usage(r) for r in _MIXED_TURNS] + ) + + some = _client_full_run_osl(events, _word_counts, performance_uuids=_MIXED_UUIDS) + none = _client_full_run_osl( + without_usage, _word_counts, performance_uuids=_MIXED_UUIDS + ) + + assert some is not None and none is not None + assert (some["partial"], some["n_missing_usage"]) == (False, 1) + assert some["output_sequence_lengths_server"]["partial"] is True + # Nothing counted on the server basis: no series at all. The empty turn + # without usage is missing usage too, as in the live count. + assert (none["partial"], none["n_missing_usage"]) == (False, 4) + assert none["output_sequence_lengths_server"] is None + assert none["n_turns_counted_server"] == 0 + + +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_full_run_other_basis_keys_follow_rule( + tmp_path: Path, primary: TokenCountSource +) -> None: + """The primary basis keeps the base names, the other basis appends its + value; ``n_errors`` exists only on the client basis and the remaining + counters are basis-independent.""" + events = _write_events(tmp_path / "events.jsonl", _MIXED_TURNS) + + block = compute_full_run_osl( + events, + _word_counts, + performance_uuids=_MIXED_UUIDS, + token_count_source=primary, + ) + + def key(base: str, basis: TokenCountSource) -> str: + return basis_key(base, basis, primary=primary) + + assert block is not None + assert set(block) == { + *( + key(base, basis) + for basis in TokenCountSource + for base in ("output_sequence_lengths", "n_turns_counted", "n_empty") + ), + key("n_errors", CLIENT), + "n_undecodable", + "n_missing", + "n_missing_usage", + "partial", + } + for basis, (total, counted, empty) in _MIXED_COUNTS.items(): + assert block[key("output_sequence_lengths", basis)]["total"] == total + assert block[key("n_turns_counted", basis)] == counted + assert block[key("n_empty", basis)] == empty + assert (block[key("n_errors", CLIENT)], block["n_missing_usage"]) == (0, 1) + # The server shortfall sets the top-level flag only when the server basis + # is primary; otherwise it marks just the server series. + assert block["partial"] is (primary is SERVER) + server_series = block[key("output_sequence_lengths", SERVER)] + assert ("partial" in server_series) is (primary is CLIENT) + assert "partial" not in block[key("output_sequence_lengths", CLIENT)] + + +def test_full_run_server_primary_missing_usage_sets_partial(tmp_path: Path) -> None: + """On the server basis a turn without usable usage, including a 0 beside + output, makes the block partial; a 0 beside empty output is empty.""" + events = _write_events( + tmp_path / "events.jsonl", + [ + _with_usage(1, "one two", 0), + _with_usage(3, "", 0), + _with_usage(5, "x", 4), + ], + ) + + block = compute_full_run_osl(events, None, token_count_source=SERVER) + + assert block is not None + assert block["output_sequence_lengths"]["total"] == 4 + assert block["n_turns_counted"] == 1 + assert (block["n_empty"], block["n_missing_usage"]) == (1, 1) + assert block["partial"] is True + # No tokenizer: the client keys are present and count nothing. + assert block["output_sequence_lengths_client"] is None + assert (block["n_turns_counted_client"], block["n_errors_client"]) == (0, 0) + + +def test_full_run_client_primary_without_tokenizer_is_none( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The client basis needs a tokenizer: without one, or with one that + cannot load, there is no block. The server basis needs none.""" + _perf_run(tmp_path, [_with_usage(1, "one", 4)]) + + assert full_run_osl_for_report(tmp_path, None, token_count_source=CLIENT) is None + server = full_run_osl_for_report(tmp_path, None, token_count_source=SERVER) + assert server is not None and server["output_sequence_lengths"]["total"] == 4 + + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", _BrokenTokenizer) + assert full_run_osl_for_report(tmp_path, "tok", token_count_source=CLIENT) is None + with pytest.raises(ValueError, match="count_batch"): + compute_full_run_osl(tmp_path / "events.jsonl", None, token_count_source=CLIENT) + + +_LOGGED_USAGE_CASES = [ + pytest.param(*case[1:], id=case[0]) for case in LOGGED_SERVER_USAGE_CASES +] + + +@pytest.mark.parametrize("primary", list(TokenCountSource)) +@pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), _LOGGED_USAGE_CASES +) +def test_full_run_missing_usage_matches_rule( + tmp_path: Path, + primary: TokenCountSource, + prompt: Any, + completion: Any, + output: TextModelOutput, + expected: ServerUsage, +) -> None: + """The full-run statistic counts each turn's usage by the shared rule.""" + events = _write_events( + tmp_path / "events.jsonl", + [ + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=output, + prompt_tokens=prompt, + completion_tokens=completion, + ) + ], + ) + + block = compute_full_run_osl( + events, _word_counts, performance_uuids={"u1"}, token_count_source=primary + ) + + def server(base: str) -> str: + return basis_key(base, SERVER, primary=primary) + + assert block is not None + assert block["n_missing_usage"] == int(expected is ServerUsage.MISSING) + assert block[server("n_turns_counted")] == int(expected is ServerUsage.USABLE) + assert block[server("n_empty")] == int(expected is ServerUsage.TRUE_ZERO) + + +@pytest.mark.parametrize( + ("tok_cls", "client_total"), [(_SpyTokenizer, 2), (_BrokenTokenizer, None)] +) +def test_for_report_server_primary_counts_client_beside_it( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + tok_cls: type, + client_total: int | None, +) -> None: + """With a tokenizer the server-primary block carries the client count; a + tokenizer that cannot load leaves the client basis uncounted, with a + warning.""" + _perf_run(tmp_path, [_with_usage(1, "one two", 3)]) + monkeypatch.setattr(full_run_osl_mod, "BatchTokenizer", tok_cls) + + with caplog.at_level(logging.WARNING): + block = full_run_osl_for_report(tmp_path, "tok", token_count_source=SERVER) + + assert block is not None and block["output_sequence_lengths"]["total"] == 3 + assert (block["output_sequence_lengths_client"] or {}).get("total") == client_total + assert ("no such tokenizer" in caplog.text) == (client_total is None) diff --git a/tests/unit/commands/test_score_accuracy.py b/tests/unit/commands/test_score_accuracy.py index 835de4a95..ca33fa1bd 100644 --- a/tests/unit/commands/test_score_accuracy.py +++ b/tests/unit/commands/test_score_accuracy.py @@ -20,15 +20,18 @@ from __future__ import annotations import sys +from pathlib import Path from types import SimpleNamespace import msgspec.json import pytest from inference_endpoint.commands.benchmark.accuracy import ( AccuracyConfiguration, + PhaseServerOsl, _accuracy_uuid_bound, _phase_osl_stats, _phase_response_counts, + _phase_server_osl_stats, score_accuracy, write_accuracy_results, ) @@ -39,12 +42,20 @@ ModelParams, ScorerMethod, ) +from inference_endpoint.core.record import EventRecord, EventType, SampleEventType +from inference_endpoint.core.types import TextModelOutput, TokenCountSource, basis_key +from inference_endpoint.evaluation import scoring as evaluation_scoring +from inference_endpoint.evaluation.scoring import Scorer +from inference_endpoint.metrics.token_count_basis import ServerUsage + +from tests.test_helpers import LOGGED_SERVER_USAGE_CASES # The module that defines score_accuracy (commands.benchmark.accuracy), resolved # via __module__ so the tests monkeypatch its module-level load_reference_backend # and score_accuracy picks up the patched one — without importing the module under # both the ``from ... import`` and ``import ...`` styles. scoring_mod = sys.modules[score_accuracy.__module__] +CLIENT, SERVER = TokenCountSource.CLIENT, TokenCountSource.SERVER class _FakeDataset: @@ -97,70 +108,45 @@ def score_breakdown(self): class _FakeOSLScorer(_FakeScorer): - """Scorer exposing the get_raw_outputs()/sample_index_map the OSL path reads.""" + """Scorer exposing the iter_complete_outputs()/sample_index_map the OSL + path reads.""" + + # Two completed samples: outputs of 2 and 4 "tokens" (words). + SAMPLE_INDEX_MAP = {"u1": 0, "u2": 1} + # (uuid, response text, server completion_tokens when the usage is usable) + OUTPUTS: list[tuple[str, str, int | None]] = [ + ("u1", "a b", None), + ("u2", "a b c d", None), + ("other", "not in this phase", None), + ] def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # Two completed samples: outputs of 2 and 4 "tokens" (words). - self.sample_index_map = {"u1": 0, "u2": 1} - - def get_raw_outputs(self, wanted_uuids=None): - import pandas as pd - - return pd.DataFrame( - [ - {"sample_uuid": "u1", "output": "a b"}, - {"sample_uuid": "u2", "output": "a b c d"}, - {"sample_uuid": "other", "output": "not in this phase"}, - ] - ) + self.sample_index_map = dict(self.SAMPLE_INDEX_MAP) + + def iter_complete_outputs(self, uuids): + yield from self.OUTPUTS class _EmptyOSLScorer(_FakeOSLScorer): """Every completion empty (all requests failed): OSL yields None.""" - def get_raw_outputs(self, wanted_uuids=None): - import pandas as pd - - return pd.DataFrame( - [ - {"sample_uuid": "u1", "output": ""}, - {"sample_uuid": "u2", "output": ""}, - ] - ) + OUTPUTS = [("u1", "", None), ("u2", "", None)] class _MixedOSLScorer(_FakeOSLScorer): """One scored, one empty (COMPLETE, blank), one missing (no COMPLETE).""" - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.sample_index_map = {"u1": 0, "u2": 1, "u3": 2} - - def get_raw_outputs(self, wanted_uuids=None): - import pandas as pd - - return pd.DataFrame( - [ - {"sample_uuid": "u1", "output": "x y"}, # scored - {"sample_uuid": "u2", "output": ""}, # empty - # u3 has no COMPLETE row -> missing - ] - ) + SAMPLE_INDEX_MAP = {"u1": 0, "u2": 1, "u3": 2} + # u3 has no COMPLETE event -> missing + OUTPUTS = [("u1", "x y", None), ("u2", "", None)] class _AllMissingScorer(_FakeOSLScorer): - """No COMPLETE rows for this phase: get_raw_outputs returns an empty but - columned frame, as the real scorer does when the bound filters everything.""" + """No COMPLETE event for this phase, as when the bound filters everything.""" - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.sample_index_map = {"u1": 0, "u2": 1, "u3": 2} - - def get_raw_outputs(self, wanted_uuids=None): - import pandas as pd - - return pd.DataFrame([], columns=["sample_uuid", "output"]) + SAMPLE_INDEX_MAP = {"u1": 0, "u2": 1, "u3": 2} + OUTPUTS = [] def _cfg(name: str, n: int, score: float, tmp, scorer=_FakeScorer, repeats: int = 1): @@ -188,17 +174,74 @@ def _ctx( tokenizer_name=None, report_dir=None, test_mode: config_schema.TestMode = config_schema.TestMode.ACC, + token_count_source: TokenCountSource = TokenCountSource.CLIENT, ): - # tokenizer_name None => OSL is skipped (fake scorers have no get_raw_outputs). + # tokenizer_name None => client OSL is skipped (and _FakeScorer has no + # iter_complete_outputs, so its entries carry no OSL block at all). # report_dir None => the uuid bound falls back to an unbounded read (no map). return SimpleNamespace( eval_configs=cfgs, tokenizer_name=tokenizer_name, report_dir=report_dir, test_mode=test_mode, + token_count_source=token_count_source, + ) + + +class _ServerOSLScorer(_FakeOSLScorer): + """u1 scored with usage, u2 scored without usage, u3 empty with usage.""" + + SAMPLE_INDEX_MAP = {"u1": 0, "u2": 1, "u3": 2} + OUTPUTS = [("u1", "a b", 3), ("u2", "a b c d", None), ("u3", "", 1), ("x", "y", 9)] + + +class _EventsScorer(_FakeScorer): + """Reads its phase's outputs from events.jsonl through the real reader.""" + + _iter_complete_records = Scorer._iter_complete_records + iter_complete_outputs = Scorer.iter_complete_outputs + + def __init__(self, name, dataset, report_dir, **kwargs): + super().__init__(name, dataset, report_dir, **kwargs) + self.report_dir = report_dir + idx_map = msgspec.json.decode((report_dir / "sample_idx_map.json").read_bytes()) + self.sample_index_map = idx_map[name] + + +_EVENT_ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + + +def _complete(uuid: str, text: str, prompt=msgspec.UNSET, completion=msgspec.UNSET): + return EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid=uuid, + data=TextModelOutput(output=text), + prompt_tokens=prompt, + completion_tokens=completion, ) +def _write_run(report_dir: Path, phases: dict[str, list[EventRecord]]) -> list[bytes]: + """Write sample_idx_map.json and events.jsonl, an ISSUED line before each + COMPLETE; return the COMPLETE lines.""" + idx_map = { + name: {r.sample_uuid: i for i, r in enumerate(records)} + for name, records in phases.items() + } + (report_dir / "sample_idx_map.json").write_bytes(msgspec.json.encode(idx_map)) + complete_lines = [] + with (report_dir / "events.jsonl").open("wb") as f: + for records in phases.values(): + for record in records: + issued = EventRecord( + event_type=SampleEventType.ISSUED, sample_uuid=record.sample_uuid + ) + f.write(_EVENT_ENCODER.encode(issued) + b"\n") + complete_lines.append(_EVENT_ENCODER.encode(record) + b"\n") + f.write(complete_lines[-1]) + return complete_lines + + class _WordBackend: """Stand-in tokenizers backend: one token per whitespace-delimited word.""" @@ -206,6 +249,25 @@ def encode_batch(self, texts, add_special_tokens=False): return [SimpleNamespace(ids=t.split()) for t in texts] +class _RaisingBackend: + """Stand-in tokenizers backend that fails on every batch.""" + + def encode_batch(self, texts, add_special_tokens=False): + raise RuntimeError("tokenizer blew up") + + +class _CountingDecoder: + """Wraps the events decoder and records every line it decodes.""" + + def __init__(self, inner): + self._inner = inner + self.lines: list[bytes] = [] + + def decode(self, line): + self.lines.append(line) + return self._inner.decode(line) + + _RESULT = SimpleNamespace(perf_results=[], phase_results=[]) @@ -326,7 +388,8 @@ def test_perf_mode_skips_only_external_scoring(self, tmp_path): def test_no_osl_without_tokenizer(self, tmp_path): # tokenizer_name None (the default) => no output_sequence_lengths attached, - # and the OSL path is never entered (fake scorers have no get_raw_outputs). + # and the OSL path stops at the read (_FakeScorer has no + # iter_complete_outputs). cfg = _cfg("aime25::gptoss", 30, 0.8, tmp_path) entry = _by_name(score_accuracy(_ctx([cfg]), _RESULT))["aime25::gptoss"] assert "output_sequence_lengths" not in entry @@ -362,14 +425,14 @@ def test_osl_skipped_for_performance_entry(self, tmp_path, monkeypatch): ] assert "output_sequence_lengths" not in entry - def test_osl_dropped_when_get_raw_outputs_raises(self, tmp_path, monkeypatch): - """A read/tokenize failure only drops OSL — scoring still succeeds.""" + def test_osl_dropped_when_events_read_raises(self, tmp_path, monkeypatch): + """A read failure only drops OSL — scoring still succeeds.""" monkeypatch.setattr( scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) class _RaisingScorer(_FakeOSLScorer): - def get_raw_outputs(self, wanted_uuids=None): + def iter_complete_outputs(self, uuids): raise RuntimeError("events.jsonl unreadable") cfg = _cfg("ds", 1, 0.8, tmp_path, scorer=_RaisingScorer) @@ -480,9 +543,9 @@ def _boom(name): assert entry["response_counts"]["scored"] == 2 assert entry["score"] == pytest.approx(0.8) - def test_uuid_bound_passed_to_get_raw_outputs(self, tmp_path, monkeypatch): + def test_uuid_bound_passed_to_events_read(self, tmp_path, monkeypatch): """End-to-end: score_accuracy computes the accuracy uuid bound from - sample_idx_map.json and passes it into get_raw_outputs.""" + sample_idx_map.json and passes it into iter_complete_outputs.""" monkeypatch.setattr( scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) @@ -492,9 +555,9 @@ def test_uuid_bound_passed_to_get_raw_outputs(self, tmp_path, monkeypatch): captured: dict = {} class _CaptureScorer(_FakeOSLScorer): - def get_raw_outputs(self, wanted_uuids=None): - captured["wanted"] = wanted_uuids - return super().get_raw_outputs(wanted_uuids) + def iter_complete_outputs(self, uuids): + captured["wanted"] = uuids + return super().iter_complete_outputs(uuids) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_CaptureScorer) score_accuracy(_ctx([cfg], tokenizer_name="fake", report_dir=tmp_path), _RESULT) @@ -507,15 +570,206 @@ def test_empty_uuid_bound_stays_bounded(self, tmp_path): captured: dict = {} class _CaptureScorer(_FakeOSLScorer): - def get_raw_outputs(self, wanted_uuids=None): - captured["wanted"] = wanted_uuids - return super().get_raw_outputs(wanted_uuids) + def iter_complete_outputs(self, uuids): + captured["wanted"] = uuids + return super().iter_complete_outputs(uuids) cfg = _cfg("swe_bench", 1, 0.8, tmp_path, scorer=_CaptureScorer) score_accuracy(_ctx([cfg], report_dir=tmp_path), _RESULT) assert captured["wanted"] == set() + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_accuracy_keys_follow_rule(self, tmp_path, monkeypatch, primary): + """Each basis's block and sample count take the base name on the + primary basis and the suffixed name on the other; missing usage is + basis-independent and marks the server block partial only when the + server basis is the other one.""" + monkeypatch.setattr( + scoring_mod, "load_reference_backend", lambda name: _WordBackend() + ) + cfg = _cfg("ds", 3, 0.8, tmp_path, scorer=_ServerOSLScorer) + + entry = _by_name( + score_accuracy( + _ctx([cfg], tokenizer_name="fake", token_count_source=primary), + _RESULT, + ) + )["ds"] + + def key(base: str, basis: TokenCountSource) -> str: + return basis_key(base, basis, primary=primary) + + assert entry["response_counts"] == { + "issued": 3, + "scored": 2, + "empty": 1, + "missing": 0, + } + # Per basis: (total tokens, samples counted) over the scored samples. + for basis, (total, n_samples) in {CLIENT: (6, 2), SERVER: (3, 1)}.items(): + assert entry[key("output_sequence_lengths", basis)]["total"] == total + assert entry[key("n_osl_samples", basis)] == n_samples + assert entry["n_missing_usage"] == 1 + server_block = entry[key("output_sequence_lengths", SERVER)] + assert ("partial" in server_block) is (primary is CLIENT) + assert "partial" not in entry[key("output_sequence_lengths", CLIENT)] + assert isinstance(entry["osl_tokenize_s"], float) + assert not [k for k in entry if k.endswith(f"_{primary.value}")] + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_server_osl_needs_no_tokenizer(self, tmp_path, primary): + cfg = _cfg("ds", 3, 0.8, tmp_path, scorer=_ServerOSLScorer) + + entry = _by_name( + score_accuracy(_ctx([cfg], token_count_source=primary), _RESULT) + )["ds"] + + def key(base: str, basis: TokenCountSource) -> str: + return basis_key(base, basis, primary=primary) + + assert entry[key("output_sequence_lengths", SERVER)]["avg"] == 3.0 + assert key("output_sequence_lengths", CLIENT) not in entry + assert "osl_tokenize_s" not in entry + assert ( + entry[key("n_osl_samples", SERVER)], + entry[key("n_osl_samples", CLIENT)], + entry["n_missing_usage"], + ) == (1, 0, 1) + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + def test_accuracy_tokenizer_failure_keeps_server_block( + self, tmp_path, monkeypatch, caplog, primary + ): + monkeypatch.setattr( + scoring_mod, "load_reference_backend", lambda name: _RaisingBackend() + ) + cfg = _cfg("ds", 3, 0.8, tmp_path, scorer=_ServerOSLScorer) + + with caplog.at_level("WARNING"): + entry = _by_name( + score_accuracy( + _ctx([cfg], tokenizer_name="fake", token_count_source=primary), + _RESULT, + ) + )["ds"] + + def key(base: str, basis: TokenCountSource) -> str: + return basis_key(base, basis, primary=primary) + + assert entry[key("output_sequence_lengths", SERVER)]["total"] == 3 + assert key("output_sequence_lengths", CLIENT) not in entry + assert "osl_tokenize_s" not in entry + assert entry["response_counts"]["scored"] == 2 + assert ( + entry[key("n_osl_samples", SERVER)], + entry[key("n_osl_samples", CLIENT)], + entry["n_missing_usage"], + ) == (1, 0, 1) + assert "tokenizer blew up" in caplog.text + + def test_accuracy_single_pass_over_events(self, tmp_path, monkeypatch): + """Two accuracy phases share one read of events.jsonl, which decodes + each COMPLETE line once and no other line.""" + monkeypatch.setattr( + scoring_mod, "load_reference_backend", lambda name: _WordBackend() + ) + complete_lines = _write_run( + tmp_path, + { + "ds1": [_complete("u1", "a b", 1, 2), _complete("u2", "c")], + "ds2": [_complete("u3", "d e f", 1, 3)], + }, + ) + decoder = _CountingDecoder(evaluation_scoring._EVENT_DECODER) + monkeypatch.setattr(evaluation_scoring, "_EVENT_DECODER", decoder) + cfgs = [ + _cfg("ds1", 2, 0.5, tmp_path, scorer=_EventsScorer), + _cfg("ds2", 1, 0.5, tmp_path, scorer=_EventsScorer), + ] + + by = _by_name( + score_accuracy( + _ctx(cfgs, tokenizer_name="fake", report_dir=tmp_path), _RESULT + ) + ) + + assert sorted(decoder.lines) == sorted(complete_lines) + assert by["ds1"]["output_sequence_lengths"]["total"] == 3 + assert by["ds2"]["output_sequence_lengths"]["total"] == 3 + assert (by["ds1"]["n_missing_usage"], by["ds2"]["n_missing_usage"]) == (1, 0) + + @pytest.mark.parametrize("primary", list(TokenCountSource)) + @pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), + [pytest.param(*case[1:], id=case[0]) for case in LOGGED_SERVER_USAGE_CASES], + ) + def test_accuracy_missing_usage_matches_rule( + self, tmp_path, primary, prompt, completion, output, expected + ): + """Accuracy OSL counts each scored sample's usage by the shared rule.""" + record = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=output, + prompt_tokens=prompt, + completion_tokens=completion, + ) + _write_run(tmp_path, {"ds": [record]}) + cfg = _cfg("ds", 1, 0.8, tmp_path, scorer=_EventsScorer) + + entry = _by_name( + score_accuracy( + _ctx([cfg], report_dir=tmp_path, token_count_source=primary), _RESULT + ) + )["ds"] + + server_key = basis_key("output_sequence_lengths", SERVER, primary=primary) + assert entry["n_missing_usage"] == int(expected is ServerUsage.MISSING) + assert (server_key in entry) is (expected is ServerUsage.USABLE) + + def test_accuracy_unsuffixed_entry_identical_with_and_without_usage_keys( + self, tmp_path, monkeypatch + ): + """On the client basis, server usage only adds keys: the shared keys + hold the same values whether or not the log carries usage.""" + monkeypatch.setattr( + scoring_mod, "load_reference_backend", lambda name: _WordBackend() + ) + with_usage = [ + _complete("u1", "a b", 1, 2), + _complete("u2", "c d e"), + _complete("u3", "", 1, 0), + _complete("u4", "f", 1, 0), + ] + without_usage = [_complete(r.sample_uuid, str(r.data)) for r in with_usage] + entries = [] + for name, records in (("with", with_usage), ("without", without_usage)): + run = tmp_path / name + run.mkdir() + _write_run(run, {"ds": records}) + cfg = _cfg("ds", 4, 0.5, run, scorer=_EventsScorer) + ctx = _ctx([cfg], tokenizer_name="fake", report_dir=run) + entries.append(_by_name(score_accuracy(ctx, _RESULT))["ds"]) + + shared = ("response_counts", "output_sequence_lengths") + assert {k: entries[0][k] for k in shared} == {k: entries[1][k] for k in shared} + assert entries[0]["output_sequence_lengths"]["total"] == 6 + + +@pytest.mark.unit +def test_phase_server_osl_counts_scored_samples_with_usable_usage(): + """Each scored sample counts its usable completion_tokens or is missing + usage; an empty sample is outside the scored population.""" + phase = _phase_server_osl_stats( + ["a", "b", "c", "d"], + {"a": "x y", "b": "z", "c": "", "d": "w"}, + {"b": 4, "c": 2}, + ) + assert isinstance(phase, PhaseServerOsl) + assert phase.block is not None and phase.block["total"] == 4 + assert (phase.n_counted, phase.n_missing_usage) == (1, 2) + @pytest.mark.unit class TestPhaseOslStats: @@ -697,7 +951,7 @@ def test_write_accuracy_results_keeps_stable_average_accuracy_key(tmp_path): {"dataset_name": "aime", "score": 0.9, "unit_samples": 30}, {"dataset_name": "lcb", "score": 0.6, "unit_samples": 1000}, ] - write_accuracy_results(tmp_path, scores) + write_accuracy_results(tmp_path, scores, token_count_source=TokenCountSource.SERVER) data = msgspec.json.decode( (tmp_path / "accuracy" / "accuracy_results.json").read_bytes() ) @@ -705,3 +959,15 @@ def test_write_accuracy_results_keeps_stable_average_accuracy_key(tmp_path): assert "samples_weighted_average_accuracy" not in data expected = (0.9 * 30 + 0.6 * 1000) / (30 + 1000) assert data["average_accuracy"] == pytest.approx(expected) + + +@pytest.mark.unit +@pytest.mark.parametrize("source", list(TokenCountSource)) +def test_write_accuracy_results_labels_the_token_basis(tmp_path, source): + scores = [{"dataset_name": "ds", "score": 0.5, "unit_samples": 2}] + write_accuracy_results(tmp_path, scores, token_count_source=source) + data = msgspec.json.decode( + (tmp_path / "accuracy" / "accuracy_results.json").read_bytes() + ) + assert data["token_count_source"] == source.value + assert data["accuracy_scores"] == scores diff --git a/tests/unit/evaluation/test_scoring.py b/tests/unit/evaluation/test_scoring.py index 6c1a89e8d..81356d07b 100644 --- a/tests/unit/evaluation/test_scoring.py +++ b/tests/unit/evaluation/test_scoring.py @@ -1616,6 +1616,93 @@ def test_empty_filter_returns_columned_frame(self, tmp_path): assert frame.empty assert list(frame.columns) == ["sample_uuid", "output"] + def test_iter_complete_outputs_reads_text_and_usage(self, tmp_path): + out = TextModelOutput(output="a b") + self._write_events( + tmp_path, + [ + EventRecord(event_type=SampleEventType.ISSUED, sample_uuid="u1"), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=out, + prompt_tokens=2, + completion_tokens=3, + ), + EventRecord( + event_type=SampleEventType.COMPLETE, sample_uuid="u2", data=out + ), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u3", + prompt_tokens=2, + completion_tokens=0, + ), + ], + ) + scorer = self._scorer(tmp_path) + + assert list(scorer.iter_complete_outputs(None)) == [ + ("u1", "a b", 3), + ("u2", "a b", None), + ("u3", "", None), + ] + bounded = scorer.iter_complete_outputs({"u1", "u3"}) + assert [uuid for uuid, *_ in bounded] == ["u1", "u3"] + # BFCL unpacks _iter_complete as (sample_uuid, data); the shape is kept. + assert [u for u, _ in scorer._iter_complete()] == ["u1", "u2", "u3"] + + def test_events_reader_decodes_only_complete_lines(self, tmp_path): + """Only COMPLETE lines are decoded, so a corrupt line of another event + type cannot fail scoring.""" + complete = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=TextModelOutput(output="a"), + ) + encoder = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + (tmp_path / "events.jsonl").write_bytes( + b'{"event_type":"sample.issued",<<>>}\n' + + encoder.encode(complete) + + b"\n" + ) + + frame = self._scorer(tmp_path).get_raw_outputs() + + assert frame.to_dict("records") == [{"sample_uuid": "u1", "output": "a"}] + + def test_non_compact_log_is_decoded(self, tmp_path): + """A line outside the compact layout is decoded whole, not dropped.""" + complete = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=TextModelOutput(output="a"), + ) + encoder = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + spaced = json.dumps(json.loads(encoder.encode(complete))).encode() + (tmp_path / "events.jsonl").write_bytes(b"\n" + spaced + b"\n") + scorer = self._scorer(tmp_path) + + assert scorer.get_raw_outputs().to_dict("records") == [ + {"sample_uuid": "u1", "output": "a"} + ] + assert list(scorer.iter_complete_outputs(None)) == [("u1", "a", None)] + + def test_compact_log_after_blank_lines_is_read(self, tmp_path): + complete = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + data=TextModelOutput(output="a"), + ) + encoder = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + (tmp_path / "events.jsonl").write_bytes( + b"\n\n" + encoder.encode(complete) + b"\n" + ) + + frame = self._scorer(tmp_path).get_raw_outputs() + + assert frame.to_dict("records") == [{"sample_uuid": "u1", "output": "a"}] + @pytest.mark.unit class TestScoreDenominator: diff --git a/tests/unit/metrics/test_archived_run_compat.py b/tests/unit/metrics/test_archived_run_compat.py index 8392fe0ac..c1e119f84 100644 --- a/tests/unit/metrics/test_archived_run_compat.py +++ b/tests/unit/metrics/test_archived_run_compat.py @@ -19,6 +19,7 @@ import pytest from inference_endpoint.commands.benchmark.full_run_osl import full_run_osl_for_report from inference_endpoint.core.record import EventRecord, EventType +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.metrics import steady_state_diagnostics as ssd from inference_endpoint.metrics.report import Report @@ -93,6 +94,8 @@ def test_report_from_archived_snapshot_reproduces_archived_summary(): _json(_RUN_DIR / "metrics/final_snapshot.json"), run_config=summary["run_config"], use_legacy_loadgen_qps_metrics=False, # concurrency run + # Archived runs counted every token with the client tokenizer. + token_count_source=TokenCountSource.CLIENT, ) rebuilt = json.loads(report.to_json()) for field in _SUMMARY_FIELDS: @@ -116,7 +119,9 @@ def test_report_from_archived_snapshot_reproduces_archived_summary(): def test_full_run_osl_recomputes_archived_block(): - block = full_run_osl_for_report(_RUN_DIR, str(_CHAR_TOKENIZER_DIR)) + block = full_run_osl_for_report( + _RUN_DIR, str(_CHAR_TOKENIZER_DIR), token_count_source=TokenCountSource.CLIENT + ) archived = _json(_SUMMARY)["output_sequence_lengths_full_run"] assert block is not None for key in ( @@ -138,7 +143,9 @@ def test_full_run_osl_recomputes_archived_block(): def test_steady_state_replay_rebuilds_archived_series(): size = ssd.superpass_size_from_events(str(_EVENTS)) assert size == 5 - series = ssd.build_super_pass_series(str(_EVENTS), size, _char_counts) + series = ssd.build_super_pass_series( + str(_EVENTS), size, _char_counts, basis=TokenCountSource.CLIENT + ).super_passes assert [ (sp.n_issued, len(sp.ttft_ns), len(sp.tpot_ns), len(sp.latency_ns)) for sp in series @@ -166,7 +173,11 @@ def test_steady_state_cli_replays_archived_verdict(tmp_path): def test_early_stopping_script_reproduces_archived_estimates(tmp_path): es = _load_es_script() - series = es.compute_series(str(_EVENTS), count_tokens=_char_counts) + series = es.compute_series( + str(_EVENTS), + count_tokens=_char_counts, + basis=TokenCountSource.CLIENT, + ) assert {k: len(v) for k, v in series.items()} == { "ttft_ns": 12, "tpot_ns": 12, diff --git a/tests/unit/metrics/test_steady_state_diagnostics.py b/tests/unit/metrics/test_steady_state_diagnostics.py index 904555d60..00b278b2b 100644 --- a/tests/unit/metrics/test_steady_state_diagnostics.py +++ b/tests/unit/metrics/test_steady_state_diagnostics.py @@ -3,22 +3,33 @@ """Tests for ``inference_endpoint.metrics.steady_state_diagnostics``. -Pins everything the detector owns: the plain-JSON event parse (wire shapes -referenced from core/record.py + core/types.py), super-pass bucketing, TTFT/TPOT -reconstruction via an injected token counter, the trend-detection algorithms, and -the CoV table. +Pins everything the detector owns: the event-log parse (wire shapes referenced +from core/record.py + core/types.py), super-pass bucketing, TTFT/TPOT +reconstruction via an injected token counter or the server usage, the +trend-detection algorithms, and the CoV table. """ import json +import msgspec import pytest +from inference_endpoint.core.record import EventRecord, EventType, SampleEventType +from inference_endpoint.core.types import TextModelOutput, TokenCountSource, basis_key from inference_endpoint.metrics import steady_state_diagnostics as mod +from inference_endpoint.metrics.token_count_basis import ServerUsage + +from tests.test_helpers import LOGGED_SERVER_USAGE_CASES pytestmark = pytest.mark.unit +def _line(record: dict) -> str: + """One event-log line, compact like the JSONL writer's.""" + return json.dumps(record, separators=(",", ":")) + + def _ev(event_type, ts, uuid="", data=None): - return json.dumps( + return _line( { "event_type": event_type, "timestamp_ns": ts, @@ -33,31 +44,70 @@ def _words(texts): return [len(t.split()) for t in texts] +def _tokenized_series(path, superpass_size): + """Super-passes replayed on the client basis with the whitespace tokenizer.""" + return mod.build_super_pass_series( + path, superpass_size, _words, basis=TokenCountSource.CLIENT + ).super_passes + + +def _tokenized_run(path, **kwargs): + """``run`` on a client-basis run, one sample per super-pass.""" + return mod.run( + path, + superpass_size=1, + count_tokens=_words, + basis=TokenCountSource.CLIENT, + primary=TokenCountSource.CLIENT, + **kwargs, + ) + + # --------------------------------------------------------------------------- # -# text_after_first_chunk — port of TextModelOutput.text_after_first_chunk over -# the parsed JSON array [tag, output, reasoning?, tool_calls?]. +# text_after_first_chunk — TextModelOutput.text_after_first_chunk without the +# tool calls; full_output_text — the whole output, tool calls as logged. # --------------------------------------------------------------------------- # def test_text_after_first_chunk_streaming_output_drops_first(): - data = ["TextModelOutput", ["hello ", "world ", "again"]] - assert mod.text_after_first_chunk(data) == "world again" + output = TextModelOutput(output=("hello ", "world ", "again")) + assert mod.text_after_first_chunk(output) == "world again" def test_text_after_first_chunk_reasoning_first_keeps_all_output(): # The first delta contains only reasoning, so content has an empty first entry. - data = ["TextModelOutput", ["", "out1 ", "out2"], ["think1 ", "think2 "]] - assert mod.text_after_first_chunk(data) == "think2 out1 out2" + output = TextModelOutput( + output=("", "out1 ", "out2"), reasoning=("think1 ", "think2 ") + ) + assert mod.text_after_first_chunk(output) == "think2 out1 out2" def test_text_after_first_chunk_joined_output_follows_streamed_reasoning(): - data = ["TextModelOutput", "out1 out2", ["think1 ", "think2 "]] - assert mod.text_after_first_chunk(data) == "think2 out1 out2" + output = TextModelOutput(output="out1 out2", reasoning=("think1 ", "think2 ")) + assert mod.text_after_first_chunk(output) == "think2 out1 out2" def test_text_after_first_chunk_non_streaming_str_has_no_first_chunk(): - data = ["TextModelOutput", "the whole answer"] - assert mod.text_after_first_chunk(data) == "" + output = TextModelOutput(output="the whole answer") + assert mod.text_after_first_chunk(output) == "" + + +def test_text_after_first_chunk_leaves_out_tool_calls(): + tool_calls = (({"index": 0, "function": {"name": "f"}},),) + output = TextModelOutput(output=("a ", "b"), tool_calls=tool_calls) + assert mod.text_after_first_chunk(output) == "b" + + +def test_full_output_text_is_reasoning_output_then_logged_tool_calls(): + tool_calls = [[{"index": 0, "function": {"name": "f", "arguments": "{}"}}]] + output = msgspec.json.decode( + msgspec.json.encode( + ["TextModelOutput", ["a ", "b"], ["r1 ", "r2 "], tool_calls] + ), + type=TextModelOutput, + ) + assert mod.full_output_text(output) == "r1 r2 a b" + json.dumps(tool_calls) + assert mod.full_output_text(TextModelOutput(output="whole")) == "whole" # --------------------------------------------------------------------------- # @@ -88,7 +138,7 @@ def test_build_super_pass_series_buckets_by_issue_order(tmp_path): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=2, count_tokens=_words) + series = _tokenized_series(path, 2) assert len(series) == 2 sp0, sp1 = series @@ -111,13 +161,13 @@ def test_build_super_pass_series_records_e2e_latency(tmp_path): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert list(series[0].latency_ns) == [2000.0] # complete(3000) - issued(1000) def test_warm_turn_ttft_excludes_cold_first_turn(tmp_path): def evt(et, ts, uuid, turn, data=None): - return json.dumps( + return _line( { "event_type": et, "timestamp_ns": ts, @@ -138,7 +188,7 @@ def evt(et, ts, uuid, turn, data=None): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert sorted(series[0].ttft_ns) == [50.0, 200.0] # all turns assert list(series[0].ttft_warm_ns) == [50.0] # turn 1 discarded @@ -153,7 +203,7 @@ def test_build_super_pass_series_ignores_events_outside_tracking(tmp_path): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert len(series) == 1 assert series[0].n_issued == 1 @@ -292,13 +342,13 @@ def _synthetic_events(tmp_path, n, ttft_ns_fn): def test_run_result_structure(tmp_path): # ttft climbs steadily -> an upward drift the scan should surface. path = _synthetic_events(tmp_path, n=8, ttft_ns_fn=lambda i: 100 + 20 * i) - result = mod.run(path, superpass_size=1, count_tokens=_words, warmup=1) - assert result["steady_state"].n_super_passes == 8 - assert result["steady_state"].warmup == 1 - assert result["drift"]["ttft_p50"]["mk_hamed_rao"] == "up" + result = _tokenized_run(path, warmup=1) + assert result.steady_state.n_super_passes == 8 + assert result.steady_state.warmup == 1 + assert result.drift["ttft_p50"]["mk_hamed_rao"] == "up" # CoV cells carry a pass/fail per bound and a gate flag. Admissibility gates on # TPOT only, so tpot_p50 is gated while ttft_p50 is now diagnostic. - cov = result["steady_state"].cov + cov = result.steady_state.cov assert cov["tpot_p50"].gated is True cell = cov["ttft_p50"] assert cell.gated is False @@ -308,9 +358,9 @@ def test_run_result_structure(tmp_path): def test_text_after_first_chunk_empty_reasoning_str_output_is_first_chunk(): - # reasoning=[] is falsy -> str output IS the sole first chunk -> excluded. - data = ["TextModelOutput", "the whole answer", []] - assert mod.text_after_first_chunk(data) == "" + # reasoning=() is falsy -> str output IS the sole first chunk -> excluded. + output = TextModelOutput(output="the whole answer", reasoning=()) + assert mod.text_after_first_chunk(output) == "" def test_retried_sample_counts_ttft_once(tmp_path): @@ -324,7 +374,7 @@ def test_retried_sample_counts_ttft_once(tmp_path): _ev("session.stop_performance_tracking", 5000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert list(series[0].ttft_ns) == [500.0] @@ -340,10 +390,63 @@ def test_valid_json_line_missing_timestamp_is_skipped(tmp_path): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert sum(sp.n_issued for sp in series) == 1 # only the well-formed sample +def test_truncated_complete_line_is_skipped(tmp_path): + # The last line of a killed run can be a cut-off COMPLETE: both replays skip it. + cut = _conv_ev("sample.complete", 3100, "B", "c2", 1, ["TextModelOutput", ["x"]]) + lines = [ + _ev("session.start_performance_tracking", 0), + _conv_ev("sample.issued", 1000, "A", "c1", 1), + _conv_ev("sample.recv_first", 1500, "A", "c1", 1), + _conv_ev( + "sample.complete", 3000, "A", "c1", 1, ["TextModelOutput", ["a ", "b"]] + ), + _conv_ev("sample.issued", 1100, "B", "c2", 1), + cut[: len(cut) // 2], + ] + path = _write_events(tmp_path, lines) + (sp,) = _tokenized_series(path, 4) + assert (sp.n_issued, list(sp.latency_ns)) == (2, [2000.0]) + natl = mod.build_trajectory_natl(path, _words, basis=TokenCountSource.CLIENT) + assert [t for t, _ in natl.pairs] == [3000] + + +def test_replays_read_a_non_compact_complete_line(tmp_path): + # A COMPLETE line outside the compact layout is decoded, not dropped: both + # replays read a spaced line as they read the compact one. + complete = _conv_ev( + "sample.complete", 3000, "A", "c1", 1, ["TextModelOutput", ["a ", "b"]] + ) + replays = [] + for name, line in ( + ("compact", complete), + ("spaced", json.dumps(json.loads(complete))), + ): + run_dir = tmp_path / name + run_dir.mkdir() + path = _write_events( + run_dir, + [ + _ev("session.start_performance_tracking", 0), + _conv_ev("sample.issued", 1000, "A", "c1", 1), + _conv_ev("sample.recv_first", 1500, "A", "c1", 1), + line, + _ev("session.stop_performance_tracking", 4000), + ], + ) + (sp,) = _tokenized_series(path, 1) + natl = mod.build_trajectory_natl(path, _words, basis=TokenCountSource.CLIENT) + replays.append(((sp.n_issued, list(sp.latency_ns)), natl.pairs)) + compact, spaced = replays + assert spaced == compact + (n_issued, latency_ns), natl_pairs = compact + assert (n_issued, latency_ns) == (1, [2000.0]) + assert [t for t, _ in natl_pairs] == [3000] + + def test_cov_pass_row_insufficient_window_is_inconclusive(): # Fewer than 2 points -> CoV undefined -> cells inconclusive (None), not True. assert mod.cov_pass_row([3.0], bounds=(0.03, 0.05)) == {0.03: None, 0.05: None} @@ -352,7 +455,7 @@ def test_cov_pass_row_insufficient_window_is_inconclusive(): def test_run_rejects_negative_warmup(tmp_path): path = _synthetic_events(tmp_path, n=6, ttft_ns_fn=lambda i: 100.0) with pytest.raises(ValueError): - mod.run(path, superpass_size=1, count_tokens=_words, warmup=-1) + _tokenized_run(path, warmup=-1) def test_super_pass_tracks_timestamps(tmp_path): @@ -367,7 +470,7 @@ def test_super_pass_tracks_timestamps(tmp_path): _ev("session.stop_performance_tracking", 4000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=4, count_tokens=_words) + series = _tokenized_series(path, 4) assert series[0].first_issue_ns == 1000 # earliest issue assert series[0].last_event_ns == 3000 # latest event (A's complete) @@ -384,7 +487,7 @@ def test_window_elapsed_and_pooling(tmp_path): _ev("session.stop_performance_tracking", 20000), ] path = _write_events(tmp_path, lines) - series = mod.build_super_pass_series(path, superpass_size=1, count_tokens=_words) + series = _tokenized_series(path, 1) assert len(series) == 2 # pooled ttft over both super-passes assert sorted(mod.pooled(series, 0, 2, "ttft_ns")) == [500.0, 500.0] @@ -669,8 +772,8 @@ def test_adaptive_warmup_capped_at_max_frac(): def test_run_auto_warmup_resolves_to_int(tmp_path): path = _synthetic_events(tmp_path, n=8, ttft_ns_fn=lambda i: 100.0) - result = mod.run(path, superpass_size=1, count_tokens=_words, warmup="auto") - assert isinstance(result["steady_state"].warmup, int) + result = _tokenized_run(path, warmup="auto") + assert isinstance(result.steady_state.warmup, int) def test_steady_state_flags_global_drift_after_first_plateau(): @@ -850,14 +953,13 @@ def test_soft_path_reports_first_plateau_even_when_short(): def test_run_result_has_steady_state_block(tmp_path): path = _synthetic_events(tmp_path, n=8, ttft_ns_fn=lambda i: 100.0) - result = mod.run(path, superpass_size=1, count_tokens=_words, warmup=1) - assert "steady_state" in result - assert result["steady_state"].anomaly is not None + result = _tokenized_run(path, warmup=1) + assert result.steady_state.anomaly is not None def test_render_text_has_section_headers(tmp_path): path = _synthetic_events(tmp_path, n=8, ttft_ns_fn=lambda i: 100 + 20 * i) - result = mod.run(path, superpass_size=1, count_tokens=_words, warmup=1) + result = _tokenized_run(path, warmup=1) text = mod.render_text(result, cov_bounds=[0.03, 0.05, 0.08]) assert "CoV" in text assert "drift (whole-run" in text @@ -960,7 +1062,7 @@ def phase(kind, turns, ts): def _conv_ev(et, ts, uuid, conv, turn=None, data=None): - return json.dumps( + return _line( { "event_type": et, "timestamp_ns": ts, @@ -1007,7 +1109,9 @@ def test_build_trajectory_natl(tmp_path): _ev("session.stop_performance_tracking", 9_000_000_000), ] path = _write_events(tmp_path, lines) - pairs = mod.build_trajectory_natl(path, _words) + trajectory = mod.build_trajectory_natl(path, _words, basis=TokenCountSource.CLIENT) + assert trajectory.n_missing_usage is None # the client basis never reads usage + pairs = trajectory.pairs # sorted by completion: c2 completes at 2e9, c1 at 3e9 assert [round(n, 3) for _, n in pairs] == [2.0, 2.5] assert [t for t, _ in pairs] == [2_000_000_000, 3_000_000_000] @@ -1029,3 +1133,404 @@ def test_build_natl_result_flat_and_varied(): ) assert r2["across_cov"] > 0.15 assert r2["found"] is False + + +# --------------------------------------------------------------------------- # +# Token-count basis: server usage from COMPLETE lines, the run-basis reader +# --------------------------------------------------------------------------- # + +_ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + + +def _logged(record: EventRecord) -> str: + """``record`` as the JSONL writer logs it.""" + return _ENCODER.encode(record).decode() + + +def _with_usage(line, completion_tokens, prompt_tokens=1): + rec = json.loads(line) + rec["prompt_tokens"] = prompt_tokens + rec["completion_tokens"] = completion_tokens + return _line(rec) + + +def _summary(run_dir, primary): + """The run's ``result_summary.json``, recording ``primary`` as its basis.""" + (run_dir / "performance").mkdir() + (run_dir / "performance" / "result_summary.json").write_text( + json.dumps({"token_count_source": primary.value}) + ) + + +def test_server_replay_uses_completion_tokens(tmp_path): + path = _write_events( + tmp_path, + [ + _ev("session.start_performance_tracking", 0), + _ev("sample.issued", 1000, "A"), + _ev("sample.recv_first", 1500, "A"), + _with_usage( + _ev("sample.complete", 3000, "A", ["TextModelOutput", ["a ", "b"]]), 4 + ), + _ev("sample.issued", 1100, "B"), + _ev("sample.recv_first", 1300, "B"), + # completion_tokens <= 1: no TPOT + _with_usage( + _ev("sample.complete", 2100, "B", ["TextModelOutput", ["x "]]), 1 + ), + _ev("sample.issued", 1200, "C"), + _ev("sample.recv_first", 1400, "C"), + # no usage: no TPOT + _ev("sample.complete", 2000, "C", ["TextModelOutput", ["p ", "q"]]), + _ev("session.stop_performance_tracking", 4000), + ], + ) + replay = mod.build_super_pass_series(path, 4, None, basis=TokenCountSource.SERVER) + (sp,) = replay.super_passes + assert (list(sp.tpot_ns), list(sp.osl), sp.out_tokens, len(sp.latency_ns)) == ( + [500.0], + [3.0], + 3, + 3, + ) + assert replay.n_missing_usage == 1 + with pytest.raises(ValueError, match="count_tokens"): + mod.build_super_pass_series(path, 4, None, basis=TokenCountSource.CLIENT) + + +def test_build_trajectory_natl_server_basis(tmp_path): + out = ["TextModelOutput", ["a"]] + path = _write_events( + tmp_path, + [ + _ev("session.start_performance_tracking", 0), + _conv_ev("sample.issued", 0, "a", "c1", 1), + _with_usage( + _conv_ev("sample.complete", 1_000_000_000, "a", "c1", 1, out), 3 + ), + _conv_ev("sample.issued", 2_000_000_000, "b", "c1", 2), + _with_usage( + _conv_ev("sample.complete", 3_000_000_000, "b", "c1", 2, out), 5 + ), + _conv_ev("sample.issued", 0, "d", "c2", 1), + # no usage: no tokens + _conv_ev("sample.complete", 2_000_000_000, "d", "c2", 1, out), + _ev("session.stop_performance_tracking", 9_000_000_000), + ], + ) + assert mod.build_trajectory_natl( + path, None, basis=TokenCountSource.SERVER + ) == mod.ReplayTrajectory( + pairs=[(3_000_000_000, 4.0)], n_missing_usage=1, n_undecodable=0 + ) + + +@pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), + [pytest.param(*case[1:], id=case[0]) for case in LOGGED_SERVER_USAGE_CASES], +) +def test_replay_missing_usage_matches_live_rule( + tmp_path, prompt, completion, output, expected +): + """Both replays read usage by ``classify_server_usage``, as the aggregator does.""" + complete = EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=3000, + sample_uuid="A", + conversation_id="c1", + data=output, + prompt_tokens=prompt, + completion_tokens=completion, + ) + path = _write_events( + tmp_path, + [ + _ev("session.start_performance_tracking", 0), + _conv_ev("sample.issued", 1000, "A", "c1"), + _conv_ev("sample.recv_first", 1500, "A", "c1"), + _logged(complete), + _ev("session.stop_performance_tracking", 4000), + ], + ) + series = mod.build_super_pass_series(path, 1, None, basis=TokenCountSource.SERVER) + natl = mod.build_trajectory_natl(path, None, basis=TokenCountSource.SERVER) + n_missing = int(expected is ServerUsage.MISSING) + assert (series.n_missing_usage, natl.n_missing_usage) == (n_missing, n_missing) + usable = expected is ServerUsage.USABLE + (sp,) = series.super_passes + assert list(sp.osl) == ([completion - 1] if usable else []) + latency_s = (3000 - 1000) / 1e9 + assert natl.pairs == ([(3000, completion / latency_s)] if usable else []) + + +def test_main_server_run_needs_no_tokenizer(tmp_path, monkeypatch): + lines = [_ev("session.start_performance_tracking", 0)] + for i in range(8): + t, u = 1000 * (i + 1), f"s{i}" + lines += [ + _ev("sample.issued", t, u), + _ev("sample.recv_first", t + 100, u), + _with_usage( + _ev("sample.complete", t + 1100, u, ["TextModelOutput", ["a ", "b"]]), + 11, + ), + ] + _write_events(tmp_path, [*lines, _ev("session.stop_performance_tracking", 20_000)]) + _summary(tmp_path, TokenCountSource.SERVER) + monkeypatch.setattr( + mod, "_make_token_counter", lambda *a, **k: pytest.fail("tokenizer loaded") + ) + argv = [ + str(tmp_path), + "--superpass-size", + "2", + "--warmup", + "0", + "--no-min-duration", + ] + assert mod.main(argv) == 0 + with pytest.raises(SystemExit): # client override: no model/tokenizer resolvable + mod.main([*argv, "--token-count-source", TokenCountSource.CLIENT.value]) + + +def _missing_usage_events(tmp_path): + """Server run: A has usage; B (streamed) and C (non-streaming) do not; D errored.""" + return _write_events( + tmp_path, + [ + _ev("sample.issued", 10, "W"), # untracked: never counted + _ev("sample.complete", 20, "W", ["TextModelOutput", ["w"]]), + _ev("session.start_performance_tracking", 0), + _ev("sample.issued", 1000, "A"), + _ev("sample.recv_first", 1500, "A"), + _with_usage( + _ev("sample.complete", 3000, "A", ["TextModelOutput", ["a ", "b"]]), 4 + ), + _ev("sample.issued", 1100, "B"), + _ev("sample.recv_first", 1300, "B"), + _ev("sample.complete", 2100, "B", ["TextModelOutput", ["x ", "y"]]), + _ev("sample.issued", 1200, "C"), + _ev("sample.complete", 2000, "C", ["TextModelOutput", "whole answer"]), + _ev("sample.issued", 1300, "D"), + _ev("sample.complete", 1900, "D", None), + _ev("session.stop_performance_tracking", 4000), + ], + ) + + +def test_run_labels_the_basis_and_counts_missing_usage(tmp_path): + path = _missing_usage_events(tmp_path) + server = mod.run( + path, + superpass_size=1, + count_tokens=None, + warmup=0, + basis=TokenCountSource.SERVER, + primary=TokenCountSource.SERVER, + ) + assert (server.basis, server.n_missing_usage) == ( + TokenCountSource.SERVER, + 2, + ) + replay = mod.build_super_pass_series(path, 1, None, basis=TokenCountSource.SERVER) + tpot = [v for sp in replay.super_passes for v in sp.tpot_ns] + assert tpot == [500.0] # missing usage is never filled in from a tokenizer + text = mod.render_text(server, cov_bounds=list(mod.DEFAULT_COV_BOUNDS)) + assert text.splitlines()[0] == "token basis: server" + client = _tokenized_run(path, warmup=0) + # The client replay never reads usage, so it reports no missing-usage count. + assert (client.basis, client.n_missing_usage) == ( + TokenCountSource.CLIENT, + None, + ) + tokenized = mod.build_super_pass_series( + path, 1, _words, basis=TokenCountSource.CLIENT + ) + assert tokenized.n_missing_usage is None + + +def test_main_json_is_labelled_with_the_basis(tmp_path, capsys): + _missing_usage_events(tmp_path) + _summary(tmp_path, TokenCountSource.SERVER) + out = tmp_path / "ss.json" + argv = [str(tmp_path), "--superpass-size", "2", "--warmup", "0"] + assert mod.main([*argv, "--no-min-duration", "--json", str(out)]) == 0 + blob = json.loads(out.read_text()) + # Top level, beside the steady_state block (which stays summary-comparable). + assert (blob["token_count_source"], blob["n_missing_usage"]) == ("server", 2) + assert "token_count_source" not in blob["steady_state"] + assert "n_undecodable" not in blob + captured = capsys.readouterr() + assert captured.out.splitlines()[0] == "token basis: server" + assert "2 tracked samples returned no server token usage" in captured.err + assert "NATL and system TPS are understated" in captured.err + assert "could not be decoded" not in captured.err + + +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_main_json_names_the_verdict_by_basis(tmp_path, monkeypatch, capsys, primary): + """The verdict keeps its name on the run's basis; the other basis suffixes it.""" + _missing_usage_events(tmp_path) + _summary(tmp_path, primary) + monkeypatch.setattr(mod, "_make_token_counter", lambda *a, **k: _words) + argv = [str(tmp_path), "--tokenizer", "words", "--superpass-size", "2"] + argv += ["--warmup", "0", "--no-min-duration"] + for override in (None, *TokenCountSource): + replayed = primary if override is None else override + flag = [] if override is None else ["--token-count-source", override.value] + out = tmp_path / f"{override}.json" + assert mod.main([*argv, *flag, "--json", str(out)]) == 0 + blob = json.loads(out.read_text()) + assert blob["token_count_source"] == replayed.value + assert basis_key("steady_state", replayed, primary=primary) in blob + assert basis_key("steady_state", replayed.other, primary=primary) not in blob + capsys.readouterr() # drain + + +def _usage_run(run_dir, *, usage: bool) -> None: + """A client-basis run whose COMPLETE lines carry server usage iff ``usage``.""" + lines = [_ev("session.start_performance_tracking", 0)] + for i in range(8): + t, u = 1000 * (i + 1), f"s{i}" + complete = _ev( + "sample.complete", t + 1100 + 10 * i, u, ["TextModelOutput", ["a ", "b c"]] + ) + lines += [ + _ev("sample.issued", t, u), + _ev("sample.recv_first", t + 100, u), + _with_usage(complete, 7) if usage else complete, + ] + run_dir.mkdir() + _write_events(run_dir, [*lines, _ev("session.stop_performance_tracking", 20_000)]) + _summary(run_dir, TokenCountSource.CLIENT) + + +def test_client_primary_replay_identical_with_and_without_usage_keys( + tmp_path, monkeypatch, capsys +): + monkeypatch.setattr(mod, "_make_token_counter", lambda *a, **k: _words) + blobs = [] + for usage in (True, False): + run_dir = tmp_path / f"usage_{usage}" + _usage_run(run_dir, usage=usage) + out = run_dir / "ss.json" + argv = [str(run_dir), "--tokenizer", "words", "--superpass-size", "2"] + argv += ["--warmup", "0", "--no-min-duration", "--json", str(out)] + assert mod.main(argv) == 0 + blobs.append(out.read_bytes()) + assert blobs[0] == blobs[1] + assert json.loads(blobs[0])["steady_state"]["n_super_passes"] == 4 + capsys.readouterr() # drain + + +def test_steady_state_cli_unreadable_summary_exits_with_hint( + tmp_path, monkeypatch, capsys +): + _missing_usage_events(tmp_path) + (tmp_path / "performance").mkdir() + (tmp_path / "performance" / "result_summary.json").write_text("{not json") + monkeypatch.setattr(mod, "_make_token_counter", lambda *a, **k: _words) + argv = [str(tmp_path), "--tokenizer", "words", "--superpass-size", "2"] + argv += ["--warmup", "0", "--no-min-duration"] + with pytest.raises(SystemExit) as exc: + mod.main(argv) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "cannot read the token-count basis" in err + assert "pass --token-count-source" in err + # The hint holds: choosing the basis replays the run. + assert mod.main([*argv, "--token-count-source", TokenCountSource.SERVER.value]) == 0 + + +def test_main_natl_json_is_labelled_with_the_basis(tmp_path, capsys): + out_text = ["TextModelOutput", ["a"]] + _write_events( + tmp_path, + [ + _ev("session.start_performance_tracking", 0), + _conv_ev("sample.issued", 0, "a", "c1", 1), + _with_usage( + _conv_ev("sample.complete", 1_000_000_000, "a", "c1", 1, out_text), 3 + ), + _conv_ev("sample.issued", 0, "d", "c2", 1), + _conv_ev("sample.complete", 2_000_000_000, "d", "c2", 1, out_text), + _ev("session.stop_performance_tracking", 9_000_000_000), + ], + ) + _summary(tmp_path, TokenCountSource.SERVER) + out = tmp_path / "natl.json" + assert mod.main([str(tmp_path), "--profile", "agentic", "--json", str(out)]) == 0 + blob = json.loads(out.read_text()) + assert (blob["token_count_source"], blob["n_missing_usage"]) == ("server", 1) + assert blob["n_trajectories"] == 1 + captured = capsys.readouterr() + assert captured.out.splitlines()[0] == "token basis: server" + assert "1 tracked samples returned no server token usage" in captured.err + + +def _undecodable(line): + """``line`` with a field of the wrong type: still a COMPLETE line, but one + the typed decode rejects.""" + rec = json.loads(line) + rec["finish_reason"] = 5 + return _line(rec) + + +def _undecodable_events(tmp_path): + """Four tracked samples with usage; D's COMPLETE does not decode.""" + lines = [_ev("session.start_performance_tracking", 0)] + for i, uuid in enumerate("ABCD"): + t = 1000 * (i + 1) + complete = _with_usage( + _conv_ev( + "sample.complete", + t + 600, + uuid, + f"c{i}", + 1, + ["TextModelOutput", ["a ", "b"]], + ), + 4, + ) + lines += [ + _conv_ev("sample.issued", t, uuid, f"c{i}", 1), + _conv_ev("sample.recv_first", t + 100, uuid, f"c{i}", 1), + _undecodable(complete) if uuid == "D" else complete, + ] + return _write_events( + tmp_path, [*lines, _ev("session.stop_performance_tracking", 9000)] + ) + + +@pytest.mark.parametrize("basis", list(TokenCountSource)) +def test_replays_count_undecodable_complete_lines(tmp_path, basis): + path = _undecodable_events(tmp_path) + count_tokens = None if basis is TokenCountSource.SERVER else _words + series = mod.build_super_pass_series(path, 4, count_tokens, basis=basis) + natl = mod.build_trajectory_natl(path, count_tokens, basis=basis) + assert (series.n_undecodable, natl.n_undecodable) == (1, 1) + (sp,) = series.super_passes + assert (len(sp.latency_ns), len(natl.pairs)) == (3, 3) + + +@pytest.mark.parametrize("profile", ["concurrency", "agentic"]) +def test_main_warns_on_undecodable_lines_and_records_them(tmp_path, capsys, profile): + _undecodable_events(tmp_path) + _summary(tmp_path, TokenCountSource.SERVER) + out = tmp_path / "ss.json" + argv = [str(tmp_path), "--profile", profile, "--superpass-size", "2"] + argv += ["--warmup", "0", "--no-min-duration", "--json", str(out)] + assert mod.main(argv) == 0 + assert json.loads(out.read_text())["n_undecodable"] == 1 + err = capsys.readouterr().err + assert "WARNING: 1 sample.complete lines could not be decoded" in err + + +def test_main_override_needs_no_archived_basis_note(tmp_path, capsys): + _missing_usage_events(tmp_path) + argv = [str(tmp_path), "--superpass-size", "2", "--warmup", "0"] + argv += ["--no-min-duration", "--token-count-source", "server"] + assert mod.main(argv) == 0 + err = capsys.readouterr().err + assert "no token-count basis recorded" not in err + assert "[token basis: server]" in err diff --git a/tests/unit/scripts/test_early_stopping_estimate_from_events.py b/tests/unit/scripts/test_early_stopping_estimate_from_events.py index d04d556e0..9ad0152f7 100644 --- a/tests/unit/scripts/test_early_stopping_estimate_from_events.py +++ b/tests/unit/scripts/test_early_stopping_estimate_from_events.py @@ -14,8 +14,13 @@ import json from pathlib import Path +import msgspec import pytest -from inference_endpoint.core.types import TextModelOutput +from inference_endpoint.core.record import EventRecord, EventType, SampleEventType +from inference_endpoint.core.types import TextModelOutput, TokenCountSource +from inference_endpoint.metrics.token_count_basis import ServerUsage + +from tests.test_helpers import LOGGED_SERVER_USAGE_CASES pytestmark = pytest.mark.unit @@ -33,8 +38,13 @@ def _load_script(): mod = _load_script() +def _line(record: dict) -> str: + """One event-log line, compact like the JSONL writer's.""" + return json.dumps(record, separators=(",", ":")) + + def _ev(event_type, ts, uuid="", data=None): - return json.dumps( + return _line( { "event_type": event_type, "timestamp_ns": ts, @@ -59,6 +69,10 @@ def _write(tmp_path, lines): return p +_CLIENT = TokenCountSource.CLIENT +_SERVER = TokenCountSource.SERVER + + def test_extract_tpot_text_gates_on_decoded_type(): # The script owns only the gate; chunk semantics are the product's # text_after_first_chunk (tested with TextModelOutput itself). @@ -99,12 +113,13 @@ def test_tracking_window_gating_and_metric_deltas(tmp_path): _ev("sample.complete", 700, "c", _tmo("nonstream full")), # no recv_first ], ) - s = mod.compute_series(events, count_tokens=_ws_counts) + s = mod.compute_series(events, count_tokens=_ws_counts, basis=_CLIENT) assert s["ttft_ns"] == [50] # a only; w untracked, c never streamed assert sorted(s["sample_latency_ns"]) == [300, 400] # a and c assert s["tpot_ns"] == [(400 - 150) / 2] # "one twothree" -> 2 ws tokens # without a tokenizer the tool degrades to ttft/latency only - assert mod.compute_series(events, count_tokens=None)["tpot_ns"] == [] + tokenless = mod.compute_series(events, count_tokens=None, basis=_CLIENT) + assert tokenless["tpot_ns"] == [] def test_imperfect_log_robustness(tmp_path, capsys): @@ -133,7 +148,7 @@ def test_imperfect_log_robustness(tmp_path, capsys): _ev("session.stop_performance_tracking", 700), ], ) - s = mod.compute_series(events, count_tokens=_ws_counts) + s = mod.compute_series(events, count_tokens=_ws_counts, basis=_CLIENT) # aggregator parity: BOTH recv_first events fire ttft (150-100, 250-200) assert sorted(s["ttft_ns"]) == [50, 50, 60] # a: complete - REFRESHED issue ts (duplicate-ISSUED parity) @@ -258,3 +273,287 @@ def test_main_cli_contract(tmp_path, capsys): with pytest.raises(SystemExit): mod.main([str(events), "--confidence", "1.0"]) capsys.readouterr() # drain + + +def _usage(line, completion_tokens, prompt_tokens=1): + rec = json.loads(line) + rec["prompt_tokens"] = prompt_tokens + rec["completion_tokens"] = completion_tokens + return _line(rec) + + +def _summary(run_dir, primary, **fields): + """The run's ``result_summary.json``, recording ``primary`` as its basis.""" + path = run_dir / "performance" / "result_summary.json" + path.parent.mkdir() + path.write_text(json.dumps({"token_count_source": primary.value, **fields})) + return path + + +def test_server_basis_tpot_from_completion_tokens(tmp_path, capsys): + events = _write( + tmp_path, + [ + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _usage(_ev("sample.complete", 400, "a", _tmo(["x", "one two"])), 3), + _ev("sample.issued", 300, "b"), + _ev("sample.recv_first", 360, "b"), + _usage( + _ev("sample.complete", 600, "b", _tmo(["x", "y"], tool_calls=[["tc"]])), + 5, + ), + _ev("sample.issued", 320, "c"), + _ev("sample.recv_first", 380, "c"), + _usage(_ev("sample.complete", 500, "c", _tmo(["x"])), 1), # <= 1: skipped + _ev("sample.issued", 340, "d"), + _ev("sample.recv_first", 390, "d"), + _ev("sample.complete", 700, "d", _tmo(["x", "y z"])), # no usage: skipped + _ev("session.stop_performance_tracking", 800), + ], + ) + s = mod.compute_series(events, count_tokens=None, basis=_SERVER) + assert sorted(s["tpot_ns"]) == [ + 60.0, + 125.0, + ] # b: 240 / 4 (tool call kept), a: 250 / 2 + assert "tool-call" not in capsys.readouterr().err + + +def test_main_reads_server_basis_from_sibling_summary(tmp_path, capsys): + events = _write( + tmp_path, + [ + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _usage(_ev("sample.complete", 400, "a", _tmo(["x", "one two"])), 3), + _ev("session.stop_performance_tracking", 500), + ], + ) + _summary(tmp_path, _SERVER) + mod.main([str(events)]) + captured = capsys.readouterr() + assert "tpot (from tpot_ns, n=1)" in captured.out + assert "TPOT skipped" not in captured.err + + +def test_server_basis_counts_missing_usage(tmp_path, capsys): + events = _write( + tmp_path, + [ + _ev("sample.issued", 5, "w"), # untracked: never counted + _ev("sample.complete", 50, "w", _tmo(["x"])), + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _usage(_ev("sample.complete", 400, "a", _tmo(["x", "one two"])), 3), + _ev("sample.issued", 300, "b"), + _ev("sample.recv_first", 360, "b"), + _ev("sample.complete", 600, "b", _tmo(["x", "y"])), # streamed, no usage + _ev("sample.issued", 320, "c"), + _ev("sample.complete", 500, "c", _tmo("whole")), # non-streaming, no usage + _ev("sample.issued", 340, "d"), + _ev("sample.complete", 700, "d", None), # error: not a missing-usage sample + _ev("session.stop_performance_tracking", 800), + ], + ) + s = mod.compute_series(events, count_tokens=None, basis=_SERVER) + assert s["tpot_ns"] == [125.0] # missing usage is never filled in from a tokenizer + err = capsys.readouterr().err + assert "WARN: 2 tracked samples returned no server token usage" in err + + +_ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + + +@pytest.mark.parametrize( + ("prompt", "completion", "output", "expected"), + [pytest.param(*case[1:], id=case[0]) for case in LOGGED_SERVER_USAGE_CASES], +) +def test_replay_missing_usage_matches_live_rule( + tmp_path, capsys, prompt, completion, output, expected +): + """The server-basis replay reads usage by ``classify_server_usage``, as live.""" + complete = EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=400, + sample_uuid="a", + data=output, + prompt_tokens=prompt, + completion_tokens=completion, + ) + events = _write( + tmp_path, + [ + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _ENCODER.encode(complete).decode(), + _ev("session.stop_performance_tracking", 500), + ], + ) + s = mod.compute_series(events, count_tokens=None, basis=_SERVER) + usable = expected is ServerUsage.USABLE + assert s["tpot_ns"] == ([250 / (completion - 1)] if usable else []) + missing = expected is ServerUsage.MISSING + err = capsys.readouterr().err + assert ("1 tracked samples returned no server token usage" in err) is missing + + +def _usage_run_events(run_dir, *, usage): + """Streamed samples whose COMPLETE lines carry server usage iff ``usage``.""" + lines = [_ev("session.start_performance_tracking", 10)] + for i in range(4): + t, u = 100 * (i + 1), f"s{i}" + complete = _ev("sample.complete", t + 70 + i, u, _tmo(["x", "one two"])) + lines += [ + _ev("sample.issued", t, u), + _ev("sample.recv_first", t + 20, u), + _usage(complete, 9) if usage else complete, + ] + run_dir.mkdir() + return _write(run_dir, [*lines, _ev("session.stop_performance_tracking", 900)]) + + +def test_client_primary_replay_identical_with_and_without_usage_keys(tmp_path): + with_usage, without_usage = ( + mod.compute_series( + _usage_run_events(tmp_path / f"usage_{usage}", usage=usage), + count_tokens=_ws_counts, + basis=_CLIENT, + ) + for usage in (True, False) + ) + assert with_usage == without_usage + assert len(with_usage["tpot_ns"]) == 4 + + +@pytest.mark.parametrize("primary", list(TokenCountSource)) +def test_es_script_refuses_json_off_primary(tmp_path, capsys, monkeypatch, primary): + # The summary holds TPOT ES maps on the run's basis. A replay on the other + # basis must not be cross-checked against them, nor written into them. + events = _write( + tmp_path, + [ + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _usage(_ev("sample.complete", 400, "a", _tmo(["x", "one two"])), 4), + _ev("session.stop_performance_tracking", 500), + ], + ) + tpot = {"percentiles": {"50.0": 83.0}, "early_stopping_percentiles": {"50.0": 1.0}} + summary = _summary(tmp_path, primary, tpot=tpot) + monkeypatch.setattr(mod, "_make_counter", lambda path: _ws_counts) + off = primary.other + argv = [str(events), "--summary", str(summary), "--tokenizer", "unused"] + argv += ["--token-count-source", off.value] + with pytest.raises(SystemExit, match=f"on the {primary.value} token basis"): + mod.main([*argv, "--json", str(tmp_path / "o.json")]) + + augmented = mod.main(argv) + out = capsys.readouterr().out + assert "tpot (from tpot_ns, n=1)" in out # the off-primary replay still prints + assert ( + f"XCHECK in-band: skipped (TPOT replayed on the {off.value} basis; " + f"the run reported {primary.value})" + ) in out + assert "MISMATCH" not in out + assert augmented["tpot"] == tpot # the run's own maps are untouched + + +def _copied_run(tmp_path, summary_basis, tpot): + """An event log copied away from its run; the run's summary stays behind.""" + (tmp_path / "copy").mkdir() + events = _write( + tmp_path / "copy", + [ + _ev("session.start_performance_tracking", 10), + _ev("sample.issued", 100, "a"), + _ev("sample.recv_first", 150, "a"), + _usage(_ev("sample.complete", 400, "a", _tmo(["x", "one two"])), 6), + _ev("session.stop_performance_tracking", 500), + ], + ) + summary = tmp_path / "run" / "result_summary.json" + summary.parent.mkdir() + summary.write_text(json.dumps({"token_count_source": summary_basis, "tpot": tpot})) + return events, summary + + +def test_es_script_summary_outside_the_run_dir_sets_the_basis( + tmp_path, capsys, monkeypatch +): + # Nothing beside the copied log records a basis; --summary records server. + tpot = {"percentiles": {"50.0": 50.0}} + events, summary = _copied_run(tmp_path, _SERVER.value, tpot) + monkeypatch.setattr(mod, "_make_counter", lambda path: _ws_counts) + argv = [str(events), "--summary", str(summary), "--tokenizer", "unused"] + off_primary = [*argv, "--token-count-source", _CLIENT.value] + with pytest.raises(SystemExit, match="on the server token basis"): + mod.main([*off_primary, "--json", str(tmp_path / "o.json")]) + + out = tmp_path / "augmented.json" + mod.main([*argv, "--json", str(out)]) + captured = capsys.readouterr() + assert "token basis: server" in captured.out + assert "no token-count basis recorded" not in captured.err + assert "tpot (from tpot_ns, n=1)" in captured.out + augmented = json.loads(out.read_text()) + assert augmented["token_count_source"] == _SERVER.value + assert augmented["tpot"]["early_stopping_percentiles"] == {"50.0": None} # n=1 + + +def test_es_script_unreadable_summary_basis_with_override_never_labels_it( + tmp_path, capsys, monkeypatch +): + events, summary = _copied_run(tmp_path, "bogus", {}) + monkeypatch.setattr(mod, "_make_counter", lambda path: _ws_counts) + argv = [str(events), "--summary", str(summary), "--tokenizer", "unused"] + with pytest.raises(SystemExit) as exc: + mod.main(argv) + assert exc.value.code == 2 + assert "pass --token-count-source" in capsys.readouterr().err + chosen = [*argv, "--token-count-source", _SERVER.value] + # The summary's own basis is unknown, so the chosen one cannot label it. + with pytest.raises(SystemExit, match="could not be read"): + mod.main([*chosen, "--json", str(tmp_path / "o.json")]) + mod.main(chosen) + assert "token basis: server" in capsys.readouterr().out + + +def test_es_script_unreadable_summary_exits_with_hint(tmp_path, capsys): + events = _write(tmp_path, [_ev("session.start_performance_tracking", 10)]) + (tmp_path / "performance").mkdir() + (tmp_path / "performance" / "result_summary.json").write_text("{not json") + with pytest.raises(SystemExit) as exc: + mod.main([str(events)]) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "cannot read the token-count basis" in err + assert "pass --token-count-source" in err + # The hint holds: choosing the basis replays the run. + mod.main([str(events), "--token-count-source", _SERVER.value]) + capsys.readouterr() # drain + + +def test_es_script_missing_summary_fails_before_the_replay(tmp_path, monkeypatch): + events = _write(tmp_path, [_ev("session.start_performance_tracking", 10)]) + monkeypatch.setattr( + mod, "compute_series", lambda *a, **k: pytest.fail("log replayed") + ) + missing = tmp_path / "nope.json" + with pytest.raises(SystemExit, match=f"FATAL: --summary {missing} not found"): + mod.main([str(events), "--summary", str(missing)]) + + +def test_es_script_override_needs_no_archived_basis_note(tmp_path, capsys): + events = _write(tmp_path, [_ev("session.start_performance_tracking", 10)]) + mod.main([str(events)]) + assert "no token-count basis recorded" in capsys.readouterr().err + mod.main([str(events), "--token-count-source", _SERVER.value]) + err = capsys.readouterr().err + assert "no token-count basis recorded" not in err + assert "NOTE: token basis: server" in err diff --git a/tests/unit/scripts/test_synth_events.py b/tests/unit/scripts/test_synth_events.py index 75a0a9789..e428256a0 100644 --- a/tests/unit/scripts/test_synth_events.py +++ b/tests/unit/scripts/test_synth_events.py @@ -10,6 +10,7 @@ the planted base values and that a flat run is detected as steady. """ +import dataclasses import importlib.util import json import statistics @@ -17,6 +18,7 @@ from pathlib import Path import pytest +from inference_endpoint.core.types import TokenCountSource from inference_endpoint.metrics import steady_state_diagnostics as diag pytestmark = pytest.mark.unit @@ -38,6 +40,13 @@ def _words(texts): return [len(t.split()) for t in texts] +def _tokenized_series(out): + """40-sample super-passes replayed on the client basis with ``_words``.""" + return diag.build_super_pass_series( + out, 40, _words, basis=TokenCountSource.CLIENT + ).super_passes + + def _args(**overrides): """Small, fast, flat-by-default run suitable for steady detection.""" base = { @@ -87,7 +96,7 @@ def _generate(tmp_path, **overrides): def test_generated_events_parse_and_recover_base_metrics(tmp_path): out, planted, gt = _generate(tmp_path) - series = diag.build_super_pass_series(out, superpass_size=40, count_tokens=_words) + series = _tokenized_series(out) assert len(series) >= 4 # Pool the back half (past the ramp) and compare medians to the planted base. @@ -220,9 +229,14 @@ def test_flat_run_detected_as_steady(tmp_path): # Isolate steadiness detection from the min-duration floor (a short synth run can't # clear 600s); this test asserts the flat plateau is *found*, not that it's long. result = diag.run( - out, superpass_size=40, count_tokens=_words, enforce_min_duration=False + out, + superpass_size=40, + count_tokens=_words, + enforce_min_duration=False, + basis=TokenCountSource.CLIENT, + primary=TokenCountSource.CLIENT, ) - ss = result["steady_state"] + ss = result.steady_state assert ss.found is True assert ss.window is not None assert ss.window.n_super_passes >= 4 @@ -234,7 +248,7 @@ def test_poisson_mode_generates_and_parses(tmp_path): assert gt["mode"] == "poisson" assert gt["concurrency"] is None assert gt["rate"] == 40.0 - series = diag.build_super_pass_series(out, superpass_size=40, count_tokens=_words) + series = _tokenized_series(out) assert sum(sp.n_issued for sp in series) == gt["n_samples"] @@ -252,3 +266,28 @@ def test_multiturn_emits_conversation_and_turn(tmp_path): assert len(convs) >= 1 assert turns <= {1, 2, 3} assert max(turns) >= 2 # multi-turn trajectories actually issued + + +def test_complete_events_carry_server_usage(tmp_path): + out, _, _ = _generate(tmp_path) + completes = [ + r + for r in map(json.loads, Path(out).read_text().splitlines()) + if r["event_type"] == "sample.complete" + ] + assert completes and all( + r["completion_tokens"] == len(r["data"][1]) for r in completes + ) + assert all( + isinstance(r["prompt_tokens"], int) and r["prompt_tokens"] >= 1 + for r in completes + ) + + +def test_server_replay_matches_whitespace_replay(tmp_path): + out, _, _ = _generate(tmp_path) + server = diag.build_super_pass_series(out, 40, None, basis=TokenCountSource.SERVER) + assert server.n_missing_usage == 0 + assert [dataclasses.asdict(sp) for sp in server.super_passes] == [ + dataclasses.asdict(sp) for sp in _tokenized_series(out) + ] diff --git a/tests/unit/validation/test_shipped_cohort_golden.py b/tests/unit/validation/test_shipped_cohort_golden.py index e599674b6..b1e10de57 100644 --- a/tests/unit/validation/test_shipped_cohort_golden.py +++ b/tests/unit/validation/test_shipped_cohort_golden.py @@ -197,6 +197,7 @@ def _accuracy(run_dir: Path) -> None: "total_samples": FLOOR, } ], + token_count_source=DEFAULT_SOURCE, ) written = run_dir / "accuracy" / "accuracy_results.json" written.replace(run_dir / "accuracy_results.json") From 6d5229eefd58f24251ec9d9aac151425d7463c01 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:19:03 -0500 Subject: [PATCH 09/26] docs: server-reported token counts, the naming rule and the usage rule Design docs (core, endpoint client, OpenAI, metrics aggregator, metrics, report, steady state, early stopping, testing), the CLI quick reference, AGENTS.md (architecture notes and module map) and the agentic examples now describe: - the primary/other basis; - the _ naming rule; - when usage is requested; - the missing-usage rule; - how to flip the default. Part of mlcommons/endpoints#540. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 161 ++++++++++++++---- docs/CLI_QUICK_REFERENCE.md | 16 ++ docs/ENDPOINT_CLIENT.md | 83 +++++---- docs/async_utils/services/DESIGN.md | 9 +- .../services/metrics_aggregator/DESIGN.md | 78 ++++++--- docs/compliance_audit_plan.md | 13 +- docs/core/DESIGN.md | 11 +- docs/early_stopping.md | 35 +++- docs/endpoint_client/DESIGN.md | 1 + docs/metrics/DESIGN.md | 76 ++++++--- docs/metrics/report_design.md | 119 ++++++++++--- docs/openai/DESIGN.md | 61 ++++++- docs/steady_state_diagnostics.md | 66 +++++-- docs/testing/DESIGN.md | 35 ++++ examples/10_Agentic_Inference/README.md | 25 ++- examples/11_Edge_Agentic_Example/README.md | 7 + .../online_edge_full_run.yaml | 17 +- .../compliance/audit_test/README.md | 6 +- 18 files changed, 639 insertions(+), 180 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e266776dd..41ee193ad 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,7 +93,7 @@ Dataset Manager --> Load Generator --> Endpoint Client --> External Endpoint | **Endpoint Client** | `src/inference_endpoint/endpoint_client/` | Multi-process HTTP workers communicating via ZMQ IPC. `HTTPEndpointClient` is the main entry point | | **Dataset Manager** | `src/inference_endpoint/dataset_manager/` | Loads JSONL, HuggingFace, CSV, JSON, Parquet datasets. `Dataset` base class with `load_sample()`/`num_samples()` interface | | **Metrics Aggregator** | `src/inference_endpoint/async_utils/services/metrics_aggregator/` | Subprocess. Subscribes to events, aggregates per-sample metrics into a `MetricsRegistry` (counters + HDR-histogram series + raw values), publishes `MetricsSnapshot` over IPC PUB at a configurable cadence (`SessionState`: `INITIALIZE` → `LIVE` → `DRAINING` → {`COMPLETE` \| `INTERRUPTED`}). Final snapshot is atomically written to `final_snapshot.json` as the **primary** Report source; the terminal pub/sub frame is a TUI "run finished" signal only. | -| **Report** | `src/inference_endpoint/metrics/report.py` | `Report.from_snapshot(dict)` builds reports from the dict form (`snapshot_to_dict`). Reads `final_snapshot.json` directly via `json.loads` (no Struct decode). Sets `complete = (state == "complete" and n_pending_tasks == 0)` and renders an explicit warning for `INTERRUPTED` runs. Keeps the steady-window verdict on `Report.steady_state` and includes it in `to_json` (unlike `accuracy`). | +| **Report** | `src/inference_endpoint/metrics/report.py` | `Report.from_snapshot(dict, token_count_source=)` builds reports from the dict form (`snapshot_to_dict`). Reads `final_snapshot.json` directly via `json.loads` (no Struct decode). `complete` depends on the token-count source: server needs `state == "complete"`, valid accounting, no dropped samples and `n_samples_missing_usage == 0`; client needs `n_pending_tasks == 0` instead of the usage check. Headline ISL/OSL/TPOT/TPS are on the primary basis (`token_count_source`); the other basis sits beside them in `_`-suffixed blocks (`Report.other_basis`) and `osl_paired`. Renders an explicit warning for `INTERRUPTED` runs. Keeps the steady-window verdict on `Report.steady_state` and includes it in `to_json` (unlike `accuracy`). | | **Config** | `src/inference_endpoint/config/`, `endpoint_client/config.py` | Pydantic-based YAML schema (`schema.py`) — the declared user surface, distinct from the resolved `runtime_settings.py`; `Timeouts` (`config/schema.py` — all give-up deadlines in one frozen model at `settings.timeouts`, incl. the whole-run `run_timeout_s` watchdog; the workload duration `runtime.max_issue_duration_ms` stays on `settings.runtime`; client worker-lifecycle timeouts stay on `settings.client`), `HTTPClientConfig` (single Pydantic model for CLI/YAML/runtime), `RuntimeSettings` | | **CLI** | `src/inference_endpoint/main.py`, `commands/benchmark/cli.py` | cyclopts-based, auto-generated from `schema.py` and `HTTPClientConfig` Pydantic models. Flat shorthands via `cyclopts.Parameter(alias=...)` | | **Async Utils** | `src/inference_endpoint/async_utils/` | `LoopManager` (uvloop + eager_task_factory), ZMQ transport layer, generic `MessageCodec[T]`-parametrized pub/sub, event publisher | @@ -142,9 +142,14 @@ working in this layer: abnormal pipeline teardown). The aggregator **ignores SIGINT**: a terminal ^C reaches the whole foreground group, and the parent's ENDED path is authoritative. Drain timeout is detected by consumers as - `state == COMPLETE and n_pending_tasks > 0` — `run_benchmark` fails the run on - it (artifacts written with `complete: false`, then non-zero exit); interrupted - runs are detected as `state == INTERRUPTED` directly. + `state == COMPLETE and n_pending_tasks > 0`. `n_pending_tasks` counts the + client-basis tokenizer work. With a client primary `run_benchmark` fails the + run on it (artifacts written with `complete: false`, then non-zero exit). With + a server primary a drain timeout leaves the run passing and marks the + client-basis blocks partial (`Report.client_token_tasks_pending`, a logged + warning); there a tracked sample without server usage fails the run instead + (see **Token count source**). Interrupted runs are detected as + `state == INTERRUPTED` directly. - **Final delivery is dual-path with separated concerns**: `publish_final` atomically writes `final_snapshot.json` (`tmp + fsync(file) + rename + fsync(parent_dir)`) — this is the **primary** @@ -172,6 +177,62 @@ working in this layer: `metrics/early_stopping.py`; post-hoc recomputation from any run's `events.jsonl` via `scripts/early_stopping_estimate_from_events.py`. See docs/early_stopping.md. +- **Token count source**: a run has one primary basis, `server` (the endpoint's + per-request `usage`) or `client` (the local tokenizer). The primary basis + keeps the base names (the headline `isl`, `osl` and `tpot_ns` series; the + `tpot`, `input_sequence_lengths`, ... report keys) and feeds TPS and the + early-stopping and steady-state inputs. The other basis is recorded beside it + under a `_` suffix from `basis_key(base, basis, *, primary)` + (`core/types.py`): `isl_server`, `osl_server`, `tpot_ns_server` beside a + client primary, `isl_client`, `osl_client`, `tpot_ns_client` beside a server + primary. `session.py` lifts + `QueryResult.metadata["prompt_tokens"/"completion_tokens"]` onto + `EventRecord.prompt_tokens`/`completion_tokens` on COMPLETE (`UNSET` when + absent, as in archived logs). The server trigger in `metrics_table.py` + (`ServerUsageTrigger`) classifies the usage once per COMPLETE + (`classify_server_usage`) and records synchronously with no tokenizer, so it + adds nothing to `n_pending_tasks`. Server TPOT is + `(complete − first chunk) / (completion_tokens − 1)`, LoadGen's rule. No + server series records a value below 1: server OSL skips 0, and server TPOT + skips `completion_tokens <= 1` and non-positive windows. + + The tokenizer triggers record into the client-basis series whenever a + tokenizer is configured (`--no-metrics-isl` drops ISL on both bases; only the + primary basis's TPOT is a tail-latency series). The `osl_paired_samples`, + `osl_paired_client_tokens` and `osl_paired_server_tokens` counters total both + OSLs over the samples counted on both bases (`Report.osl_paired`). On either + primary a tracked, successful sample whose COMPLETE carries no usable usage + (no usage, or 0 `completion_tokens` beside non-empty output) increments + `tracked_samples_missing_usage` (and + `tracked_samples_missing_usage_empty_output` when it has no usage and empty + output). The post-run counts (full-run OSL, accuracy OSL, the standalone + replays) apply the same rule. With a server primary any such sample makes + `Report.complete` false, and `run_benchmark` raises `ExecutionError` (exit 4) + with the cause-aware message of `Report.incomplete_reason()`; with a client + primary it only marks the server-basis blocks partial. + + `--token-count-source {server,client}` (`settings.token_count_source`, + `TokenCountSource` in `core/types.py`) is resolved once by + `commands/benchmark/pipeline.py:resolve_token_count_source`, the only reader + of the setting; api_type `videogen` resolves to `client`. The result lives on + `BenchmarkContext.token_count_source`, goes to the aggregator as + `--token-count-source`, and is recorded as the top-level `token_count_source` + of `result_summary.json` and `accuracy/accuracy_results.json`. Asking for + usage is a run setting, `settings.client.stream_include_usage` (on by default; + `--no-stream-include-usage` opts out), passed to every adapter as + `encode_query(query, *, include_usage=)`; the OpenAI adapters send + `stream_options: {"include_usage": true}` only on streaming requests with at + most one choice. A server primary rejects the opt-out on a streaming OpenAI + run and fails setup on a row with `n > 1` (a client primary logs a warning + instead). Standalone post-run tools (the steady-state replay CLI and the + early-stopping script) take the primary from + `metrics/token_count_basis.py:token_count_source_for_run(run_dir, *, summary_path=None)`: + the top-level `token_count_source` of the summary (`summary_path`, by default + `performance/result_summary.json` under `run_dir`), else `config.yaml`; a run + that records neither (archived runs) means `client`. Both go through + `replay_basis(run_dir, *, override, summary_path=None)`, which pairs that + primary with the basis to replay on (`--token-count-source` when given). + - **Steady-state detection (opt-in, off by default)**: Enable with `settings.steady_state.enabled` / `--steady-state`. The parent resolves a workload profile (`commands/benchmark/pipeline.py:steady_state_profile`) and @@ -180,37 +241,45 @@ working in this layer: The aggregator feeds a `SuperPassCollector` (`metrics/steady_state_diagnostics.py`) from the sample events it already - routes. It reuses `TpotTrigger`'s token counts, so nothing is tokenized twice. - Per-sample series are `array("d")`, ~34 bytes/sample. Bucket size arrives on - `PHASE_START` (`PhaseData.num_turns`) and is latched while the series is - empty. + routes. It reuses the headline TPOT token counts, so nothing is tokenized + twice: in server mode it is fed by `ServerUsageTrigger` with count + `completion_tokens − 1`, in client mode by `TpotTrigger`. Per-sample series + are `array("d")`, ~34 bytes/sample. Bucket size arrives on `PHASE_START` + (`PhaseData.num_turns`) and is latched while the series is empty. The collector holds no per-sample state. The super-pass index lives on - `SampleRow.sp_index`; the standalone parse keeps its own `_PendingRow` because - it has no table. This avoids holding the same in-flight sample in two uuid - maps. + `SampleRow.superpass_index`; the standalone parse keeps its own `_PendingRow` + because it has no table. This avoids holding the same in-flight sample in two + uuid maps. Two ordering rules preserve the index. First, the router reads the row **before** `table.set_field`, because COMPLETE drops it from the table and RECV_FIRST overwrites `recv_first_ns`. That overwrite is how a retry's second - first-chunk is told apart. Second, `TpotTrigger` reads `row.sp_index` **at - fire time** and carries the int in the recorder closure. The count arrives at - the drain flush, long after the row is gone. + first-chunk is told apart. Second, the client-mode `TpotTrigger` reads + `row.superpass_index` **at fire time** and carries the int in the recorder + closure. The count arrives at the drain flush, long after the row is gone. + `ServerUsageTrigger` has no such lag: it runs inside `set_field(COMPLETE)` + while the row is still in the table, reads `row.superpass_index` directly and + feeds the collector synchronously. `ISSUED` is also gated on `table.is_tracking`, because `set_field` returns early outside the window without dropping the row. The verdict is computed after token drain. It rides `MetricsSnapshot.steady_state` into the single `final_snapshot.json` write, then onto `Report.steady_state`. - The series describes only a run with terminal state `COMPLETE` and - `n_pending_tasks == 0`. This path is best-effort: no failure in steady-state - detection can fail a run. + The series describes only a run with terminal state `COMPLETE` whose TPOT and + OSL counts are whole. In server mode the verdict is withheld when any tracked + sample lacks usage (`tracked_samples_missing_usage > 0`); a pending client + tokenizer drain does not withhold it, since those counts never feed the + series. In client mode it is withheld on `n_pending_tasks > 0`. This path is + best-effort: no failure in steady-state detection can fail a run. Eligibility is read from `Profile.supported`; offline and agentic are `False`, and an unclassified load pattern resolves to no profile. This keeps the parent's gate and the child's profile selection in sync. The gate also - requires streaming, because without `TpotTrigger` every verdict would be - `found: false`; a resolved tokenizer; and a non-accuracy-only run. + requires streaming, because without a TPOT trigger every verdict would be + `found: false`; a resolved tokenizer **or** server token-count mode (server + TPOT needs none); and a non-accuracy-only run. Agentic is refused twice: by load-pattern name in the parent gate and by `Profile.supported` in the aggregator. Agentic runs do not get a steady-state @@ -314,8 +383,9 @@ docs/compliance_audit_plan.md ("Run ordering"). `run_audit`: 2. Runs each `AuditRunSpec` phase (from `AuditTest.plan_runs`) back-to-back under its own `/