refactor(benchmark): split execute.py into execute/profiling/accuracy/pipeline

[arekay-nv]

Refactor: split commands/benchmark/execute.py into cohesive modules

Branch: refactor/execute-split · Scope: strictly behavior-preserving · Status: ready for review (3 review rounds complete, all findings addressed)

Why

src/inference_endpoint/commands/benchmark/execute.py had grown to 1691 lines and
concentrated almost all CLI benchmark orchestration in one file. Two functions dominated:
_run_benchmark_async (~403 lines) interleaved ZMQ setup, two subprocess launches, HTTP-client
construction, agentic wiring, profiling triggers, the session run, and a ~110-line finally
doing five unrelated teardown jobs; _score_accuracy (~172 lines) mixed tokenizer loading, uuid
bounding, scoring, response counts, and OSL. Duplicated tmpfs-salvage / AccuracyConfiguration
construction / profile start-stop blocks and best-effort except-sprawl compounded it.

The goal was to split the file along its natural seams without changing any observable
behavior — same exception types (and therefore exit codes), same on-disk artifacts, same ZMQ
connect-before-bind and teardown ordering, same QPS/report semantics, same tmpfs ownership.

What changed

execute.py split into four modules under src/inference_endpoint/commands/benchmark/:

Module	Lines	Responsibility
execute.py	1691 → 1076	Thin orchestrator: setup_benchmark / run_benchmark_async / finalize_benchmark / run_benchmark, _run_benchmark_async, _build_phases, _load_datasets, _resolve_accuracy_components, _PerfPhaseTimeout, BenchmarkContext/BenchmarkResult/ResponseCollector, and new helpers (_create_issuer, _build_agentic_strategy, _wire_on_sample_complete, _write_report_artifacts, _summarize_and_log_metrics)
profiling.py (new)	222	Profiler-trigger protocol (_derive_profile_urls, _post_profile, _render_profile_status, _write_profiling_section) + new ProfileController (owns URL derivation + start/stop/payload lifecycle)
accuracy.py (new)	405	AccuracyConfiguration, _phase_osl_stats, _phase_response_counts, _accuracy_uuid_bound, _score_accuracy, new _load_osl_backend, new write_accuracy_results
pipeline.py (new)	345	new MetricsPipeline (ZMQ + publisher + subscriber + aggregator/event-logger subprocess lifecycle) + _build_aggregator_args, _build_event_logger_args, _build_report_from_snapshot, _load_final_snapshot_from_disk

Two new abstractions

• ProfileController (profiling.py) — collapses the profile URL pre-derivation, the
  /start_profile (PERFORMANCE-phase-only) fire, the /stop_profile (only for status==200
  starts) fire, and the {engine, starts, stops} payload build into one object. Disabled and
  inert when engine is None; raises up-front (fail-before-run) when an engine is set but
  endpoints are empty.
• MetricsPipeline (pipeline.py) — owns the ZMQ context, event publisher, snapshot
  subscriber, and the two service subprocesses via explicit lifecycle methods — deliberately
  not an async context manager (an early design that a review rejected). The run has three
  distinct teardown paths a single __aexit__ can't express cleanly:
  • start() — bring up, unwinding partial resources if service launch fails (no ZMQ leak);
  • drain_and_build_report() — graceful drain (publisher close → wait for services → source the
    final snapshot → build the Report), run on both clean-finish and session-failure;
  • abort() — connect-failure fast path: kill services without a graceful drain.
    close() exits the ZMQ scope idempotently.

Import contracts preserved

Everything commands/audit.py imports from execute (BenchmarkResult, TestMode,
_salvage_tmpfs, finalize_benchmark, run_benchmark_async, setup_benchmark) and cli.py's
imports (resolve_report_dir, run_benchmark) stay in execute — audit.py, cli.py, the
package __init__, and the TEST04 audit tests are unchanged. Runtime import direction is
one-way (execute → {profiling, accuracy, pipeline}); the sibling modules reference
BenchmarkContext only under TYPE_CHECKING, so there is no cycle.

Three unit test files were repointed to the new module paths (imports / mock.patch targets
only — no assertion changes): test_benchmark.py, test_score_accuracy.py,
test_benchmark_final_snapshot.py.

Review process (3 rounds, Codex + Claude in parallel each round)

Every round ran two independent reviewers against the diff vs main; each finding was fixed and
re-verified before the next round.

#	Finding	Sev	Reviewer	Fix
F1	Publisher double-closed on drain path (safe — close() is idempotent — but violated the null-after-drain invariant)	Low	Claude R1	Null publisher/subscriber after closing in drain_and_build_report
F2	ProfileController built before the run try, so a misconfig ValueError bypassed tmpfs/pbar cleanup	Low	Claude R1	Construct it as the first statement inside the try (also closes a latent service-subprocess leak that existed in the original)
F3	tmpfs-salvage vs pbar/zmq ordering on the failure path	Info	Claude R1	Confirmed benign — no change
F4	tmpfs/event-log/metrics mkdirs ran before the cleanup try; a mkdir failure after tmpfs_dir existed would leak it	Med	Codex R1	Compute paths before the try, move the mkdirs inside it
F5	Flat finally: pbar.close(); pipe.close() — if pbar.close() raised, ZMQ leaked and tmpfs wasn't salvaged (the original guaranteed both via the ZMQ with)	Med	Codex R2	Nested try: <body> finally: (try: pbar.close() finally: pipe.close()) wrapped by the outer except BaseException salvage
—	Stale execute_mod comment/var name in test_score_accuracy.py	Nit	Claude R3	Renamed to scoring_mod, comment corrected

Round 3 verdict (both reviewers): clean / ready to commit. Both traced the ZMQ scope exiting
exactly once on all five teardown paths (clean success, session ExecutionError, connect
SetupError→abort(), launcher.launch raising, mid-run KeyboardInterrupt), confirmed
pipe.close() survives a pbar.close() failure, that cleanup exceptions route to tmpfs salvage,
and that no result/profiler/publisher/http_client is possibly-unbound at any reachable use.

Verification

This was developed on macOS, where a chunk of the suite fails for
platform reasons unrelated to this change (pin_loadgen() raises
UnsupportedPlatformError on darwin; the metrics-aggregator/sched_* paths are
Linux-only). To separate refactor regressions from environment noise, every suite
was run both with the change and against a clean main tree (git stash) — the
pass/fail profile is identical, so the refactor adds zero failures.

• Affected unit tests (test_benchmark, test_score_accuracy,
  test_benchmark_final_snapshot, compliance test_output_caching): 233 passed —
  these all run on macOS and cover the moved code directly.
• Full unit suite: main 58 failed / 1542 passed → this branch 58 failed / 1542
  passed (identical). All 58 failures are in metrics_aggregator/ (44) and
  endpoint_client/test_cpu_affinity.py (14) — untouched modules that fail on macOS
  only; none is in commands/benchmark/.
• Integration commands/ (drives run_benchmark end-to-end): main 29 failed / 8
  passed → this branch 29 failed / 8 passed (identical). The failures are the
  pin_loadgen()-requires-Linux path, which the original setup_benchmark invoked the
  same way.
• End-to-end smoke (echo server, --no-cpu-affinity): the full refactored path runs
  clean and writes every artifact — config.yaml, events.jsonl, sample_idx_map.json,
  report.txt, metrics/final_snapshot.json, performance/result_summary.json
  (qps≈10630, n_samples_issued=7000, 7000/7000 successful, state=complete,
  complete=True). Confirms MetricsPipeline bring-up + drain (final_snapshot.json
  sourcing → Report), tmpfs event-log salvage, and clean teardown.
• mypy: clean on all four files. (The pre-commit mypy hook reports 3 errors in
  cpu_affinity.py / token_metrics.py — the same macOS-only sched_setaffinity /
  sched_getaffinity false positives in untouched files; they pass in CI/Linux, and are
  why the hook needs --no-verify locally on macOS.)
• ruff / ruff-format / license headers / whitespace: all pass.

Note on CI: these platform failures are macOS-only. The full suite should be run on
Linux (CI) to confirm a green board; locally on macOS the baseline-diff above is the
meaningful signal.

Reproduce

# The moved code, all macOS-runnable:
uv run pytest tests/unit/commands/test_benchmark.py tests/unit/commands/test_score_accuracy.py \
  tests/unit/commands/test_benchmark_final_snapshot.py tests/unit/compliance/test_output_caching.py  # 233 pass
# Full suite on Linux/CI for a green board; on macOS compare against a stashed `main`
# tree instead (identical failure profile = no regression):
uv run pytest                                              # Linux/CI
uv run pre-commit run --files src/inference_endpoint/commands/benchmark/{execute,profiling,accuracy,pipeline}.py
# echo-server smoke — report dir gets config.yaml, events.jsonl, sample_idx_map.json,
# metrics/final_snapshot.json, performance/result_summary.json, report.txt.
# On macOS pass --no-cpu-affinity (pin_loadgen is Linux-only); dummy_1k.jsonl needs parser.prompt.
uv run python -m inference_endpoint.testing.echo_server --port 8765 &
uv run inference-endpoint benchmark offline --endpoints http://localhost:8765 --model test-model \
  --dataset "tests/assets/datasets/dummy_1k.jsonl,samples=40,parser.prompt=text_input" \
  --workers 2 --no-cpu-affinity --max-connections 16

What to focus on in review

1. MetricsPipeline teardown (pipeline.py) and its three call sites in
   execute._run_benchmark_async — the ZMQ-scope-exactly-once invariant across all five paths.
2. The nested cleanup structure in _run_benchmark_async (F5) — that pipe.close() always
   runs and cleanup failures still salvage tmpfs.
3. accuracy._score_accuracy / write_accuracy_results equivalence to the original inline
   scoring + accuracy_results.json block (OSL fast-backend gating, lazy uuid_to_text, numpy
   coercion, finalize write order: report artifacts before accuracy results).

---

🤖 Generated with Claude Code

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request refactors the benchmark execution logic by modularizing cohesive sub-concerns from execute.py into sibling modules: accuracy.py for scoring, pipeline.py for managing the metrics and event-log service pipeline, and profiling.py for handling profiler triggers. Unit tests have been updated accordingly to reflect these new module boundaries. The review feedback highlights three robust error-handling improvements: ensuring background subprocesses are killed if the pipeline is closed with an active publisher, wrapping the _salvage_tmpfs call in a try-except block to prevent masking original exceptions during cleanup, and adding exception handling around self.subscriber.close() during the graceful drain process.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If the pipeline is closed without having successfully completed drain_and_build_report (for example, due to an initialization error or an unhandled exception during the run), self.publisher will still be active. In this case, close() is called with kill=False, which prevents self._launcher.kill_all() from running and leaks the background service subprocesses. Forcing kill = True when self.publisher is not None ensures that active subprocesses are always cleaned up on abnormal termination.

Suggested change

	if self.publisher is not None:
	try:
	self.publisher.close()
	except Exception as e: # noqa: BLE001 — teardown best-effort
	logger.warning("Publisher close error: %s", e)
	if kill and self._launcher is not None:
	self._launcher.kill_all()
	if self.publisher is not None:
	kill = True
	try:
	self.publisher.close()
	except Exception as e: # noqa: BLE001 — teardown best-effort
	logger.warning("Publisher close error: %s", e)
	if kill and self._launcher is not None:
	self._launcher.kill_all()

[gemini-code-assist]

During exception handling in the except BaseException: block, if _salvage_tmpfs raises an exception (e.g., due to disk full or permission issues during shutil.copy2), it will propagate and mask the original exception, while also skipping the shutil.rmtree cleanup. Wrapping the salvage call in a try...except block ensures that cleanup always completes and the original exception is preserved.

            try:
                _salvage_tmpfs(ctx.report_dir, tmpfs_dir)
            except Exception as e:
                logger.warning(f"Failed to salvage tmpfs: {e}")
            shutil.rmtree(tmpfs_dir, ignore_errors=True)

[gemini-code-assist]

In drain_and_build_report, self.subscriber.close() is called without exception handling. If closing the subscriber raises an exception, it will abort the graceful drain and prevent the benchmark report from being built. Wrapping it in a try...except block (similar to the implementation in _teardown) ensures robust finalization.

Suggested change

	# re-close them (close() is idempotent, but this keeps the invariant clean).
	self.publisher = None
	if self.subscriber is not None:
	if self.subscriber is not None:
	try:
	self.subscriber.close()
	except Exception as e:
	logger.warning("Subscriber close error during drain: %s", e)
	self.subscriber = None

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.