feat: report OpenAI finish reason distribution (#422)

* feat: report OpenAI finish reasons

* refactor: enumerate finish reason counters

Author: hvagadia

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

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:

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
+                        ),
                     )
                 )
 

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.
 
@@ -204,6 +217,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 +339,13 @@ 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 = {
+            reason.value.removeprefix(finish_reason_prefix): int(
+                counters.get(reason.value, 0)
+            )
+            for reason in _FINISH_REASON_COUNTERS
+        }
 
         return cls(
             version=str(version_info.get("version", "unknown")),
@@ -344,6 +365,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,
         )
 

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

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)

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)

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
 

tests/unit/metrics/test_report_builder.py

@@ -257,6 +257,24 @@ def test_failed_uses_tracked_counter(self):
         report = _build_report(registry)
         assert report.n_samples_failed == 1
 
+    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)
+
+        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)
         report = _build_report(registry, state=SessionState.COMPLETE, n_pending_tasks=0)

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):