Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Comment thread
hvagadia marked this conversation as resolved.
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"
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/inference_endpoint/core/record.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
8 changes: 8 additions & 0 deletions src/inference_endpoint/load_generator/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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,
Expand All @@ -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
),
)
)

Expand Down
22 changes: 22 additions & 0 deletions src/inference_endpoint/metrics/report.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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")),
Expand All @@ -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,
)

Expand Down
6 changes: 5 additions & 1 deletion src/inference_endpoint/openai/completions_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
8 changes: 7 additions & 1 deletion tests/unit/load_generator/test_async_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
18 changes: 18 additions & 0 deletions tests/unit/metrics/test_report_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions tests/unit/openai/test_completions_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading