refactor(benchmark): split execute.py into execute/profiling/accuracy/pipeline - #421
refactor(benchmark): split execute.py into execute/profiling/accuracy/pipeline#421arekay-nv wants to merge 13 commits into
Conversation
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| _salvage_tmpfs(ctx.report_dir, tmpfs_dir) | ||
| shutil.rmtree(tmpfs_dir, ignore_errors=True) |
There was a problem hiding this comment.
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)
…/pipeline
execute.py concentrated all CLI benchmark orchestration; the large
_run_benchmark_async and _score_accuracy dominated its complexity. Split along
natural seams into four cohesive modules:
- profiling.py (new): profile-trigger protocol (vLLM /start_profile,/stop_profile)
+ ProfileController (URL derivation + start/stop/payload lifecycle)
- accuracy.py (new): AccuracyConfiguration, _score_accuracy, _load_osl_backend,
write_accuracy_results
- pipeline.py (new): MetricsPipeline — ZMQ + metrics-aggregator/event-logger
subprocess lifecycle (start/drain_and_build_report/abort/close) + snapshot→Report
- execute.py: thin orchestrator
Everything audit.py/cli.py import stays in execute (TestMode, BenchmarkResult,
_salvage_tmpfs, setup/run/finalize, run_benchmark, resolve_report_dir); one-way
import graph (execute → {profiling, accuracy, pipeline}). Three unit test files
repointed.
Rebased onto main after the SWE-bench scorer (#342). #342's execute.py accuracy
changes are routed into their post-split homes: AccuracyConfiguration
model_params/endpoint_config fields, _effective_external_sample_count,
_accuracy_uuid_bound (None-on-unavailable), and SKIP_ENDPOINT_PHASE handling in
_score_accuracy land in accuracy.py; _validate_accuracy_config_for_scorer, the
external-scorer skips/logging, and the dataset-loader/preflight wiring land in
execute.py.
Also restores the early-stopping aggregator flag the split had dropped:
"--early-stopping" is now threaded through pipeline._build_aggregator_args (gated
on settings.early_stopping.enabled), replacing an orphaned append to a variable
that no longer existed in _run_benchmark_async after the split.
Verified: unit suite green except pre-existing macOS-only failures
(sched_setaffinity / IPC-socket-path in cpu_affinity/token_metrics/metrics_aggregator,
none in commands/benchmark). ruff/ruff-format/prettier/license clean; mypy clean on
the changed files (remaining mypy errors are the macOS-only os.sched_* attr-defined
false-positives in untouched files).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
85c7f5d to
57d9a66
Compare
A setup failure between MetricsPipeline.start() (which launches the aggregator + event-logger subprocesses) and session.run — e.g. _build_phases rejecting an agentic accuracy dataset, or BenchmarkSession construction raising — never reached the graceful drain, and the finally called only pipe.close() (_teardown(kill=False)). The launched subprocesses were thus neither drained nor killed; with the aggregator drain-timeout defaulting to unlimited they could linger indefinitely waiting for an ENDED that never arrives. Track whether the drain ran; if not, the finally now aborts (kill=True) so the services are killed. The connect-failure path already did this via pipe.abort() for SetupError; this extends the guarantee to every other post-start() setup failure. _teardown is idempotent (_closed guard), so the existing SetupError->abort path double-calls abort() harmlessly. Pre-existing latent gap (the monolith's ZMQ `with` __exit__ likewise never killed the launcher), surfaced by the review council on the execute.py split. Adds a regression test that drives a post-launch setup failure and asserts kill_all() runs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The setup-failure abort added in the previous commit runs `pipe.abort()` → `_teardown(kill=True)` from a `finally` while an exception is propagating. `kill_all()` sat unguarded between the try/except-wrapped publisher and subscriber closes; a raise there (e.g. a child exiting between `poll()` and `kill()`) would mask the original setup error and skip the ZMQ `__exit__` — and since `_teardown` already set `_closed`, the trailing `pipe.close()` would no-op, leaking the ZMQ scope (violating the finally's "close must always run" invariant). Wrap `kill_all()` in the same best-effort try/except as the neighbouring closes so teardown always reaches the ZMQ `__exit__` and never masks the in-flight exception. Surfaced by the review council (Cursor + Claude) on the prior fix. Also adds a companion test: a SetupError after launch hits both abort() sites (explicit except + finally); asserts kill_all runs exactly once, locking in the `_closed` idempotency guard the finally-abort depends on. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
nv-alicheng
left a comment
nv-alicheng
left a comment
There was a problem hiding this comment.
Review Council — Multi-AI Code Review
Reviewed by: Claude + Code-Quality | Depth: thorough
See the summary comment for the tiered breakdown. codex was unavailable in this environment, so this is a Claude + Code-Quality review.
| ) | ||
| self._starts.append(rec) | ||
|
|
||
| def stop(self, completed_normally: bool) -> None: |
There was a problem hiding this comment.
[Claude] medium (testing): ProfileController.start/stop have no coverage. Every _run_benchmark_async test in test_benchmark.py raises before session.run completes, so profiler.start() (execute.py:897) and profiler.stop() (execute.py:912) never execute. The refactor moved real branching into stop() — the non-200 / i >= len(self._stop_urls) skip (line 195) and the stop_reason abort-vs-phase_end selection (line 193) — none of it exercised. Add a direct ProfileController unit test covering the start→stop index mapping.
| and load_pattern.use_legacy_loadgen_qps_metrics | ||
| ), | ||
| ) | ||
| if not report.complete: |
There was a problem hiding this comment.
[Claude] low (testing): _build_report_from_snapshot moved two warning branches out of the monolith that are now untested: the incomplete-report warning (if not report.complete:, 170-175) and the legacy-loadgen-QPS deprecation warning (176-183), plus the swallow-to-None except at 185. test_benchmark_final_snapshot.py only tests _load_final_snapshot_from_disk; _build_report_from_snapshot is never called from tests.
| # Wall-clock of just this phase's tokenization (seconds); | ||
| # summed across datasets for the accuracy report's total. | ||
| entry["osl_tokenize_s"] = round(time.perf_counter() - t0, 3) | ||
| except Exception as e: # noqa: BLE001 - optional blocks; never fail scoring |
There was a problem hiding this comment.
[Claude] low (testing): In the moved _score_accuracy, the get_raw_outputs build-once branch (342-355) and the broad except Exception that drops response-counts/OSL without failing scoring (370-375) are load-bearing. test_score_accuracy.py drives _score_accuracy with fake backends but no test forces get_raw_outputs/tokenize to raise and asserts the entry is still appended with score present and the counts/OSL keys simply absent — the docstring's 'a read failure only drops these blocks, never fails scoring' invariant is unverified.
Review Council — Multi-AI Code ReviewReviewed by: Claude + Code-Quality | Depth: thorough Found 12 issues across 4 files. The refactor is largely faithful — imports, call sites, operation ordering, and error paths are preserved across the split. No critical/high defects; findings are testing gaps, maintainability debt on the new module boundaries, and one dead accessor. 11 posted inline; 1 (execute.py:1082) is summary-only (outside the diff hunk). 🔴 Must Fix (critical/high)None. 🟡 Should Fix (medium)
🔵 Consider (low)
Cross-cutting note: #1/#2 share a root cause — the split left helpers that are now cross-module API with private (
|
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
…eanups Address PR #421 review comments (renames + one-liners; no behavior change): - Rename cross-module helpers public (drop the `_` prefix now that the execute.py split imports them across a package boundary): effective_external_sample_count, score_accuracy, write_profiling_section - profiling: type `action` as Literal["start", "stop"]; name the profile-POST timeout constant (_PROFILE_POST_TIMEOUT_S); drop the dead ProfileController.enabled property - execute: annotate _wire_on_sample_complete `publisher` as EventPublisherService (TYPE_CHECKING import) - accuracy: mark AccuracyConfiguration frozen=True (build-once/read-many) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…rors In run_benchmark_async's teardown finally, an exception raised by pbar.close() (e.g. BrokenPipeError on a closed stderr) would replace the exception being unwound — and on an otherwise-clean run would propagate to the outer except → tmpfs salvage, turning a successful benchmark into a failure that never returns its BenchmarkResult. Catch and log it; the surrounding finally still guarantees pipe.abort()/pipe.close() run. Addresses PR #421 Codex review comment. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the hand-rolled start()/abort()/close()/_teardown() lifecycle with a real async context manager, addressing PR #421 review (r3677374272 CM smell + the _teardown BaseException-leak it hides). pipeline.py: - __aenter__ = start(); __aexit__ releases via a contextlib.ExitStack that owns the ZMQ scope + publisher/subscriber closers, so every step runs even under a BaseException (Ctrl-C/SystemExit) mid-teardown and exceptions chain — no more skipped ZMQ __exit__ or leaked service children. No hand-called __enter__/__exit__. - drain_and_build_report() nulls self.publisher on success; __aexit__ reads that as the drained-cleanly signal to release-only, else kills the services (they'd otherwise linger on the aggregator's unlimited drain-timeout). - Removed abort()/close()/_teardown()/_zmq_cm/_closed; added _close_publisher/ _close_subscriber ExitStack callbacks. execute.py: - Consumer uses `async with pipe:`; drops the `drained` flag, the explicit `except SetupError: pipe.abort()`, and the `if not drained: abort(); close()` finally. Setup/connect/session failures skip the drain, so __aexit__ kills. Closes the hvagadia _teardown/BaseException high, the Gemini subscriber.close-in- drain and publisher-leak items, and the CM-smell design comment. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…filer/pipeline gaps Address the remaining PR #421 review items. Correctness (execute.py): - _salvage_tmpfs is now best-effort at both sites (the run_benchmark_async except-BaseException and the run_benchmark finally): a raising salvage no longer masks the exception being unwound or skips the rmtree cleanup. - profiling.json is written after the report artifacts and guarded (OSError → warn), so it can no longer be first-and-unguarded and sink the built perf report on an IO error. Tests: - ProfileController.start/stop: start→stop index mapping, non-200 start skips its stop, abort vs phase_end stop_reason, disabled no-op, stop-without-start. - _build_report_from_snapshot: incomplete-report and legacy-QPS warnings, and the swallow-to-None on a malformed snapshot. - drain_and_build_report snapshot fallback: disk-missing → subscriber.latest, and no-snapshot → None (the SIGKILL/OOM recovery path). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
…ver __aexit__ policy Address review-council findings on the execute.py/pipeline.py split. - execute.py: bind http_client at function scope and shut it down in the outer finally (which always runs) instead of only in the session finally. A setup error after _create_issuer (e.g. _build_agentic_strategy/_build_phases raising) previously skipped the session finally and leaked the client's worker subprocesses; shutdown_async() is idempotent so the clean-path call is a no-op. - execute.py: make the end-of-run drain best-effort so a teardown error (publisher.close / wait_for_exit) can't replace the run's in-flight exception on the session-failure path (matches the sibling client-shutdown guard). - AGENTS.md: MetricsPipeline is an async context manager now — drop the stale abort/close from its Code Organization entry. Tests: - HTTP client is shut down when setup raises after _create_issuer. - __aexit__ kill policy: clean drain (publisher nulled) does NOT kill the services; never-drained (publisher still set) kills exactly once. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Review-council follow-ups on the execute.py split. - execute.py: the end-of-run drain except is now conditional. On a clean run (session_completed_normally) a drain / report-build failure propagates so the run fails loudly — an unconditional swallow would exit 0 with report=None and no perf artifacts. On the session-failure path it still swallows, so a teardown error can't replace the run's in-flight exception. (Refines the earlier best-effort-drain change, which over-corrected.) - test_benchmark.py: a clean session with a failing drain propagates the error. - test_benchmark_final_snapshot.py: drop the duplicate `import ... as pipeline_mod` (kept the `from ... import` style) and patch via string targets — resolves the dual-import lint on the file. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fail clean benchmark sessions when metrics draining yields no usable report instead of exiting successfully without performance artifacts. Run service termination through the pipeline ExitStack so publisher, subscriber, and ZMQ cleanup still execute when teardown is interrupted. Add regressions for missing reports and repeated interrupts during both normal and partial-start teardown.
3 participants
Refactor: split
commands/benchmark/execute.pyinto cohesive modulesBranch:
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.pyhad grown to 1691 lines andconcentrated almost all CLI benchmark orchestration in one file. Two functions dominated:
_run_benchmark_async(~403 lines) interleaved ZMQ setup, two subprocess launches, HTTP-clientconstruction, agentic wiring, profiling triggers, the session run, and a ~110-line
finallydoing five unrelated teardown jobs;
_score_accuracy(~172 lines) mixed tokenizer loading, uuidbounding, scoring, response counts, and OSL. Duplicated tmpfs-salvage /
AccuracyConfigurationconstruction / 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.pysplit into four modules undersrc/inference_endpoint/commands/benchmark/:execute.pysetup_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)_derive_profile_urls,_post_profile,_render_profile_status,_write_profiling_section) + newProfileController(owns URL derivation + start/stop/payload lifecycle)accuracy.py(new)AccuracyConfiguration,_phase_osl_stats,_phase_response_counts,_accuracy_uuid_bound,_score_accuracy, new_load_osl_backend, newwrite_accuracy_resultspipeline.py(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_diskTwo new abstractions
ProfileController(profiling.py) — collapses the profile URL pre-derivation, the/start_profile(PERFORMANCE-phase-only) fire, the/stop_profile(only forstatus==200starts) fire, and the
{engine, starts, stops}payload build into one object. Disabled andinert when
engine is None; raises up-front (fail-before-run) when an engine is set butendpoints are empty.
MetricsPipeline(pipeline.py) — owns the ZMQ context, event publisher, snapshotsubscriber, 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 thefinal 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.pyimports fromexecute(BenchmarkResult,TestMode,_salvage_tmpfs,finalize_benchmark,run_benchmark_async,setup_benchmark) andcli.py'simports (
resolve_report_dir,run_benchmark) stay inexecute—audit.py,cli.py, thepackage
__init__, and the TEST04 audit tests are unchanged. Runtime import direction isone-way (
execute → {profiling, accuracy, pipeline}); the sibling modules referenceBenchmarkContextonly underTYPE_CHECKING, so there is no cycle.Three unit test files were repointed to the new module paths (imports /
mock.patchtargetsonly — 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 andre-verified before the next round.
close()is idempotent — but violated the null-after-drain invariant)publisher/subscriberafter closing indrain_and_build_reportProfileControllerbuilt before the runtry, so a misconfigValueErrorbypassed tmpfs/pbar cleanuptry(also closes a latent service-subprocess leak that existed in the original)mkdirs ran before the cleanuptry; a mkdir failure aftertmpfs_direxisted would leak ittry, move themkdirs inside itfinally: pbar.close(); pipe.close()— ifpbar.close()raised, ZMQ leaked and tmpfs wasn't salvaged (the original guaranteed both via the ZMQwith)try: <body> finally: (try: pbar.close() finally: pipe.close())wrapped by the outerexcept BaseExceptionsalvageexecute_modcomment/var name intest_score_accuracy.pyscoring_mod, comment correctedRound 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, connectSetupError→abort(),launcher.launchraising, mid-runKeyboardInterrupt), confirmedpipe.close()survives apbar.close()failure, that cleanup exceptions route to tmpfs salvage,and that no
result/profiler/publisher/http_clientis possibly-unbound at any reachable use.Verification
test_benchmark,test_score_accuracy,test_benchmark_final_snapshot, compliancetest_output_caching): 233 passed —these all run on macOS and cover the moved code directly.
main58 failed / 1542 passed → this branch 58 failed / 1542passed (identical). All 58 failures are in
metrics_aggregator/(44) andendpoint_client/test_cpu_affinity.py(14) — untouched modules that fail on macOSonly; none is in
commands/benchmark/.commands/(drivesrun_benchmarkend-to-end):main29 failed / 8passed → this branch 29 failed / 8 passed (identical). The failures are the
pin_loadgen()-requires-Linux path, which the originalsetup_benchmarkinvoked thesame way.
--no-cpu-affinity): the full refactored path runsclean 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). ConfirmsMetricsPipelinebring-up + drain (final_snapshot.jsonsourcing → Report), tmpfs event-log salvage, and clean teardown.
mypyhook reports 3 errors incpu_affinity.py/token_metrics.py— the same macOS-onlysched_setaffinity/sched_getaffinityfalse positives in untouched files; they pass in CI/Linux, and arewhy the hook needs
--no-verifylocally on macOS.)Reproduce
What to focus on in review
MetricsPipelineteardown (pipeline.py) and its three call sites inexecute._run_benchmark_async— the ZMQ-scope-exactly-once invariant across all five paths._run_benchmark_async(F5) — thatpipe.close()alwaysruns and cleanup failures still salvage tmpfs.
accuracy._score_accuracy/write_accuracy_resultsequivalence to the original inlinescoring +
accuracy_results.jsonblock (OSL fast-backend gating, lazyuuid_to_text, numpycoercion, finalize write order: report artifacts before accuracy results).
🤖 Generated with Claude Code