Skip to content

Refactor commands/benchmark/execute.py — split the benchmark God-module #384

Description

@viraatc

Problem

src/inference_endpoint/commands/benchmark/execute.py has grown to ~1270 lines and concentrates most of the run lifecycle's responsibilities in one module. _run_benchmark_async alone is ~370 lines (L663–L1030) and interleaves unrelated concerns, which makes it hard to read, test in isolation, and extend without touching the hot path.

Distinct responsibilities currently living in this one file:

  • Setup & context (setup_benchmark, BenchmarkContext, _default_report_path)
  • Tokenizer/accuracy resolution (_check_tokenizer_exists, _resolve_accuracy_components)
  • Dataset loading (_load_datasets)
  • Phase construction (_build_phases)
  • Profiling triggers (_derive_profile_urls, _post_profile, _render_profile_status, _write_profiling_section)
  • Power monitoring (_build_power_payload) — added in DRAFT: feat(power): vendor-agnostic power/energy monitoring #383
  • Snapshot/IO (_load_final_snapshot_from_disk, _salvage_tmpfs, _write_scoring_artifacts)
  • Async run orchestration (_run_benchmark_async, run_benchmark_async)
  • Finalization/reporting (finalize_benchmark, run_benchmark)

Why now

#383 (power monitoring) had to thread start/stop hooks and a payload builder into an already-large _run_benchmark_async + finalize_benchmark. Profiling did the same earlier. Each optional sidecar feature (profiling, power, and future ones) adds more interleaved branches to the same long functions. This is a recurring smell — the module wants an extension seam.

Suggested direction (non-prescriptive)

  • Extract the optional "around-the-run" features (profiling, power, future telemetry) behind a small uniform run-hook seam (e.g. on_run_start / on_phase_start / on_run_stop / contribute_to_finalize) so _run_benchmark_async stops growing per feature.
  • Move profiling helpers to a profiling/-adjacent module and power helpers already live in power/execute.py should orchestrate, not implement, each sidecar.
  • Split finalization (scoring artifacts, sibling-JSON writers, report.txt sections) into a finalize/reporting module.
  • Break _run_benchmark_async into named stages (service launch → warmup → measure → drain → collect-results).

Behavior should be unchanged; this is structure-only. Good incremental targets: (1) hook seam + move profiling/power render out, (2) split finalize, (3) decompose _run_benchmark_async.

Acceptance

  • No behavior change; result_summary.json / report.txt / profiling.json / power.json outputs identical.
  • execute.py orchestrates; feature implementations live in their own modules.
  • Existing unit/integration tests pass; long functions decomposed into independently testable units.

Filed as follow-up context from #383.


Also in scope: split config/schema.py + redundancy audit

src/inference_endpoint/config/schema.py is ~1210 lines with ~30 classes/enums in one module (enums, OSL distribution, model params, datasets, accuracy, runtime, load-pattern, warmup, drain, profiling, power, settings, endpoint, the BenchmarkConfig discriminated union, …). It's the second God-module next to execute.py.

Split (behavior-preserving)

Break into a config/schema/ package by concern, e.g.:

  • enums.py (LoadPatternType, OSLDistributionType, DatasetType, EvalMethod, ScorerMethod, TestMode, StreamingMode, TestType, TEMPLATE_TYPE_MAP)
  • model.py (ModelParams, OSLDistribution)
  • dataset.py (Dataset, AccuracyConfig)
  • runtime.py (RuntimeConfig, LoadPattern, WarmupConfig, DrainConfig)
  • sidecars.py (ProfilingConfig, PowerConfig) — or keep each feature's config beside its package
  • settings.py (Settings, Offline/OnlineSettings, EndpointConfig)
  • benchmark.py (BenchmarkConfig, Offline/OnlineBenchmarkConfig, discriminator)
  • config/schema.py → thin re-export shim (from .schema.enums import *, etc.) so the ~many from inference_endpoint.config.schema import X call sites keep working unchanged — maximal backward compatibility, zero churn for importers.

Redundancy audit (DRY)

Audit and de-duplicate while splitting:

  • Repeated model_config = ConfigDict(extra="forbid", frozen=True) on nearly every model → a shared frozen base model.
  • The @model_serializer(mode="wrap") "drop-when-disabled/irrelevant" pattern now appears in 2+ places (RuntimeConfig's use_legacy_loadgen_qps_metrics, PowerConfig's disabled-collapse) → factor a helper.
  • Repeated cyclopts.Parameter(alias=...) + duplicated help/description strings (help and description are often identical) → derive one from the other.
  • Overlapping validation idioms across _resolve_and_validate / per-model validators.
  • Any dead fields / unused imports surfaced by the move.

Acceptance

  • No public import breaks (config.schema re-export shim), no behavior change, identical generated templates + result_summary.json.
  • ruff/mypy/pytest green; the audit removes duplication rather than relocating it.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions