From ab0ff15cf7000c2b518c50b5a606c4394aad2366 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Fri, 17 Jul 2026 14:03:21 -0700 Subject: [PATCH 1/2] feat: report OpenAI finish reasons --- .../services/metrics_aggregator/aggregator.py | 21 ++++++++++++ src/inference_endpoint/core/record.py | 1 + .../load_generator/session.py | 8 +++++ src/inference_endpoint/metrics/report.py | 8 +++++ .../openai/completions_adapter.py | 6 +++- .../services/event_logger/test_file_writer.py | 15 +++++--- .../metrics_aggregator/test_aggregator.py | 34 +++++++++++++++++++ .../unit/load_generator/test_async_session.py | 8 ++++- tests/unit/metrics/test_report_builder.py | 13 +++++++ tests/unit/openai/test_completions_adapter.py | 1 + 10 files changed, 109 insertions(+), 6 deletions(-) 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 8a8f1af4e..cd34a88f2 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py @@ -69,6 +69,12 @@ class MetricCounterKey(str, Enum): # 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" @@ -80,6 +86,15 @@ class MetricCounterKey(str, Enum): 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, + "tool_calls": MetricCounterKey.TRACKED_FINISH_REASON_TOOL_CALLS, + "content_filter": MetricCounterKey.TRACKED_FINISH_REASON_CONTENT_FILTER, + "function_call": MetricCounterKey.TRACKED_FINISH_REASON_FUNCTION_CALL, +} + + _TRACKED_SAMPLE_EVENTS = frozenset( { SampleEventType.ISSUED, @@ -397,6 +412,12 @@ async def process(self, records: list[EventRecord]) -> None: registry.increment(MetricCounterKey.TOTAL_SAMPLES_COMPLETED.value) if is_tracked: registry.increment(MetricCounterKey.TRACKED_SAMPLES_COMPLETED.value) + if isinstance(record.finish_reason, str): + counter = _FINISH_REASON_COUNTERS.get( + record.finish_reason, + MetricCounterKey.TRACKED_FINISH_REASON_OTHER, + ) + registry.increment(counter.value) if saw_shutdown: # ENDED has been observed; transition to DRAINING so any tick diff --git a/src/inference_endpoint/core/record.py b/src/inference_endpoint/core/record.py index 9ac60e1b6..4eb689b37 100644 --- a/src/inference_endpoint/core/record.py +++ b/src/inference_endpoint/core/record.py @@ -156,6 +156,7 @@ class EventRecord(msgspec.Struct, kw_only=True, frozen=True, gc=False): # type: conversation_id: str = "" turn: int | None = None data: OUTPUT_TYPE | PromptData | ErrorData | None = None + finish_reason: str | msgspec.UnsetType = msgspec.UNSET class EventRecordCodec: diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index 0d991ab1f..1b7e9a828 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -30,6 +30,8 @@ from enum import Enum from typing import Any, Protocol +import msgspec + from ..config.runtime_settings import RuntimeSettings from ..config.schema import LoadPatternType from ..core.record import ( @@ -601,6 +603,7 @@ def _handle_response(self, resp: QueryResult | StreamChunk) -> None: ) ) if self._current_phase_type != PhaseType.WARMUP: + finish_reason = resp.metadata.get("finish_reason") self._publisher.publish( EventRecord( event_type=SampleEventType.COMPLETE, @@ -611,6 +614,11 @@ def _handle_response(self, resp: QueryResult | StreamChunk) -> None: conversation_id=conv_id_str, turn=turn_num, data=resp.response_output, + finish_reason=( + finish_reason + if isinstance(finish_reason, str) + else msgspec.UNSET + ), ) ) diff --git a/src/inference_endpoint/metrics/report.py b/src/inference_endpoint/metrics/report.py index ac5615700..9ac63f39b 100644 --- a/src/inference_endpoint/metrics/report.py +++ b/src/inference_endpoint/metrics/report.py @@ -204,6 +204,7 @@ class Report(msgspec.Struct, frozen=True): # type: ignore[call-arg] # tokenizer unavailable). qps: float | None = None tps: float | None = None + finish_reason_counts: dict[str, int] = msgspec.field(default_factory=dict) # Run configuration (load_pattern, warmup, and the scheduler/dataloader RNG # seeds), from config. Carried so result_summary.json is self-describing and a @@ -325,6 +326,12 @@ def _series_dict(key: str) -> dict[str, Any]: # 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 = { + name.removeprefix(finish_reason_prefix): int(value) + for name, value in counters.items() + if name.startswith(finish_reason_prefix) and value + } return cls( version=str(version_info.get("version", "unknown")), @@ -344,6 +351,7 @@ def _series_dict(key: str) -> dict[str, Any]: legacy_loadgen_window_duration_ns=legacy_loadgen_window_duration_ns, qps=qps, tps=tps, + finish_reason_counts=finish_reason_counts, run_config=run_config, ) diff --git a/src/inference_endpoint/openai/completions_adapter.py b/src/inference_endpoint/openai/completions_adapter.py index d7dbd8c86..6c0124ee3 100644 --- a/src/inference_endpoint/openai/completions_adapter.py +++ b/src/inference_endpoint/openai/completions_adapter.py @@ -104,9 +104,13 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> QueryResult: resp = cls._response_decoder.decode(response_bytes) if not resp.choices: raise ValueError("Response must contain at least one choice") + choice = resp.choices[0] return QueryResult( id=query_id, - response_output=TextModelOutput(output=resp.choices[0].text), + response_output=TextModelOutput(output=choice.text), + metadata=( + {"finish_reason": choice.finish_reason} if choice.finish_reason else {} + ), ) @classmethod 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 c01079f82..9e9ae6f96 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 @@ -56,15 +56,22 @@ def test_writes_valid_jsonl(self, tmp_path): writer = JSONLWriter(tmp_path / "events", flush_interval=1) try: writer.write(_record(SampleEventType.ISSUED, uuid="s1", ts=1000)) - writer.write(_record(SampleEventType.COMPLETE, uuid="s1", ts=2000)) + writer.write( + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=2000, + sample_uuid="s1", + finish_reason="tool_calls", + ) + ) finally: writer.close() lines = (tmp_path / "events.jsonl").read_text().strip().split("\n") assert len(lines) == 2 - for line in lines: - parsed = json.loads(line) - assert isinstance(parsed, dict) + records = [json.loads(line) for line in lines] + assert "finish_reason" not in records[0] + assert records[1]["finish_reason"] == "tool_calls" def test_record_roundtrip_fields(self, tmp_path): writer = JSONLWriter(tmp_path / "events", flush_interval=1) 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 4510462cb..ca2aa2c99 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -808,6 +808,40 @@ async def test_total_vs_tracked_counters(self, tmp_path): finally: agg.close() + @pytest.mark.asyncio + 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") + try: + await agg.process( + [ + session_event( + SessionEventType.START_PERFORMANCE_TRACKING, ts=0 + ), + sample_event(SampleEventType.ISSUED, "s1", ts=1), + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=2, + sample_uuid="s1", + finish_reason="stop", + ), + sample_event(SampleEventType.ISSUED, "s2", ts=3), + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=4, + sample_uuid="s2", + finish_reason="future_reason", + ), + ] + ) + + counters = snapshot_counters(registry) + assert counters[MetricCounterKey.TRACKED_FINISH_REASON_STOP.value] == 1 + assert counters[MetricCounterKey.TRACKED_FINISH_REASON_OTHER.value] == 1 + finally: + agg.close() + # --------------------------------------------------------------------------- # Token trigger tests (with mock BatchTokenizer and real event loop) diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 278589c17..21a2bc2b7 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -725,10 +725,16 @@ async def test_handle_response_stamps_conversation_id_and_turn(self): phase_issuer.uuid_to_conv_info["q-ok"] = ("conv-9", 5) phase_issuer.inflight = 1 session._handle_response( - QueryResult(id="q-ok", response_output="ok", completed_at=12345) + QueryResult( + id="q-ok", + response_output="ok", + metadata={"finish_reason": "stop"}, + completed_at=12345, + ) ) complete = publisher.events_of_type(SampleEventType.COMPLETE) assert [(e.conversation_id, e.turn) for e in complete] == [("conv-9", 5)] + assert complete[0].finish_reason == "stop" assert "q-ok" not in phase_issuer.uuid_to_conv_info assert "q-ok" not in phase_issuer.completed_uuids diff --git a/tests/unit/metrics/test_report_builder.py b/tests/unit/metrics/test_report_builder.py index 0ed22c322..2b0e170a3 100644 --- a/tests/unit/metrics/test_report_builder.py +++ b/tests/unit/metrics/test_report_builder.py @@ -257,6 +257,19 @@ def test_failed_uses_tracked_counter(self): report = _build_report(registry) assert report.n_samples_failed == 1 + def test_finish_reason_counts(self): + registry = _make_registry(n_samples=2) + registry.increment(MetricCounterKey.TRACKED_FINISH_REASON_STOP.value, 1) + registry.increment(MetricCounterKey.TRACKED_FINISH_REASON_LENGTH.value, 1) + + report = _build_report(registry) + + assert report.finish_reason_counts == {"stop": 1, "length": 1} + assert json.loads(report.to_json())["finish_reason_counts"] == { + "stop": 1, + "length": 1, + } + def test_complete_flag_true_when_state_complete_and_no_pending(self): registry = _make_registry(n_samples=5) report = _build_report(registry, state=SessionState.COMPLETE, n_pending_tasks=0) diff --git a/tests/unit/openai/test_completions_adapter.py b/tests/unit/openai/test_completions_adapter.py index 8cb89b5b7..391ab88db 100644 --- a/tests/unit/openai/test_completions_adapter.py +++ b/tests/unit/openai/test_completions_adapter.py @@ -140,6 +140,7 @@ def test_decode_non_streaming(self): assert result.id == "qid-1" assert isinstance(result.response_output, TextModelOutput) assert result.response_output.output == "hello world" + assert result.metadata["finish_reason"] == "stop" @pytest.mark.unit def test_decode_empty_choices_raises(self): From ef15278267b72457ace688ef48f74576dafdd092 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Fri, 17 Jul 2026 19:03:43 -0700 Subject: [PATCH 2/2] refactor: enumerate finish reason counters --- src/inference_endpoint/metrics/report.py | 20 +++++++++++++++++--- tests/unit/metrics/test_report_builder.py | 11 ++++++++--- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/inference_endpoint/metrics/report.py b/src/inference_endpoint/metrics/report.py index 9ac63f39b..cd6924fa8 100644 --- a/src/inference_endpoint/metrics/report.py +++ b/src/inference_endpoint/metrics/report.py @@ -26,6 +26,9 @@ import msgspec.json import msgspec.structs +from inference_endpoint.async_utils.services.metrics_aggregator.aggregator import ( + MetricCounterKey, +) from inference_endpoint.async_utils.services.metrics_aggregator.registry import ( build_token_series_dict, ) @@ -66,6 +69,16 @@ def place_early_stopping_percentiles( return out +_FINISH_REASON_COUNTERS = ( + MetricCounterKey.TRACKED_FINISH_REASON_STOP, + MetricCounterKey.TRACKED_FINISH_REASON_LENGTH, + MetricCounterKey.TRACKED_FINISH_REASON_TOOL_CALLS, + MetricCounterKey.TRACKED_FINISH_REASON_CONTENT_FILTER, + MetricCounterKey.TRACKED_FINISH_REASON_FUNCTION_CALL, + MetricCounterKey.TRACKED_FINISH_REASON_OTHER, +) + + def _series_to_metric_dict(stat: dict[str, Any]) -> dict[str, Any]: """Convert a series-stat dict into the shape ``display()`` expects. @@ -328,9 +341,10 @@ def _series_dict(key: str) -> dict[str, Any]: n_pending_tasks = snap.get("n_pending_tasks", 0) finish_reason_prefix = "tracked_finish_reason_" finish_reason_counts = { - name.removeprefix(finish_reason_prefix): int(value) - for name, value in counters.items() - if name.startswith(finish_reason_prefix) and value + reason.value.removeprefix(finish_reason_prefix): int( + counters.get(reason.value, 0) + ) + for reason in _FINISH_REASON_COUNTERS } return cls( diff --git a/tests/unit/metrics/test_report_builder.py b/tests/unit/metrics/test_report_builder.py index 2b0e170a3..3c4bb29ed 100644 --- a/tests/unit/metrics/test_report_builder.py +++ b/tests/unit/metrics/test_report_builder.py @@ -257,18 +257,23 @@ def test_failed_uses_tracked_counter(self): report = _build_report(registry) assert report.n_samples_failed == 1 - def test_finish_reason_counts(self): + def test_finish_reason_counts_include_zeros(self): registry = _make_registry(n_samples=2) registry.increment(MetricCounterKey.TRACKED_FINISH_REASON_STOP.value, 1) registry.increment(MetricCounterKey.TRACKED_FINISH_REASON_LENGTH.value, 1) report = _build_report(registry) - assert report.finish_reason_counts == {"stop": 1, "length": 1} - assert json.loads(report.to_json())["finish_reason_counts"] == { + expected_counts = { "stop": 1, "length": 1, + "tool_calls": 0, + "content_filter": 0, + "function_call": 0, + "other": 0, } + assert report.finish_reason_counts == expected_counts + assert json.loads(report.to_json())["finish_reason_counts"] == expected_counts def test_complete_flag_true_when_state_complete_and_no_pending(self): registry = _make_registry(n_samples=5)