AIPerf Code Patterns

1
# aiperf/cli.py — register with lazy import strings
2
app.command("aiperf.cli_commands.profile:app", name="profile")

1
# aiperf/cli_commands/profile.py — thin command definition
2
from cyclopts import App
3
from aiperf.config.flags import CLIConfig
4
5
app = App(name="profile")
6
7
@app.default
8
def profile(*, cli_config: CLIConfig) -> None:
9
    """Run the Profile subcommand."""
10
    from aiperf.cli_utils import exit_on_error
11
    from aiperf.config.loader.errors import ConfigurationError
12
13
    with exit_on_error(title="Error Running AIPerf System", show_traceback=False):
14
        from aiperf.config.flags.resolver import resolve_config
15
        from aiperf.config.loader import build_benchmark_plan
16
17
        config = resolve_config(cli_config, cli_config.config_file)
18
        plan = build_benchmark_plan(config)
19
20
    with exit_on_error(
21
        title="Error Running AIPerf System",
22
        quiet_for=(ConfigurationError,),
23
    ):
24
        from aiperf.cli_runner import run_benchmark  # heavy import deferred
25
26
        run_benchmark(plan)

1
# src/aiperf/config/flags/cli_config.py — add the field in its section block
2
my_new_flag: Annotated[
3
    int | None,
4
    Field(
5
        ge=1,
6
        description="One-line user-facing description rendered in --help "
7
        "and docs/cli-options.md. Mention units, defaults, and obvious "
8
        "interactions with other flags.",
9
    ),
10
    CLIParameter(
11
        name=("--my-new-flag",),       # CLI flag name (independent of attr name)
12
        group=Groups.LOAD_GENERATOR,   # cyclopts --help group; pick from list below
13
    ),
14
] = None

1
class MyService(BaseComponentService):
2
    @on_message(MessageType.MY_MSG)
3
    async def _handle(self, msg: MyMsg) -> None:
4
        await self.publish(ResponseMsg(data=msg.data))

1
from pydantic import Field
2
from aiperf.common.models import AIPerfBaseModel
3
4
class Record(AIPerfBaseModel):
5
    ts_ns: int = Field(description="Timestamp in nanoseconds")
6
    value: float = Field(description="Measured value")

1
from aiperf.common.messages import Message, MessageTypeT
2
from aiperf.common.hooks import on_message
3
4
class MyMsg(Message):
5
    message_type: MessageTypeT = MessageType.MY_MSG
6
    data: list[Record] = Field(description="Records to process")
7
8
# In service class:
9
@on_message(MessageType.MY_MSG)
10
async def _handle(self, msg: MyMsg) -> None:
11
    await self.publish(OtherMsg(data=msg.data))

1
from aiperf.plugin import plugins
2
from aiperf.plugin.enums import PluginType
3
4
EndpointClass = plugins.get_class(PluginType.ENDPOINT, 'chat')

1
try:
2
    await risky_operation()
3
except Exception as e:
4
    self.error(f"Operation failed: {e!r}")
5
    await self.publish(ResultMsg(error=ErrorDetails.from_exception(e)))

1
# Expensive - lambda defers evaluation
2
self.debug(lambda: f"Processing {len(self._items())} items")
3
4
# Cheap - direct string is fine
5
self.info("Starting service")

1
from pydantic import Field
2
from aiperf.common.finite import FiniteFloat
3
from aiperf.common.models import AIPerfBaseModel
4
5
class MetricSummary(AIPerfBaseModel):
6
    mean: FiniteFloat = Field(description="Sample mean (must be finite)")
7
    std: FiniteFloat | None = Field(
8
        default=None,
9
        description="Sample stddev; None means insufficient samples",
10
    )
11
    p99: FiniteFloat | None = Field(
12
        default=None,
13
        description="99th percentile latency in ms; None means no samples",
14
    )

1
import orjson
2
from aiperf.common.finite import scrub_non_finite
3
4
def export_records_json(records: list[Record], out_path: Path) -> None:
5
    payload = {"records": [r.model_dump() for r in records]}
6
    out_path.write_bytes(orjson.dumps(scrub_non_finite(payload)))

1
from aiperf.common.finite import is_finite_value
2
3
def maybe_record_throughput(value: float) -> None:
4
    if not is_finite_value(value):
5
        self.warning(lambda: f"Skipping non-finite throughput: {value!r}")
6
        return
7
    self._records.append(value)

1
from aiperf.common.finite import nan_safe_mean, nan_safe_std
2
3
# Partial-failure samples may contain NaN; np.mean would propagate.
4
samples = [r.latency_ms for r in records]  # may contain NaN
5
mean = nan_safe_mean(samples)               # None if no finite values
6
std = nan_safe_std(samples, ddof=1)         # None if < 2 finite values

1
# WRONG: raw float field accepts NaN silently
2
class BadSummary(AIPerfBaseModel):
3
    p99: float = Field(description="99th percentile latency")  # accepts NaN
4
5
# WRONG: orjson silently coerces NaN/inf to JSON null
6
out_path.write_bytes(orjson.dumps({"p99": float("nan")}))
7
# Result on disk: {"p99": null}  -- indistinguishable from "missing"
8
9
# WRONG: np.mean propagates NaN through Pareto/BO downstream
10
import numpy as np
11
mean = float(np.mean([1.0, 2.0, float("nan")]))  # NaN, poisons callers

1
from aiperf.common.path_safety import safe_read_template_path
2
3
body = safe_read_template_path(user_string)
4
if body is None:
5
    # safety check failed — caller picks the fallback semantic
6
    body = user_string          # template "path or inline" idiom
7
    # or: raise ValueError(f"Template file not readable: {user_string!r}")

1
import pytest
2
from aiperf.plugin import plugins
3
from aiperf.plugin.enums import PluginType
4
from tests.harness import mock_plugin
5
6
@pytest.mark.asyncio
7
async def test_async_operation():
8
    result = await some_async_func()
9
    assert result.status == "ok"
10
11
@pytest.mark.parametrize("input,expected",
12
    [
13
        ("a", 1),
14
        ("b", 2),
15
    ]
16
)  # fmt: skip
17
def test_with_params(input, expected):
18
    assert process(input) == expected
19
20
def test_with_mock_plugin():
21
    with mock_plugin(PluginType.ENDPOINT, "test", MockClass):
22
        assert plugins.get_class(PluginType.ENDPOINT, "test") == MockClass

1
# src/aiperf/exporters/internal_metrics_console_exporter.py — gated single-table
2
class ConsoleInternalMetricsExporter(ConsoleMetricsExporter):
3
    """Console exporter for INTERNAL framework metrics, gated on dev mode."""
4
5
    title = "[yellow]NVIDIA AIPerf | Internal Metrics[/yellow]"
6
    require_flags = MetricFlags.INTERNAL    # records must have this flag
7
    exclude_flags = MetricFlags.ERROR_ONLY  # records with this flag are hidden
8
    console_groups = None                   # single combined table; ignore groups
9
10
    def _check_enabled(self, exporter_config: ExporterConfig) -> None:
11
        if not (Environment.DEV.MODE and Environment.DEV.SHOW_INTERNAL_METRICS):
12
            raise ConsoleExporterDisabled("Internal metrics are not enabled, ...")

1
from aiperf.plot.models.uncertainty import BenchmarkPoint, LatencyThroughputUncertaintyData
2
3
point = BenchmarkPoint(
4
    x_mean=10.0, y_mean=100.0,
5
    x_ci_low=8.0, x_ci_high=12.0,
6
    y_ci_low=90.0, y_ci_high=110.0,
7
    cov_xy=5.0,  # enables rotated ellipses; None for axis-aligned
8
    label="concurrency=4",
9
)
10
data = LatencyThroughputUncertaintyData(
11
    points=[point],
12
    confidence_level=0.95,
13
    title="Latency vs Throughput",
14
    x_label="Latency (ms)",
15
    y_label="Throughput (tok/s)",
16
)

1
from aiperf.plot.models.uncertainty import (
2
    BenchmarkPoint, LatencyThroughputUncertaintyData, UncertaintySeries,
3
)
4
5
# One series per experiment variant (e.g., request_count=20 vs 50).
6
# When `series` is non-empty it overrides `points`; see get_series().
7
data = LatencyThroughputUncertaintyData(
8
    series=[
9
        UncertaintySeries(name="request_count=20", points=[
10
            BenchmarkPoint(x_mean=5.0, y_mean=50.0, x_ci_low=4.0, x_ci_high=6.0,
11
                           y_ci_low=45.0, y_ci_high=55.0, label="c=2", n_runs=10),
12
            BenchmarkPoint(x_mean=15.0, y_mean=120.0, x_ci_low=13.0, x_ci_high=17.0,
13
                           y_ci_low=110.0, y_ci_high=130.0, label="c=10", n_runs=8),
14
        ]),
15
        UncertaintySeries(name="request_count=50", points=[
16
            BenchmarkPoint(x_mean=6.0, y_mean=48.0, x_ci_low=4.5, x_ci_high=7.5,
17
                           y_ci_low=42.0, y_ci_high=54.0, label="c=2", n_runs=10),
18
            BenchmarkPoint(x_mean=18.0, y_mean=110.0, x_ci_low=15.0, x_ci_high=21.0,
19
                           y_ci_low=100.0, y_ci_high=120.0, label="c=10", n_runs=10),
20
        ]),
21
    ],
22
    confidence_level=0.95,
23
    title="Latency vs Throughput by Request Count",
24
    x_label="Latency (ms)",
25
    y_label="Throughput (tok/s)",
26
)

1
from aiperf.plot.core.plot_generator import PlotGenerator
2
3
pg = PlotGenerator()
4
fig = pg.create_uncertainty_plot(data)
5
fig.write_image("output.png")  # Kaleido export

1
from aiperf.plot.exporters import export_uncertainty_matplotlib
2
from pathlib import Path
3
4
export_uncertainty_matplotlib(data, Path("output.png"))

1
from aiperf.plot.geometry import compute_ellipse_vertices, compute_axis_aligned_ellipse_vertices
2
import numpy as np
3
4
cov = np.array([[4.0, 1.0], [1.0, 9.0]])
5
vertices = compute_ellipse_vertices(cov, center=(10.0, 100.0), confidence_level=0.95)
6
# Returns list of (x, y) tuples forming a closed polygon

1
# src/aiperf/common/validators/orchestrator_v1.py - gate convention
2
raise NotImplementedError(
3
    f"conversation '{conv.conversation_id}' turn {idx}: "
4
    f"prerequisite kind '{prereq.kind}' not supported by v1 orchestrator"
5
)

1
# src/aiperf/endpoints/raw_endpoint.py - composing a mixin
2
class RawEndpoint(JMESPathResponseMixin, BaseEndpoint):
3
    def __init__(self, model_endpoint: ModelEndpointInfo, **kwargs: Any) -> None:
4
        super().__init__(model_endpoint, **kwargs)
5
        self._init_response_parser()

1
from typing import Protocol, runtime_checkable
2
3
from aiperf.common.messages.inference_messages import MetricRecordsData
4
from aiperf.common.models import CreditPhaseStats
5
6
OTelResultData = MetricRecordsData | CreditPhaseStats
7
8
9
@runtime_checkable
10
class OTelResultsStrategyProtocol(Protocol):
11
    """Public extension point for new streamed OTel result domains.
12
13
    A strategy owns exactly one ``OTelResultData`` variant and emits its
14
    telemetry via ``OTelStrategyContextProtocol``. Strategies MUST NOT touch
15
    OTel instruments, the fanout queue, or the MLflow client directly — the
16
    context owns instrument lifecycle and cross-strategy state so fanout
17
    stays consistent across strategies.
18
    """
19
20
    def supports(self, record_data: OTelResultData) -> bool:
21
        """Return True iff ``record_data`` is the variant this strategy consumes.
22
23
        Implementations use ``isinstance`` against a single concrete type —
24
        strategies are mutually exclusive by record type.
25
        """
26
        ...
27
28
    async def process(self, record_data: OTelResultData) -> None:
29
        """Emit telemetry for ``record_data`` without blocking the hot path.
30
31
        Instrument access goes through the context's ``get_or_create_*``
32
        factories, which enqueue fanout events rather than touching the OTel
33
        SDK inline. Raising is permitted; the processor is best-effort, so
34
        the records manager logs and swallows the failure.
35
        """
36
        ...

1
from aiperf.post_processors.strategies.core import (
2
    OTelResultData,
3
    OTelResultsStrategyProtocol,
4
    OTelStrategyContextProtocol,
5
)
6
7
8
class MetricResultsStrategy(OTelResultsStrategyProtocol):
9
    """Streams per-request metric records as histogram observations."""
10
11
    def __init__(self, context: OTelStrategyContextProtocol) -> None:
12
        self._context = context
13
14
    def supports(self, record_data: OTelResultData) -> bool:
15
        return isinstance(record_data, MetricRecordsData)
16
17
    async def process(self, record_data: OTelResultData) -> None:
18
        # Emit histogram observations for each metric in the record.
19
        ...
20
21
22
class TimingResultsStrategy(OTelResultsStrategyProtocol):
23
    """Streams phase-level timing snapshots using counters and gauges."""
24
25
    def __init__(self, context: OTelStrategyContextProtocol) -> None:
26
        self._context = context
27
28
    def supports(self, record_data: OTelResultData) -> bool:
29
        return isinstance(record_data, CreditPhaseStats)
30
31
    async def process(self, record_data: OTelResultData) -> None:
32
        # Emit counter deltas and gauge snapshots for timing data.
33
        ...

1
for strategy in self._strategies:
2
    if strategy.supports(record_data):
3
        await strategy.process(record_data)

1
import multiprocessing as mp
2
from aiperf.common.environment import Environment
3
4
event_queue = mp.Queue(maxsize=Environment.OTEL.MAX_BUFFERED_RECORDS)  # default 10 000

1
from queue import Empty, Full
2
3
def _queue_fanout_event(self, event_type: str, payload: dict[str, Any]) -> None:
4
    """Enqueue streaming event for the fanout process without blocking the event loop."""
5
    if self._fanout_queue is None:
6
        return
7
8
    event = {"type": event_type, "payload": payload}
9
    try:
10
        self._fanout_queue.put_nowait(event)
11
        self._fanout_sent_events += 1
12
    except Full:
13
        if self._drop_oldest_fanout_event():
14
            try:
15
                self._fanout_queue.put_nowait(event)
16
                self._fanout_sent_events += 1
17
                return
18
            except Full:
19
                pass
20
        self._record_fanout_drop(
21
            "OTel fanout queue remained full; dropping newest event"
22
        )
23
    except Exception as exc:
24
        self.warning(f"Failed to enqueue OTel fanout event: {exc!r}")