feat: enable multiple warmup and profiling phases in a single run#1150
feat: enable multiple warmup and profiling phases in a single run#1150ilana-n wants to merge 6 commits into
Conversation
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@cadd9242f8c23eb636f365a4f4de49ef84cc0cceRecommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@cadd9242f8c23eb636f365a4f4de49ef84cc0cceLast updated for commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughChangesThe PR makes adaptive scale YAML-only, introduces explicit phase kind/name/index identity, supports repeated warmup and profiling phases, routes runtime state by phase instance, and adds MMLU-Pro accuracy benchmarking with phase-scoped artifacts and manifests. Phase configuration and resolution
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
This comment has been minimized.
This comment has been minimized.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/aiperf/records/records_manager.py (2)
1852-1888: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winScope
output_tokens_per_jouleto the same phase window as energy.records_resultsis built withoutphase_index, sooutput_tokens_per_joulecan use tokens aggregated across all profiling phases while_apply_gpu_efficiency_metrics()computes energy from only the final phase window. Thread a phase-scoped token source into the efficiency helper, or derive the numerator fromphase_records_resultsinstead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/records/records_manager.py` around lines 1852 - 1888, Update _apply_gpu_efficiency_metrics so its output_tokens_per_joule numerator is scoped to the current phase, matching the energy window used for the final phase. Pass a phase-indexed token source into the helper or derive tokens from phase_records_results, and avoid using the unscoped records_results token total.
1155-1175: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMark the currently active profiling phase as cancelled. Hardcoding
_final_results_phase_indexleaves the interrupted phase’s tracker unmarked, so_summarize_per_phase_metric_records()can emit a per-phase artifact withwas_cancelled=Falsefor the phase that actually received Ctrl+C. Pass the active profilingphase_indexhere instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/records/records_manager.py` around lines 1155 - 1175, Update _on_profile_cancel_command to use the currently active profiling phase_index when calling mark_phase_cancelled and _process_results, instead of _final_results_phase_index. Preserve the existing cancellation handling while ensuring the interrupted phase is marked and summarized as cancelled.
🧹 Nitpick comments (8)
src/aiperf/credit/issuer.py (1)
96-98: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_phase_key(phase-runtime-key) computation duplicated acrosscredit/issuer.pyandtiming/phase/runner.py. Both independently computephase_index if phase_index is not None else phaseto key concurrency-manager and slot operations; the same pattern also appears (per graph evidence) incredit/callback_handler.pyandrecords_manager.py. A single shared helper (e.g.phase_runtime_key(phase, phase_index) -> PhaseRuntimeKey) would remove the risk of these diverging as multi-phase concurrency isolation evolves — this key is central to correctness (wrong key = concurrency limits leaking across phases).
src/aiperf/credit/issuer.py#L96-L98: replace the inline_phase_keyproperty body with a call to the shared helper.src/aiperf/timing/phase/runner.py#L189-L195: replace the inline_phase_keyproperty body with a call to the same shared helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/credit/issuer.py` around lines 96 - 98, The phase runtime-key computation is duplicated across the credit issuer and phase runner. Add or reuse a shared phase_runtime_key helper, then update _phase_key in src/aiperf/credit/issuer.py at lines 96-98 and src/aiperf/timing/phase/runner.py at lines 189-195 to delegate to it, preserving the phase_index-when-present, otherwise phase behavior.src/aiperf/config/phases.py (1)
255-269: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLegacy canonical-name →
kindinference rule duplicated acrossconfig/phases.pyandconfig/flags/resolver.py. Both implement the same "ifkindis omitted andnameiswarmup/profiling, inferkindfromname" rule independently (a third copy also exists inconfig/loader/normalizers.py::_infer_phase_kind, outside this review batch). Any future change to the canonical-name set or inference semantics would need to be kept in sync by hand across all copies.
src/aiperf/config/phases.py#L255-L269: extract the inference check (lines 258-259) into a shared helper (e.g.infer_legacy_phase_kind) usable by both the Pydantic validator and the raw-dict pre-merge pass.src/aiperf/config/flags/resolver.py#L459-L487: replace the inline inference block (lines 470-473) with a call to the same shared helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/phases.py` around lines 255 - 269, Extract the canonical-name inference logic from PhaseConfig._validate_phase_constraints into a shared infer_legacy_phase_kind helper, preserving the existing warmup/profiling behavior and None handling; update src/aiperf/config/phases.py lines 255-269 to use it, and replace the duplicated inline inference in src/aiperf/config/flags/resolver.py lines 459-487 with the same helper so both validation and raw-dict pre-merge paths share one implementation.src/aiperf/records/records_manager.py (1)
1690-1732: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftPer-phase artifacts never receive GPU efficiency metrics.
_summarize_per_phase_metric_recordsbuildsPhaseProfileResultsfor every configured warmup/profiling phase entry but never calls_apply_gpu_efficiency_metricsfor any of them —total_gpu_power/total_gpu_energy/output_tokens_per_joule/energy_per_useronly ever land on the rootrecords. For a feature explicitly designed to compare named phases with different concurrency/config (e.g. a low-vs-storm profiling comparison), missing per-phase GPU efficiency numbers limits the usefulness of the newphases/<phase_name>/artifacts for GPU-telemetry users.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/records/records_manager.py` around lines 1690 - 1732, Update _summarize_per_phase_metric_records to apply _apply_gpu_efficiency_metrics to each constructed PhaseProfileResults before appending or returning it, ensuring total_gpu_power, total_gpu_energy, output_tokens_per_joule, and energy_per_user are populated for every phase artifact as they are for the root records.tests/unit/common/models/test_credit_models.py (1)
87-90: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the repository’s required
param(..., id=...)parametrization style.As per coding guidelines, use “the
param(..., id=...)style with# fmt: skipon the closing parenthesis line.”Proposed update
+from pytest import param + - `@pytest.mark.parametrize`("field", ["phase_index", "profiling_index"]) + `@pytest.mark.parametrize`( + "field", + [ + param("phase_index", id="phase-index"), + param("profiling_index", id="profiling-index"), + ], + ) # fmt: skip🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/common/models/test_credit_models.py` around lines 87 - 90, Update the test_phase_indexes_must_be_non_negative parametrization to use pytest.param entries with explicit id= values, and add # fmt: skip to the closing parenthesis line as required by repository style. Preserve both phase_index and profiling_index cases and the existing validation assertion.Source: Coding guidelines
tests/unit/credit/test_callback_handler.py (1)
165-172: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename the test to include the function under test.
test_registers_same_kind_phases_by_runtime_indexdoes not follow the requiredtest_<function>_<scenario>_<expected>convention. Consider a name such astest_register_phase_same_kind_phases_by_runtime_index.As per coding guidelines, tests must be named
test_<function>_<scenario>_<expected>.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/credit/test_callback_handler.py` around lines 165 - 172, Rename the test method test_registers_same_kind_phases_by_runtime_index to follow the required test_<function>_<scenario>_<expected> convention, using the function under test register_phase while preserving the existing scenario and expected behavior.Source: Coding guidelines
tests/unit/config/test_converter_profiling_phase_routes.py (1)
260-269: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the test constant as immutable class data.
Ruff reports RUF012 for the mutable class attribute. Use
ClassVarandfrozenset.+from typing import ClassVar + class TestAdaptiveScaleCliRemoval: - REMOVED_FIELDS = { + REMOVED_FIELDS: ClassVar[frozenset[str]] = frozenset({ ... - } + })🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/config/test_converter_profiling_phase_routes.py` around lines 260 - 269, Update the REMOVED_FIELDS class constant to be annotated with ClassVar and initialized as a frozenset, preserving all existing field names while satisfying Ruff RUF012.Source: Linters/SAST tools
tests/unit/exporters/test_exporter_manager.py (1)
87-89: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new tests. The changed test functions omit the required
-> Noneannotation.
tests/unit/exporters/test_exporter_manager.py#L87-L89: add-> Nonetotest_export_writes_phase_metric_artifacts.tests/unit/records/test_dag_metadata_tagging.py#L29-L29,L109-L111,L126-L126,L172-L172: add-> Noneto all four new test methods.As per coding guidelines, annotate all functions with parameter and return types.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/exporters/test_exporter_manager.py` around lines 87 - 89, Add the required -> None return annotation to test_export_writes_phase_metric_artifacts in tests/unit/exporters/test_exporter_manager.py and to all four new test methods in tests/unit/records/test_dag_metadata_tagging.py, preserving their existing parameters and bodies.Source: Coding guidelines
src/aiperf/common/models/record_models.py (1)
198-200: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winUse a shared
PhaseKindtype forphase_kind. These fields still accept arbitrary strings, so thewarmup/profilingcontract can drift across the record and timing models. Makephase_kinduse one shared enum/typed alias everywhere it is serialized.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/common/models/record_models.py` around lines 198 - 200, Replace the free-form str annotation on phase_kind in the record model with the shared PhaseKind enum or typed alias, and update the corresponding phase_kind fields in the timing models to use the same type wherever they are serialized. Preserve the existing optional default of None and the warmup/profiling values.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/common/models/record_models.py`:
- Around line 275-277: Update the profiling_index Field declaration to enforce a
minimum value of zero by adding the same non-negative validation used by
MetricRecordMetadata and RecordContext, while preserving its optional None
default and existing description.
In `@src/aiperf/exporters/exporter_manager.py`:
- Line 104: Wrap the _export_phase_metric_artifacts call with handling for
expected filesystem/path errors from directory creation and manifest writes,
then continue into the deferred exporter loop when those errors occur. Apply the
same protection to the corresponding calls at the other referenced locations
without suppressing unrelated exceptions.
In `@src/aiperf/records/records_tracker.py`:
- Around line 291-315: Update was_phase_cancelled, mark_phase_cancelled, and
check_and_set_all_records_received_for_phase to resolve phase_index through the
same deterministic latest-phase resolver used by create_stats_for_phase before
calling _get_phase_tracker. Preserve explicitly supplied indices, and ensure
index-less named-phase calls target the existing latest indexed tracker rather
than creating an enum-keyed tracker.
In `@src/aiperf/timing/config.py`:
- Around line 87-131: Update TimingConfig.from_run so default_cancellation
always remains a fresh RequestCancellationConfig() and is not overwritten from
any profiling phase. Preserve phase-local cancellation handling through
_build_phase_config and ensure phases omitting cancellation inherit the clean
default configuration.
In `@src/aiperf/workers/worker.py`:
- Around line 735-738: Update the dataset-manager failure path’s fallback
RequestInfo constructor to propagate phase_index, profiling_index, phase_name,
and phase_kind from the current credit, matching the metadata fields used in the
surrounding RequestInfo construction. Preserve the existing fallback error
behavior while ensuring later named-phase retrieval errors retain their phase
identity.
In `@tests/component_integration/timing/test_adaptive_scale.py`:
- Around line 72-75: Quote the config_path interpolation in the cli.run_sync
command so paths containing spaces remain a single --config argument. Update the
command construction around cli.run_sync without changing the surrounding
profiling options.
In `@tests/unit/timing/strategies/test_adaptive_scale.py`:
- Line 109: Annotate the tmp_path parameters with Path in both
test_setup_phase_writes_phase_scoped_manifest and the function at
tests/unit/timing/strategies/test_adaptive_scale.py lines 132-132, preserving
their existing return annotations and behavior.
---
Outside diff comments:
In `@src/aiperf/records/records_manager.py`:
- Around line 1852-1888: Update _apply_gpu_efficiency_metrics so its
output_tokens_per_joule numerator is scoped to the current phase, matching the
energy window used for the final phase. Pass a phase-indexed token source into
the helper or derive tokens from phase_records_results, and avoid using the
unscoped records_results token total.
- Around line 1155-1175: Update _on_profile_cancel_command to use the currently
active profiling phase_index when calling mark_phase_cancelled and
_process_results, instead of _final_results_phase_index. Preserve the existing
cancellation handling while ensuring the interrupted phase is marked and
summarized as cancelled.
---
Nitpick comments:
In `@src/aiperf/common/models/record_models.py`:
- Around line 198-200: Replace the free-form str annotation on phase_kind in the
record model with the shared PhaseKind enum or typed alias, and update the
corresponding phase_kind fields in the timing models to use the same type
wherever they are serialized. Preserve the existing optional default of None and
the warmup/profiling values.
In `@src/aiperf/config/phases.py`:
- Around line 255-269: Extract the canonical-name inference logic from
PhaseConfig._validate_phase_constraints into a shared infer_legacy_phase_kind
helper, preserving the existing warmup/profiling behavior and None handling;
update src/aiperf/config/phases.py lines 255-269 to use it, and replace the
duplicated inline inference in src/aiperf/config/flags/resolver.py lines 459-487
with the same helper so both validation and raw-dict pre-merge paths share one
implementation.
In `@src/aiperf/credit/issuer.py`:
- Around line 96-98: The phase runtime-key computation is duplicated across the
credit issuer and phase runner. Add or reuse a shared phase_runtime_key helper,
then update _phase_key in src/aiperf/credit/issuer.py at lines 96-98 and
src/aiperf/timing/phase/runner.py at lines 189-195 to delegate to it, preserving
the phase_index-when-present, otherwise phase behavior.
In `@src/aiperf/records/records_manager.py`:
- Around line 1690-1732: Update _summarize_per_phase_metric_records to apply
_apply_gpu_efficiency_metrics to each constructed PhaseProfileResults before
appending or returning it, ensuring total_gpu_power, total_gpu_energy,
output_tokens_per_joule, and energy_per_user are populated for every phase
artifact as they are for the root records.
In `@tests/unit/common/models/test_credit_models.py`:
- Around line 87-90: Update the test_phase_indexes_must_be_non_negative
parametrization to use pytest.param entries with explicit id= values, and add #
fmt: skip to the closing parenthesis line as required by repository style.
Preserve both phase_index and profiling_index cases and the existing validation
assertion.
In `@tests/unit/config/test_converter_profiling_phase_routes.py`:
- Around line 260-269: Update the REMOVED_FIELDS class constant to be annotated
with ClassVar and initialized as a frozenset, preserving all existing field
names while satisfying Ruff RUF012.
In `@tests/unit/credit/test_callback_handler.py`:
- Around line 165-172: Rename the test method
test_registers_same_kind_phases_by_runtime_index to follow the required
test_<function>_<scenario>_<expected> convention, using the function under test
register_phase while preserving the existing scenario and expected behavior.
In `@tests/unit/exporters/test_exporter_manager.py`:
- Around line 87-89: Add the required -> None return annotation to
test_export_writes_phase_metric_artifacts in
tests/unit/exporters/test_exporter_manager.py and to all four new test methods
in tests/unit/records/test_dag_metadata_tagging.py, preserving their existing
parameters and bodies.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f2b0a1f4-c9c2-48eb-8eea-a33f255ab796
📒 Files selected for processing (56)
docs/benchmark-modes/timing-modes-reference.mddocs/cli-options.mddocs/dev/patterns.mddocs/tutorials/adaptive-scale.mddocs/tutorials/yaml-config.mdsrc/aiperf/common/accumulator_protocols.pysrc/aiperf/common/models/__init__.pysrc/aiperf/common/models/credit_models.pysrc/aiperf/common/models/record_models.pysrc/aiperf/config/config.pysrc/aiperf/config/flags/_adaptive_control_cli.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/flags/_resolver_adaptive.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/flags/converter.pysrc/aiperf/config/flags/resolver.pysrc/aiperf/config/loader/helpers.pysrc/aiperf/config/loader/normalizers.pysrc/aiperf/config/phases.pysrc/aiperf/config/sweep/expand.pysrc/aiperf/credit/callback_handler.pysrc/aiperf/credit/issuer.pysrc/aiperf/credit/messages.pysrc/aiperf/credit/structs.pysrc/aiperf/exporters/exporter_manager.pysrc/aiperf/gpu_telemetry/accumulator.pysrc/aiperf/gpu_telemetry/protocols.pysrc/aiperf/metrics/accumulator.pysrc/aiperf/post_processors/otel_metrics_results_processor.pysrc/aiperf/records/record_processor_service.pysrc/aiperf/records/records_manager.pysrc/aiperf/records/records_tracker.pysrc/aiperf/timing/concurrency.pysrc/aiperf/timing/config.pysrc/aiperf/timing/phase/progress_tracker.pysrc/aiperf/timing/phase/runner.pysrc/aiperf/timing/phase_orchestrator.pysrc/aiperf/timing/strategies/adaptive_scale.pysrc/aiperf/timing/strategies/adaptive_scale_artifacts.pysrc/aiperf/timing/strategies/adaptive_scale_backends.pysrc/aiperf/timing/strategies/adaptive_scale_runtime.pysrc/aiperf/workers/worker.pytests/component_integration/timing/test_adaptive_scale.pytests/unit/common/models/test_credit_models.pytests/unit/config/test_converter_profiling_phase_routes.pytests/unit/config/test_named_phase_kind.pytests/unit/config/test_v1_resolver_streaming_override.pytests/unit/credit/test_callback_handler.pytests/unit/exporters/test_exporter_manager.pytests/unit/gpu_telemetry/test_accumulator.pytests/unit/post_processors/test_metrics_accumulator.pytests/unit/records/test_dag_metadata_tagging.pytests/unit/records/test_records_manager.pytests/unit/timing/strategies/test_adaptive_scale.pytests/unit/timing/test_race_conditions.py
💤 Files with no reviewable changes (7)
- src/aiperf/config/flags/_adaptive_control_cli.py
- src/aiperf/config/flags/_section_fields.py
- src/aiperf/config/flags/_converter_profiling.py
- src/aiperf/config/flags/cli_config.py
- docs/cli-options.md
- tests/unit/config/test_v1_resolver_streaming_override.py
- src/aiperf/config/flags/_resolver_adaptive.py
ilana-n
force-pushed
the
ilana/aip-1004-1006-named-phase-kind-workflows
branch
from
ca81100 to
5865129
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/aiperf/config/flags/_converter_profiling.py (1)
328-337: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUnhandled
AttributeErrorif a JSONL line is valid JSON but not an object.
load_json_str(stripped)only raisesValueError/TypeErroron parse failure, but a well-formed JSON line that isn't an object (e.g. a bare array or number) will pass parsing and then crash ondata.get("timestamp")since non-dict JSON values don't support.get(). This exception isn't caught by the surroundingexcept OSError, so it propagates and would abort CLI config resolution instead of just returningFalse.🐛 Proposed fix
try: data = load_json_str(stripped) except (ValueError, TypeError): return False - return data.get("timestamp") is not None + return isinstance(data, dict) and data.get("timestamp") is not None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/flags/_converter_profiling.py` around lines 328 - 337, Update the JSONL validation logic around load_json_str so parsed values are verified to be mapping/object data before calling data.get("timestamp"). Return False for valid non-object JSON values, while preserving the existing behavior for parse failures and timestamp checks.src/aiperf/dataset/composer/custom.py (1)
127-150: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
UnicodeDecodeErrorneeds to wrap the file read, not just JSON parsing.
for line in f:can raiseUnicodeDecodeErrorbeforeload_json_str()runs, and that exception is aValueError, so the outer handler logs and re-raises instead of falling back to filename-based detection. Catch the decode/read step too.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/dataset/composer/custom.py` around lines 127 - 150, Update the file-reading logic in the dataset type inference method around the with open block so UnicodeDecodeError raised while iterating through f is handled by the filename-based fallback, not the outer ValueError handler. Wrap the line-reading iteration together with load_json_str processing in the existing decode-error handling, preserving the current behavior for valid JSON and other non-JSON files.src/aiperf/config/schema/aiperf-config.schema.json (1)
3358-3414: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMirror the templated-string variants into the singular
dataset:FileDatasetschema.traceSessionSampleRatio,interTurnDelayCapSeconds,maxIdleGapCapSeconds, andreplaySpeedupaccept Jinja2/env-var strings in the canonicalFileDatasetcopy, but the singular shorthand still exposes only numeric/null variants, so the same config shape is treated differently depending on whether it usesdataset:ordatasets:.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/schema/aiperf-config.schema.json` around lines 3358 - 3414, Update the singular dataset FileDataset schema definitions for traceSessionSampleRatio, interTurnDelayCapSeconds, maxIdleGapCapSeconds, and replaySpeedup to match the templated-string variants already present in the canonical FileDataset schema. Preserve the existing numeric/null validation and defaults while adding the corresponding Jinja2/env-var string alternatives so dataset: and datasets: accept the same config shapes.
🧹 Nitpick comments (5)
src/aiperf/post_processors/metric_results_processor.py (1)
57-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
setdefaultdefault is always constructed and discarded.
_instances_mapis already populated for every registry tag before this comprehension runs (lines 50-52), and_setup_metrics(MetricType.DERIVED)only returns tags that are a subset ofMetricRegistry.all_tags(). Sometric.tagalways already exists in_instances_map, meaningtype(metric)()is evaluated and thrown away on every derived metric, every time a processor is constructed —setdefault's default argument is evaluated eagerly regardless of whether the key exists.♻️ Proposed simplification
self.derive_funcs: dict[ MetricTagT, Callable[[MetricResultsDict], MetricValueTypeT] ] = { - metric.tag: self._instances_map.setdefault( - metric.tag, type(metric)() - ).derive_value # type: ignore + metric.tag: self._instances_map[metric.tag].derive_value # type: ignore for metric in self._setup_metrics(MetricType.DERIVED) if metric.type == MetricType.DERIVED }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/post_processors/metric_results_processor.py` around lines 57 - 65, Update the derived-metric comprehension initializing self.derive_funcs to retrieve each existing instance directly from _instances_map by metric.tag instead of calling setdefault with type(metric)(). Preserve the existing derive_value callback mapping and filtering for MetricType.DERIVED.tests/unit/post_processors/test_derived_metric_instance_scope.py (1)
24-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
fixed_schedule_runfixture across two test files. Both files independently define the sameBenchmarkRunconstruction for a single fixed-schedule profiling phase; the shared root cause is a missing common fixture for this scenario.
tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46: keep this definition (or move it) and import it as the shared fixture.tests/unit/post_processors/test_base_metrics_processor.py#L349-L374: remove this local copy and reuse the shared fixture (e.g. viatests/unit/post_processors/conftest.py) instead of redefining it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/post_processors/test_derived_metric_instance_scope.py` around lines 24 - 46, The fixed_schedule_run fixture is duplicated across both test modules. In tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46, retain or move the existing fixture into tests/unit/post_processors/conftest.py as the shared fixture; in tests/unit/post_processors/test_base_metrics_processor.py#L349-L374, remove the local definition and reuse the shared fixed_schedule_run fixture.tests/unit/dataset/loader/test_baseten_offline_suite.py (1)
164-178: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: use
itertools.pairwise()for successive-pair iteration.Static analysis (RUF007) flags
zip(ts, ts[1:], strict=False)at lines 169 and 177;itertools.pairwise(ts)is the more idiomatic equivalent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/dataset/loader/test_baseten_offline_suite.py` around lines 164 - 178, The successive-timestamp gap calculations in test_global_idle_gap_capped and test_gap_cap_is_opt_in trigger RUF007; import and use itertools.pairwise(ts) instead of zip(ts, ts[1:], strict=False), preserving the existing gap assertions.Source: Linters/SAST tools
tests/unit/dataset/loader/test_baseten_replay_timemodel.py (1)
61-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
itertools.pairwiseoverzip(..., strict=False).Ruff RUF007 flags this pattern;
itertools.pairwise(in_time_order)is clearer for consecutive-pair iteration.♻️ Suggested diff
+from itertools import pairwise + def test_monotonic_nondecreasing_in_time_order(self): ts = [0, 100, 99999, 100050, 500000] out = reflow_idle_gaps(ts, 2000) in_time_order = [out[i] for i in sorted(range(len(ts)), key=lambda i: ts[i])] assert in_time_order == sorted(in_time_order) # every consecutive gap <= cap - gaps = [b - a for a, b in zip(in_time_order, in_time_order[1:], strict=False)] + gaps = [b - a for a, b in pairwise(in_time_order)] assert all(g <= 2000 for g in gaps)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/dataset/loader/test_baseten_replay_timemodel.py` around lines 61 - 68, Update test_monotonic_nondecreasing_in_time_order to import and use itertools.pairwise for consecutive gap calculation instead of zip with strict=False, preserving the existing gap assertion.Source: Linters/SAST tools
src/aiperf/config/dataset/resolver.py (1)
23-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the baseten-only replay fields into a shared constant
The resolver warning list duplicates the same baseten-only fields already guarded in
_converter_dataset.py; centralize them so the warning and reject paths stay in sync.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/dataset/resolver.py` around lines 23 - 58, Move the baseten_trace-only field tuple from the resolver into a shared constant module or existing shared location, then import and reuse it in both _warn_ignored_baseten_only_fields and the validation logic in _converter_dataset.py. Remove the duplicate field list and preserve the current warning and rejection behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/metrics-reference.md`:
- Around line 85-89: Update the “Individual Flags” table in
docs/metrics-reference.md to include MetricFlags.FIXED_SCHEDULE_ONLY, alongside
the existing metric flag entries, while preserving the table’s established
format and descriptions.
In `@src/aiperf/dataset/loader/base_trace_loader.py`:
- Around line 289-301: Update the `extra_body` assignment in the trace-to-`Turn`
conversion to distinguish an explicitly provided empty `request_body` from an
absent value: use the declared-field lookup result whenever it is not `None`,
and fall back to `trace.extra` only when no declared `request_body` is
available.
In `@tests/unit/config/test_baseten_replay_flags.py`:
- Around line 23-24: Guard the module-level pyarrow imports in the test module
so unavailable Windows ARM64 environments skip only the parquet-dependent
coverage instead of failing collection. Use pytest.importorskip("pyarrow")
before accessing pyarrow, or move both imports into trace_parquet while
preserving existing parquet behavior.
---
Outside diff comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 328-337: Update the JSONL validation logic around load_json_str so
parsed values are verified to be mapping/object data before calling
data.get("timestamp"). Return False for valid non-object JSON values, while
preserving the existing behavior for parse failures and timestamp checks.
In `@src/aiperf/config/schema/aiperf-config.schema.json`:
- Around line 3358-3414: Update the singular dataset FileDataset schema
definitions for traceSessionSampleRatio, interTurnDelayCapSeconds,
maxIdleGapCapSeconds, and replaySpeedup to match the templated-string variants
already present in the canonical FileDataset schema. Preserve the existing
numeric/null validation and defaults while adding the corresponding
Jinja2/env-var string alternatives so dataset: and datasets: accept the same
config shapes.
In `@src/aiperf/dataset/composer/custom.py`:
- Around line 127-150: Update the file-reading logic in the dataset type
inference method around the with open block so UnicodeDecodeError raised while
iterating through f is handled by the filename-based fallback, not the outer
ValueError handler. Wrap the line-reading iteration together with load_json_str
processing in the existing decode-error handling, preserving the current
behavior for valid JSON and other non-JSON files.
---
Nitpick comments:
In `@src/aiperf/config/dataset/resolver.py`:
- Around line 23-58: Move the baseten_trace-only field tuple from the resolver
into a shared constant module or existing shared location, then import and reuse
it in both _warn_ignored_baseten_only_fields and the validation logic in
_converter_dataset.py. Remove the duplicate field list and preserve the current
warning and rejection behavior.
In `@src/aiperf/post_processors/metric_results_processor.py`:
- Around line 57-65: Update the derived-metric comprehension initializing
self.derive_funcs to retrieve each existing instance directly from
_instances_map by metric.tag instead of calling setdefault with type(metric)().
Preserve the existing derive_value callback mapping and filtering for
MetricType.DERIVED.
In `@tests/unit/dataset/loader/test_baseten_offline_suite.py`:
- Around line 164-178: The successive-timestamp gap calculations in
test_global_idle_gap_capped and test_gap_cap_is_opt_in trigger RUF007; import
and use itertools.pairwise(ts) instead of zip(ts, ts[1:], strict=False),
preserving the existing gap assertions.
In `@tests/unit/dataset/loader/test_baseten_replay_timemodel.py`:
- Around line 61-68: Update test_monotonic_nondecreasing_in_time_order to import
and use itertools.pairwise for consecutive gap calculation instead of zip with
strict=False, preserving the existing gap assertion.
In `@tests/unit/post_processors/test_derived_metric_instance_scope.py`:
- Around line 24-46: The fixed_schedule_run fixture is duplicated across both
test modules. In
tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46, retain
or move the existing fixture into tests/unit/post_processors/conftest.py as the
shared fixture; in
tests/unit/post_processors/test_base_metrics_processor.py#L349-L374, remove the
local definition and reuse the shared fixed_schedule_run fixture.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fc6971fc-d85d-42b4-87d2-cd10d8b1adcd
📒 Files selected for processing (54)
README.mddocs/benchmark-datasets.mddocs/cli-options.mddocs/genai-perf-feature-comparison.mddocs/index.ymldocs/metrics-reference.mddocs/tutorials/baseten-trace.mddocs/tutorials/custom-dataset.mdsrc/aiperf/common/enums/metric_enums.pysrc/aiperf/common/models/dataset_models.pysrc/aiperf/common/models/record_models.pysrc/aiperf/config/dataset/config.pysrc/aiperf/config/dataset/defaults.pysrc/aiperf/config/dataset/resolver.pysrc/aiperf/config/flags/_converter_dataset.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/dataset/agentic_code_gen/datasets/1k_sessions_200k_ctx/README.mdsrc/aiperf/dataset/composer/custom.pysrc/aiperf/dataset/generator/video.pysrc/aiperf/dataset/loader/_baseten_replay_timemodel.pysrc/aiperf/dataset/loader/base_trace_loader.pysrc/aiperf/dataset/loader/baseten_trace.pysrc/aiperf/dataset/loader/models.pysrc/aiperf/endpoints/openai_completions.pysrc/aiperf/exporters/exporter_manager.pysrc/aiperf/metrics/base_derived_metric.pysrc/aiperf/metrics/types/replay_sched_lag_metrics.pysrc/aiperf/plugin/enums.pysrc/aiperf/plugin/plugins.yamlsrc/aiperf/post_processors/base_metrics_processor.pysrc/aiperf/post_processors/metric_results_processor.pysrc/aiperf/post_processors/timeslice_metric_results_processor.pytests/component_integration/dataset/test_baseten_trace_replay.pytests/integration/test_video.pytests/integration/utils.pytests/unit/common/config/test_input_config.pytests/unit/config/test_baseten_replay_flags.pytests/unit/dataset/generator/test_video_generator.pytests/unit/dataset/loader/test_bailian_trace.pytests/unit/dataset/loader/test_baseten_offline_suite.pytests/unit/dataset/loader/test_baseten_replay_timemodel.pytests/unit/dataset/loader/test_baseten_trace.pytests/unit/dataset/loader/test_baseten_trace_optional_pyarrow.pytests/unit/dataset/loader/test_can_load.pytests/unit/endpoints/test_chat_endpoint.pytests/unit/endpoints/test_completions_endpoint.pytests/unit/endpoints/test_openai_completions.pytests/unit/metrics/test_replay_sched_lag_metrics.pytests/unit/post_processors/test_base_metrics_processor.pytests/unit/post_processors/test_derived_metric_instance_scope.pytests/unit/post_processors/test_timeslice_metric_results_processor.pytests/unit/records/test_records_manager_process_results.py
💤 Files with no reviewable changes (1)
- src/aiperf/dataset/agentic_code_gen/datasets/1k_sessions_200k_ctx/README.md
🚧 Files skipped from review as they are similar to previous changes (2)
- src/aiperf/exporters/exporter_manager.py
- src/aiperf/common/models/record_models.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/aiperf/config/flags/_converter_profiling.py (1)
328-337: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUnhandled
AttributeErrorif a JSONL line is valid JSON but not an object.
load_json_str(stripped)only raisesValueError/TypeErroron parse failure, but a well-formed JSON line that isn't an object (e.g. a bare array or number) will pass parsing and then crash ondata.get("timestamp")since non-dict JSON values don't support.get(). This exception isn't caught by the surroundingexcept OSError, so it propagates and would abort CLI config resolution instead of just returningFalse.🐛 Proposed fix
try: data = load_json_str(stripped) except (ValueError, TypeError): return False - return data.get("timestamp") is not None + return isinstance(data, dict) and data.get("timestamp") is not None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/flags/_converter_profiling.py` around lines 328 - 337, Update the JSONL validation logic around load_json_str so parsed values are verified to be mapping/object data before calling data.get("timestamp"). Return False for valid non-object JSON values, while preserving the existing behavior for parse failures and timestamp checks.src/aiperf/dataset/composer/custom.py (1)
127-150: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
UnicodeDecodeErrorneeds to wrap the file read, not just JSON parsing.
for line in f:can raiseUnicodeDecodeErrorbeforeload_json_str()runs, and that exception is aValueError, so the outer handler logs and re-raises instead of falling back to filename-based detection. Catch the decode/read step too.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/dataset/composer/custom.py` around lines 127 - 150, Update the file-reading logic in the dataset type inference method around the with open block so UnicodeDecodeError raised while iterating through f is handled by the filename-based fallback, not the outer ValueError handler. Wrap the line-reading iteration together with load_json_str processing in the existing decode-error handling, preserving the current behavior for valid JSON and other non-JSON files.src/aiperf/config/schema/aiperf-config.schema.json (1)
3358-3414: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMirror the templated-string variants into the singular
dataset:FileDatasetschema.traceSessionSampleRatio,interTurnDelayCapSeconds,maxIdleGapCapSeconds, andreplaySpeedupaccept Jinja2/env-var strings in the canonicalFileDatasetcopy, but the singular shorthand still exposes only numeric/null variants, so the same config shape is treated differently depending on whether it usesdataset:ordatasets:.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/schema/aiperf-config.schema.json` around lines 3358 - 3414, Update the singular dataset FileDataset schema definitions for traceSessionSampleRatio, interTurnDelayCapSeconds, maxIdleGapCapSeconds, and replaySpeedup to match the templated-string variants already present in the canonical FileDataset schema. Preserve the existing numeric/null validation and defaults while adding the corresponding Jinja2/env-var string alternatives so dataset: and datasets: accept the same config shapes.
🧹 Nitpick comments (5)
src/aiperf/post_processors/metric_results_processor.py (1)
57-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
setdefaultdefault is always constructed and discarded.
_instances_mapis already populated for every registry tag before this comprehension runs (lines 50-52), and_setup_metrics(MetricType.DERIVED)only returns tags that are a subset ofMetricRegistry.all_tags(). Sometric.tagalways already exists in_instances_map, meaningtype(metric)()is evaluated and thrown away on every derived metric, every time a processor is constructed —setdefault's default argument is evaluated eagerly regardless of whether the key exists.♻️ Proposed simplification
self.derive_funcs: dict[ MetricTagT, Callable[[MetricResultsDict], MetricValueTypeT] ] = { - metric.tag: self._instances_map.setdefault( - metric.tag, type(metric)() - ).derive_value # type: ignore + metric.tag: self._instances_map[metric.tag].derive_value # type: ignore for metric in self._setup_metrics(MetricType.DERIVED) if metric.type == MetricType.DERIVED }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/post_processors/metric_results_processor.py` around lines 57 - 65, Update the derived-metric comprehension initializing self.derive_funcs to retrieve each existing instance directly from _instances_map by metric.tag instead of calling setdefault with type(metric)(). Preserve the existing derive_value callback mapping and filtering for MetricType.DERIVED.tests/unit/post_processors/test_derived_metric_instance_scope.py (1)
24-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
fixed_schedule_runfixture across two test files. Both files independently define the sameBenchmarkRunconstruction for a single fixed-schedule profiling phase; the shared root cause is a missing common fixture for this scenario.
tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46: keep this definition (or move it) and import it as the shared fixture.tests/unit/post_processors/test_base_metrics_processor.py#L349-L374: remove this local copy and reuse the shared fixture (e.g. viatests/unit/post_processors/conftest.py) instead of redefining it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/post_processors/test_derived_metric_instance_scope.py` around lines 24 - 46, The fixed_schedule_run fixture is duplicated across both test modules. In tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46, retain or move the existing fixture into tests/unit/post_processors/conftest.py as the shared fixture; in tests/unit/post_processors/test_base_metrics_processor.py#L349-L374, remove the local definition and reuse the shared fixed_schedule_run fixture.tests/unit/dataset/loader/test_baseten_offline_suite.py (1)
164-178: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: use
itertools.pairwise()for successive-pair iteration.Static analysis (RUF007) flags
zip(ts, ts[1:], strict=False)at lines 169 and 177;itertools.pairwise(ts)is the more idiomatic equivalent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/dataset/loader/test_baseten_offline_suite.py` around lines 164 - 178, The successive-timestamp gap calculations in test_global_idle_gap_capped and test_gap_cap_is_opt_in trigger RUF007; import and use itertools.pairwise(ts) instead of zip(ts, ts[1:], strict=False), preserving the existing gap assertions.Source: Linters/SAST tools
tests/unit/dataset/loader/test_baseten_replay_timemodel.py (1)
61-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
itertools.pairwiseoverzip(..., strict=False).Ruff RUF007 flags this pattern;
itertools.pairwise(in_time_order)is clearer for consecutive-pair iteration.♻️ Suggested diff
+from itertools import pairwise + def test_monotonic_nondecreasing_in_time_order(self): ts = [0, 100, 99999, 100050, 500000] out = reflow_idle_gaps(ts, 2000) in_time_order = [out[i] for i in sorted(range(len(ts)), key=lambda i: ts[i])] assert in_time_order == sorted(in_time_order) # every consecutive gap <= cap - gaps = [b - a for a, b in zip(in_time_order, in_time_order[1:], strict=False)] + gaps = [b - a for a, b in pairwise(in_time_order)] assert all(g <= 2000 for g in gaps)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/dataset/loader/test_baseten_replay_timemodel.py` around lines 61 - 68, Update test_monotonic_nondecreasing_in_time_order to import and use itertools.pairwise for consecutive gap calculation instead of zip with strict=False, preserving the existing gap assertion.Source: Linters/SAST tools
src/aiperf/config/dataset/resolver.py (1)
23-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the baseten-only replay fields into a shared constant
The resolver warning list duplicates the same baseten-only fields already guarded in
_converter_dataset.py; centralize them so the warning and reject paths stay in sync.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/dataset/resolver.py` around lines 23 - 58, Move the baseten_trace-only field tuple from the resolver into a shared constant module or existing shared location, then import and reuse it in both _warn_ignored_baseten_only_fields and the validation logic in _converter_dataset.py. Remove the duplicate field list and preserve the current warning and rejection behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/metrics-reference.md`:
- Around line 85-89: Update the “Individual Flags” table in
docs/metrics-reference.md to include MetricFlags.FIXED_SCHEDULE_ONLY, alongside
the existing metric flag entries, while preserving the table’s established
format and descriptions.
In `@src/aiperf/dataset/loader/base_trace_loader.py`:
- Around line 289-301: Update the `extra_body` assignment in the trace-to-`Turn`
conversion to distinguish an explicitly provided empty `request_body` from an
absent value: use the declared-field lookup result whenever it is not `None`,
and fall back to `trace.extra` only when no declared `request_body` is
available.
In `@tests/unit/config/test_baseten_replay_flags.py`:
- Around line 23-24: Guard the module-level pyarrow imports in the test module
so unavailable Windows ARM64 environments skip only the parquet-dependent
coverage instead of failing collection. Use pytest.importorskip("pyarrow")
before accessing pyarrow, or move both imports into trace_parquet while
preserving existing parquet behavior.
---
Outside diff comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 328-337: Update the JSONL validation logic around load_json_str so
parsed values are verified to be mapping/object data before calling
data.get("timestamp"). Return False for valid non-object JSON values, while
preserving the existing behavior for parse failures and timestamp checks.
In `@src/aiperf/config/schema/aiperf-config.schema.json`:
- Around line 3358-3414: Update the singular dataset FileDataset schema
definitions for traceSessionSampleRatio, interTurnDelayCapSeconds,
maxIdleGapCapSeconds, and replaySpeedup to match the templated-string variants
already present in the canonical FileDataset schema. Preserve the existing
numeric/null validation and defaults while adding the corresponding
Jinja2/env-var string alternatives so dataset: and datasets: accept the same
config shapes.
In `@src/aiperf/dataset/composer/custom.py`:
- Around line 127-150: Update the file-reading logic in the dataset type
inference method around the with open block so UnicodeDecodeError raised while
iterating through f is handled by the filename-based fallback, not the outer
ValueError handler. Wrap the line-reading iteration together with load_json_str
processing in the existing decode-error handling, preserving the current
behavior for valid JSON and other non-JSON files.
---
Nitpick comments:
In `@src/aiperf/config/dataset/resolver.py`:
- Around line 23-58: Move the baseten_trace-only field tuple from the resolver
into a shared constant module or existing shared location, then import and reuse
it in both _warn_ignored_baseten_only_fields and the validation logic in
_converter_dataset.py. Remove the duplicate field list and preserve the current
warning and rejection behavior.
In `@src/aiperf/post_processors/metric_results_processor.py`:
- Around line 57-65: Update the derived-metric comprehension initializing
self.derive_funcs to retrieve each existing instance directly from
_instances_map by metric.tag instead of calling setdefault with type(metric)().
Preserve the existing derive_value callback mapping and filtering for
MetricType.DERIVED.
In `@tests/unit/dataset/loader/test_baseten_offline_suite.py`:
- Around line 164-178: The successive-timestamp gap calculations in
test_global_idle_gap_capped and test_gap_cap_is_opt_in trigger RUF007; import
and use itertools.pairwise(ts) instead of zip(ts, ts[1:], strict=False),
preserving the existing gap assertions.
In `@tests/unit/dataset/loader/test_baseten_replay_timemodel.py`:
- Around line 61-68: Update test_monotonic_nondecreasing_in_time_order to import
and use itertools.pairwise for consecutive gap calculation instead of zip with
strict=False, preserving the existing gap assertion.
In `@tests/unit/post_processors/test_derived_metric_instance_scope.py`:
- Around line 24-46: The fixed_schedule_run fixture is duplicated across both
test modules. In
tests/unit/post_processors/test_derived_metric_instance_scope.py#L24-L46, retain
or move the existing fixture into tests/unit/post_processors/conftest.py as the
shared fixture; in
tests/unit/post_processors/test_base_metrics_processor.py#L349-L374, remove the
local definition and reuse the shared fixed_schedule_run fixture.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fc6971fc-d85d-42b4-87d2-cd10d8b1adcd
📒 Files selected for processing (54)
README.mddocs/benchmark-datasets.mddocs/cli-options.mddocs/genai-perf-feature-comparison.mddocs/index.ymldocs/metrics-reference.mddocs/tutorials/baseten-trace.mddocs/tutorials/custom-dataset.mdsrc/aiperf/common/enums/metric_enums.pysrc/aiperf/common/models/dataset_models.pysrc/aiperf/common/models/record_models.pysrc/aiperf/config/dataset/config.pysrc/aiperf/config/dataset/defaults.pysrc/aiperf/config/dataset/resolver.pysrc/aiperf/config/flags/_converter_dataset.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/dataset/agentic_code_gen/datasets/1k_sessions_200k_ctx/README.mdsrc/aiperf/dataset/composer/custom.pysrc/aiperf/dataset/generator/video.pysrc/aiperf/dataset/loader/_baseten_replay_timemodel.pysrc/aiperf/dataset/loader/base_trace_loader.pysrc/aiperf/dataset/loader/baseten_trace.pysrc/aiperf/dataset/loader/models.pysrc/aiperf/endpoints/openai_completions.pysrc/aiperf/exporters/exporter_manager.pysrc/aiperf/metrics/base_derived_metric.pysrc/aiperf/metrics/types/replay_sched_lag_metrics.pysrc/aiperf/plugin/enums.pysrc/aiperf/plugin/plugins.yamlsrc/aiperf/post_processors/base_metrics_processor.pysrc/aiperf/post_processors/metric_results_processor.pysrc/aiperf/post_processors/timeslice_metric_results_processor.pytests/component_integration/dataset/test_baseten_trace_replay.pytests/integration/test_video.pytests/integration/utils.pytests/unit/common/config/test_input_config.pytests/unit/config/test_baseten_replay_flags.pytests/unit/dataset/generator/test_video_generator.pytests/unit/dataset/loader/test_bailian_trace.pytests/unit/dataset/loader/test_baseten_offline_suite.pytests/unit/dataset/loader/test_baseten_replay_timemodel.pytests/unit/dataset/loader/test_baseten_trace.pytests/unit/dataset/loader/test_baseten_trace_optional_pyarrow.pytests/unit/dataset/loader/test_can_load.pytests/unit/endpoints/test_chat_endpoint.pytests/unit/endpoints/test_completions_endpoint.pytests/unit/endpoints/test_openai_completions.pytests/unit/metrics/test_replay_sched_lag_metrics.pytests/unit/post_processors/test_base_metrics_processor.pytests/unit/post_processors/test_derived_metric_instance_scope.pytests/unit/post_processors/test_timeslice_metric_results_processor.pytests/unit/records/test_records_manager_process_results.py
💤 Files with no reviewable changes (1)
- src/aiperf/dataset/agentic_code_gen/datasets/1k_sessions_200k_ctx/README.md
🚧 Files skipped from review as they are similar to previous changes (2)
- src/aiperf/exporters/exporter_manager.py
- src/aiperf/common/models/record_models.py
🛑 Comments failed to post (3)
docs/metrics-reference.md (1)
85-89: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash rg -n -A5 'FIXED_SCHEDULE_ONLY|Metric Flags Reference' docs/metrics-reference.mdRepository: ai-dynamo/aiperf
Length of output: 1976
🏁 Script executed:
#!/bin/bash set -euo pipefail sed -n '2005,2085p' docs/metrics-reference.md printf '\n---\n' rg -n 'FIXED_SCHEDULE_ONLY' -A3 -B3 .Repository: ai-dynamo/aiperf
Length of output: 18104
Add
FIXED_SCHEDULE_ONLYto the Metric Flags Reference
docs/metrics-reference.mdshould listMetricFlags.FIXED_SCHEDULE_ONLYin the “Individual Flags” table, since the new replay schedule metrics already rely on it and the section currently omits it.docs/metrics-reference.md#L2005-L2030🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/metrics-reference.md` around lines 85 - 89, Update the “Individual Flags” table in docs/metrics-reference.md to include MetricFlags.FIXED_SCHEDULE_ONLY, alongside the existing metric flag entries, while preserving the table’s established format and descriptions.src/aiperf/dataset/loader/base_trace_loader.py (1)
289-301: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Falsy check lets an empty
request_bodyfall back toextra.
request_body or getattr(trace, "extra", None)treats an explicitly-set-but-emptyrequest_body(e.g.{}fromBasetenTraceDatasetLoader._set_request_bodywhenforce_min_tokens=Falseandomit_kv_hints=True) the same as "not declared", silently falling back totrace.extrainstead of the intended empty dict. Currently benign forBasetenTrace(noextrafield), but fragile for any future trace type that declares both fields.🐛 Proposed fix: check identity, not truthiness
return Turn( timestamp=getattr(trace, "timestamp", None), delay=getattr(trace, "delay", None), texts=[Text(name="text", contents=[prompt])], max_tokens=getattr(trace, "output_length", None), - extra_body=request_body or getattr(trace, "extra", None), + extra_body=( + request_body + if request_body is not None + else getattr(trace, "extra", None) + ), )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.# extra="allow" surfaces stray row keys as attributes; only a DECLARED # request_body field (e.g. BasetenTrace) may override the row's extra. request_body = ( getattr(trace, "request_body", None) if "request_body" in type(trace).model_fields else None ) return Turn( timestamp=getattr(trace, "timestamp", None), delay=getattr(trace, "delay", None), texts=[Text(name="text", contents=[prompt])], max_tokens=getattr(trace, "output_length", None), extra_body=( request_body if request_body is not None else getattr(trace, "extra", None) ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/dataset/loader/base_trace_loader.py` around lines 289 - 301, Update the `extra_body` assignment in the trace-to-`Turn` conversion to distinguish an explicitly provided empty `request_body` from an absent value: use the declared-field lookup result whenever it is not `None`, and fall back to `trace.extra` only when no declared `request_body` is available.tests/unit/config/test_baseten_replay_flags.py (1)
23-24: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash rg -n 'pyarrow' pyproject.tomlRepository: ai-dynamo/aiperf
Length of output: 405
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the test file and the nearby optional-pyarrow test for context. sed -n '1,120p' tests/unit/config/test_baseten_replay_flags.py printf '\n--- optional pyarrow test ---\n' sed -n '1,200p' tests/unit/dataset/loader/test_baseten_trace_optional_pyarrow.py printf '\n--- pyproject pyarrow context ---\n' sed -n '20,70p' pyproject.tomlRepository: ai-dynamo/aiperf
Length of output: 8573
🏁 Script executed:
#!/bin/bash set -euo pipefail # Check whether this test module is the only place importing pyarrow at module scope. rg -n '^\s*(import pyarrow|from pyarrow)' tests/unit -g '*.py' # Find any skips/importorskip patterns used around pyarrow. rg -n 'importorskip|pyarrow' tests/unit -g '*.py'Repository: ai-dynamo/aiperf
Length of output: 9256
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the file size to decide whether to read more surrounding lines if needed. wc -l tests/unit/config/test_baseten_replay_flags.pyRepository: ai-dynamo/aiperf
Length of output: 205
Guard the module-level
pyarrowimport.pyarrowis intentionally absent on Windows ARM64, so importing it at module scope can make this whole test module fail to collect there and skip the unrelated CLI-plumbing coverage. Move the import behindpytest.importorskip("pyarrow")or intotrace_parquet.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/config/test_baseten_replay_flags.py` around lines 23 - 24, Guard the module-level pyarrow imports in the test module so unavailable Windows ARM64 environments skip only the parquet-dependent coverage instead of failing collection. Use pytest.importorskip("pyarrow") before accessing pyarrow, or move both imports into trace_parquet while preserving existing parquet behavior.
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/aiperf/config/flags/_converter_profiling.py (1)
328-339: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winUnhandled
AttributeErroron non-dict JSON line.
data.get("timestamp")runs outside theexcept (ValueError, TypeError)guard. If the first non-empty JSONL line is valid JSON but not an object (e.g. a list/scalar), this raises an uncaughtAttributeErrorinstead of returningFalse, crashing config resolution on a malformed-but-parseable trace file.🐛 Proposed fix
try: data = load_json_str(stripped) except (ValueError, TypeError): return False - return data.get("timestamp") is not None + return isinstance(data, dict) and data.get("timestamp") is not None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/flags/_converter_profiling.py` around lines 328 - 339, Update the JSONL validation logic around load_json_str in the profiling converter so a successfully parsed non-dict value returns False instead of raising when accessing data.get("timestamp"). Validate the parsed data type before the timestamp lookup, while preserving the existing behavior for malformed JSON, missing timestamps, and file access errors.
🧹 Nitpick comments (1)
src/aiperf/exporters/exporter_manager.py (1)
169-204: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReaching into exporter internals via
_generate_content().
_write_phase_exportcalls the leading-underscoreexporter._generate_content()across a class boundary. Consider exposing a small public API (e.g.generate_content()) on the exporters for this cross-class use rather than depending on a private method.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/exporters/exporter_manager.py` around lines 169 - 204, Add a public content-generation API on MetricsJsonExporter and MetricsCsvExporter, such as generate_content(), and update _write_phase_export to call it instead of the private _generate_content() method. Preserve the existing generated content and file-writing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 343-371: Update _count_dataset_records so errors reading or
decoding an individual JSONL file inside the directory loop are handled per
file, skipping only that file while preserving the accumulated total from
successfully processed files. Keep the existing behavior for directory discovery
and non-directory file handling unchanged.
In `@src/aiperf/exporters/exporter_manager.py`:
- Around line 120-163: Update _export_phase_metric_artifacts to isolate each
phase’s directory creation and export operations within a per-phase exception
boundary, including phase_dir.mkdir and both _write_phase_export calls. When one
phase raises an OSError or export error, record the failure consistently and
continue processing subsequent phase_records instead of aborting the loop.
In `@tests/unit/common/models/test_credit_models.py`:
- Around line 87-96: Update test_phase_indexes_must_be_non_negative to import
param directly from pytest and replace pytest.param calls with param. Move the
fmt: skip comment from the parameter list’s closing bracket to the closing
parenthesis of the `@pytest.mark.parametrize` call.
In `@tests/unit/post_processors/test_metrics_accumulator.py`:
- Around line 204-207: Update
test_export_results_filters_same_kind_records_by_phase_index with explicit type
annotations for every fixture parameter, including mock_metric_registry and
mock_run, and retain the async test’s None return annotation.
---
Outside diff comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 328-339: Update the JSONL validation logic around load_json_str in
the profiling converter so a successfully parsed non-dict value returns False
instead of raising when accessing data.get("timestamp"). Validate the parsed
data type before the timestamp lookup, while preserving the existing behavior
for malformed JSON, missing timestamps, and file access errors.
---
Nitpick comments:
In `@src/aiperf/exporters/exporter_manager.py`:
- Around line 169-204: Add a public content-generation API on
MetricsJsonExporter and MetricsCsvExporter, such as generate_content(), and
update _write_phase_export to call it instead of the private _generate_content()
method. Preserve the existing generated content and file-writing behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fe0daed4-c94a-42ba-9043-cdc30216d187
📒 Files selected for processing (59)
docs/benchmark-modes/timing-modes-reference.mddocs/cli-options.mddocs/dev/patterns.mddocs/tutorials/adaptive-scale.mddocs/tutorials/yaml-config.mdsrc/aiperf/common/accumulator_protocols.pysrc/aiperf/common/models/__init__.pysrc/aiperf/common/models/credit_models.pysrc/aiperf/common/models/record_models.pysrc/aiperf/common/phase.pysrc/aiperf/common/types.pysrc/aiperf/config/config.pysrc/aiperf/config/flags/_adaptive_control_cli.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/flags/_resolver_adaptive.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/flags/converter.pysrc/aiperf/config/flags/resolver.pysrc/aiperf/config/loader/helpers.pysrc/aiperf/config/loader/normalizers.pysrc/aiperf/config/phases.pysrc/aiperf/config/sweep/expand.pysrc/aiperf/credit/callback_handler.pysrc/aiperf/credit/issuer.pysrc/aiperf/credit/messages.pysrc/aiperf/credit/structs.pysrc/aiperf/exporters/exporter_manager.pysrc/aiperf/gpu_telemetry/accumulator.pysrc/aiperf/gpu_telemetry/protocols.pysrc/aiperf/metrics/accumulator.pysrc/aiperf/post_processors/otel_metrics_results_processor.pysrc/aiperf/records/record_processor_service.pysrc/aiperf/records/records_manager.pysrc/aiperf/records/records_tracker.pysrc/aiperf/timing/concurrency.pysrc/aiperf/timing/config.pysrc/aiperf/timing/phase/progress_tracker.pysrc/aiperf/timing/phase/runner.pysrc/aiperf/timing/phase_orchestrator.pysrc/aiperf/timing/strategies/adaptive_scale.pysrc/aiperf/timing/strategies/adaptive_scale_artifacts.pysrc/aiperf/timing/strategies/adaptive_scale_backends.pysrc/aiperf/timing/strategies/adaptive_scale_runtime.pysrc/aiperf/workers/worker.pytests/component_integration/timing/test_adaptive_scale.pytests/unit/common/models/test_credit_models.pytests/unit/config/test_converter_profiling_phase_routes.pytests/unit/config/test_named_phase_kind.pytests/unit/config/test_v1_resolver_streaming_override.pytests/unit/credit/test_callback_handler.pytests/unit/exporters/test_exporter_manager.pytests/unit/gpu_telemetry/test_accumulator.pytests/unit/post_processors/test_metrics_accumulator.pytests/unit/records/test_dag_metadata_tagging.pytests/unit/records/test_records_manager.pytests/unit/records/test_records_manager_process_results.pytests/unit/timing/strategies/test_adaptive_scale.pytests/unit/timing/test_race_conditions.py
💤 Files with no reviewable changes (4)
- src/aiperf/config/flags/_section_fields.py
- src/aiperf/config/flags/_resolver_adaptive.py
- src/aiperf/config/flags/_adaptive_control_cli.py
- tests/unit/config/test_v1_resolver_streaming_override.py
🚧 Files skipped from review as they are similar to previous changes (39)
- docs/benchmark-modes/timing-modes-reference.md
- docs/tutorials/adaptive-scale.md
- src/aiperf/records/record_processor_service.py
- src/aiperf/workers/worker.py
- src/aiperf/common/models/init.py
- src/aiperf/credit/messages.py
- src/aiperf/credit/structs.py
- src/aiperf/post_processors/otel_metrics_results_processor.py
- src/aiperf/gpu_telemetry/protocols.py
- docs/dev/patterns.md
- src/aiperf/config/loader/helpers.py
- src/aiperf/timing/phase_orchestrator.py
- tests/unit/timing/test_race_conditions.py
- src/aiperf/timing/strategies/adaptive_scale_runtime.py
- src/aiperf/timing/strategies/adaptive_scale.py
- src/aiperf/common/models/credit_models.py
- src/aiperf/common/accumulator_protocols.py
- tests/unit/exporters/test_exporter_manager.py
- src/aiperf/timing/strategies/adaptive_scale_artifacts.py
- src/aiperf/config/sweep/expand.py
- tests/unit/credit/test_callback_handler.py
- tests/component_integration/timing/test_adaptive_scale.py
- src/aiperf/config/phases.py
- src/aiperf/timing/phase/progress_tracker.py
- src/aiperf/timing/strategies/adaptive_scale_backends.py
- src/aiperf/gpu_telemetry/accumulator.py
- src/aiperf/config/config.py
- tests/unit/records/test_records_manager_process_results.py
- src/aiperf/config/loader/normalizers.py
- src/aiperf/credit/issuer.py
- src/aiperf/timing/phase/runner.py
- tests/unit/config/test_converter_profiling_phase_routes.py
- tests/unit/records/test_dag_metadata_tagging.py
- src/aiperf/records/records_tracker.py
- src/aiperf/records/records_manager.py
- src/aiperf/config/flags/resolver.py
- src/aiperf/credit/callback_handler.py
- tests/unit/records/test_records_manager.py
- tests/unit/timing/strategies/test_adaptive_scale.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/aiperf/config/flags/_converter_profiling.py (1)
328-339: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winUnhandled
AttributeErroron non-dict JSON line.
data.get("timestamp")runs outside theexcept (ValueError, TypeError)guard. If the first non-empty JSONL line is valid JSON but not an object (e.g. a list/scalar), this raises an uncaughtAttributeErrorinstead of returningFalse, crashing config resolution on a malformed-but-parseable trace file.🐛 Proposed fix
try: data = load_json_str(stripped) except (ValueError, TypeError): return False - return data.get("timestamp") is not None + return isinstance(data, dict) and data.get("timestamp") is not None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/flags/_converter_profiling.py` around lines 328 - 339, Update the JSONL validation logic around load_json_str in the profiling converter so a successfully parsed non-dict value returns False instead of raising when accessing data.get("timestamp"). Validate the parsed data type before the timestamp lookup, while preserving the existing behavior for malformed JSON, missing timestamps, and file access errors.
🧹 Nitpick comments (1)
src/aiperf/exporters/exporter_manager.py (1)
169-204: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReaching into exporter internals via
_generate_content().
_write_phase_exportcalls the leading-underscoreexporter._generate_content()across a class boundary. Consider exposing a small public API (e.g.generate_content()) on the exporters for this cross-class use rather than depending on a private method.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/exporters/exporter_manager.py` around lines 169 - 204, Add a public content-generation API on MetricsJsonExporter and MetricsCsvExporter, such as generate_content(), and update _write_phase_export to call it instead of the private _generate_content() method. Preserve the existing generated content and file-writing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 343-371: Update _count_dataset_records so errors reading or
decoding an individual JSONL file inside the directory loop are handled per
file, skipping only that file while preserving the accumulated total from
successfully processed files. Keep the existing behavior for directory discovery
and non-directory file handling unchanged.
In `@src/aiperf/exporters/exporter_manager.py`:
- Around line 120-163: Update _export_phase_metric_artifacts to isolate each
phase’s directory creation and export operations within a per-phase exception
boundary, including phase_dir.mkdir and both _write_phase_export calls. When one
phase raises an OSError or export error, record the failure consistently and
continue processing subsequent phase_records instead of aborting the loop.
In `@tests/unit/common/models/test_credit_models.py`:
- Around line 87-96: Update test_phase_indexes_must_be_non_negative to import
param directly from pytest and replace pytest.param calls with param. Move the
fmt: skip comment from the parameter list’s closing bracket to the closing
parenthesis of the `@pytest.mark.parametrize` call.
In `@tests/unit/post_processors/test_metrics_accumulator.py`:
- Around line 204-207: Update
test_export_results_filters_same_kind_records_by_phase_index with explicit type
annotations for every fixture parameter, including mock_metric_registry and
mock_run, and retain the async test’s None return annotation.
---
Outside diff comments:
In `@src/aiperf/config/flags/_converter_profiling.py`:
- Around line 328-339: Update the JSONL validation logic around load_json_str in
the profiling converter so a successfully parsed non-dict value returns False
instead of raising when accessing data.get("timestamp"). Validate the parsed
data type before the timestamp lookup, while preserving the existing behavior
for malformed JSON, missing timestamps, and file access errors.
---
Nitpick comments:
In `@src/aiperf/exporters/exporter_manager.py`:
- Around line 169-204: Add a public content-generation API on
MetricsJsonExporter and MetricsCsvExporter, such as generate_content(), and
update _write_phase_export to call it instead of the private _generate_content()
method. Preserve the existing generated content and file-writing behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fe0daed4-c94a-42ba-9043-cdc30216d187
📒 Files selected for processing (59)
docs/benchmark-modes/timing-modes-reference.mddocs/cli-options.mddocs/dev/patterns.mddocs/tutorials/adaptive-scale.mddocs/tutorials/yaml-config.mdsrc/aiperf/common/accumulator_protocols.pysrc/aiperf/common/models/__init__.pysrc/aiperf/common/models/credit_models.pysrc/aiperf/common/models/record_models.pysrc/aiperf/common/phase.pysrc/aiperf/common/types.pysrc/aiperf/config/config.pysrc/aiperf/config/flags/_adaptive_control_cli.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/flags/_resolver_adaptive.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/flags/converter.pysrc/aiperf/config/flags/resolver.pysrc/aiperf/config/loader/helpers.pysrc/aiperf/config/loader/normalizers.pysrc/aiperf/config/phases.pysrc/aiperf/config/sweep/expand.pysrc/aiperf/credit/callback_handler.pysrc/aiperf/credit/issuer.pysrc/aiperf/credit/messages.pysrc/aiperf/credit/structs.pysrc/aiperf/exporters/exporter_manager.pysrc/aiperf/gpu_telemetry/accumulator.pysrc/aiperf/gpu_telemetry/protocols.pysrc/aiperf/metrics/accumulator.pysrc/aiperf/post_processors/otel_metrics_results_processor.pysrc/aiperf/records/record_processor_service.pysrc/aiperf/records/records_manager.pysrc/aiperf/records/records_tracker.pysrc/aiperf/timing/concurrency.pysrc/aiperf/timing/config.pysrc/aiperf/timing/phase/progress_tracker.pysrc/aiperf/timing/phase/runner.pysrc/aiperf/timing/phase_orchestrator.pysrc/aiperf/timing/strategies/adaptive_scale.pysrc/aiperf/timing/strategies/adaptive_scale_artifacts.pysrc/aiperf/timing/strategies/adaptive_scale_backends.pysrc/aiperf/timing/strategies/adaptive_scale_runtime.pysrc/aiperf/workers/worker.pytests/component_integration/timing/test_adaptive_scale.pytests/unit/common/models/test_credit_models.pytests/unit/config/test_converter_profiling_phase_routes.pytests/unit/config/test_named_phase_kind.pytests/unit/config/test_v1_resolver_streaming_override.pytests/unit/credit/test_callback_handler.pytests/unit/exporters/test_exporter_manager.pytests/unit/gpu_telemetry/test_accumulator.pytests/unit/post_processors/test_metrics_accumulator.pytests/unit/records/test_dag_metadata_tagging.pytests/unit/records/test_records_manager.pytests/unit/records/test_records_manager_process_results.pytests/unit/timing/strategies/test_adaptive_scale.pytests/unit/timing/test_race_conditions.py
💤 Files with no reviewable changes (4)
- src/aiperf/config/flags/_section_fields.py
- src/aiperf/config/flags/_resolver_adaptive.py
- src/aiperf/config/flags/_adaptive_control_cli.py
- tests/unit/config/test_v1_resolver_streaming_override.py
🚧 Files skipped from review as they are similar to previous changes (39)
- docs/benchmark-modes/timing-modes-reference.md
- docs/tutorials/adaptive-scale.md
- src/aiperf/records/record_processor_service.py
- src/aiperf/workers/worker.py
- src/aiperf/common/models/init.py
- src/aiperf/credit/messages.py
- src/aiperf/credit/structs.py
- src/aiperf/post_processors/otel_metrics_results_processor.py
- src/aiperf/gpu_telemetry/protocols.py
- docs/dev/patterns.md
- src/aiperf/config/loader/helpers.py
- src/aiperf/timing/phase_orchestrator.py
- tests/unit/timing/test_race_conditions.py
- src/aiperf/timing/strategies/adaptive_scale_runtime.py
- src/aiperf/timing/strategies/adaptive_scale.py
- src/aiperf/common/models/credit_models.py
- src/aiperf/common/accumulator_protocols.py
- tests/unit/exporters/test_exporter_manager.py
- src/aiperf/timing/strategies/adaptive_scale_artifacts.py
- src/aiperf/config/sweep/expand.py
- tests/unit/credit/test_callback_handler.py
- tests/component_integration/timing/test_adaptive_scale.py
- src/aiperf/config/phases.py
- src/aiperf/timing/phase/progress_tracker.py
- src/aiperf/timing/strategies/adaptive_scale_backends.py
- src/aiperf/gpu_telemetry/accumulator.py
- src/aiperf/config/config.py
- tests/unit/records/test_records_manager_process_results.py
- src/aiperf/config/loader/normalizers.py
- src/aiperf/credit/issuer.py
- src/aiperf/timing/phase/runner.py
- tests/unit/config/test_converter_profiling_phase_routes.py
- tests/unit/records/test_dag_metadata_tagging.py
- src/aiperf/records/records_tracker.py
- src/aiperf/records/records_manager.py
- src/aiperf/config/flags/resolver.py
- src/aiperf/credit/callback_handler.py
- tests/unit/records/test_records_manager.py
- tests/unit/timing/strategies/test_adaptive_scale.py
🛑 Comments failed to post (4)
src/aiperf/config/flags/_converter_profiling.py (1)
343-371: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
One bad JSONL file zeroes the whole directory count.
The single outer
try/except (OSError, UnicodeDecodeError): return 0spans the entirerglob("*.jsonl")loop. A decode error in any one file among many discards thetotalalready accumulated from prior files and returns0for the whole directory, understating the dataset size rather than just skipping the offending file.🐛 Proposed fix — isolate per-file failures
if path.is_dir(): total = 0 for jsonl in path.rglob("*.jsonl"): - with open(jsonl, encoding="utf-8") as f: - total += sum(1 for line in f if line.strip()) + try: + with open(jsonl, encoding="utf-8") as f: + total += sum(1 for line in f if line.strip()) + except (OSError, UnicodeDecodeError): + continue return total📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.def _count_dataset_records(file_path: object) -> int: """Count records across a JSONL file/directory or Parquet trace file.""" from pathlib import Path path = Path(file_path) try: if path.is_dir(): total = 0 for jsonl in path.rglob("*.jsonl"): try: with open(jsonl, encoding="utf-8") as f: total += sum(1 for line in f if line.strip()) except (OSError, UnicodeDecodeError): continue return total if path.suffix.lower() == ".parquet" and path.is_file(): try: import pyarrow as pa import pyarrow.parquet as pq except ImportError: return 0 try: return pq.ParquetFile(path).metadata.num_rows except (OSError, pa.ArrowException): return 0 if path.is_file(): with open(path, encoding="utf-8") as f: return sum(1 for line in f if line.strip()) except (OSError, UnicodeDecodeError): return 0 return 0🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 351-351: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(jsonl, encoding="utf-8")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').(open-filename-from-request)
[warning] 366-366: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(path, encoding="utf-8")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/flags/_converter_profiling.py` around lines 343 - 371, Update _count_dataset_records so errors reading or decoding an individual JSONL file inside the directory loop are handled per file, skipping only that file while preserving the accumulated total from successfully processed files. Keep the existing behavior for directory discovery and non-directory file handling unchanged.src/aiperf/exporters/exporter_manager.py (1)
120-163: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Isolate per-phase failures inside the loop, not just around the whole method.
phase_dir.mkdir(...)(line 127) is unguarded; anOSErrorfor one phase's directory aborts the entirefor phase_result in phase_recordsloop, so every phase after the failing one silently loses its metric artifacts even though the outer try/except (lines 105-108) prevents a hard crash.🛡️ Proposed fix
for phase_result in phase_records: phase_dir = self._run.cfg.artifacts.dir / "phases" / phase_result.phase_name - await asyncio.to_thread(phase_dir.mkdir, parents=True, exist_ok=True) + try: + await asyncio.to_thread(phase_dir.mkdir, parents=True, exist_ok=True) + except OSError as exc: + self.warning( + f"Failed to create phase artifact dir {phase_dir}: {exc}" + ) + continue completed = (📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.async def _export_phase_metric_artifacts(self) -> None: phase_records = getattr(self._results, "phase_records", None) or [] if not phase_records: return manifest_entries: list[dict[str, Any]] = [] for phase_result in phase_records: phase_dir = self._run.cfg.artifacts.dir / "phases" / phase_result.phase_name try: await asyncio.to_thread(phase_dir.mkdir, parents=True, exist_ok=True) except OSError as exc: self.warning( f"Failed to create phase artifact dir {phase_dir}: {exc}" ) continue completed = ( phase_result.successful_request_count + phase_result.error_request_count ) phase_profile = ProfileResults( records=phase_result.records, completed=completed, start_ns=phase_result.start_ns or self._results.start_ns, end_ns=phase_result.end_ns or self._results.end_ns, was_cancelled=phase_result.was_cancelled, successful_request_count=phase_result.successful_request_count, error_request_count=phase_result.error_request_count, error_summary=[], ) entry: dict[str, Any] = { "phase_index": phase_result.phase_index, "profiling_index": phase_result.profiling_index, "phase_name": phase_result.phase_name, "phase_kind": phase_result.phase_kind, } await self._write_phase_export( exporter_cls=MetricsJsonExporter, phase_profile=phase_profile, file_path=phase_dir / self._run.cfg.artifacts.profile_export_json_file.name, manifest_entry=entry, manifest_key="metrics_json", ) await self._write_phase_export( exporter_cls=MetricsCsvExporter, phase_profile=phase_profile, file_path=phase_dir / self._run.cfg.artifacts.profile_export_csv_file.name, manifest_entry=entry, manifest_key="metrics_csv", ) manifest_entries.append(entry)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/exporters/exporter_manager.py` around lines 120 - 163, Update _export_phase_metric_artifacts to isolate each phase’s directory creation and export operations within a per-phase exception boundary, including phase_dir.mkdir and both _write_phase_export calls. When one phase raises an OSError or export error, record the failure consistently and continue processing subsequent phase_records instead of aborting the loop.tests/unit/common/models/test_credit_models.py (1)
87-96: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
pytest.paramimport style and misplaced# fmt: skip.Uses
pytest.param(...)instead of importingparam, and places# fmt: skipon the closing]of the list (line 92) instead of the closing)of@pytest.mark.parametrize(...)(line 93) — Ruff RUF028 confirms the comment is ineffective there.As per path instructions, "When using pytest.param, import it with from pytest import param and place # fmt: skip on the closing ) line of the parametrization block."🔧 Proposed fix
+from pytest import param + `@pytest.mark.parametrize`( "field", [ - pytest.param("phase_index", id="phase-index"), - pytest.param("profiling_index", id="profiling-index"), - ], # fmt: skip - ) + param("phase_index", id="phase-index"), + param("profiling_index", id="profiling-index"), + ], + ) # fmt: skip def test_phase_indexes_must_be_non_negative(self, field: str) -> None:🧰 Tools
🪛 Ruff (0.15.21)
[warning] 92-92: This suppression comment is invalid because it cannot be in an expression, pattern, argument list, or other non-statement
Remove this comment
(RUF028)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/common/models/test_credit_models.py` around lines 87 - 96, Update test_phase_indexes_must_be_non_negative to import param directly from pytest and replace pytest.param calls with param. Move the fmt: skip comment from the parameter list’s closing bracket to the closing parenthesis of the `@pytest.mark.parametrize` call.Sources: Path instructions, Linters/SAST tools
tests/unit/post_processors/test_metrics_accumulator.py (1)
204-207: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the required test annotations.
Annotate the fixture parameters and return type for this new async test. As per coding guidelines, “Annotate all functions with type hints for every parameter and return value.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/post_processors/test_metrics_accumulator.py` around lines 204 - 207, Update test_export_results_filters_same_kind_records_by_phase_index with explicit type annotations for every fixture parameter, including mock_metric_registry and mock_run, and retain the async test’s None return annotation.Source: Coding guidelines
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/aiperf/config/schema/aiperf-config.schema.json (1)
1043-1050: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUpdate the phase schema to allow named workflows
src/aiperf/config/schema/aiperf-config.schema.jsonstill constrains phasenametowarmup/profilingand doesn’t exposekind, so schema-based validation will reject the new arbitrary phase names. Regenerate or patch the schema to matchsrc/aiperf/config/phases.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/config/schema/aiperf-config.schema.json` around lines 1043 - 1050, Update the phase definition schema in aiperf-config.schema.json to match phases.py: allow arbitrary string values for phase name instead of restricting it to warmup/profiling, and add the phase kind property with the same validation and requiredness as the Python model. Regenerate the schema if supported; otherwise patch the corresponding phase schema directly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/post_processors/conftest.py`:
- Around line 71-76: Move uuid, BenchmarkConfig, and BenchmarkRun imports from
fixed_schedule_run to the module-level import block, then annotate
fixed_schedule_run with -> BenchmarkRun while preserving its existing fixture
behavior.
---
Outside diff comments:
In `@src/aiperf/config/schema/aiperf-config.schema.json`:
- Around line 1043-1050: Update the phase definition schema in
aiperf-config.schema.json to match phases.py: allow arbitrary string values for
phase name instead of restricting it to warmup/profiling, and add the phase kind
property with the same validation and requiredness as the Python model.
Regenerate the schema if supported; otherwise patch the corresponding phase
schema directly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 87470b2e-642e-4c3c-b1f0-d845c577e384
📒 Files selected for processing (16)
docs/metrics-reference.mdsrc/aiperf/config/dataset/constants.pysrc/aiperf/config/dataset/resolver.pysrc/aiperf/config/flags/_converter_dataset.pysrc/aiperf/config/flags/_converter_profiling.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/config/sweep/expand.pysrc/aiperf/dataset/composer/custom.pysrc/aiperf/dataset/loader/base_trace_loader.pysrc/aiperf/post_processors/metric_results_processor.pytests/unit/config/test_baseten_replay_flags.pytests/unit/dataset/loader/test_baseten_offline_suite.pytests/unit/dataset/loader/test_baseten_replay_timemodel.pytests/unit/post_processors/conftest.pytests/unit/post_processors/test_base_metrics_processor.pytests/unit/post_processors/test_derived_metric_instance_scope.py
💤 Files with no reviewable changes (2)
- tests/unit/post_processors/test_base_metrics_processor.py
- tests/unit/post_processors/test_derived_metric_instance_scope.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/aiperf/config/sweep/expand.py
- src/aiperf/config/flags/_converter_profiling.py
ilana-n
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
src/aiperf/accuracy/benchmarks/mmlu.py (1)
271-280: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated CoT instruction-building logic.
The base-instruction + conditional CoT clause + trailing blank line block is duplicated verbatim between
_format_promptand_build_chat_messages. Extracting a shared helper (e.g._build_instruction(subject, enable_cot)) would avoid the two copies drifting if the CoT wording is tweaked later.♻️ Proposed refactor
+ def _build_instruction(self, subject: str, enable_cot: bool) -> str: + instruction = ( + "The following are multiple choice questions (with answers) " + f"about {subject.replace('_', ' ')}." + ) + if enable_cot: + instruction += ( + " Think step by step and then output the answer in the format of " + '"The answer is (X)" at the end.' + ) + instruction += "\n\n" + return instructionThen in
_format_promptand_build_chat_messages, replace the inline block withinstruction = self._build_instruction(subject, enable_cot).Also applies to: 315-324
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/accuracy/benchmarks/mmlu.py` around lines 271 - 280, Extract the duplicated instruction construction from _format_prompt and _build_chat_messages into a shared _build_instruction(subject, enable_cot) helper, preserving the subject formatting, conditional CoT wording, and trailing blank line. Replace both inline blocks with calls to this helper so future wording changes remain centralized.src/aiperf/plugin/plugins.yaml (1)
1515-1521: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd metadata for the new grader plugin.
mmlu_prodefinesclassanddescriptionbut omitsmetadata. Add a metadata mapping consistent with the plugin registry contract. As per coding guidelines, “For new plugins, add an entry withclass,description, andmetadata.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/plugin/plugins.yaml` around lines 1515 - 1521, Add a metadata mapping to the mmlu_pro plugin entry alongside its existing class and description fields, following the structure and values required by the surrounding plugin registry entries and contract. Preserve the current class and description unchanged.Source: Coding guidelines
tests/unit/accuracy/test_mmlu_pro_grader.py (1)
24-63: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the required test-name structure.
Rename the new tests to
test_<function>_<scenario>_<expected>so failures identify the exercised method and intended outcome.
tests/unit/accuracy/test_mmlu_pro_grader.py#L24-L63: rename the newgrade()tests, e.g.test_grade_tier1_correct_clean.tests/unit/accuracy/test_multiple_choice_grader.py#L127-L154: rename the CoT fallback tests with an explicit expected outcome, e.g.test_grade_cot_answer_marker_clean.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/accuracy/test_mmlu_pro_grader.py` around lines 24 - 63, Rename the new grade() tests in tests/unit/accuracy/test_mmlu_pro_grader.py lines 24-63 to the required test_grade_<scenario>_<expected> structure, using names that identify the exercised scenario and outcome. Also rename the CoT fallback tests in tests/unit/accuracy/test_multiple_choice_grader.py lines 127-154 to include the grade method, scenario, and explicit expected result; change test names only and preserve test behavior.Source: Coding guidelines
src/aiperf/accuracy/graders/mmlu_pro.py (1)
36-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid annotating dynamic
**kwargsasAny.These methods do not consume a defined keyword contract; leave
**kwargsuntyped or replace it with explicit parameters. Based on learnings, “avoid adding type annotations to**kwargslike**kwargs: Any.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/accuracy/graders/mmlu_pro.py` around lines 36 - 42, Remove the Any annotation from the dynamic **kwargs parameters in extract_answer and grade, or replace them with explicit keyword parameters if a defined contract exists. Keep the existing method behavior and return annotations unchanged.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/integration/test_adaptive_scale.py`:
- Around line 122-127: Strengthen the assertions in the adaptive-scale artifact
test so adaptive_phase["events_path"] and adaptive_phase["summary_path"] are
verified to reside under the phases/profiling/ directory, preferably by
asserting their exact expected relative paths. Retain the existing existence
checks and top-level artifact absence checks.
- Around line 112-115: Update the async test’s artifact checks around
manifest_path and the later existence checks to avoid synchronous Path.exists()
and read_bytes() calls. Use the project’s async file API, or await
asyncio.to_thread for each filesystem operation, while preserving the existing
manifest parsing and schema_version assertion.
---
Nitpick comments:
In `@src/aiperf/accuracy/benchmarks/mmlu.py`:
- Around line 271-280: Extract the duplicated instruction construction from
_format_prompt and _build_chat_messages into a shared
_build_instruction(subject, enable_cot) helper, preserving the subject
formatting, conditional CoT wording, and trailing blank line. Replace both
inline blocks with calls to this helper so future wording changes remain
centralized.
In `@src/aiperf/accuracy/graders/mmlu_pro.py`:
- Around line 36-42: Remove the Any annotation from the dynamic **kwargs
parameters in extract_answer and grade, or replace them with explicit keyword
parameters if a defined contract exists. Keep the existing method behavior and
return annotations unchanged.
In `@src/aiperf/plugin/plugins.yaml`:
- Around line 1515-1521: Add a metadata mapping to the mmlu_pro plugin entry
alongside its existing class and description fields, following the structure and
values required by the surrounding plugin registry entries and contract.
Preserve the current class and description unchanged.
In `@tests/unit/accuracy/test_mmlu_pro_grader.py`:
- Around line 24-63: Rename the new grade() tests in
tests/unit/accuracy/test_mmlu_pro_grader.py lines 24-63 to the required
test_grade_<scenario>_<expected> structure, using names that identify the
exercised scenario and outcome. Also rename the CoT fallback tests in
tests/unit/accuracy/test_multiple_choice_grader.py lines 127-154 to include the
grade method, scenario, and explicit expected result; change test names only and
preserve test behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: afd366d2-a912-4439-91ee-33a0681b0dfc
📒 Files selected for processing (36)
docs/accuracy/accuracy-benchmarking.mddocs/cli-options.mddocs/tutorials/adaptive-scale.mddocs/tutorials/yaml-config.mdsrc/aiperf/accuracy/accuracy_record_processor.pysrc/aiperf/accuracy/accuracy_results_processor.pysrc/aiperf/accuracy/benchmarks/mmlu.pysrc/aiperf/accuracy/benchmarks/mmlu_pro.pysrc/aiperf/accuracy/graders/_choice_extract.pysrc/aiperf/accuracy/graders/mmlu_pro.pysrc/aiperf/accuracy/graders/multiple_choice.pysrc/aiperf/accuracy/models.pysrc/aiperf/common/enums/metric_enums.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/dataset/loader/accuracy_dataset_loader.pysrc/aiperf/metrics/types/accuracy_metrics.pysrc/aiperf/metrics/types/osl_mismatch_metrics.pysrc/aiperf/plugin/enums.pysrc/aiperf/plugin/plugins.yamlsrc/aiperf/post_processors/base_metrics_processor.pysrc/aiperf/records/records_manager.pysrc/aiperf/timing/config.pytests/integration/test_adaptive_scale.pytests/unit/accuracy/test_accuracy_record_processor.pytests/unit/accuracy/test_choice_extract.pytests/unit/accuracy/test_hf_benchmark_datasets.pytests/unit/accuracy/test_mmlu_cot.pytests/unit/accuracy/test_mmlu_pro_benchmark.pytests/unit/accuracy/test_mmlu_pro_grader.pytests/unit/accuracy/test_multiple_choice_grader.pytests/unit/config/test_accuracy_enable_cot_flag.pytests/unit/post_processors/conftest.pytests/unit/records/test_records_manager.pytests/unit/records/test_records_manager_process_results.pytests/unit/timing/test_timing_config.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/unit/records/test_records_manager_process_results.py
- src/aiperf/config/flags/cli_config.py
- docs/tutorials/adaptive-scale.md
- tests/unit/records/test_records_manager.py
- src/aiperf/timing/config.py
- src/aiperf/records/records_manager.py
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
ilana-n
force-pushed
the
ilana/aip-1004-1006-named-phase-kind-workflows
branch
from
df3e1d1 to
ed07d76
Compare
Signed-off-by: inguyen <inguyen@nvidia.com>
ilana-n
force-pushed
the
ilana/aip-1004-1006-named-phase-kind-workflows
branch
from
ed07d76 to
7a96498
Compare
Signed-off-by: inguyen <inguyen@nvidia.com>
Signed-off-by: inguyen <inguyen@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: inguyen <inguyen@nvidia.com>
ilana-n
requested review from
FrankD412,
ajcasagrande,
debermudez,
ganeshku1,
lkomali and
matthewkotila
as code owners
Signed-off-by: inguyen <inguyen@nvidia.com>
ajcasagrande
left a comment
ajcasagrande
left a comment
There was a problem hiding this comment.
AIPerf multi-phase review — solid, coherent feature; all 2089 targeted unit tests pass and the per-phase artifact pipeline works end-to-end (verified against the in-repo mock server with a warmup + 3-profiling-phase config). Two real correctness defects found, both confined to the new multi-phase surface — single-profiling-phase and CLI paths are unaffected. Both reproduced at runtime.
Fix order:
- (High) Root/top-level export silently reflects only the last profiling phase's metric distributions while counts are summed across all phases → internally inconsistent primary artifact.
- (Medium) Per-phase cancellation doesn't inherit the run default the docstring promises; the run-level value and
PhaseOrchestrator._cancellation_policyare dead code.
Working well: the name/kind split and validation; concurrency/credit slot keying round-trip (no leak/over-release); adaptive backend keying; final-phase deferral (no stranded results); NaN-safe accumulator masking; clean adaptive-scale CLI removal (no dangling refs); path-safe phase-scoped artifacts.
| start_ns=phase_stats.start_ns, | ||
| end_ns=phase_stats.requests_end_ns, | ||
| phase=phase, | ||
| phase_index=phase_stats.phase_index, |
There was a problem hiding this comment.
High: root export collapses to only the last profiling phase. _process_results (line 1920) passes phase_index=None for the PROFILING summary intending to aggregate all profiling phases via the categorical benchmark_phase mask. But here create_stats_for_phase(phase, phase_index=None) resolves via _latest_tracker_for_phase to the last profiling tracker, so ctx.phase_index becomes that concrete int (not None), and accumulator._mask_for_export_context filters to only that phase's records.
Reproduced (warmup + low/storm/recovery, 20 req each): root profile_export_aiperf.json reports request_count=20 and is byte-identical to recovery, while completed (from _aggregate_records_stats_for_phase) sums to 60 — an inconsistent summary. The same defect collapses the warmup aggregate. Suggest ExportContext(..., phase_index=phase_index, ...) (honor the requested arg) plus a regression test that exercises the real column-store masking with ≥2 profiling phases.
There was a problem hiding this comment.
ah, this is important thanks for catching it - thought I had fixed it but I see not entirely!
There was a problem hiding this comment.
gonna wait on #1144 since that will be involved with the fix here
| _build_phase_config( | ||
| phase, | ||
| artifact_dir=artifact_dir, | ||
| default_cancellation=default_cancellation, |
There was a problem hiding this comment.
Medium: omitted per-phase cancellation gets NO cancellation, not the documented run default. The threaded default_cancellation (line 101) is an empty RequestCancellationConfig(); the real run default run_cancellation (line 102) is only stored on TimingConfig.request_cancellation, which PhaseOrchestrator no longer reads (it builds a fresh per-phase simulator), so run_cancellation and _cancellation_policy are dead. The from_run docstring (line 95) says "omitted per-phase cancellation inherits the run default."
Reproduced (two profiling phases: cancel_phase rate 100, inherit_phase omitted): cancel_phase cancels 20/20, inherit_phase cancels 0/20 — inheritance never happens. Also _default_cancellation_config now scans warmup phases (ordering change vs get_profiling_phases()). Please pick one semantic and make code + docstring agree, and drop the dead run-level policy.
There was a problem hiding this comment.
thank you! will address
2 participants
Pull Request Summary
Adds named phase support to AIPerf so benchmark configs can model Rednote-style workflows with multiple warmup/profiling phases, phase-specific runtime metadata, cancellation behavior, adaptive behavior, and phase-scoped artifacts. This PR separates a phase’s workflow
namefrom its semantickind, threads phase identity through timing, credit, records, metrics, and export paths, and updates docs/tests for multi-phase targeting.It also intentionally removes the unreleased adaptive-scale CLI surface. Adaptive scale is now a YAML-only phase feature because the CLI form becomes ambiguous once configs can contain multiple profiling/adaptive phases; YAML keeps the adaptive controller attached to the exact phase it controls.
Changes
PhaseKind = Literal["warmup", "profiling"]and separated phasenamefrom phase semantickind.phase_index,profiling_index,phase_name, andphase_kind.MetricRecordMetadata,ColumnStore, andExportContext.phase_indexbefore semantic phase, so two profiling-kind phases do not collapse into the same export.CreditPhase.PROFILINGrecords with differentphase_indexvalues export separately.phases/<phase_name>/...plus an adaptive scale manifest.phases.profiling.*only resolves when unambiguous.kindguidance, multi-phase sweep paths, phase-scoped adaptive artifacts, and phase-aware export behavior.--adaptive-scale--adaptive-sustain-duration--adaptive-assessment-period--adaptive-scale-control--adaptive-control-*--adaptive-scale-slaAdaptive Scale CLI Removal
Adaptive scale was initially exposed through CLI flags for the simple one-phase case, but named phases make that interface ambiguous. A config can now have multiple profiling phases, and more than one of them may be adaptive. In that world, a command like
--adaptive-sustain-duration 45has no natural target unless we add another phase-selection mini-language to the CLI.Rather than adding that complexity before release, this PR removes the adaptive-scale CLI surface entirely. Adaptive scale is more natural as YAML because the controller settings live directly on the phase they affect:
This avoids ambiguous targeting, keeps multi-phase configs readable, and leaves normal CLI overrides available for unambiguous non-adaptive loadgen knobs.
Export Behavior
This PR extends the export pipeline so phase identity survives beyond runtime execution. Metric records now carry concrete phase metadata, and the accumulator stores that metadata alongside metric columns. During export,
ExportContextincludesphase_index,phase_name,phase_kind, andprofiling_index, allowing summaries and metric outputs to isolate a specific phase instead of grouping by semantic warmup/profiling alone.That matters for named workflows where several phases are all
kind: profiling. For example,low,storm, andrecoveryshould not merge into one profiling export. The accumulator now prefersphase_indexwhen present, then falls back to the legacy semantic phase filter for older/programmatic records.Why This Is Important
Rednote-style workflows need more than a single canonical
warmupplusprofilingpair. They need ordered phases like setup, low-cancel, storm, recovery, and repeated profiling windows, each with distinct cancellation/adaptive behavior and separate runtime/result identity.Without concrete phase metadata, several paths collapse all profiling work together: callback state can overwrite another profiling phase, records can lose which profiling phase produced a metric, and exports can merge distinct profiling windows into one result set. This PR makes phase identity explicit across config, runtime, records, metrics, artifacts, and export outputs.
The
name/kindsplit keeps user workflows expressive while preserving the existing semantic distinction the credit and result pipelines actually need: warmup versus profiling.How To Test
Run focused config and named phase tests:
Run adaptive scale strategy tests:
Run records/metrics accumulator coverage:
Run callback and timing race coverage:
Run adaptive component integration with the marker enabled:
Expected Output
Multi-phase runs preserve concrete phase metadata:
phase_indexprofiling_indexphase_namephase_kindMetrics exports use concrete phase identity when available, so multiple profiling-kind phases export independently instead of merging into one profiling bucket.
Adaptive scale artifacts are phase-scoped:
Related Issues
Summary by CodeRabbit
adaptive_scale_manifest.json.mmlu_pro, with new--accuracy-enable-cot/--accuracy-no-enable-cotbehavior.FIXED_SCHEDULE_ONLYmetric flag.kindinference and case-insensitive phase-name uniqueness.