Repository navigation
Refactor commands/benchmark/execute.py — split the benchmark God-module #384
Description
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_asyncstops growing per feature. - Move profiling helpers to a
profiling/-adjacent module and power helpers already live inpower/—execute.pyshould orchestrate, not implement, each sidecar. - Split finalization (scoring artifacts, sibling-JSON writers, report.txt sections) into a
finalize/reporting module. - Break
_run_benchmark_asyncinto 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.jsonoutputs identical. execute.pyorchestrates; 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 packagesettings.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 ~manyfrom inference_endpoint.config.schema import Xcall 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'suse_legacy_loadgen_qps_metrics, PowerConfig's disabled-collapse) → factor a helper. - Repeated
cyclopts.Parameter(alias=...)+ duplicated help/description strings (help anddescriptionare 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.schemare-export shim), no behavior change, identical generated templates +result_summary.json. - ruff/mypy/pytest green; the audit removes duplication rather than relocating it.
Activity
arekay-nv commentedon Aug 25, 2026
closed via #421