Skip to content

feat: add ProfileCallback for per-invocation profiling#91

Open
acere wants to merge 1 commit into
awslabs:mainfrom
acere:feat/profile-callback
Open

feat: add ProfileCallback for per-invocation profiling#91
acere wants to merge 1 commit into
awslabs:mainfrom
acere:feat/profile-callback

Conversation

@acere

@acere acere commented Jun 19, 2026

Copy link
Copy Markdown
Contributor

Closes #90

Summary

Adds a built-in ProfileCallback that captures per-invocation timing data during a Runner run and produces a structured profiling report.

Changes

New: llmeter/callbacks/profiling.py

  • InvocationProfile dataclass — per-request record with timing phases, token counts, cache info, retries
  • ProfileReport dataclass — aggregated summary with time accounting, cache hit rate, retry stats, percentile statistics
  • ProfileCallback — lifecycle hooks that collect data and generate the report in after_run
  • Saves profile_report.json + profile_invocations.jsonl alongside run results

New: llmeter/plotting/profiling.py

  • plot_phase_breakdown — stacked bar (TTFT vs generation per request)
  • plot_request_timeline — Gantt-style concurrency view
  • plot_throughput_over_time — tok/s scatter with rolling average
  • plot_time_accounting — pie chart (server prefill vs decode vs client overhead)
  • plot_tpot_distribution — TPOT histogram
  • plot_ttft_vs_input_tokens — scatter with cache hits highlighted
  • plot_profile_summary — 2×2 multi-panel dashboard

Supporting changes

  • llmeter/callbacks/__init__.py — exports ProfileCallback
  • llmeter/plotting/__init__.py — exports profiling plot functions
  • mkdocs.yml — nav entries for new API reference pages
  • docs/reference/callbacks/profiling.md + docs/reference/plotting/profiling.md
  • examples/Profiling.ipynb — full walkthrough notebook

Tests

  • tests/unit/callbacks/test_profiling.py — 29 tests covering serialization, lifecycle, report computation, stats contribution, file saving, and edge cases

Terminology

Clear server/client separation in naming:

  • TTFT = server prefill + network round-trip
  • Generation time = server-side decode (TTLT − TTFT)
  • Runner overhead = client-side (queue, token counting, callbacks)

Testing

uv run pytest tests/unit/callbacks/test_profiling.py -v  # 29 passed
uv run ruff check llmeter/callbacks/profiling.py llmeter/plotting/profiling.py

Adds a new ProfileCallback that instruments each request during a Runner
run to capture detailed timing phases, token counts, cache hit info,
retries, and throughput metrics.

Key features:
- Per-invocation profiles with TTFT, generation time (server decode),
  TPOT, tokens/second, cache hit ratio, reasoning tokens, and retries
- Aggregated ProfileReport with time accounting (server vs client),
  cache hit rate, retry stats, and percentile statistics
- Saves profile_report.json and profile_invocations.jsonl alongside
  run results for post-hoc analysis
- Contributes stats to result.stats under the profile_ prefix

Also adds plotting utilities in llmeter.plotting for profiling data:
- plot_phase_breakdown: stacked bar (TTFT vs generation per request)
- plot_request_timeline: Gantt-style concurrency view
- plot_throughput_over_time: tok/s scatter with rolling average
- plot_time_accounting: pie chart of time split
- plot_tpot_distribution: TPOT histogram
- plot_ttft_vs_input_tokens: scatter with cache hits highlighted
- plot_profile_summary: 2x2 multi-panel dashboard

Includes 29 unit tests, API reference docs, and example notebook.
@athewsey
athewsey force-pushed the feat/profile-callback branch from dacb97f to dc2be88 Compare June 24, 2026 08:32
@athewsey

Copy link
Copy Markdown
Collaborator

(Note, seems like this PR depends on plotting.defaults introduced in #89)

Just seen a weird intermittent error while testing the example notebook here... Requests failing with:

Endpoint invocation failed: 'dict' object has no attribute 'id'
Traceback (most recent call last):
  File "{...}/llmeter/llmeter/endpoints/base.py", line 259, in wrapper
    self.process_raw_response(raw_response, start_t, response)
  File "{...}/llmeter/llmeter/endpoints/openai_response.py", line 281, in process_raw_response
    response.id = event.response.id
                  ^^^^^^^^^^^^^^^^^
AttributeError: 'dict' object has no attribute 'id'

The same happens with google.gemma-4-31b on the same new /openai/v1 endpoint, but interestingly not with openai.gpt-oss-120b on the old /v1 endpoint. Claude couldn't figure yet what path through the OpenAI Python SDK could possibly be allowing this object to pass through un-parsed.

Will come back to it later

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add a built-in profiling callback for per-invocation timing analysis

2 participants