AIPerf Plugin System

The AIPerf plugin system provides a flexible, extensible architecture for customizing benchmark behavior. It uses YAML-based configuration with lazy loading, priority-based conflict resolution, and dynamic enum generation.

Table of Contents

• Overview
  • Terminology
  • Key Components
• Architecture
• Plugin Categories
• Using Plugins
• Creating Custom Plugins
• Plugin Configuration
• CLI Commands
• Advanced Topics

Overview

The plugin system enables:

• Extensibility: Add custom endpoints, exporters, and timing strategies without modifying core code
• Lazy Loading: Classes load on first access, avoiding circular imports
• Conflict Resolution: Higher priority plugins override lower priority ones
• Type Safety: Auto-generated enums provide IDE autocomplete
• Validation: Validate plugins without importing them

Terminology

Hierarchy:

Registry (singleton)
└── Package (1+) ─── discovered via entry points
    └── Manifest (1+ per package) ─── plugins.yaml files
        └── Category (1+)
            └── Entry (1+) ─── PluginEntry
                ├── Class ─── lazy-loaded Python class
                └── Metadata ─── optional typed config

Key Components

Architecture

Discovery Flow

Entry Points → plugins.yaml → Pydantic Validation → Registry
                                                      ↓
                              get_class() → Import Module → Cache

Registry Singleton Pattern

The plugin registry follows the singleton pattern with module-level exports:

1
from aiperf.plugin import plugins
2
from aiperf.plugin.enums import PluginType
3
4
# Get a plugin class by name
5
EndpointClass = plugins.get_class(PluginType.ENDPOINT, "chat")
6
7
# Iterate all plugins in a category
8
for entry, cls in plugins.iter_all(PluginType.ENDPOINT):
9
    print(f"{entry.name}: {entry.description}")

Plugin Categories

AIPerf supports 33 plugin categories organized by function, including api_router and public_dataset_loader:

Timing Categories

Dataset Categories

Endpoint and Transport Categories

Processing Categories

Accuracy Categories

UI and Selection Categories

Service Categories

Visualization and Telemetry Categories

Infrastructure Categories (Internal)

Sweep / Adaptive Search Categories

Using Plugins

1
from aiperf.plugin import plugins
2
from aiperf.plugin.enums import PluginType, EndpointType
3
4
# Get class by name, enum, or full path
5
ChatEndpoint = plugins.get_class(PluginType.ENDPOINT, "chat")
6
ChatEndpoint = plugins.get_class(PluginType.ENDPOINT, EndpointType.CHAT)
7
ChatEndpoint = plugins.get_class(PluginType.ENDPOINT, "aiperf.endpoints.openai_chat:ChatEndpoint")
8
9
# Iterate plugins
10
for entry, cls in plugins.iter_all(PluginType.ENDPOINT):
11
    print(f"{entry.name}: {entry.class_path}")
12
13
# Get metadata (raw dict or typed)
14
metadata = plugins.get_metadata("endpoint", "chat")
15
endpoint_meta = plugins.get_endpoint_metadata("chat")  # Returns EndpointMetadata

Creating Custom Plugins

Quick Start (4 steps):

Minimal Endpoint Example

1
# my_package/endpoints/custom_endpoint.py
2
class MyCustomEndpoint(BaseEndpoint):
3
    def format_payload(self, request_info: RequestInfo) -> dict[str, Any]:
4
        turn = request_info.turns[-1]
5
        texts = [content for text in turn.texts for content in text.contents if content]
6
        return {"prompt": texts[0] if texts else ""}
7
8
    def parse_response(self, response: InferenceServerResponse) -> ParsedResponse | None:
9
        if json_obj := response.get_json():
10
            return ParsedResponse(perf_ns=response.perf_ns, data=TextResponseData(text=json_obj.get("text", "")))
11
        return None

1
# yaml-language-server: $schema=https://raw.githubusercontent.com/ai-dynamo/aiperf/refs/heads/main/src/aiperf/plugin/schema/plugins.schema.json
2
# my_package/plugins.yaml
3
schema_version: "1.0"
4
endpoint:
5
  my_custom:
6
    class: my_package.endpoints.custom_endpoint:MyCustomEndpoint
7
    description: Custom endpoint for my API.
8
    metadata: { endpoint_path: /v1/generate, supports_streaming: true, produces_tokens: true, tokenizes_input: true, metrics_title: My Custom Metrics }

Plugin Configuration

categories.yaml Schema

Defines plugin categories with their protocols and metadata schemas:

1
# yaml-language-server: $schema=https://raw.githubusercontent.com/ai-dynamo/aiperf/refs/heads/main/src/aiperf/plugin/schema/categories.schema.json
2
schema_version: "1.0"
3
4
endpoint:
5
  protocol: aiperf.endpoints.protocols:EndpointProtocol
6
  metadata_class: aiperf.plugin.schema.schemas:EndpointMetadata
7
  enum: EndpointType
8
  description: |
9
    Endpoints define how to format requests and parse responses for different APIs.
10
  internal: false  # Set to true for infrastructure categories

plugins.yaml Schema

Registers plugin implementations:

1
# yaml-language-server: $schema=https://raw.githubusercontent.com/ai-dynamo/aiperf/refs/heads/main/src/aiperf/plugin/schema/plugins.schema.json
2
schema_version: "1.0"
3
4
endpoint:
5
  chat:
6
    class: aiperf.endpoints.openai_chat:ChatEndpoint
7
    description: OpenAI Chat Completions endpoint.
8
    priority: 0  # Higher priority wins conflicts
9
    metadata:
10
      endpoint_path: /v1/chat/completions
11
      supports_streaming: true
12
      produces_tokens: true
13
      tokenizes_input: true
14
      metrics_title: LLM Metrics

Metadata Schemas

Category-specific metadata is validated against Pydantic models in aiperf.plugin.schema.schemas:

CLI Commands

$
$ aiperf plugins endpoint chat
$
╭───────────────── endpoint:chat ─────────────────╮
$
│ Type: chat                                      │
$
│ Category: endpoint                              │
$
│ Package: aiperf                                 │
$
│ Class: aiperf.endpoints.openai_chat:ChatEndpoint│
$
│                                                 │
$
│ OpenAI Chat Completions endpoint. Supports      │
$
│ multi-modal inputs and streaming responses.     │
$
╰─────────────────────────────────────────────────╯

Advanced Topics

Conflict Resolution

API Reference

1
# Runtime registration (testing)
2
plugins.register("endpoint", "test", TestEndpoint, priority=10)
3
plugins.reset_registry()  # Reset to initial state
4
5
# Dynamic enum generation
6
MyEndpointType = plugins.create_enum(PluginType.ENDPOINT, "MyEndpointType", module=__name__)
7
8
# Validation without importing
9
errors = plugins.validate_all(check_class=True)  # {category: [(name, error), ...]}
10
11
# Reverse lookup
12
name = plugins.find_registered_name(PluginType.ENDPOINT, ChatEndpoint)  # "chat"
13
14
# Package metadata
15
pkg = plugins.get_package_metadata("aiperf")  # PackageInfo(version, author, ...)

Type Safety: get_class() returns typed results (e.g., type[EndpointProtocol]) with IDE autocomplete.

Built-in Plugins Reference

Endpoints

Timing Strategies

Arrival Patterns

Dataset Composers

UI Types

Accuracy Benchmarks

Accuracy Graders

Troubleshooting

Plugin Not Found

TypeNotFoundError: Type 'my_plugin' not found for category 'endpoint'.

Solutions:

1. Verify the plugin is registered in plugins.yaml
2. Check the entry point is defined in pyproject.toml
3. Reinstall the package in the active environment: uv pip install -e .
4. Run aiperf plugins --validate to check for errors

Module Import Errors

ImportError: Failed to import module for endpoint:my_plugin

Solutions:

1. Verify the class path format: module.path:ClassName
2. Check all dependencies are installed
3. Verify the module is importable: python -c "import module.path"

Class Not Found

AttributeError: Class 'MyClass' not found

Solutions:

1. Verify the class name matches exactly (case-sensitive)
2. Ensure the class is exported from the module
3. Run aiperf plugins --validate for detailed error

Conflict Resolution Issues

If your plugin is being shadowed by another:

1. Use higher priority: priority: 10 in plugins.yaml
2. Access by full class path: plugins.get_class("endpoint", "my_pkg.endpoints:MyEndpoint")
3. Check aiperf plugins to see which packages are loaded