diff --git a/docs/reference/callbacks/profiling.md b/docs/reference/callbacks/profiling.md new file mode 100644 index 0000000..844594e --- /dev/null +++ b/docs/reference/callbacks/profiling.md @@ -0,0 +1 @@ +::: llmeter.callbacks.profiling diff --git a/docs/reference/plotting/profiling.md b/docs/reference/plotting/profiling.md new file mode 100644 index 0000000..3cb631d --- /dev/null +++ b/docs/reference/plotting/profiling.md @@ -0,0 +1 @@ +::: llmeter.plotting.profiling diff --git a/examples/Profiling.ipynb b/examples/Profiling.ipynb new file mode 100644 index 0000000..dc4fcae --- /dev/null +++ b/examples/Profiling.ipynb @@ -0,0 +1,428 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Profiling LLMeter Runs\n", + "\n", + "This notebook demonstrates the `ProfileCallback` — a built-in callback that\n", + "collects per-request timing data and generates a profiling report.\n", + "\n", + "## What you'll learn\n", + "\n", + "1. How to attach the profiler to a Runner\n", + "2. Reading the profile report (time accounting, cache hits, throughput)\n", + "3. Inspecting per-invocation profiles\n", + "4. Generating profiling plots (phase breakdown, timeline, throughput)\n", + "5. Loading saved profiles from disk for post-hoc analysis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llmeter.callbacks import ProfileCallback\n", + "from llmeter.endpoints import OpenAIResponseStreamEndpoint\n", + "from llmeter.runner import Runner" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup\n", + "\n", + "Create an endpoint and attach the `ProfileCallback` to a Runner.\n", + "The profiler works with any endpoint — streaming endpoints give the richest data.\n", + "\n", + "Here we use the OpenAI Responses API on Bedrock Mantle with GPT-5.4." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from aws_bedrock_token_generator import provide_token\n", + "\n", + "AWS_REGION = os.environ.get(\"AWS_REGION\", \"us-east-1\")\n", + "MODEL_ID = \"openai.gpt-5.4\"\n", + "\n", + "token = provide_token(region=AWS_REGION)\n", + "\n", + "endpoint = OpenAIResponseStreamEndpoint(\n", + " model_id=MODEL_ID,\n", + " endpoint_name=\"GPT-5.4-mantle\",\n", + " provider=\"bedrock-mantle\",\n", + " api_key=token,\n", + " base_url=f\"https://bedrock-mantle.{AWS_REGION}.api.aws/openai/v1\",\n", + ")\n", + "\n", + "profiler = ProfileCallback(\n", + " print_report=True, # Print summary when the run finishes\n", + " save_report=True, # Save profile_report.json + profile_invocations.jsonl\n", + ")\n", + "\n", + "runner = Runner(\n", + " endpoint=endpoint,\n", + " output_path=\"./outputs/profiling-demo\",\n", + " callbacks=[profiler],\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Run the Benchmark\n", + "\n", + "Run a small batch. The profiler will automatically collect per-request\n", + "timing data and produce a report at the end." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "payload = endpoint.create_payload(\n", + " \"Explain the difference between amortized and average-case analysis \"\n", + " \"in algorithm design. Give one example of each. Keep it under 150 words.\"\n", + ")\n", + "\n", + "result = await runner.run(\n", + " payload=payload,\n", + " n_requests=10,\n", + " clients=2,\n", + " run_name=\"profiling-demo\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Inspect the Report\n", + "\n", + "The report is printed automatically (since `print_report=True`), but\n", + "you can also access it programmatically." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "report = profiler.report\n", + "\n", + "print(f\"Wall clock: {report.total_wall_clock:.1f}s\")\n", + "print(f\"API time: {report.total_api_time:.1f}s ({report.api_time_fraction:.0%})\")\n", + "print(f\"Runner overhead: {report.runner_overhead:.1f}s ({report.runner_overhead_fraction:.0%})\")\n", + "print(f\"Cache hit rate: {report.cache_hit_rate:.0%}\")\n", + "print(f\"Total retries: {report.total_retries}\")\n", + "print(f\"Tokens in/out: {report.total_tokens_input:,} / {report.total_tokens_output:,}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Per-Invocation Data\n", + "\n", + "Each request's profile is available as an `InvocationProfile` object.\n", + "This is the same data saved to `profile_invocations.jsonl`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "# Convert to DataFrame for easy analysis\n", + "from dataclasses import asdict\n", + "\n", + "df = pd.DataFrame([asdict(p) for p in profiler.invocation_profiles])\n", + "df[[\"sequence\", \"ttft\", \"generation_time\", \"output_tokens_per_second\",\n", + " \"tokens_input\", \"tokens_output\", \"cache_hit\", \"error\"]]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Profiling Plots\n", + "\n", + "The `llmeter.plotting` module provides ready-made visualizations for profiling data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llmeter.plotting import (\n", + " plot_phase_breakdown,\n", + " plot_request_timeline,\n", + " plot_throughput_over_time,\n", + " plot_time_accounting,\n", + " plot_tpot_distribution,\n", + " plot_ttft_vs_input_tokens,\n", + " plot_profile_summary,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Phase Breakdown\n", + "\n", + "Stacked bar showing TTFT vs generation time for each request." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_phase_breakdown(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Request Timeline\n", + "\n", + "Gantt-style view showing request concurrency. Overlapping bars = parallel requests." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_request_timeline(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Throughput Over Time\n", + "\n", + "How token generation speed varies as the run progresses." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_throughput_over_time(profiler, window=3).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Time Accounting\n", + "\n", + "Where does wall-clock time go? TTFT (prefill), stream (decode), or runner overhead?" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_time_accounting(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### TPOT Distribution\n", + "\n", + "Histogram of per-token latency across all requests." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_tpot_distribution(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### TTFT vs Input Tokens\n", + "\n", + "Does prefill time scale with prompt length? Cache hits are highlighted." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_ttft_vs_input_tokens(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### All-in-One Summary\n", + "\n", + "A 2×2 multi-panel combining phase breakdown, throughput, TPOT, and time accounting." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_profile_summary(profiler).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Loading Saved Profiles\n", + "\n", + "Profile data is saved alongside the run results. You can reload it later\n", + "without needing the original callback or live endpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from pathlib import Path\n", + "from llmeter.callbacks.profiling import InvocationProfile\n", + "\n", + "# Load per-invocation profiles from saved JSONL\n", + "profile_path = Path(\"./outputs/profiling-demo/profiling-demo/profile_invocations.jsonl\")\n", + "\n", + "if profile_path.exists():\n", + " with open(profile_path) as f:\n", + " profiles = [InvocationProfile(**json.loads(line)) for line in f]\n", + " \n", + " print(f\"Loaded {len(profiles)} invocation profiles\")\n", + " \n", + " # These can be passed directly to plotting functions\n", + " plot_phase_breakdown(profiles, title=\"Loaded from disk\").show()\n", + "else:\n", + " print(f\"No saved profiles found at {profile_path}\")\n", + " print(\"(Run the benchmark cells above first)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Load the summary report\n", + "report_path = Path(\"./outputs/profiling-demo/profiling-demo/profile_report.json\")\n", + "\n", + "if report_path.exists():\n", + " with open(report_path) as f:\n", + " report_data = json.load(f)\n", + " \n", + " print(\"Saved profile report:\")\n", + " for key, value in report_data.items():\n", + " if not isinstance(value, dict):\n", + " print(f\" {key}: {value}\")\n", + "else:\n", + " print(f\"No saved report found at {report_path}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Combining with Other Callbacks\n", + "\n", + "The profiler plays well with other callbacks like `SystemMetricsMonitor` and `CostModel`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llmeter.callbacks import CostModel, ProfileCallback\n", + "from llmeter.callbacks.system_metrics import SystemMetricsMonitor\n", + "from llmeter.callbacks.cost.dimensions import InputTokens, OutputTokens\n", + "\n", + "# Combine profiling with cost estimation and system monitoring\n", + "profiler = ProfileCallback()\n", + "monitor = SystemMetricsMonitor(sample_interval=0.5)\n", + "cost_model = CostModel(\n", + " request_dims=[\n", + " InputTokens(price_per_million=3.0),\n", + " OutputTokens(price_per_million=15.0),\n", + " ]\n", + ")\n", + "\n", + "runner = Runner(\n", + " endpoint=endpoint,\n", + " output_path=\"./outputs/profiling-full\",\n", + " callbacks=[profiler, monitor, cost_model],\n", + ")\n", + "\n", + "result = await runner.run(payload=payload, n_requests=10, run_name=\"full-profile\")\n", + "print(\"Uncomment the line above to run with all callbacks\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/llmeter/callbacks/__init__.py b/llmeter/callbacks/__init__.py index 42db055..3915a45 100644 --- a/llmeter/callbacks/__init__.py +++ b/llmeter/callbacks/__init__.py @@ -16,6 +16,7 @@ from .base import Callback # noqa: F401 from .cost import CostModel # noqa: F401 +from .profiling import ProfileCallback # noqa: F401 spec = importlib.util.find_spec("mlflow") if spec: diff --git a/llmeter/callbacks/profiling.py b/llmeter/callbacks/profiling.py new file mode 100644 index 0000000..225cb4a --- /dev/null +++ b/llmeter/callbacks/profiling.py @@ -0,0 +1,583 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Profiling callback for LLMeter runs. + +Collects per-request phase timings and produces a summary report after the run, +breaking down where time is spent: server-side prefill (TTFT), server-side +generation (decode), client-side runner overhead, and callback overhead. + +Optionally saves: +- ``profile_report.json`` — aggregated summary with statistics +- ``profile_invocations.jsonl`` — per-invocation detail (one JSON object per line) +""" + +import json +import logging +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime + +from upath.types import ReadablePathLike, WritablePathLike + +from ..endpoints.base import InvocationResponse +from ..results import Result +from ..runner import _RunConfig +from ..utils import ensure_path, summary_stats_from_list +from .base import Callback + +logger = logging.getLogger(__name__) + + +@dataclass +class InvocationProfile: + """Per-invocation timing and metadata captured during a run. + + This is the per-request record saved to ``profile_invocations.jsonl``. + """ + + #: Sequence number within the run (0-indexed, in order of completion). + sequence: int = 0 + #: Unique request ID (from InvocationResponse.id), if available. + request_id: str | None = None + #: Wall-clock time when the request was sent (ISO format). + request_time: str | None = None + #: Offset in seconds from the start of the run to when this request completed. + offset_from_run_start: float = 0.0 + + # --- Timing phases --- + #: Time to first token (seconds). Server-side prefill + network round-trip. + ttft: float | None = None + #: Time to last token (seconds). Total response time. + ttlt: float | None = None + #: Generation time = TTLT - TTFT (seconds). Server-side decode (token generation). + generation_time: float | None = None + #: Time per output token (seconds/token) = generation_time / tokens_output. + time_per_output_token: float | None = None + #: Output token throughput (tokens/second) = tokens_output / generation_time. + output_tokens_per_second: float | None = None + + # --- Token counts --- + #: Number of input tokens. + tokens_input: int | None = None + #: Number of output tokens. + tokens_output: int | None = None + #: Number of input tokens served from cache (prompt caching). + tokens_input_cached: int | None = None + #: Whether this request had a cache hit (any tokens served from cache). + cache_hit: bool = False + #: Cache hit ratio (cached_tokens / total_input_tokens). + cache_hit_ratio: float | None = None + #: Number of reasoning/thinking tokens (subset of output tokens). + tokens_output_reasoning: int | None = None + + # --- Error/retry info --- + #: Error message, or None if successful. + error: str | None = None + #: Number of retries for this request. + retries: int | None = None + + # --- Overhead --- + #: Time spent in after_invoke callbacks (seconds). Measures profiling overhead. + callback_overhead: float = 0.0 + + +@dataclass +class ProfileReport: + """Aggregated profiling report produced after a run. + + Contains both timing breakdowns and richer metadata like cache hit rates, + retry statistics, and throughput analysis. + """ + + # --- Time accounting --- + total_wall_clock: float = 0.0 + total_api_time: float = 0.0 + runner_overhead: float = 0.0 + api_time_fraction: float = 0.0 + runner_overhead_fraction: float = 0.0 + + # --- Request counts --- + total_requests: int = 0 + successful_requests: int = 0 + failed_requests: int = 0 + retried_requests: int = 0 + total_retries: int = 0 + + # --- Cache stats --- + cache_hits: int = 0 + cache_hit_rate: float = 0.0 + total_tokens_cached: int = 0 + + # --- Token totals --- + total_tokens_input: int = 0 + total_tokens_output: int = 0 + total_tokens_reasoning: int = 0 + + # --- Timing statistics --- + ttft_stats: dict[str, float] = field(default_factory=dict) + generation_time_stats: dict[str, float] = field(default_factory=dict) + tpot_stats: dict[str, float] = field(default_factory=dict) + tokens_per_second_stats: dict[str, float] = field(default_factory=dict) + callback_overhead_stats: dict[str, float] = field(default_factory=dict) + + def __str__(self) -> str: + lines = [ + "╭─── LLMeter Profile Report ───╮", + f" Wall clock: {self.total_wall_clock:.3f}s", + f" API time: {self.total_api_time:.3f}s ({self.api_time_fraction:.1%})", + f" Runner overhead: {self.runner_overhead:.3f}s" + f" ({self.runner_overhead_fraction:.1%})", + f" Requests: {self.successful_requests} ok," + f" {self.failed_requests} failed", + ] + + if self.retried_requests > 0: + lines.append( + f" Retries: {self.total_retries}" + f" ({self.retried_requests} requests)" + ) + + if self.total_tokens_input > 0: + lines.append("") + lines.append( + f" Tokens in/out: {self.total_tokens_input:,}" + f" / {self.total_tokens_output:,}" + ) + if self.total_tokens_reasoning > 0: + lines.append( + f" Reasoning tokens: {self.total_tokens_reasoning:,}" + f" ({self.total_tokens_reasoning / max(self.total_tokens_output, 1):.0%}" + " of output)" + ) + if self.cache_hits > 0: + lines.append( + f" Cache hits: {self.cache_hits}" + f" ({self.cache_hit_rate:.0%})," + f" {self.total_tokens_cached:,} tokens cached" + ) + + lines.append("") + if self.ttft_stats: + lines.append(" TTFT (server prefill + network):") + lines.append( + f" avg={self.ttft_stats.get('average', 0):.3f}s " + f"p50={self.ttft_stats.get('p50', 0):.3f}s " + f"p90={self.ttft_stats.get('p90', 0):.3f}s " + f"p99={self.ttft_stats.get('p99', 0):.3f}s" + ) + if self.generation_time_stats: + lines.append(" Generation (server decode):") + lines.append( + f" avg={self.generation_time_stats.get('average', 0):.3f}s " + f"p50={self.generation_time_stats.get('p50', 0):.3f}s " + f"p90={self.generation_time_stats.get('p90', 0):.3f}s " + f"p99={self.generation_time_stats.get('p99', 0):.3f}s" + ) + if self.tpot_stats: + lines.append(" Time per output token:") + lines.append( + f" avg={self.tpot_stats.get('average', 0) * 1000:.1f}ms " + f"p50={self.tpot_stats.get('p50', 0) * 1000:.1f}ms " + f"p90={self.tpot_stats.get('p90', 0) * 1000:.1f}ms" + ) + if self.tokens_per_second_stats: + lines.append(" Output throughput:") + lines.append( + f" avg={self.tokens_per_second_stats.get('average', 0):.1f} tok/s " + f"p50={self.tokens_per_second_stats.get('p50', 0):.1f} tok/s" + ) + if self.callback_overhead_stats: + avg_ms = self.callback_overhead_stats.get("average", 0) * 1000 + p99_ms = self.callback_overhead_stats.get("p99", 0) * 1000 + lines.append( + f" Callback overhead: avg={avg_ms:.2f}ms p99={p99_ms:.2f}ms" + ) + + lines.append("╰──────────────────────────────╯") + return "\n".join(lines) + + def __repr__(self) -> str: + return ( + f"ProfileReport(wall_clock={self.total_wall_clock:.3f}s, " + f"api_time={self.total_api_time:.3f}s ({self.api_time_fraction:.1%}), " + f"overhead={self.runner_overhead:.3f}s ({self.runner_overhead_fraction:.1%}), " + f"requests={self.successful_requests}ok/{self.failed_requests}fail)" + ) + + +@dataclass +class ProfileCallback(Callback): + """Profiling callback that collects per-invocation timing data and produces a report. + + This callback instruments each request to capture detailed phase timings, + token counts, cache hit information, retries, and throughput metrics. + + After the run completes: + 1. A `ProfileReport` summary is generated and contributed to ``result.stats`` + 2. Optionally printed to stdout (``print_report=True``) + 3. Optionally saved to disk (``save_report=True``): + - ``profile_report.json`` — aggregated summary + - ``profile_invocations.jsonl`` — per-invocation detail records + + Args: + print_report: If True (default), print the report summary to stdout after the run. + save_report: If True (default), save report files to the result's output directory. + Only takes effect when the result has an ``output_path`` set. + + Example:: + + from llmeter.callbacks import ProfileCallback + from llmeter.runner import Runner + + profiler = ProfileCallback() + runner = Runner(endpoint=endpoint, callbacks=[profiler]) + result = await runner.run(payload=payload, n_requests=20) + + # Access the report + print(profiler.report) + + # Inspect per-invocation data + for inv in profiler.invocation_profiles: + if inv.cache_hit: + print(f"Request {inv.sequence}: cached {inv.tokens_input_cached} tokens") + + # Per-invocation data is also saved as profile_invocations.jsonl + + Contributed stats (prefixed with ``profile_``): + + - ``profile_total_wall_clock``: Total run wall-clock time (seconds) + - ``profile_total_api_time``: Sum of TTLT across successful requests (seconds) + - ``profile_runner_overhead``: Wall clock minus API time (seconds) + - ``profile_api_time_fraction``: Fraction of wall clock in API calls + - ``profile_runner_overhead_fraction``: Fraction of wall clock in overhead + - ``profile_cache_hit_rate``: Fraction of requests with cached input tokens + - ``profile_total_retries``: Total retry count across all requests + - ``profile_ttft-average``, ``-p50``, ``-p90``, ``-p99``: TTFT statistics + - ``profile_generation_time-average``, ``-p50``, ``-p90``, ``-p99``: Generation time stats + - ``profile_tpot-average``, ``-p50``, ``-p90``, ``-p99``: Time per output token stats + - ``profile_callback_overhead-average``, ``-p50``, ``-p90``, ``-p99``: Callback overhead + - ``profile_tokens_per_second-average``, ``-p50``, ``-p90``, ``-p99``: Token throughput + """ + + print_report: bool = True + save_report: bool = True + + def __post_init__(self): + self._invocation_profiles: list[InvocationProfile] = [] + self._pending: dict = {} + self._sequence: int = 0 + self._run_start_time: float = 0.0 + self._run_end_time: float = 0.0 + self._report: ProfileReport | None = None + + def __getstate__(self): + """Support pickling/deepcopy by returning only serializable state.""" + return { + "print_report": self.print_report, + "save_report": self.save_report, + "_invocation_profiles": self._invocation_profiles, + "_sequence": self._sequence, + "_run_start_time": self._run_start_time, + "_run_end_time": self._run_end_time, + } + + def __setstate__(self, state): + """Restore from pickle/deepcopy.""" + self.print_report = state["print_report"] + self.save_report = state.get("save_report", True) + self._invocation_profiles = state["_invocation_profiles"] + self._sequence = state["_sequence"] + self._run_start_time = state["_run_start_time"] + self._run_end_time = state["_run_end_time"] + self._pending = {} + self._report = None + + @property + def invocation_profiles(self) -> list[InvocationProfile]: + """Access the raw per-invocation profile data.""" + return self._invocation_profiles + + @property + def report(self) -> ProfileReport | None: + """The generated profile report (available after the run completes).""" + return self._report + + async def before_run(self, run_config: _RunConfig) -> None: + """Reset state and record run start time.""" + self._invocation_profiles = [] + self._pending = {} + self._sequence = 0 + self._report = None + self._run_start_time = time.perf_counter() + + async def before_invoke(self, payload: dict) -> None: + """Record the timestamp just before the endpoint is called.""" + self._pending = {"before_invoke_time": time.perf_counter()} + + async def after_invoke(self, response: InvocationResponse) -> None: + """Capture full invocation profile from the response.""" + t_enter = time.perf_counter() + + # Compute derived timing fields + ttft = response.time_to_first_token + ttlt = response.time_to_last_token + generation_time = None + tpot = None + tps = None + + if ttft is not None and ttlt is not None and ttlt > ttft: + generation_time = ttlt - ttft + if response.num_tokens_output and response.num_tokens_output > 0: + tpot = generation_time / response.num_tokens_output + tps = response.num_tokens_output / generation_time + + # Cache info + cached = response.num_tokens_input_cached or 0 + has_cache_hit = cached > 0 + cache_ratio = None + if ( + has_cache_hit + and response.num_tokens_input + and response.num_tokens_input > 0 + ): + cache_ratio = cached / response.num_tokens_input + + # Request time formatting + req_time_str = None + if response.request_time is not None: + req_time_str = ( + response.request_time.isoformat() + if isinstance(response.request_time, datetime) + else str(response.request_time) + ) + + cb_overhead = time.perf_counter() - t_enter + + profile = InvocationProfile( + sequence=self._sequence, + request_id=response.id, + request_time=req_time_str, + offset_from_run_start=t_enter - self._run_start_time, + ttft=ttft, + ttlt=ttlt, + generation_time=generation_time, + time_per_output_token=tpot, + output_tokens_per_second=tps, + tokens_input=response.num_tokens_input, + tokens_output=response.num_tokens_output, + tokens_input_cached=response.num_tokens_input_cached, + cache_hit=has_cache_hit, + cache_hit_ratio=cache_ratio, + tokens_output_reasoning=response.num_tokens_output_reasoning, + error=response.error, + retries=response.retries, + callback_overhead=cb_overhead, + ) + + self._invocation_profiles.append(profile) + self._sequence += 1 + self._pending = {} + + async def after_run(self, result: Result) -> None: + """Generate the profile report, contribute stats, and optionally save to disk.""" + self._run_end_time = time.perf_counter() + self._report = self._compute_report() + + # Contribute stats to result + stats = self._report_to_stats(self._report) + result._update_contributed_stats(stats) + + if self.print_report: + print(self._report) + + if self.save_report and result.output_path: + self._save_report(result.output_path) + self._save_invocations(result.output_path) + + def _compute_report(self) -> ProfileReport: + """Compute the profile report from collected invocation profiles.""" + report = ProfileReport() + report.total_wall_clock = self._run_end_time - self._run_start_time + report.total_requests = len(self._invocation_profiles) + + successful = [p for p in self._invocation_profiles if p.error is None] + failed = [p for p in self._invocation_profiles if p.error is not None] + report.successful_requests = len(successful) + report.failed_requests = len(failed) + + # Retry stats (across all requests, including failed) + retried = [p for p in self._invocation_profiles if p.retries and p.retries > 0] + report.retried_requests = len(retried) + report.total_retries = sum(p.retries for p in retried if p.retries) + + if not successful: + return report + + # Cache stats + cache_hits = [p for p in successful if p.cache_hit] + report.cache_hits = len(cache_hits) + report.cache_hit_rate = len(cache_hits) / len(successful) if successful else 0.0 + report.total_tokens_cached = sum(p.tokens_input_cached or 0 for p in successful) + + # Token totals + report.total_tokens_input = sum(p.tokens_input or 0 for p in successful) + report.total_tokens_output = sum(p.tokens_output or 0 for p in successful) + report.total_tokens_reasoning = sum( + p.tokens_output_reasoning or 0 for p in successful + ) + + # Timing metrics + ttft_values: list[float] = [] + generation_times: list[float] = [] + tpot_values: list[float] = [] + tps_values: list[float] = [] + cb_overhead_values: list[float] = [] + + for p in successful: + if p.ttft is not None: + ttft_values.append(p.ttft) + if p.generation_time is not None: + generation_times.append(p.generation_time) + if p.time_per_output_token is not None: + tpot_values.append(p.time_per_output_token) + if p.output_tokens_per_second is not None: + tps_values.append(p.output_tokens_per_second) + cb_overhead_values.append(p.callback_overhead) + + # Total API time = sum of TTLT for successful requests + api_times = [p.ttlt for p in successful if p.ttlt is not None] + report.total_api_time = sum(api_times) + report.runner_overhead = max( + 0.0, report.total_wall_clock - report.total_api_time + ) + + if report.total_wall_clock > 0: + report.api_time_fraction = report.total_api_time / report.total_wall_clock + report.runner_overhead_fraction = ( + report.runner_overhead / report.total_wall_clock + ) + + # Compute summary statistics for each metric + if ttft_values: + report.ttft_stats = summary_stats_from_list(ttft_values) + if generation_times: + report.generation_time_stats = summary_stats_from_list(generation_times) + if tpot_values: + report.tpot_stats = summary_stats_from_list(tpot_values) + if tps_values: + report.tokens_per_second_stats = summary_stats_from_list(tps_values) + if cb_overhead_values: + report.callback_overhead_stats = summary_stats_from_list(cb_overhead_values) + + return report + + @staticmethod + def _report_to_stats(report: ProfileReport) -> dict[str, float | int]: + """Convert a ProfileReport into a flat dict suitable for result.stats.""" + stats: dict[str, float | int] = { + "profile_total_wall_clock": report.total_wall_clock, + "profile_total_api_time": report.total_api_time, + "profile_runner_overhead": report.runner_overhead, + "profile_api_time_fraction": report.api_time_fraction, + "profile_runner_overhead_fraction": report.runner_overhead_fraction, + "profile_successful_requests": report.successful_requests, + "profile_failed_requests": report.failed_requests, + "profile_retried_requests": report.retried_requests, + "profile_total_retries": report.total_retries, + "profile_cache_hits": report.cache_hits, + "profile_cache_hit_rate": report.cache_hit_rate, + "profile_total_tokens_cached": report.total_tokens_cached, + "profile_total_tokens_input": report.total_tokens_input, + "profile_total_tokens_output": report.total_tokens_output, + "profile_total_tokens_reasoning": report.total_tokens_reasoning, + } + + for metric_name, metric_stats in [ + ("profile_ttft", report.ttft_stats), + ("profile_generation_time", report.generation_time_stats), + ("profile_tpot", report.tpot_stats), + ("profile_callback_overhead", report.callback_overhead_stats), + ("profile_tokens_per_second", report.tokens_per_second_stats), + ]: + for agg, value in metric_stats.items(): + stats[f"{metric_name}-{agg}"] = value + + return stats + + def _save_report( + self, + output_path: WritablePathLike, + file_name: str = "profile_report.json", + ) -> None: + """Save the profile report as JSON to the result output directory. + + Args: + output_path: Directory where the report file will be written. + file_name: Name of the report file. Defaults to ``profile_report.json``. + """ + if self._report is None: + return + + out_dir = ensure_path(output_path) + out_dir.mkdir(parents=True, exist_ok=True) + report_path = out_dir / file_name + + report_data = asdict(self._report) + with report_path.open("w") as f: + json.dump(report_data, f, indent=2) + + logger.info("Profile report saved to %s", report_path) + + def _save_invocations( + self, + output_path: WritablePathLike, + file_name: str = "profile_invocations.jsonl", + ) -> None: + """Save per-invocation profiles as JSONL to the result output directory. + + Each line is a JSON object representing one invocation's profile data, + suitable for loading into pandas or custom analysis tools. + + Args: + output_path: Directory where the file will be written. + file_name: Name of the file. Defaults to ``profile_invocations.jsonl``. + """ + if not self._invocation_profiles: + return + + out_dir = ensure_path(output_path) + out_dir.mkdir(parents=True, exist_ok=True) + invocations_path = out_dir / file_name + + with invocations_path.open("w") as f: + for profile in self._invocation_profiles: + f.write(json.dumps(asdict(profile)) + "\n") + + logger.info( + "Profile invocations (%d records) saved to %s", + len(self._invocation_profiles), + invocations_path, + ) + + def save_to_file(self, path: WritablePathLike) -> None: + """Save this ProfileCallback configuration to file.""" + out_path = ensure_path(path) + config = { + "type": "ProfileCallback", + "print_report": self.print_report, + "save_report": self.save_report, + } + with out_path.open("w") as f: + json.dump(config, f, indent=2) + + @classmethod + def _load_from_file(cls, path: ReadablePathLike) -> "ProfileCallback": + """Load a ProfileCallback from file.""" + in_path = ensure_path(path) + with in_path.open("r") as f: + config = json.load(f) + + return cls( + print_report=config.get("print_report", True), + save_report=config.get("save_report", True), + ) diff --git a/llmeter/plotting/__init__.py b/llmeter/plotting/__init__.py index 340f7d6..060fc3a 100644 --- a/llmeter/plotting/__init__.py +++ b/llmeter/plotting/__init__.py @@ -14,12 +14,28 @@ plot_load_test_results, scatter_histogram_2d, ) +from .profiling import ( + plot_phase_breakdown, + plot_profile_summary, + plot_request_timeline, + plot_throughput_over_time, + plot_time_accounting, + plot_tpot_distribution, + plot_ttft_vs_input_tokens, +) __all__ = [ - "plot_heatmap", - "scatter_histogram_2d", "boxplot_by_dimension", - "plot_load_test_results", - "histogram_by_dimension", "color_sequences", + "histogram_by_dimension", + "plot_heatmap", + "plot_load_test_results", + "plot_phase_breakdown", + "plot_profile_summary", + "plot_request_timeline", + "plot_throughput_over_time", + "plot_time_accounting", + "plot_tpot_distribution", + "plot_ttft_vs_input_tokens", + "scatter_histogram_2d", ] diff --git a/llmeter/plotting/profiling.py b/llmeter/plotting/profiling.py new file mode 100644 index 0000000..f41b62e --- /dev/null +++ b/llmeter/plotting/profiling.py @@ -0,0 +1,651 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plotting utilities for ProfileCallback data. + +Provides ready-made Plotly figures for visualizing profiling results: +- Phase breakdown (stacked bar chart showing TTFT vs generation time per request) +- Request timeline (Gantt-style view of request concurrency) +- Throughput over time (tokens/second as the run progresses) +- Time accounting (pie/sunburst showing where wall-clock time goes) +- TPOT distribution (histogram of per-token latency) + +All functions accept either a `ProfileCallback` instance (with data still in memory) +or a list of `InvocationProfile` records (e.g. loaded from `profile_invocations.jsonl`). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from ..utils import DeferredError +from .defaults import DEFAULT_TEMPLATE + +if not TYPE_CHECKING: + try: + import plotly.graph_objects as go + from plotly.subplots import make_subplots + except ImportError as e: + go = DeferredError(e) + make_subplots = DeferredError(e) +else: + import plotly.graph_objects as go + from plotly.subplots import make_subplots + +if TYPE_CHECKING: + from ..callbacks.profiling import InvocationProfile, ProfileCallback + + +def _get_profiles( + source: ProfileCallback | list[InvocationProfile], +) -> list[InvocationProfile]: + """Extract InvocationProfile list from either a callback or a raw list.""" + if isinstance(source, list): + return source + return source.invocation_profiles + + +def plot_phase_breakdown( + source: ProfileCallback | list[InvocationProfile], + *, + title: str = "Per-Request Phase Breakdown", + show_overhead: bool = False, + **layout_kwargs: Any, +) -> go.Figure: + """Stacked bar chart showing TTFT vs generation time for each request. + + Each bar represents one request, split into: + - TTFT (prefill + network latency) + - Generation (server-side decode) + - Callback overhead (optional, usually negligible) + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + title: Chart title. + show_overhead: If True, include callback overhead as a third segment. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with the stacked bar chart. + """ + profiles = _get_profiles(source) + successful = [p for p in profiles if p.error is None] + + x_labels = list(range(1, len(successful) + 1)) + ttft_values = [p.ttft or 0 for p in successful] + generation_values = [p.generation_time or 0 for p in successful] + + fig = go.Figure() + + fig.add_trace( + go.Bar( + x=x_labels, + y=ttft_values, + name="TTFT (prefill + network)", + marker_color="#636EFA", + ) + ) + fig.add_trace( + go.Bar( + x=x_labels, + y=generation_values, + name="Generation (server decode)", + marker_color="#EF553B", + ) + ) + + if show_overhead: + overhead_values = [p.callback_overhead for p in successful] + fig.add_trace( + go.Bar( + x=x_labels, + y=overhead_values, + name="Callback overhead", + marker_color="#00CC96", + ) + ) + + fig.update_layout( + barmode="stack", + title=title, + xaxis_title="Request #", + yaxis_title="Time (seconds)", + template=DEFAULT_TEMPLATE, + **layout_kwargs, + ) + + return fig + + +def plot_request_timeline( + source: ProfileCallback | list[InvocationProfile], + *, + title: str = "Request Timeline", + color_by: str = "phase", + **layout_kwargs: Any, +) -> go.Figure: + """Gantt-style timeline showing when each request was active. + + Visualizes request concurrency — overlapping bars indicate parallel requests. + Each request is shown as a horizontal bar from its start offset to completion, + split into TTFT and stream phases. + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + title: Chart title. + color_by: How to color bars. "phase" splits TTFT/stream with different colors. + "throughput" colors by output tokens/second. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with the timeline. + """ + profiles = _get_profiles(source) + successful = [p for p in profiles if p.error is None] + + if not successful: + fig = go.Figure() + fig.update_layout(title=title, template=DEFAULT_TEMPLATE) + return fig + + fig = go.Figure() + + # Sort by offset for visual clarity + sorted_profiles = sorted(successful, key=lambda p: p.offset_from_run_start) + + y_labels = [f"#{p.sequence}" for p in sorted_profiles] + + if color_by == "phase": + # TTFT phase + ttft_starts = [] + ttft_durations = [] + stream_starts = [] + generation_durations = [] + + for p in sorted_profiles: + ttlt = p.ttlt or 0 + ttft = p.ttft or 0 + # Request started at offset - ttlt (it completed at offset) + req_start = p.offset_from_run_start - ttlt + ttft_starts.append(req_start) + ttft_durations.append(ttft) + stream_starts.append(req_start + ttft) + generation_durations.append(p.generation_time or 0) + + fig.add_trace( + go.Bar( + y=y_labels, + x=ttft_durations, + base=ttft_starts, + orientation="h", + name="TTFT", + marker_color="#636EFA", + opacity=0.8, + ) + ) + fig.add_trace( + go.Bar( + y=y_labels, + x=generation_durations, + base=stream_starts, + orientation="h", + name="Generation", + marker_color="#EF553B", + opacity=0.8, + ) + ) + else: + # Color by throughput + tps_values = [p.output_tokens_per_second or 0 for p in sorted_profiles] + starts = [] + durations = [] + + for p in sorted_profiles: + ttlt = p.ttlt or 0 + req_start = p.offset_from_run_start - ttlt + starts.append(req_start) + durations.append(ttlt) + + fig.add_trace( + go.Bar( + y=y_labels, + x=durations, + base=starts, + orientation="h", + name="Request", + marker=dict( + color=tps_values, + colorscale="Viridis", + colorbar=dict(title="tok/s"), + ), + opacity=0.8, + ) + ) + + fig.update_layout( + barmode="stack", + title=title, + xaxis_title="Time since run start (seconds)", + yaxis_title="Request", + template=DEFAULT_TEMPLATE, + height=max(300, len(successful) * 25 + 100), + **layout_kwargs, + ) + + return fig + + +def plot_throughput_over_time( + source: ProfileCallback | list[InvocationProfile], + *, + title: str = "Output Throughput Over Time", + window: int = 5, + **layout_kwargs: Any, +) -> go.Figure: + """Scatter plot of per-request output throughput (tok/s) over the run timeline. + + Shows how throughput varies as the run progresses, with an optional rolling + average to smooth noise. + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + title: Chart title. + window: Rolling average window size (number of requests). Set to 1 to disable. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with throughput over time. + """ + profiles = _get_profiles(source) + successful = [ + p + for p in profiles + if p.error is None and p.output_tokens_per_second is not None + ] + + if not successful: + fig = go.Figure() + fig.update_layout(title=title, template=DEFAULT_TEMPLATE) + return fig + + # Sort by completion time + sorted_profiles = sorted(successful, key=lambda p: p.offset_from_run_start) + offsets = [p.offset_from_run_start for p in sorted_profiles] + tps_values = [p.output_tokens_per_second or 0 for p in sorted_profiles] + + fig = go.Figure() + + # Individual points + fig.add_trace( + go.Scatter( + x=offsets, + y=tps_values, + mode="markers", + name="Per-request", + marker=dict(size=6, opacity=0.5, color="#636EFA"), + ) + ) + + # Rolling average + if window > 1 and len(tps_values) >= window: + rolling_avg = [] + for i in range(len(tps_values)): + start = max(0, i - window + 1) + rolling_avg.append(sum(tps_values[start : i + 1]) / (i - start + 1)) + + fig.add_trace( + go.Scatter( + x=offsets, + y=rolling_avg, + mode="lines", + name=f"Rolling avg (window={window})", + line=dict(color="#EF553B", width=2), + ) + ) + + fig.update_layout( + title=title, + xaxis_title="Time since run start (seconds)", + yaxis_title="Output tokens/second", + template=DEFAULT_TEMPLATE, + **layout_kwargs, + ) + + return fig + + +def plot_time_accounting( + source: ProfileCallback | list[InvocationProfile], + *, + total_wall_clock: float | None = None, + title: str = "Time Accounting", + **layout_kwargs: Any, +) -> go.Figure: + """Pie chart showing how wall-clock time is split between API time and overhead. + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + total_wall_clock: Total run wall-clock time in seconds. If source is a + ProfileCallback with a computed report, this is extracted automatically. + title: Chart title. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with the pie chart. + """ + from ..callbacks.profiling import ProfileCallback as _PC + + profiles = _get_profiles(source) + successful = [p for p in profiles if p.error is None] + + api_time = sum(p.ttlt or 0 for p in successful) + + # Try to get wall clock from report + if total_wall_clock is None and isinstance(source, _PC) and source.report: + total_wall_clock = source.report.total_wall_clock + + if total_wall_clock is None or total_wall_clock <= 0: + # Can't compute overhead without wall clock + total_wall_clock = api_time # fallback: show API time only + + overhead = max(0.0, total_wall_clock - api_time) + + # Further break down API time into TTFT vs stream + total_ttft = sum(p.ttft or 0 for p in successful) + total_generation = sum(p.generation_time or 0 for p in successful) + + labels = ["TTFT (prefill)", "Generation (decode)", "Runner overhead"] + values = [total_ttft, total_generation, overhead] + colors = ["#636EFA", "#EF553B", "#AB63FA"] + + fig = go.Figure( + data=[ + go.Pie( + labels=labels, + values=values, + marker=dict(colors=colors), + textinfo="label+percent", + textposition="inside", + hole=0.3, + ) + ] + ) + + fig.update_layout( + title=title, + template=DEFAULT_TEMPLATE, + **layout_kwargs, + ) + + return fig + + +def plot_tpot_distribution( + source: ProfileCallback | list[InvocationProfile], + *, + title: str = "Time Per Output Token Distribution", + bin_size_ms: float | None = None, + **layout_kwargs: Any, +) -> go.Figure: + """Histogram of time-per-output-token (TPOT) values in milliseconds. + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + title: Chart title. + bin_size_ms: Histogram bin size in milliseconds. Auto if None. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with the TPOT histogram. + """ + profiles = _get_profiles(source) + tpot_ms = [ + p.time_per_output_token * 1000 + for p in profiles + if p.error is None and p.time_per_output_token is not None + ] + + if not tpot_ms: + fig = go.Figure() + fig.update_layout(title=title, template=DEFAULT_TEMPLATE) + return fig + + xbins = dict(size=bin_size_ms) if bin_size_ms else None + + fig = go.Figure( + data=[ + go.Histogram( + x=tpot_ms, + name="TPOT", + marker_color="#636EFA", + opacity=0.8, + xbins=xbins, + ) + ] + ) + + fig.update_layout( + title=title, + xaxis_title="Time per output token (ms)", + yaxis_title="Count", + template=DEFAULT_TEMPLATE, + **layout_kwargs, + ) + + return fig + + +def plot_ttft_vs_input_tokens( + source: ProfileCallback | list[InvocationProfile], + *, + title: str = "TTFT vs Input Tokens", + **layout_kwargs: Any, +) -> go.Figure: + """Scatter plot of TTFT against input token count. + + Useful for understanding how prefill time scales with prompt length. + Cache hits are highlighted with a different marker. + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + title: Chart title. + **layout_kwargs: Additional kwargs passed to fig.update_layout(). + + Returns: + A Plotly Figure with the scatter plot. + """ + profiles = _get_profiles(source) + successful = [ + p + for p in profiles + if p.error is None and p.ttft is not None and p.tokens_input is not None + ] + + if not successful: + fig = go.Figure() + fig.update_layout(title=title, template=DEFAULT_TEMPLATE) + return fig + + # Split into cached vs non-cached + non_cached = [p for p in successful if not p.cache_hit] + cached = [p for p in successful if p.cache_hit] + + fig = go.Figure() + + if non_cached: + fig.add_trace( + go.Scatter( + x=[p.tokens_input for p in non_cached], + y=[p.ttft for p in non_cached], + mode="markers", + name="No cache", + marker=dict(size=8, color="#636EFA", opacity=0.7), + ) + ) + + if cached: + fig.add_trace( + go.Scatter( + x=[p.tokens_input for p in cached], + y=[p.ttft for p in cached], + mode="markers", + name="Cache hit", + marker=dict(size=10, color="#00CC96", opacity=0.8, symbol="diamond"), + ) + ) + + fig.update_layout( + title=title, + xaxis_title="Input tokens", + yaxis_title="TTFT (seconds)", + template=DEFAULT_TEMPLATE, + **layout_kwargs, + ) + + return fig + + +def plot_profile_summary( + source: ProfileCallback | list[InvocationProfile], + *, + total_wall_clock: float | None = None, + title: str = "Profile Summary", +) -> go.Figure: + """Multi-panel figure combining the key profiling visualizations. + + Creates a 2x2 subplot with: + - Phase breakdown (stacked bar) + - Throughput over time (scatter + rolling avg) + - TPOT distribution (histogram) + - Time accounting (pie) + + Args: + source: A ProfileCallback instance or list of InvocationProfile records. + total_wall_clock: Total run wall-clock time (auto-detected from ProfileCallback). + title: Overall figure title. + + Returns: + A Plotly Figure with the 2x2 subplot layout. + """ + from ..callbacks.profiling import ProfileCallback as _PC + + profiles = _get_profiles(source) + successful = [p for p in profiles if p.error is None] + + if isinstance(source, _PC) and source.report: + if total_wall_clock is None: + total_wall_clock = source.report.total_wall_clock + + fig = make_subplots( + rows=2, + cols=2, + subplot_titles=( + "Phase Breakdown", + "Throughput Over Time", + "TPOT Distribution", + "Time Accounting", + ), + specs=[ + [{"type": "bar"}, {"type": "scatter"}], + [{"type": "histogram"}, {"type": "pie"}], + ], + ) + + # Panel 1: Phase breakdown + x_labels = list(range(1, len(successful) + 1)) + fig.add_trace( + go.Bar( + x=x_labels, + y=[p.ttft or 0 for p in successful], + name="TTFT", + marker_color="#636EFA", + showlegend=True, + ), + row=1, + col=1, + ) + fig.add_trace( + go.Bar( + x=x_labels, + y=[p.generation_time or 0 for p in successful], + name="Generation", + marker_color="#EF553B", + showlegend=True, + ), + row=1, + col=1, + ) + + # Panel 2: Throughput over time + sorted_by_time = sorted(successful, key=lambda p: p.offset_from_run_start) + offsets = [p.offset_from_run_start for p in sorted_by_time] + tps_values = [p.output_tokens_per_second or 0 for p in sorted_by_time] + + fig.add_trace( + go.Scatter( + x=offsets, + y=tps_values, + mode="markers+lines", + name="tok/s", + marker=dict(size=5, color="#00CC96"), + line=dict(color="#00CC96", width=1), + showlegend=True, + ), + row=1, + col=2, + ) + + # Panel 3: TPOT histogram + tpot_ms = [ + p.time_per_output_token * 1000 + for p in successful + if p.time_per_output_token is not None + ] + if tpot_ms: + fig.add_trace( + go.Histogram( + x=tpot_ms, + name="TPOT (ms)", + marker_color="#AB63FA", + opacity=0.8, + showlegend=True, + ), + row=2, + col=1, + ) + + # Panel 4: Time accounting pie + api_time = sum(p.ttlt or 0 for p in successful) + total_ttft = sum(p.ttft or 0 for p in successful) + total_generation = sum(p.generation_time or 0 for p in successful) + overhead = max(0.0, (total_wall_clock or api_time) - api_time) + + fig.add_trace( + go.Pie( + labels=["TTFT", "Generation", "Overhead"], + values=[total_ttft, total_generation, overhead], + marker=dict(colors=["#636EFA", "#EF553B", "#AB63FA"]), + textinfo="label+percent", + hole=0.3, + showlegend=False, + ), + row=2, + col=2, + ) + + fig.update_layout( + barmode="stack", + title_text=title, + template=DEFAULT_TEMPLATE, + height=700, + width=1100, + ) + + # Axis labels for subplots + fig.update_xaxes(title_text="Request #", row=1, col=1) + fig.update_yaxes(title_text="Time (s)", row=1, col=1) + fig.update_xaxes(title_text="Run time (s)", row=1, col=2) + fig.update_yaxes(title_text="tok/s", row=1, col=2) + fig.update_xaxes(title_text="TPOT (ms)", row=2, col=1) + fig.update_yaxes(title_text="Count", row=2, col=1) + + return fig diff --git a/mkdocs.yml b/mkdocs.yml index 27fb90b..80687a7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -57,6 +57,8 @@ nav: - results: reference/callbacks/cost/results.md - serde: reference/callbacks/cost/serde.md - mlflow: reference/callbacks/mlflow.md + - profiling: reference/callbacks/profiling.md + - system_metrics: reference/callbacks/system_metrics.md - endpoints: - reference/endpoints/index.md - anthropic_messages: reference/endpoints/anthropic_messages.md @@ -71,6 +73,9 @@ nav: - json_utils: reference/json_utils.md - plotting: - reference/plotting/index.md + - defaults: reference/plotting/defaults.md + - percentage: reference/plotting/percentage.md + - profiling: reference/plotting/profiling.md - prompt_utils: reference/prompt_utils.md - results: reference/results.md - runner: reference/runner.md diff --git a/tests/unit/callbacks/test_profiling.py b/tests/unit/callbacks/test_profiling.py new file mode 100644 index 0000000..fa370d3 --- /dev/null +++ b/tests/unit/callbacks/test_profiling.py @@ -0,0 +1,475 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import copy +import json +import pickle +import tempfile +from dataclasses import asdict +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from llmeter.callbacks.profiling import ( + ProfileCallback, +) +from llmeter.endpoints.base import InvocationResponse +from llmeter.results import Result + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +class FakeRunConfig: + pass + + +def _make_response( + *, + idx: int = 0, + ttft: float = 0.25, + ttlt: float = 3.0, + tokens_in: int = 100, + tokens_out: int = 80, + cached: int | None = None, + reasoning: int | None = None, + error: str | None = None, + retries: int | None = None, +) -> InvocationResponse: + return InvocationResponse( + response_text="output" if not error else None, + id=f"req-{idx:03d}", + request_time=datetime(2025, 6, 19, 10, 0, idx, tzinfo=timezone.utc), + time_to_first_token=ttft if not error else None, + time_to_last_token=ttlt if not error else None, + num_tokens_input=tokens_in if not error else None, + num_tokens_output=tokens_out if not error else None, + num_tokens_input_cached=cached, + num_tokens_output_reasoning=reasoning, + error=error, + retries=retries, + ) + + +@pytest.fixture +def profiler(): + return ProfileCallback(print_report=False, save_report=False) + + +@pytest.fixture +def profiler_with_data(): + """Returns a profiler that has run through a simulated run.""" + + async def _setup(): + cb = ProfileCallback(print_report=False, save_report=True) + await cb.before_run(FakeRunConfig()) + + responses = [ + _make_response(idx=0, ttft=0.2, ttlt=2.5, tokens_in=100, tokens_out=80), + _make_response( + idx=1, ttft=0.08, ttlt=2.0, tokens_in=100, tokens_out=60, cached=70 + ), + _make_response( + idx=2, + ttft=0.5, + ttlt=6.0, + tokens_in=200, + tokens_out=300, + reasoning=150, + ), + _make_response( + idx=3, ttft=0.9, ttlt=4.0, tokens_in=80, tokens_out=50, retries=2 + ), + _make_response(idx=4, error="ThrottlingException", retries=3), + ] + + for resp in responses: + await cb.before_invoke({"messages": [{"role": "user", "content": "test"}]}) + await cb.after_invoke(resp) + + result = Result( + responses=[], + total_requests=5, + clients=1, + n_requests=5, + total_test_time=20.0, + ) + await cb.after_run(result) + return cb, result + + return asyncio.run(_setup()) + + +# --------------------------------------------------------------------------- +# Serialization / deepcopy tests +# --------------------------------------------------------------------------- + + +class TestProfileCallbackSerialization: + def test_deepcopy(self, profiler): + copied = copy.deepcopy(profiler) + assert copied.print_report == profiler.print_report + assert copied.save_report == profiler.save_report + + def test_pickle_roundtrip(self, profiler): + data = pickle.dumps(profiler) + restored = pickle.loads(data) + assert restored.print_report is False + assert restored.save_report is False + + def test_asdict(self, profiler): + d = asdict(profiler) + assert d == {"print_report": False, "save_report": False} + + def test_save_and_load_config(self, profiler): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "profiler.json" + profiler.save_to_file(path) + + with open(path) as f: + config = json.load(f) + assert config["type"] == "ProfileCallback" + assert config["print_report"] is False + assert config["save_report"] is False + + restored = ProfileCallback._load_from_file(path) + assert restored.print_report is False + assert restored.save_report is False + + +# --------------------------------------------------------------------------- +# Lifecycle tests +# --------------------------------------------------------------------------- + + +class TestProfileCallbackLifecycle: + def test_before_run_resets_state(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + assert profiler._invocation_profiles == [] + assert profiler._sequence == 0 + assert profiler._report is None + assert profiler._run_start_time > 0 + + asyncio.run(_test()) + + def test_captures_invocation_data(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke( + _make_response(idx=0, ttft=0.3, ttlt=2.5, tokens_in=100, tokens_out=80) + ) + + assert len(profiler.invocation_profiles) == 1 + p = profiler.invocation_profiles[0] + assert p.sequence == 0 + assert p.request_id == "req-000" + assert p.ttft == 0.3 + assert p.ttlt == 2.5 + assert p.generation_time == pytest.approx(2.2) + assert p.tokens_input == 100 + assert p.tokens_output == 80 + assert p.error is None + assert p.callback_overhead >= 0 + + asyncio.run(_test()) + + def test_captures_cache_info(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke( + _make_response(idx=0, cached=60, tokens_in=100, tokens_out=50) + ) + + p = profiler.invocation_profiles[0] + assert p.cache_hit is True + assert p.tokens_input_cached == 60 + assert p.cache_hit_ratio == pytest.approx(0.6) + + asyncio.run(_test()) + + def test_captures_reasoning_tokens(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke( + _make_response(idx=0, tokens_out=200, reasoning=120) + ) + + p = profiler.invocation_profiles[0] + assert p.tokens_output_reasoning == 120 + + asyncio.run(_test()) + + def test_captures_errors_and_retries(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke( + _make_response(idx=0, error="timeout", retries=3) + ) + + p = profiler.invocation_profiles[0] + assert p.error == "timeout" + assert p.retries == 3 + assert p.ttft is None + assert p.generation_time is None + + asyncio.run(_test()) + + def test_sequence_increments(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + for i in range(5): + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke(_make_response(idx=i)) + + sequences = [p.sequence for p in profiler.invocation_profiles] + assert sequences == [0, 1, 2, 3, 4] + + asyncio.run(_test()) + + def test_reuse_across_runs(self, profiler): + """A profiler should reset between runs.""" + + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke(_make_response(idx=0)) + result1 = Result(responses=[], total_requests=1) + await profiler.after_run(result1) + + assert len(profiler.invocation_profiles) == 1 + + # Second run + await profiler.before_run(FakeRunConfig()) + assert len(profiler.invocation_profiles) == 0 + assert profiler.report is None + + asyncio.run(_test()) + + +# --------------------------------------------------------------------------- +# Report computation tests +# --------------------------------------------------------------------------- + + +class TestProfileReport: + def test_report_computed(self, profiler_with_data): + cb, result = profiler_with_data + report = cb.report + assert report is not None + assert report.total_requests == 5 + assert report.successful_requests == 4 + assert report.failed_requests == 1 + + def test_report_cache_stats(self, profiler_with_data): + cb, _ = profiler_with_data + report = cb.report + assert report.cache_hits == 1 + assert report.cache_hit_rate == pytest.approx(0.25) + assert report.total_tokens_cached == 70 + + def test_report_retry_stats(self, profiler_with_data): + cb, _ = profiler_with_data + report = cb.report + assert report.retried_requests == 2 # idx=3 and idx=4 + assert report.total_retries == 5 # 2 + 3 + + def test_report_token_totals(self, profiler_with_data): + cb, _ = profiler_with_data + report = cb.report + assert report.total_tokens_input == 480 # 100+100+200+80 + assert report.total_tokens_output == 490 # 80+60+300+50 + assert report.total_tokens_reasoning == 150 + + def test_report_timing_stats(self, profiler_with_data): + cb, _ = profiler_with_data + report = cb.report + assert "average" in report.ttft_stats + assert "p50" in report.ttft_stats + assert "average" in report.generation_time_stats + assert "average" in report.tpot_stats + assert "average" in report.tokens_per_second_stats + + def test_report_str(self, profiler_with_data): + cb, _ = profiler_with_data + text = str(cb.report) + assert "Profile Report" in text + assert "4 ok" in text + assert "1 failed" in text + assert "Cache hits" in text + assert "Reasoning tokens" in text + assert "Retries" in text + + def test_report_repr(self, profiler_with_data): + cb, _ = profiler_with_data + text = repr(cb.report) + assert "ProfileReport" in text + + +# --------------------------------------------------------------------------- +# Stats contribution tests +# --------------------------------------------------------------------------- + + +class TestProfileStatsContribution: + def test_contributes_to_result_stats(self, profiler_with_data): + _, result = profiler_with_data + stats = result.stats + assert "profile_total_wall_clock" in stats + assert "profile_total_api_time" in stats + assert "profile_runner_overhead" in stats + assert "profile_api_time_fraction" in stats + assert "profile_successful_requests" in stats + assert stats["profile_successful_requests"] == 4 + assert stats["profile_failed_requests"] == 1 + + def test_contributes_cache_stats(self, profiler_with_data): + _, result = profiler_with_data + assert result.stats["profile_cache_hit_rate"] == pytest.approx(0.25) + assert result.stats["profile_total_tokens_cached"] == 70 + + def test_contributes_timing_stats(self, profiler_with_data): + _, result = profiler_with_data + assert "profile_ttft-average" in result.stats + assert "profile_ttft-p50" in result.stats + assert "profile_generation_time-average" in result.stats + assert "profile_tpot-average" in result.stats + assert "profile_tokens_per_second-average" in result.stats + + +# --------------------------------------------------------------------------- +# File saving tests +# --------------------------------------------------------------------------- + + +class TestProfileSaving: + def test_saves_report_json(self, profiler_with_data): + cb, _ = profiler_with_data + with tempfile.TemporaryDirectory() as tmpdir: + cb._save_report(tmpdir) + report_path = Path(tmpdir) / "profile_report.json" + assert report_path.exists() + + with open(report_path) as f: + data = json.load(f) + assert data["successful_requests"] == 4 + assert data["failed_requests"] == 1 + assert "ttft_stats" in data + + def test_saves_invocations_jsonl(self, profiler_with_data): + cb, _ = profiler_with_data + with tempfile.TemporaryDirectory() as tmpdir: + cb._save_invocations(tmpdir) + inv_path = Path(tmpdir) / "profile_invocations.jsonl" + assert inv_path.exists() + + with open(inv_path) as f: + records = [json.loads(line) for line in f] + assert len(records) == 5 + assert records[0]["sequence"] == 0 + assert records[0]["request_id"] == "req-000" + assert records[1]["cache_hit"] is True + assert records[4]["error"] == "ThrottlingException" + + def test_save_report_false_skips(self): + async def _test(): + cb = ProfileCallback(print_report=False, save_report=False) + await cb.before_run(FakeRunConfig()) + await cb.before_invoke({"messages": []}) + await cb.after_invoke(_make_response(idx=0)) + + with tempfile.TemporaryDirectory() as tmpdir: + result = Result(responses=[], total_requests=1, output_path=tmpdir) + await cb.after_run(result) + assert not (Path(tmpdir) / "profile_report.json").exists() + assert not (Path(tmpdir) / "profile_invocations.jsonl").exists() + + asyncio.run(_test()) + + def test_no_output_path_no_error(self): + async def _test(): + cb = ProfileCallback(print_report=False, save_report=True) + await cb.before_run(FakeRunConfig()) + await cb.before_invoke({"messages": []}) + await cb.after_invoke(_make_response(idx=0)) + + result = Result(responses=[], total_requests=1, output_path=None) + await cb.after_run(result) # Should not raise + + asyncio.run(_test()) + + def test_stats_persist_save_load(self, profiler_with_data): + _, result = profiler_with_data + with tempfile.TemporaryDirectory() as tmpdir: + result.output_path = tmpdir + result.save() + + loaded = Result.load(tmpdir, load_responses=False) + assert "profile_total_wall_clock" in loaded.stats + assert "profile_ttft-average" in loaded.stats + assert loaded.stats["profile_cache_hit_rate"] == pytest.approx(0.25) + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +class TestProfileEdgeCases: + def test_all_requests_fail(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + for i in range(3): + await profiler.before_invoke({"messages": []}) + await profiler.after_invoke(_make_response(idx=i, error="fail")) + + result = Result(responses=[], total_requests=3) + await profiler.after_run(result) + + assert profiler.report.successful_requests == 0 + assert profiler.report.failed_requests == 3 + assert profiler.report.ttft_stats == {} + + asyncio.run(_test()) + + def test_no_requests(self, profiler): + async def _test(): + await profiler.before_run(FakeRunConfig()) + result = Result(responses=[], total_requests=0) + await profiler.after_run(result) + + assert profiler.report.total_requests == 0 + assert profiler.report.successful_requests == 0 + + asyncio.run(_test()) + + def test_missing_ttft(self, profiler): + """Requests with None TTFT should not break computation.""" + + async def _test(): + await profiler.before_run(FakeRunConfig()) + await profiler.before_invoke({"messages": []}) + resp = InvocationResponse( + response_text="ok", + time_to_first_token=None, + time_to_last_token=2.0, + num_tokens_input=50, + num_tokens_output=30, + error=None, + ) + await profiler.after_invoke(resp) + + result = Result(responses=[], total_requests=1) + await profiler.after_run(result) + + # Should not crash, TTFT stats should be empty + assert profiler.report.ttft_stats == {} + + asyncio.run(_test())