Code Patterns | NVIDIA AIPerf Documentation

Code examples for common development tasks. Referenced from CLAUDE.md.

CLI Command Pattern

Commands live in src/aiperf/cli_commands/, one file per command. They are lazily loaded via import strings in aiperf.cli — modules are only imported when their command is invoked:

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.common.config import ServiceConfig, UserConfig
4 
5 app = App(name="profile")
6 
7 @app.default
8 def profile(user_config: UserConfig, service_config: ServiceConfig | None = None) -> None:
9     """Run the Profile subcommand."""
10     from aiperf.cli_runner import run_system_controller  # heavy import deferred
11 
12     run_system_controller(user_config, service_config)

Conventions:

Export a single App named app.
Hyphenate multi-word commands: App(name="analyze-trace").
Keep module-level imports minimal; heavy deps go inside the function body.
Heavy implementation logic lives in a cli.py inside the owning domain package (e.g. aiperf/plugin/cli.py), lazily imported at call time.

Service Pattern

Services run in separate processes via bootstrap.py:

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 service:
2   my_service:
3     class: aiperf.my_module.my_service:MyService
4     description: My custom service
5     metadata:
6       required: true
7       auto_start: true

Config types:

ServiceConfig: infrastructure (ZMQ ports, logging level)
UserConfig: benchmark params (endpoints, loadgen settings)

Model Pattern

Use AIPerfBaseModel for data, BaseConfig for configuration:

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")

Message Pattern

Messages require message_type field and handler decorator:

1 from aiperf.common.messages import Message
2 from aiperf.common.hooks import on_message
3 
4 class MyMsg(Message):
5     message_type: MessageType = 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))

Auto-subscription happens during @on_init phase.

Plugin System Pattern

YAML-based registry with lazy-loading:

1 # plugins.yaml
2 endpoint:
3   chat:
4     class: aiperf.endpoints.openai_chat:ChatEndpoint
5     description: OpenAI Chat Completions endpoint
6     metadata:
7       endpoint_path: /v1/chat/completions
8       supports_streaming: true
9       produces_tokens: true
10       tokenizes_input: true
11       supports_audio: true
12       supports_images: true
13       supports_videos: true
14       metrics_title: LLM Metrics

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

Error Handling Pattern

Log errors and publish ErrorDetails in messages:

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)))

Logging Pattern

Use lambda for expensive log messages:

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")

Testing Pattern

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

Auto-fixtures (always active): asyncio.sleep runs instantly, RNG=42, singletons reset.