Skip to content

AGENTS.md

Latest commit

 

History

History
516 lines (398 loc) · 54.3 KB

AGENTS.md

File metadata and controls

516 lines (398 loc) · 54.3 KB
Raw

AGENTS.md

Guidelines for AI coding agents working in the MLPerf Inference Endpoint Benchmarking System repository.

Project Overview

High-performance benchmarking tool for LLM inference endpoints targeting 50k+ QPS. Python 3.12+, Apache 2.0 licensed.

Common Commands

# Development setup
uv sync --extra dev --extra test
uv run pre-commit install

# Testing
uv run pytest                                        # All tests (excludes slow/performance)
uv run pytest -m unit                                # Unit tests only
uv run pytest -m integration                         # Integration tests only
uv run pytest --cov=src --cov-report=html            # With coverage
uv run pytest -xvs tests/unit/path/to/test_file.py  # Single test file

# Code quality (run before commits)
uv run pre-commit run --all-files

# Local testing with echo server
uv run python -m inference_endpoint.testing.echo_server --port 8765
uv run inference-endpoint probe --endpoints http://localhost:8765 --model test-model

# CLI usage
uv run inference-endpoint benchmark offline --endpoints URL --model NAME --dataset PATH
uv run inference-endpoint benchmark online --endpoints URL --model NAME --dataset PATH --load-pattern poisson --target-qps 100
uv run inference-endpoint benchmark from-config --config config.yaml

Backward-compatible setup (pip + venv)

Does not use uv.lock — dependency versions may differ from the lockfile.

python3.12 -m venv venv && source venv/bin/activate
pip install -e ".[dev,test]"
pre-commit install

# After activating the venv, commands run without the `uv run` prefix:
pytest                                        # All tests (excludes slow/performance)
pytest -m unit                                # Unit tests only
pytest -m integration                         # Integration tests only
pytest --cov=src --cov-report=html            # With coverage
pytest -xvs tests/unit/path/to/test_file.py  # Single test file

# Code quality — MUST run before every commit, no exceptions
pre-commit run --all-files

# Local testing with echo server
python -m inference_endpoint.testing.echo_server --port 8765
inference-endpoint probe --endpoints http://localhost:8765 --model test-model

# CLI usage
inference-endpoint benchmark offline --endpoints URL --model NAME --dataset PATH
inference-endpoint benchmark online --endpoints URL --model NAME --dataset PATH --load-pattern poisson --target-qps 100
inference-endpoint benchmark from-config --config config.yaml

Architecture

Core Data Flow

Dataset Manager --> Load Generator --> Endpoint Client --> External Endpoint
                        |
                   EventPublisher (events PUB)
                        |
        +---------------+---------------+
        |                               |
   EventLoggerService            MetricsAggregatorService
   (events.jsonl)                (registry → publisher)
                                        |
                                  metrics PUB
                                        |
                                Main process SUB
                                        |
                                Report.from_snapshot

Key Components

Component Location Purpose
Load Generator src/inference_endpoint/load_generator/ Central orchestrator: BenchmarkSession owns the lifecycle, PhaseIssuer drives per-phase execution, TimedIssueStrategy/BurstStrategy/ConcurrencyStrategy control timing. Emits ERROR before COMPLETE for failed queries (metrics aggregator depends on this order).
Endpoint Client src/inference_endpoint/endpoint_client/ Multi-process HTTP workers communicating via ZMQ IPC. HTTPEndpointClient is the main entry point
Dataset Manager src/inference_endpoint/dataset_manager/ Loads JSONL, HuggingFace, CSV, JSON, Parquet datasets. Dataset base class with load_sample()/num_samples() interface
Metrics Aggregator src/inference_endpoint/async_utils/services/metrics_aggregator/ Subprocess. Subscribes to events, aggregates per-sample metrics into a MetricsRegistry (counters + HDR-histogram series + raw values), publishes MetricsSnapshot over IPC PUB at a configurable cadence (SessionState: INITIALIZE → LIVE → DRAINING → {COMPLETE | INTERRUPTED}). Final snapshot is atomically written to final_snapshot.json as the primary Report source; the terminal pub/sub frame is a TUI "run finished" signal only.
Report src/inference_endpoint/metrics/report.py Report.from_snapshot(dict) — pure-function builder consuming the dict form (snapshot_to_dict). Reads final_snapshot.json directly via json.loads (no Struct decode). Plumbs complete = (state == "complete" and n_pending_tasks == 0); renders an explicit warning for INTERRUPTED runs.
Config src/inference_endpoint/config/, endpoint_client/config.py Pydantic-based YAML schema (schema.py), HTTPClientConfig (single Pydantic model for CLI/YAML/runtime), RuntimeSettings
CLI src/inference_endpoint/main.py, commands/benchmark/cli.py cyclopts-based, auto-generated from schema.py and HTTPClientConfig Pydantic models. Flat shorthands via cyclopts.Parameter(alias=...)
Async Utils src/inference_endpoint/async_utils/ LoopManager (uvloop + eager_task_factory), ZMQ transport layer, generic MessageCodec[T]-parametrized pub/sub, event publisher
OpenAI/SGLang src/inference_endpoint/openai/, sglang/ Protocol adapters and response accumulators for different API formats. openai_completions adapter (completions_adapter.py) sends pre-tokenized token IDs to /v1/completions, bypassing the server chat template — required for gpt-oss-120b on vLLM. sglang adapter sends to /generate via input_ids. Both apply Harmonize() client-side.
DeepSeek-R1 (MLPerf) src/inference_endpoint/evaluation/scoring.py (LegacyMLPerfDeepSeekR1Scorer), examples/07_DeepSeekR1_Example/ MLPerf DeepSeek-R1 accuracy. TensorRT-LLM is OpenAI-compatible, so it is served via api_type: openai / openai_completions (no dedicated trtllm adapter). The combined multi-subset eval (math500/aime/gpqa/mmlu_pro/livecodebench) is the official MLCommons eval_accuracy.py, run out-of-process via uv run --project against the isolated subproject at src/inference_endpoint/evaluation/legacy_mlperf_deepseek_r1/ (a uv subproject excluded from the parent wheel; mirrors the VBench pattern). The example feeds the exact MLPerf prompt via pre-tokenized input_tokens to /v1/completions.
VideoGen src/inference_endpoint/videogen/ Adapter for video-generation endpoints (e.g. trtllm-serve POST /v1/videos/generations, used by MLPerf WAN2.2-T2V-A14B). Defaults to response_format=video_path (server saves video to shared storage and returns path) to avoid large byte payloads. Accuracy mode also runs on video_path: the adapter mirrors the path into response_output so the event log carries it to VBenchScorer (see evaluation/scoring.py), which scores videos via VBench from a sibling uv subproject at examples/09_Wan22_VideoGen_Example/accuracy/ (vbench's transformers==4.33.2 + numpy<2 pins are incompatible with the parent env, so it runs out-of-process via uv run --project). Dataset is ingested via the generic JSONL loader.
SWE-bench src/inference_endpoint/dataset_manager/predefined/swe_bench/, src/inference_endpoint/evaluation/swe_bench_scorer.py, src/inference_endpoint/evaluation/swebench_service/ SWEBench predefined dataset (HuggingFace princeton-nlp/SWE-bench_Verified or _Lite; ACCURACY_ONLY=True). SWEBenchScorer sets SKIP_ENDPOINT_PHASE=True and bypasses the built-in accuracy phase entirely: it delegates agent execution and grading to the configured SWE-bench service via accuracy_config.extras.swebench_service_url. The service is an isolated uv subproject; its host owns Docker/runtime execution, artifacts, and credentials, while the benchmark client remains the report-producing entrypoint.
Compliance (submission checker) src/inference_endpoint/compliance/checker.py, scripts/check_compliance.py Validates a completed run's report directory against a registered ruleset. check_submission(report_dir, ruleset, model) reads the resolved config.yaml plus scorer output (accuracy/accuracy_results.json for accuracy, scores.json for the agentic perf run) and runs config-lock (deterministic + single-stream), the accuracy gate (score >= factor x reference, factor 0.97 for Edge-Agentic), and run validity (0 dropped turns). Server-side launch flags (--reasoning off, --ctx-size) aren't in client artifacts, so they're surfaced as manual attestations. CLI: scripts/check_compliance.py REPORT_DIR (exit 0 = pass).
Compliance (audit tests) src/inference_endpoint/compliance/, commands/audit.py MLPerf compliance audits. AuditTest protocol + AuditRunSpec/AuditRunArtifacts + registry (compliance/__init__.py); OutputCachingAudit (compliance/audit_test/output_caching_test.py, which also owns the QPS-specific AuditRunStats) implements MLPerf TEST04 output-caching detection — reference phase (distinct samples) vs. fixed-sample audit phase, comparing QPS against threshold. commands/audit.py:run_audit runs phases via AuditTest.plan_runs/validate, writing audit_result.json/verify_<TEST>.txt atomically via compliance/result.py. Enabled by the audit: YAML block; cli._run runs it after the main benchmark (upstream MLPerf order: perf run, then TEST04), or standalone with audit.only: true. Perf-only by default (a phase may opt into accuracy via AuditRunSpec.test_mode, but this is unused today).

Hot-Path Architecture

Multi-process, event-loop design optimized for throughput:

  • BenchmarkSession thread schedules samples with busy-wait timing
  • Worker processes (N instances) handle HTTP requests via ZMQ IPC
  • Uses eager_task_factory and uvloop for minimal async overhead
  • CPU affinity support (cpu_affinity.py) for performance tuning
  • Custom HTTP connection pooling (http.py) with httptools parser

Metrics Aggregator subprocess (pub/sub)

The aggregator is a separate process (python -m inference_endpoint.async_utils.services.metrics_aggregator) that subscribes to events and publishes MetricsSnapshot messages. Key facts for working in this layer:

  • Series storage: each SeriesSampler keeps three parallel views: O(1) cheap rollups (count/total/min/max/sum_sq, exact), an HDR Histogram (cheap live percentiles), and an in-memory array.array of raw values (for exact percentiles in the COMPLETE snapshot). Hot path is registry.record(name, value) — no allocation, no I/O.
  • Counter API: registry.increment(name, delta=1) for sample-event counters. registry.set_counter(name, value) only for the three derived-duration counters (total_duration_ns max-of-elapsed, tracked_duration_ns sum-of-blocks, legacy_loadgen_window_duration_ns first-issue→last-issued-completion span for LoadGen-parity QPS/TPS).
  • Lifecycle: INITIALIZE (constructed, awaiting first STARTED) → LIVE (run in progress, ticking every --publish-interval seconds) → DRAINING (set on ENDED; tick continues; bounded by the --drain-timeout budget — schema default 0 = unlimited) → terminal: COMPLETE (clean end via publish_final, exact stats) or INTERRUPTED (signal-handler-triggered final via SIGTERM/SIGINT; best-effort partial stats). Drain timeout detected by consumers as state == COMPLETE and n_pending_tasks > 0; interrupted runs are detected as state == INTERRUPTED directly.
  • Final delivery is dual-path with separated concerns: publish_final atomically writes final_snapshot.json (tmp + fsync(file) + rename + fsync(parent_dir)) — this is the primary Report source — AND emits the terminal-state snapshot over pub/sub as a TUI shutdown signal. Each path is wrapped in its own try/except so one failure cannot suppress the other. Main process consumer reads final_snapshot.json (via json.loads to dict, no Struct decode); falls back to the subscriber's latest live snapshot only if the file is missing (e.g. SIGKILL / OOM before the signal handler ran). The dict form is the canonical consumer contract (see snapshot_to_dict).
  • Early stopping (on by default): series registered with register_series(..., tail_latency=True) (today ttft/tpot/latency) get MLPerf early-stopping percentile estimates on the COMPLETE (exact) snapshot — a compact early_stopping_percentiles map in result_summary.json whose keys mirror the percentiles grid (≥ p50) with estimate-or-null values; rich detail is INFO-logged. On by default (cold-path only; the exact path shares one in-place sort between the percentile grid and the estimates); settings.early_stopping.enabled: false / --no-early-stopping opts out. Confidence/tolerance are LoadGen constants. Pure math in metrics/early_stopping.py; post-hoc recomputation from any run's events.jsonl via scripts/early_stopping_estimate_from_events.py. See docs/early_stopping.md.
  • Histogram bucket edges are dynamic per snapshot: log-spaced over the observed [min, max]. Bucket count is fixed at construction; consumers MUST re-render from the snapshot's (lo, hi, count) triples each frame and MUST NOT track bucket-by-index across snapshots.

CLI Modes

CLI is auto-generated from config/schema.py Pydantic models via cyclopts. Fields annotated with cyclopts.Parameter(alias="--flag") get flat shorthands; all other fields get auto-generated dotted flags (kebab-case).

  • CLI mode (offline/online): cyclopts constructs OfflineBenchmarkConfig/OnlineBenchmarkConfig (subclasses in config/schema.py) directly from CLI args. Type locked via Literal. --dataset is repeatable with TOML-style format [perf|acc:]<path>[,key=value...] (e.g. --dataset data.csv,samples=500,parser.prompt=article). Full accuracy support via accuracy_config.eval_method=pass_at_1 etc.
  • YAML mode (from-config): BenchmarkConfig.from_yaml_file() loads YAML, resolves env vars, and auto-selects the right subclass via Pydantic discriminated union. Optional --timeout/--mode overrides via config.with_updates().
  • eval: Not yet implemented (raises CLIError with a tracking issue link)

Config Construction & Validation

Both CLI and YAML produce the same subclass via Pydantic discriminated union on type:

CLI offline/online:  cyclopts → OfflineBenchmarkConfig/OnlineBenchmarkConfig → with_updates(datasets) → run_benchmark
YAML from-config:    from_yaml_file(path) → discriminated union → same subclass → run_benchmark

OfflineBenchmarkConfig and OnlineBenchmarkConfig (in config/schema.py) inherit BenchmarkConfig:

  • type: locked via Literal[TestType.OFFLINE] / Literal[TestType.ONLINE]
  • settings: OfflineSettings (hides load pattern) / OnlineSettings
  • submission_ref, benchmark_mode: show=False on base class

Validation is layered:

  1. Field-level (Pydantic): Field(ge=0) on durations, Field(ge=-1) on workers, Literal on benchmark_mode
  2. Field validators: workers != 0 check
  3. Model validator (_resolve_and_validate): streaming AUTO resolution, model name from submission_ref, load pattern vs test type, cross-field duration check, duplicate datasets

Load Patterns

  • max_throughput: Offline burst (all queries at t=0)
  • poisson: Fixed QPS with Poisson arrival distribution
  • concurrency: Fixed concurrent requests

Compliance Audits

Orthogonal to the main run: a YAML-only audit: block (show=False, no CLI flag) on BenchmarkConfig selects an AuditTest. If audit: is set, cli._run freezes a shared report_dir, runs the main benchmark, then calls commands/audit.py:run_audit — the upstream MLPerf order (perf run, then TEST04). Unlike upstream there is no SUT reset between stages; see docs/compliance_audit_plan.md ("Run ordering"). run_audit:

  1. Builds a per-phase config that is performance-only by default (drops accuracy datasets so no phase re-issues or re-scores them), overridable per-AuditRunSpec via test_mode (ACC/BOTH keeps accuracy datasets — supported but currently unused; every registered audit runs perf-only), and, after the first phase's dataset loads, calls AuditTest.validate(cfg, dataset_size, load_pattern) to bounds-check the test's sample counts/indices and load pattern before any phase issues load (reuses the loaded dataset — no extra load).
  2. Runs each AuditRunSpec phase (from AuditTest.plan_runs) back-to-back under its own <report_dir>/<label>/ subdir via setup_benchmark/run_benchmark_async. A phase whose Report.complete is False (drain timeout / interrupt) aborts with ExecutionError — no result on partial data.
  3. Calls AuditTest.verify(...) and atomically writes audit_result.json (durable record) then verify_<TEST>.txt (validator marker).

run_audit returns an AuditResult; cli._run maps passed to the process exit code (0 PASS / 1 FAIL; setup/IO errors → non-zero via the standard error path). A crashed or interrupted main run skips the audit entirely (as upstream: TEST04 only runs once a perf result exists); a FAIL doesn't cost the submission its perf report — cli._run raises only after both stages have run. The output-caching audit (MLPerf TEST04) is the only registered audit today (AuditTestId.OUTPUT_CACHING_TEST); add new audits by implementing the AuditTest protocol (including validate) and adding the instance to the AUDIT_TESTS map in compliance/__init__.py.

Code Organization

src/inference_endpoint/
├── main.py                    # Entry point + CLI app: cyclopts app, commands, error formatter, run()
├── exceptions.py              # CLIError, ExecutionError, InputValidationError, SetupError
├── commands/                  # Command execution logic
│   ├── benchmark/
│   │   ├── __init__.py
│   │   ├── cli.py             # benchmark_app: offline, online, from-config subcommands
│   │   ├── execute.py         # Phased orchestration: setup_benchmark/run_benchmark_async/finalize_benchmark + BenchmarkContext; run_benchmark runs the main benchmark (cli._run dispatches run_audit when audit: is set)
│   │   ├── profiling.py       # Profiler-trigger protocol (vLLM /start_profile,/stop_profile) + ProfileController (URL derivation + start/stop/payload lifecycle)
│   │   ├── accuracy.py        # AccuracyConfiguration + per-dataset scoring (_score_accuracy, OSL/response-count rollups, write_accuracy_results)
│   │   └── pipeline.py        # MetricsPipeline: async context manager for the ZMQ + metrics-aggregator/event-logger subprocess lifecycle (__aenter__/__aexit__/start/drain_and_build_report) + snapshot→Report
│   ├── audit.py               # run_audit() — compliance audit orchestrator (phases → verify → result)
│   ├── probe.py               # ProbeConfig + execute_probe()
│   ├── info.py                # execute_info()
│   ├── validate.py            # execute_validate()
│   └── init.py                # execute_init()
├── compliance/                # MLPerf compliance audits
│   ├── __init__.py            # AuditTest protocol + AuditRunSpec/AuditRunArtifacts + test registry
│   ├── result.py              # AuditResult + atomic write_result (audit_result.json + verify_<TEST>.txt)
│   └── audit_test/             # audit tests, added to the AUDIT_TESTS map in compliance/__init__.py
│       ├── __init__.py         # package marker; registry wiring lives in compliance/__init__.py
│       ├── output_caching_test.py  # OutputCachingAudit (MLPerf TEST04): caching detection (reference vs fixed-sample QPS)
│       └── README.md           # TEST04 output-caching audit usage (WAN 2.2 example)
├── core/
│   ├── types.py               # APIType, Query, QueryResult, StreamChunk, QueryStatus (msgspec Structs)
│   └── record.py              # EventRecord — transport record used by event logger and ZMQ transport
├── load_generator/
│   ├── session.py             # BenchmarkSession, PhaseIssuer, PhaseConfig, PhaseResult, SessionResult
│   ├── strategy.py            # TimedIssueStrategy, BurstStrategy, ConcurrencyStrategy, LoadStrategy
│   ├── agentic_inference_strategy.py # AgenticInferenceStrategy
│   ├── conversation_manager.py # ConversationManager, ConversationState
│   ├── sample_order.py        # SampleOrder, WithoutReplacement/WithReplacement/SingleSampleOrder, create_sample_order
│   └── delay.py               # poisson_delay_fn, make_delay_fn
├── endpoint_client/
│   ├── http_client.py         # HTTPEndpointClient - main client interface
│   ├── worker.py              # Worker process implementation
│   ├── worker_manager.py      # Manages worker lifecycle
│   ├── http.py                # ConnectionPool, HttpRequestTemplate, raw HTTP
│   ├── http_sample_issuer.py  # Bridges load generator to HTTP client
│   ├── config.py              # HTTPClientConfig (single Pydantic model — CLI/YAML/runtime)
│   ├── adapter_protocol.py    # HttpRequestAdapter protocol
│   ├── accumulator_protocol.py # Response accumulation protocol
│   ├── cpu_affinity.py        # CPU pinning
│   └── utils.py               # Port range helpers
├── async_utils/
│   ├── loop_manager.py        # LoopManager (uvloop + eager_task_factory)
│   ├── runner.py              # run_async() — uvloop + eager_task_factory entry point for CLI commands
│   ├── event_publisher.py     # Async event pub/sub
│   ├── services/
│   │   ├── event_logger/      # EventLoggerService: writes EventRecords to JSONL/SQLite
│   │   └── metrics_aggregator/  # MetricsAggregatorService: subscribes to events, publishes MetricsSnapshot
│   │       ├── __main__.py     # Subprocess entry: --metrics-socket, --metrics-output-dir, --publish-interval, --hdr-sig-figs, --n-histogram-buckets
│   │       ├── aggregator.py   # MetricsAggregatorService (event router); SessionState lifecycle; tracked_samples_failed
│   │       ├── snapshot.py     # MetricsSnapshot wire schema + SessionState enum + msgpack codec
│   │       ├── registry.py     # MetricsRegistry, CounterSampler, SeriesSampler (HDR + raw array.array + cheap rollups)
│   │       ├── publisher.py    # MetricsPublisher (tick task + atomic disk fallback)
│   │       ├── subscriber.py   # MetricsSnapshotSubscriber (latest + COMPLETE snapshot capture)
│   │       ├── metrics_table.py # In-flight sample rows + trigger dispatch (TTFT/TPOT/ISL/OSL)
│   │       └── token_metrics.py # BatchTokenizer (live thread lane + drain-only sharded pool) + TokenBatchQueue (defer-to-flush buffer, owns the live flush loop) for ISL/OSL/TPOT
│   └── transport/             # ZMQ-based IPC transport layer
│       ├── protocol.py        # Transport protocols + TransportConfig + MessageCodec[T]
│       └── zmq/               # ZMQ implementation (context, pubsub, transport, ZMQTransportConfig)
├── dataset_manager/
│   ├── dataset.py             # Dataset base class, DatasetFormat enum
│   ├── factory.py             # Dataset factory
│   ├── transforms.py          # ColumnRemap and other transforms
│   └── predefined/            # Built-in datasets (aime25, cnndailymail, gpqa, etc.)
├── metrics/
│   ├── report.py              # Report.from_snapshot(MetricsSnapshot); display + JSON serialization
│   ├── metric.py              # Metric types (Throughput, etc.)
│   ├── early_stopping.py      # MLPerf LoadGen early-stopping percentile estimates (pure math; see docs/early_stopping.md)
│   └── results_plots.py       # Standardized run-artifact plots (matplotlib-guarded); CLI: scripts/plot_results.py
├── config/
│   ├── schema.py              # Single source of truth: Pydantic models + cyclopts annotations
│   ├── runtime_settings.py    # RuntimeSettings + SampleOrderSpec dataclasses
│   ├── ruleset_base.py        # BenchmarkSuiteRuleset base
│   ├── ruleset_registry.py    # Ruleset registry
│   ├── user_config.py         # UserConfig dataclass for ruleset user overrides
│   ├── rulesets/mlcommons/    # MLCommons-specific rules, datasets, models
│   └── templates/             # YAML config templates (_template.yaml minimal, _template_full.yaml all defaults)
├── openai/                    # OpenAI-compatible API types and adapters
│   ├── types.py               # OpenAI response types (chat + text completion)
│   ├── openai_adapter.py      # Chat completions adapter (/v1/chat/completions)
│   ├── openai_msgspec_adapter.py  # msgspec-based chat completions adapter (fast path)
│   ├── completions_adapter.py # Text completions adapter (/v1/completions, pre-tokenized input)
│   ├── accumulator.py         # Streaming response accumulator (shared by chat + completions)
│   └── harmony.py             # openai_harmony integration
├── sglang/                    # SGLang API adapter (/generate with input_ids)
├── videogen/                  # Video generation adapter (e.g. WAN2.2 T2V workload)
│   ├── __init__.py
│   ├── types.py               # Pydantic: VideoPathRequest, VideoPathResponse, VideoPayloadResponse
│   └── adapter.py             # VideoGenAdapter (HttpRequestAdapter) + VideoGenAccumulator (no-op)
├── evaluation/                # Accuracy evaluation (extractor, scoring, livecodebench)
│   └── swebench_service/      # Isolated uv service for Docker-backed SWE-bench runs
├── compliance/                # Submission compliance checks (config-lock, accuracy gate, run validity)
│   ├── __init__.py
│   └── checker.py             # check_submission() + Check/ComplianceReport (Edge-Agentic ruleset)
├── plugins/                   # Plugin system
├── profiling/                 # line_profiler integration, pytest plugin
├── testing/
│   ├── echo_server.py         # Local echo server for testing
│   ├── max_throughput_server.py # Max throughput test server
│   └── docker_server.py       # Docker-based server management
└── utils/
    ├── logging.py             # Logging setup
    ├── version.py             # Version info
    ├── dataset_utils.py       # Dataset utilities
    └── benchmark_httpclient.py # HTTP client throughput benchmarking utility

tests/
├── conftest.py                # Shared fixtures (echo/oracle servers, datasets, settings)
├── test_helpers.py            # Test utility functions
├── unit/                      # Unit tests (mirror src/ structure)
├── integration/               # Integration tests (real servers, end-to-end)
│   ├── endpoint_client/       # HTTP client integration tests
│   └── commands/              # CLI command integration tests
├── performance/               # Performance benchmarks (pytest-benchmark)
└── datasets/                  # Test data (dummy_1k.jsonl, squad_pruned/)

Development Standards

Code Style and Pre-commit Hooks

  • Formatter/Linter: ruff (line-length 88, target Python 3.12)
  • Type checking: mypy (via pre-commit)
  • Formatting: ruff-format (double quotes, space indent)
  • License headers: Required on all Python files (enforced by pre-commit hook scripts/add_license_header.py)
  • Conventional commits: feat:, fix:, docs:, test:, chore:

Pre-commit Hooks

All of these run automatically on commit:

  • trailing-whitespace, end-of-file-fixer, check-yaml, check-merge-conflict, debug-statements
  • ruff (lint + autofix) and ruff-format
  • mypy type checking
  • prettier for YAML/JSON/Markdown
  • License header enforcement
  • regenerate-templates: auto-regenerates YAML config templates from schema defaults when schema.py, config.py, or regenerate_templates.py change

IMPORTANT: Always run pre-commit run --all-files before every commit. Hooks may modify files (prettier, ruff-format, license headers). If files are modified, stage the changes and commit once. Never commit without running pre-commit first.

See Development Guide for full setup and workflow details.

Data Types & Serialization

  • Core types (Query, QueryResult, StreamChunk): msgspec.Struct with frozen=True, array_like=True, gc=False, omit_defaults=True
  • Config types: pydantic.BaseModel for validation
  • Enums: str, Enum pattern for serializable enums (e.g., LoadPatternType, APIType)
  • Serialization: msgspec.json for hot-path (ZMQ transport), pydantic for config

Testing

Coverage target: >90% for all new code.

Test markers:

@pytest.mark.unit           # Unit tests
@pytest.mark.integration    # Integration tests (may need servers)
@pytest.mark.slow           # Skip in CI
@pytest.mark.performance    # No timeout, skip in CI
@pytest.mark.run_explicitly # Only run when explicitly selected

Async tests: Use @pytest.mark.asyncio — strict mode is configured globally in pyproject.toml (asyncio_mode = "strict"). Do NOT pass mode="strict" to the marker — it's not a valid argument.

Key fixtures (defined in tests/conftest.py):

  • mock_http_echo_server — real HTTP echo server on dynamic port
  • mock_http_oracle_server — dataset-driven response server
  • dummy_dataset — in-memory test dataset
  • hf_squad_dataset — HuggingFace squad dataset
  • max_throughput_runtime_settings, poisson_runtime_settings, concurrency_runtime_settings — preset configs

Test data: tests/assets/datasets/dummy_1k.jsonl (1000 samples), tests/assets/datasets/squad_pruned/

Performance Guidelines

These apply especially to code in the hot path (load generator, endpoint client, transport):

  • No match statements in hot paths — use dict dispatch instead
  • Use dataclass(slots=True) for frequently instantiated classes (or msgspec.Struct)
  • Prefer generators over list comprehensions for large datasets
  • Minimize async suspends in hot path code
  • Use msgspec over json/pydantic for serialization in the data path
  • Connection pooling: The HTTP client uses custom ConnectionPool with httptools parser — not aiohttp/requests
  • Event loop: uvloop with eager_task_factory via LoopManager
  • IPC: ZMQ-based transport between main process and worker processes

Key Dependencies

Package Purpose
uvloop Performance-optimized event loop
httptools Fast HTTP parser for custom connection pool
msgspec Fast serialization for core types, ZMQ transport, MetricsSnapshot wire
pyzmq ZMQ IPC between main process and workers / metrics aggregator
hdrhistogram HDR Histogram for live percentiles in metrics aggregator (C-backed)
pydantic Configuration validation
cyclopts CLI framework — auto-generates flags from Pydantic
duckdb Data aggregation
transformers Tokenization for OSL reporting

Files to NOT Modify

  • src/inference_endpoint/openai/openai_types_gen.py — auto-generated, excluded from ruff/pre-commit
  • src/inference_endpoint/openai/openapi.yaml — OpenAI API spec, excluded from pre-commit

Documentation references — no local-only artifacts

Code, comments, docstrings, tests, and committed Markdown MUST NOT reference paths that aren't in the repository. This includes anything under .gitignored directories (e.g. .cursor_artifacts/, design scratch dirs, untracked working notes), absolute paths to a contributor's workstation, build outputs, or unmerged branch artifacts. A reviewer fetching the PR should be able to follow every reference cited in the diff.

Why: stale references compound — See foo.md §3 is meaningless once foo.md is gone, renamed, or never existed in the merged tree, and rotting cross-references are how docs stop being trusted. AI agents reading the codebase later treat dangling pointers as ground truth and propagate confusion.

Allowed:

  • Paths to files committed to the repo (docs/..., src/..., tests/..., README.md, etc.).
  • External URLs (issue trackers, PRs, RFCs, vendor docs).
  • Generic references to environment/setup that the reader is expected to create themselves (e.g. source .venv/bin/activate in a setup README, where .venv is the user's local venv).

Disallowed examples:

  • See .cursor_artifacts/foo_design.md §2 — .cursor_artifacts/ is gitignored.
  • See ~/work/notes/architecture.txt — contributor-local.
  • Tracked in metrics_pubsub_design_v5.md test impact section — same gitignored doc.

If a design doc is worth referencing from the source tree, commit it to docs/ or inline the relevant content into the code comment / docstring. For one-off rationale that won't survive the conversation, prefer a self-contained explanation in the comment itself rather than a pointer to ephemera.

Comments and docstrings — describe current state, not development history

Don't write comments or docstrings that narrate iteration on the codebase. Pointers to abandoned approaches, prior implementations, or design pivots belong in the PR description and git log, not in the source tree. They rot quickly: the prior implementation is gone, the reader has no way to evaluate the comparison, and the scaffolding accumulates with every iteration. Future readers — humans and AI agents alike — treat the comment as if it describes load-bearing context when it's actually historical clutter.

This applies especially to AI-assisted development, where it's tempting to leave a paper trail of "we tried X first, then switched to Y" inside the source. That paper trail belongs in the PR description.

Disallowed patterns:

  • # Originally used X, but switched to Y for ...
  • # An earlier implementation did X — this version does Y
  • # Removed the foo parameter / # Replaced bar with baz
  • # Note: this used to be sync but is now async
  • # Regression: an earlier shape did X — even in regression-test docstrings, drop the narrative framing.
  • # An alternative design considered ... but was rejected because ... (unless the rejected alternative is a common path a future contributor might re-attempt — in that case, frame it as "don't do X because Y", not as developer history).

Allowed:

  • Current rationale: # Uses dict dispatch — hot path measured at sub-ms (describes why the current design exists; no history).
  • Regression context that doesn't narrate the prior bug's discovery: # Without this check, value > hdr_high silently corrupts the histogram total (describes the bug being prevented, framed as a current invariant — not "we used to have a bug here").
  • Inline TODO/FIXME pointing at a tracking issue (URL or issue number, not "we plan to do X eventually").

Rule of thumb: if removing the comment would leave the code's intent unchanged for someone seeing it for the first time, the comment is fine. If the comment only makes sense to someone who saw the prior version, delete it.

Comments and docstrings — no line-of-code estimates

Don't reference line counts in comments or docstrings. Phrasing like "one ~20-line block", "this 50-LOC walker", "(~30 LOC)" rots the moment the code is refactored, conveys no actionable meaning, and adds maintenance burden — every edit nearby risks invalidating the count, and there is no warning when it does.

Disallowed patterns:

  • # Manual mapping (one ~20-line block) is the source of truth
  • # ~50 LOC of mirror types below
  • # This function is 30 lines; consider splitting if it grows past 50

Allowed:

  • Describe what the code does and why, not how big it is.
  • If size is genuinely the point (e.g. a perf comment about an inlined hot path), name the property that matters: "kept inlined to avoid the call overhead measured at X µs", not "this is 12 lines".

Why: the value of a comment is the invariant it pins. A line count isn't an invariant — it's an accident of formatting and current scope. Future readers will trust the number; it will be wrong; you've now created a misleading comment instead of an absent one.

Keeping AGENTS.md Up to Date

This file is the source of truth for AI agents working in this repo. If it is stale or wrong, every AI-assisted session starts from a broken foundation.

When to Update

Update AGENTS.md as part of any PR that includes a significant refactor, meaning:

  • Moved, renamed, or deleted modules/packages — update the Code Organization tree and Key Components table
  • Changed architectural boundaries (e.g., new IPC transport, replaced pydantic with msgspec for config) — update Architecture and Data Types sections
  • Added or removed CLI commands/subcommands — update CLI Modes and Common Commands
  • Changed test infrastructure (new fixtures, changed markers, new test directories) — update Testing section
  • Added or removed key dependencies — update Key Dependencies table
  • Changed build/tooling (new pre-commit hooks, changed ruff config, new CI steps) — update docs/DEVELOPMENT.md
  • Changed hot-path patterns (new transport, changed serialization, new performance constraints) — update Performance Guidelines

How to Update

  1. Treat AGENTS.md changes as part of the refactor itself — include them in the same PR, not as a follow-up
  2. Keep descriptions factual and concise — what exists and where, not aspirational design docs
  3. If you add a new top-level module under src/inference_endpoint/, add it to both the Key Components table and the Code Organization tree
  4. If you remove something, remove it from AGENTS.md — stale entries are worse than missing ones

Reviewer Checklist

When reviewing PRs with significant structural changes, verify:

  • Code Organization tree matches the actual directory structure post-merge
  • Key Components table reflects any moved/renamed/new components
  • No references to deleted files, classes, or modules remain

Common AI Coding Pitfalls

Known failure modes when AI tools generate code for this project. Reference these during code review of AI-assisted PRs.

Architecture & Design

  • Inventing abstractions that don't exist: AI may introduce new base classes, registries, or factory patterns that don't match the existing architecture. This project uses concrete types and explicit wiring — check that new code follows existing patterns rather than imposing unfamiliar frameworks.
  • Misunderstanding the multi-process boundary: The endpoint client uses separate worker processes (not threads) communicating over ZMQ. AI-generated code often assumes shared memory, passes unpicklable objects across processes, or adds synchronization primitives (locks, semaphores) that don't work cross-process.
  • Confusing hot-path vs. cold-path: AI tends to treat all code uniformly. Code in load_generator/, endpoint_client/worker.py, and async_utils/transport/ is latency-critical. Pydantic validation, excessive logging, or try/except blocks in these paths will degrade throughput.

Serialization & Types

  • Using the wrong serialization library: This project uses msgspec for hot-path data and pydantic for config. AI frequently defaults to json.dumps/json.loads, dataclasses, or applies pydantic where msgspec is required. If the type is a msgspec.Struct, encode/decode with msgspec.json, not stdlib json.
  • Breaking msgspec Struct constraints: Core types (Query, QueryResult, StreamChunk) are frozen=True with array_like=True. AI may try to mutate fields directly (must use force_setattr), add mutable default fields, or assume dict-like serialization when the wire format is array-based.
  • Adding dataclass where msgspec.Struct is expected: If neighboring types use msgspec, new types in the same module should too. AI defaults to @dataclass out of habit.

Async & Concurrency

  • Mixing sync and async incorrectly: AI may await in a sync context, call blocking I/O inside an async def, or use asyncio.run() when a loop is already running (this project manages its own loops via LoopManager).
  • Creating new event loops: Workers and the main process already have managed event loops. AI may call asyncio.new_event_loop() or asyncio.run() which conflicts with the existing LoopManager/uvloop setup.
  • Ignoring eager_task_factory: The project uses Python 3.12's eager_task_factory for performance. AI-generated code that creates coroutines expecting lazy scheduling will behave differently than expected.

Testing

  • Generating mock-heavy tests for integration scenarios: This project has real echo/oracle server fixtures. AI tends to mock HTTP calls even when mock_http_echo_server or mock_http_oracle_server fixtures exist and should be used.
  • Missing test markers: Every test function needs @pytest.mark.unit, @pytest.mark.integration, or another marker. AI-generated tests almost always omit markers, which breaks CI filtering.
  • Wrong asyncio marker: Tests must use bare @pytest.mark.asyncio — strict mode is configured globally in pyproject.toml. Do NOT pass mode="strict" to the marker (it's not a valid argument and will cause errors). AI sometimes hallucinates this parameter.
  • Fabricating fixture names: AI may invent fixtures that don't exist in conftest.py. Always check that referenced fixtures actually exist before using them.

Code Style & Repo Conventions

  • Missing license headers: Every Python file needs the Apache 2.0 SPDX header. AI never generates these — the pre-commit hook will add them, but be aware of this when reviewing diffs.
  • Importing removed or renamed modules: After refactors, AI (working from stale context) may import old module paths. Always verify imports resolve to actual files.
  • Over-documenting: AI generates verbose docstrings, inline comments explaining obvious code, and type annotations on trivial variables. This project prefers minimal comments — only where the why isn't obvious from the code.
  • Adding backwards-compatibility shims: If something was renamed or removed, AI may add re-exports, aliases, or deprecation wrappers. In this project, just delete the old thing and update all call sites.
  • Empty except blocks: Every except block must contain either a comment explaining why the exception is ignored, or a logging statement. Bare except: pass without explanation is disallowed. AI often generates empty handlers — always add the reason.
  • No lazy imports: All imports must be at the top of the file. Imports inside functions, methods, or conditional blocks (other than TYPE_CHECKING) are disallowed. The only exceptions are: (1) circular import avoidance with a documenting comment, (2) optional dependencies with a top-level try/except that sets the import to None, (3) security sandboxing code that intentionally restricts imports.

Dependency & Environment

  • Adding new dependencies without justification: AI may pip install or add imports for packages not in pyproject.toml. Any new runtime, dev, or test dependency must be justified, added to the correct optional group, and pinned to an exact version (==). After adding a dependency, run pip-audit (included in dev extras) to verify it has no known vulnerabilities. When adding dependencies, use uv add <package>==<version> to update both pyproject.toml and uv.lock atomically, then run uv run pip-audit to check for vulnerabilities. Note: [build-system] requires is also pinned to exact versions for reproducibility.
  • Using requests/aiohttp for HTTP: This project has its own HTTP client (endpoint_client/http.py) using httptools. AI defaults to requests or aiohttp — these should not appear in production code (test dependencies are fine).