fix(dataset): harden inputs.json and raw-payload loader input validation#1120
fix(dataset): harden inputs.json and raw-payload loader input validation#1120ajcasagrande wants to merge 1 commit into
Conversation
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@63749f2c204c147935c703a4f6aab65e1455a452Recommended 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@63749f2c204c147935c703a4f6aab65e1455a452Last updated for commit: |
WalkthroughThis PR adds required-key validation to InputsJsonPayloadLoader.load_dataset, raising ValueError when session_id or payloads are missing, and tightens RawPayloadDatasetLoader.can_load to reject non-dict data and to stop silently swallowing OSError during directory probing. Two new adversarial test modules cover these behaviors. ChangesLoader validation fixes
Estimated code review effort: 2 (Simple) | ~12 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/unit/dataset/loader/test_inputs_json_adversarial.py (1)
32-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueParametrize missing
pytest.param/# fmt: skip.As per path instructions,
tests/**/*.pyshould "Usefrom pytest import paramand put# fmt: skipon the)line for@pytest.mark.parametrize."🧹 Proposed fix
-import orjson -import pytest +import orjson +import pytest +from pytest import param from pydantic import ValidationError ... - `@pytest.mark.parametrize`("bad_data", [[], "s", 123]) + `@pytest.mark.parametrize`( + "bad_data", [param([]), param("s"), param(123)] + ) # 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/dataset/loader/test_inputs_json_adversarial.py` at line 32, The `@pytest.mark.parametrize` in the test module should follow the project’s pytest formatting convention. Update the parametrization in the test function to use pytest.param entries from from pytest import param, and add # fmt: skip on the closing parenthesis line of the decorator so formatting is preserved. Use the parametrized test definition around the bad_data cases as the target for this change.Source: Path instructions
src/aiperf/dataset/loader/raw_payload.py (1)
164-168: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant exception type in tuple.
orjson.JSONDecodeErroris already a subclass ofValueError, so catchingValueErroralone would produce identical behavior. Not incorrect, just redundant — the explicit inclusion oforjson.JSONDecodeErrorcould be seen as documenting intent.🤖 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/raw_payload.py` around lines 164 - 168, The exception tuple in the raw payload loader is redundant because `orjson.JSONDecodeError` is already covered by `ValueError`. Simplify the `except` in the JSON parsing path of `raw_payload.py` by keeping only `ValueError` in the `can_load`/payload-loading logic, or otherwise make the choice consistent with the intended intent if you want the explicit type as documentation.tests/unit/dataset/loader/test_raw_payload_adversarial.py (1)
31-31: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueParametrize missing
pytest.param/# fmt: skip.As per path instructions,
tests/**/*.pyshould "Usefrom pytest import paramand put# fmt: skipon the)line for@pytest.mark.parametrize."🧹 Proposed fix
-import orjson -import pytest +import orjson +import pytest +from pytest import param ... - `@pytest.mark.parametrize`("bad_data", [[], "string", 123]) + `@pytest.mark.parametrize`( + "bad_data", [param([]), param("string"), param(123)] + ) # 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/dataset/loader/test_raw_payload_adversarial.py` at line 31, The parametrized test in the `test_raw_payload_adversarial` module should follow the repository’s formatting rule for `tests/**/*.py`: import `param` from `pytest`, use `param(...)` entries instead of bare literals in `pytest.mark.parametrize`, and place `# fmt: skip` on the closing parenthesis line of the decorator. Update the `@pytest.mark.parametrize` declaration accordingly so it matches the style used by other tests and avoids formatter churn.Source: Path instructions
🤖 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/dataset/loader/inputs_json.py`:
- Around line 74-80: The validation loop in inputs_json loading should reject
non-dict entries before checking required keys. Update the entry validation in
the data_list iteration so the code first verifies each item is a mapping (for
example in the loop that checks session_id and payloads) and raises the same
ValueError path for invalid types, then perform the existing key presence checks
on valid dict entries only.
In `@tests/unit/dataset/loader/test_inputs_json_adversarial.py`:
- Around line 3-9: The module docstring in TestWave2FixForwardCompatibility is
stale because this file does not contain any xfail-strict tests. Update the
docstring to describe the actual coverage in test_inputs_json_adversarial.py:
the adversarial can_load/load_dataset behavior and the direct assertions for the
Wave 2 fixes, without mentioning xfail markers or tests that do not exist here.
---
Nitpick comments:
In `@src/aiperf/dataset/loader/raw_payload.py`:
- Around line 164-168: The exception tuple in the raw payload loader is
redundant because `orjson.JSONDecodeError` is already covered by `ValueError`.
Simplify the `except` in the JSON parsing path of `raw_payload.py` by keeping
only `ValueError` in the `can_load`/payload-loading logic, or otherwise make the
choice consistent with the intended intent if you want the explicit type as
documentation.
In `@tests/unit/dataset/loader/test_inputs_json_adversarial.py`:
- Line 32: The `@pytest.mark.parametrize` in the test module should follow the
project’s pytest formatting convention. Update the parametrization in the test
function to use pytest.param entries from from pytest import param, and add #
fmt: skip on the closing parenthesis line of the decorator so formatting is
preserved. Use the parametrized test definition around the bad_data cases as the
target for this change.
In `@tests/unit/dataset/loader/test_raw_payload_adversarial.py`:
- Line 31: The parametrized test in the `test_raw_payload_adversarial` module
should follow the repository’s formatting rule for `tests/**/*.py`: import
`param` from `pytest`, use `param(...)` entries instead of bare literals in
`pytest.mark.parametrize`, and place `# fmt: skip` on the closing parenthesis
line of the decorator. Update the `@pytest.mark.parametrize` declaration
accordingly so it matches the style used by other tests and avoids formatter
churn.
🪄 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: 65751dfd-d6aa-4a0f-a083-c66d7ae8757b
📒 Files selected for processing (4)
src/aiperf/dataset/loader/inputs_json.pysrc/aiperf/dataset/loader/raw_payload.pytests/unit/dataset/loader/test_inputs_json_adversarial.pytests/unit/dataset/loader/test_raw_payload_adversarial.py
| for idx, entry in enumerate(data_list): | ||
| if "session_id" not in entry: | ||
| raise ValueError( | ||
| f"{path}: data[{idx}] missing required key 'session_id'" | ||
| ) | ||
| if "payloads" not in entry: | ||
| raise ValueError(f"{path}: data[{idx}] missing required key 'payloads'") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file around the reported lines.
sed -n '1,180p' src/aiperf/dataset/loader/inputs_json.py
# Find any surrounding validation for entry types or JSON shape in the loader.
rg -n "session_id|payloads|isinstance\\(entry|data_list|ValueError" src/aiperf/dataset/loader -nRepository: ai-dynamo/aiperf
Length of output: 26717
Validation should reject non-object entries first.
"session_id" not in entry assumes each data[] item is a mapping. A null, number, or list will raise an unhandled TypeError here instead of the intended ValueError. Add an isinstance(entry, dict) guard before the key checks.
🤖 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/inputs_json.py` around lines 74 - 80, The
validation loop in inputs_json loading should reject non-dict entries before
checking required keys. Update the entry validation in the data_list iteration
so the code first verifies each item is a mapping (for example in the loop that
checks session_id and payloads) and raises the same ValueError path for invalid
types, then perform the existing key presence checks on valid dict entries only.
| """Adversarial coverage for InputsJsonPayloadLoader. | ||
|
|
||
| Pins current behavior for edge cases in `can_load` and `load_dataset`, | ||
| and marks two xfail-strict tests that will flip to pass when the known | ||
| bugs (duplicate session_id overwrite, bare KeyError on missing keys) are | ||
| fixed in Wave 2. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale docstring: no xfail-strict tests exist in this file.
The module docstring says it "marks two xfail-strict tests that will flip to pass," but TestWave2FixForwardCompatibility tests assert the fixed behavior directly with pytest.raises(...), with no @pytest.mark.xfail anywhere in the file. This looks like a copy-paste holdover from the raw_payload adversarial suite (which does still have such a marker) and will confuse future readers.
🤖 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_inputs_json_adversarial.py` around lines 3 -
9, The module docstring in TestWave2FixForwardCompatibility is stale because
this file does not contain any xfail-strict tests. Update the docstring to
describe the actual coverage in test_inputs_json_adversarial.py: the adversarial
can_load/load_dataset behavior and the direct assertions for the Wave 2 fixes,
without mentioning xfail markers or tests that do not exist here.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
Summary
Hardens input validation in two dataset loaders so malformed user inputs fail with actionable errors instead of bare tracebacks or silent misclassification. Carved from the
ajc/agentx-on-mainadversarial-hardening work; includes the two new adversarial test suites that pin the fixed behavior.Crashes fixed (with pre-fix exceptions)
1.
InputsJsonPayloadLoader.load_dataset— missing required keys (src/aiperf/dataset/loader/inputs_json.py)An
inputs.jsonentry missingsession_idorpayloadscrashed with a bare, context-free exception:Now raises a path- and index-rich error:
2.
RawPayloadDatasetLoader.can_load— non-dict probe data (src/aiperf/dataset/loader/raw_payload.py)Type-dispatch auto-detection feeds arbitrary first-record shapes into
can_load. A non-dict record (list, string, scalar) crashed the detection chain mid-walk:Now guarded with
isinstance(data, dict)so detection falls through cleanly (can_loadreturnsFalse).Deliberate behavior change: OSError now propagates from
_dir_has_raw_payload_jsonlPreviously
_dir_has_raw_payload_jsonlcaught(orjson.JSONDecodeError, OSError)and skipped to the next file. That meant an unreadable.jsonlfile (e.g. permission-denied) was silently swallowed and the directory was misclassified as not-raw-payload, sending the run down the wrong loader path with no signal.The except clause is narrowed to
(orjson.JSONDecodeError, ValueError): only malformed JSON is skipped;OSError/PermissionErrorfrom an unreadable file now propagates out ofcan_loadso misconfigured inputs fail loudly at detection time. This is an intentional loud-failure change, pinned bytest_can_load_directory_with_unreadable_jsonl_raises_post_fix.Tests
tests/unit/dataset/loader/test_inputs_json_adversarial.py(new, 142 lines) —can_loadboundary shapes, Pydanticmin_lengthenforcement, session_id passthrough, one-turn-per-payload, duplicate/missing-key error pins.tests/unit/dataset/loader/test_raw_payload_adversarial.py(new, 218 lines) — non-dict probe data, directory-peek discrimination (empty/null/array first lines, wrong extensions, malformed JSONL), JSONL line edge cases (blank lines, trailing newline, verbatimraw_payload), unreadable-file loud-failure pin.Revert-verification (pin discipline)
Reverse-applied the two src fixes onto
mainand ran the adversarial suites — 5 tests failed, reproducing the exact pre-fix crashes:Re-applied the fixes: all green.
Evidence
uv run pytest -n auto tests/unit/dataset/loader/test_inputs_json_adversarial.py tests/unit/dataset/loader/test_raw_payload_adversarial.py tests/unit/dataset/loader/test_inputs_json_payload.py tests/unit/dataset/loader/test_raw_payload.py— 47 passedpytest --collect-only tests/unit/dataset/loader/— clean (931 collected)ruff format+ruff checkon all four touched files — cleanmake check-ruff-baselined— OK (118 total, 118 baselined, 0 new)make check-ergonomics— OK (58 total, 58 baselined, 0 new)🤖 Generated with Claude Code
Summary by CodeRabbit