From 63749f2c204c147935c703a4f6aab65e1455a452 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Mon, 6 Jul 2026 20:38:43 -0700 Subject: [PATCH] fix(dataset): harden inputs.json and raw-payload loader input validation Co-Authored-By: Claude Fable 5 Signed-off-by: Anthony Casagrande --- src/aiperf/dataset/loader/inputs_json.py | 6 + src/aiperf/dataset/loader/raw_payload.py | 10 +- .../loader/test_inputs_json_adversarial.py | 142 ++++++++++++ .../loader/test_raw_payload_adversarial.py | 218 ++++++++++++++++++ 4 files changed, 375 insertions(+), 1 deletion(-) create mode 100644 tests/unit/dataset/loader/test_inputs_json_adversarial.py create mode 100644 tests/unit/dataset/loader/test_raw_payload_adversarial.py diff --git a/src/aiperf/dataset/loader/inputs_json.py b/src/aiperf/dataset/loader/inputs_json.py index c7019005e3..2f40ebce6d 100644 --- a/src/aiperf/dataset/loader/inputs_json.py +++ b/src/aiperf/dataset/loader/inputs_json.py @@ -72,6 +72,12 @@ def load_dataset(self) -> dict[str, list[InputsJsonSession]]: result: dict[str, list[InputsJsonSession]] = {} 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'") session = InputsJsonSession( session_id=entry["session_id"], payloads=entry["payloads"], diff --git a/src/aiperf/dataset/loader/raw_payload.py b/src/aiperf/dataset/loader/raw_payload.py index f6a6b78362..c43079e9cd 100644 --- a/src/aiperf/dataset/loader/raw_payload.py +++ b/src/aiperf/dataset/loader/raw_payload.py @@ -49,6 +49,11 @@ def can_load( InputsFile structures (``data`` key holding a list). """ if data is not None: + # Type-dispatch plugins feed arbitrary first-record shapes here; + # guard against non-dict inputs (list, string, scalar) so + # auto-detection falls through cleanly instead of AttributeError. + if not isinstance(data, dict): + return False if is_speed_bench_row(data): return False if not isinstance(data.get("messages"), list): @@ -156,7 +161,10 @@ def _dir_has_raw_payload_jsonl(directory: Path) -> bool: return isinstance(record, dict) and isinstance( record.get("messages"), list ) - except (orjson.JSONDecodeError, OSError): + except (orjson.JSONDecodeError, ValueError): + # Narrow: only malformed JSON is skipped. OSError/PermissionError + # from an unreadable file must propagate so can_load fails loud + # instead of silently misclassifying the directory. continue return False diff --git a/tests/unit/dataset/loader/test_inputs_json_adversarial.py b/tests/unit/dataset/loader/test_inputs_json_adversarial.py new file mode 100644 index 0000000000..2ce5353188 --- /dev/null +++ b/tests/unit/dataset/loader/test_inputs_json_adversarial.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import orjson +import pytest +from pydantic import ValidationError + +from aiperf.dataset.loader.inputs_json import InputsJsonPayloadLoader +from aiperf.dataset.loader.models import InputsJsonSession + + +def _make_loader(filename): + loader = InputsJsonPayloadLoader.__new__(InputsJsonPayloadLoader) + loader.filename = str(filename) + loader.info = MagicMock() + loader.debug = MagicMock() + return loader + + +class TestCanLoadAdversarial: + @pytest.mark.parametrize("bad_data", [[], "s", 123]) + def test_can_load_non_dict_data_returns_false(self, bad_data): + """`can_load` with non-dict `data` must return False, not raise.""" + assert InputsJsonPayloadLoader.can_load(data=bad_data) is False + + def test_can_load_dict_without_data_key_returns_false(self): + """Dict without a top-level ``data`` key must return False.""" + assert InputsJsonPayloadLoader.can_load(data={"not_data": []}) is False + + def test_can_load_data_not_a_list_returns_false(self): + """Top-level ``data`` value that isn't a list must return False.""" + assert InputsJsonPayloadLoader.can_load(data={"data": "str"}) is False + + def test_can_load_file_with_non_json_extension_returns_false(self, tmp_path): + """Files with a non-``.json`` suffix must return False even if content is valid JSON.""" + path = tmp_path / "inputs.txt" + path.write_bytes( + orjson.dumps({"data": [{"session_id": "s", "payloads": [{"model": "m"}]}]}) + ) + assert InputsJsonPayloadLoader.can_load(filename=path) is False + + def test_can_load_zero_byte_file_returns_false(self, tmp_path): + """Empty files must return False (orjson raises, caught).""" + path = tmp_path / "empty.json" + path.write_bytes(b"") + assert InputsJsonPayloadLoader.can_load(filename=path) is False + + def test_can_load_file_with_json_array_top_level_returns_false(self, tmp_path): + """Files whose root JSON value is an array must return False.""" + path = tmp_path / "array.json" + path.write_bytes(orjson.dumps([{"session_id": "x", "payloads": [{"m": 1}]}])) + assert InputsJsonPayloadLoader.can_load(filename=path) is False + + +class TestLoadDatasetAdversarial: + def test_load_dataset_entry_with_empty_payloads_list_rejected_by_pydantic( + self, tmp_path + ): + """``InputsJsonSession`` has ``min_length=1`` on payloads; empty list raises ValidationError.""" + path = tmp_path / "inputs.json" + path.write_bytes(orjson.dumps({"data": [{"session_id": "s", "payloads": []}]})) + loader = _make_loader(path) + with pytest.raises(ValidationError): + loader.load_dataset() + + def test_convert_to_conversations_session_id_passthrough(self, tmp_path): + """The ``session_id`` from the file must appear verbatim on the resulting Conversation.""" + data = { + "data": [ + { + "session_id": "custom-id", + "payloads": [{"model": "m", "messages": []}], + } + ] + } + path = tmp_path / "inputs.json" + path.write_bytes(orjson.dumps(data)) + loader = _make_loader(path) + conversations = loader.convert_to_conversations(loader.load_dataset()) + assert len(conversations) == 1 + assert conversations[0].session_id == "custom-id" + + def test_convert_to_conversations_emits_one_turn_per_payload(self, tmp_path): + """Three payloads must produce three Turns, each with ``raw_payload`` set.""" + payloads = [ + {"model": "m", "turn": 1}, + {"model": "m", "turn": 2}, + {"model": "m", "turn": 3}, + ] + data = {"data": [{"session_id": "s", "payloads": payloads}]} + path = tmp_path / "inputs.json" + path.write_bytes(orjson.dumps(data)) + loader = _make_loader(path) + conversations = loader.convert_to_conversations(loader.load_dataset()) + assert len(conversations) == 1 + turns = conversations[0].turns + assert len(turns) == 3 + for turn, expected in zip(turns, payloads, strict=True): + assert turn.raw_payload == expected + + +class TestWave2FixForwardCompatibility: + """Tests that pin the post-fix behavior after Wave 2 bug fixes landed.""" + + def test_load_dataset_duplicate_session_id_rejected_post_fix(self, tmp_path): + data = { + "data": [ + {"session_id": "dup", "payloads": [{"model": "m1"}]}, + {"session_id": "dup", "payloads": [{"model": "m2"}]}, + ] + } + path = tmp_path / "inputs.json" + path.write_bytes(orjson.dumps(data)) + loader = _make_loader(path) + with pytest.raises(ValueError, match="duplicate"): + loader.load_dataset() + + def test_load_dataset_missing_required_key_raises_value_error_post_fix( + self, tmp_path + ): + path = tmp_path / "inputs.json" + path.write_bytes(orjson.dumps({"data": [{"payloads": [{"x": 1}]}]})) + loader = _make_loader(path) + with pytest.raises(ValueError, match="session_id"): + loader.load_dataset() + + +def test_inputs_json_session_model_rejects_empty_payloads_directly(): + """Sanity check that the Pydantic constraint is on the model, not the loader.""" + with pytest.raises(ValidationError): + InputsJsonSession(session_id="s", payloads=[]) diff --git a/tests/unit/dataset/loader/test_raw_payload_adversarial.py b/tests/unit/dataset/loader/test_raw_payload_adversarial.py new file mode 100644 index 0000000000..85e1977f58 --- /dev/null +++ b/tests/unit/dataset/loader/test_raw_payload_adversarial.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adversarial coverage for RawPayloadDatasetLoader. + +Exercises can_load boundary inputs (non-dict data, malformed JSONL, directory +discrimination rules) and load_dataset line-parsing edge cases that the +shipped unit tests do not cover. Task 1 of the raw-payload adversarial pass. +""" + +from unittest.mock import MagicMock + +import orjson +import pytest + +from aiperf.dataset.loader.raw_payload import RawPayloadDatasetLoader + + +def _make_loader(filename): + """Construct a loader bypassing __init__ to avoid UserConfig wiring.""" + loader = RawPayloadDatasetLoader.__new__(RawPayloadDatasetLoader) + loader.filename = str(filename) + loader.session_id_generator = MagicMock() + loader.session_id_generator.next.side_effect = [f"s{i}" for i in range(100)] + loader.info = MagicMock() + loader.debug = MagicMock() + return loader + + +class TestCanLoadDataShape: + @pytest.mark.parametrize("bad_data", [[], "string", 123]) + def test_can_load_non_dict_data_returns_false(self, bad_data): + """can_load guards against non-dict inputs and returns False cleanly. + Auto-detection plugins feed arbitrary first-record shapes here; prior + to the defensive guard, non-dict data raised AttributeError at + ``data.get()`` and broke the detection chain mid-walk. + """ + assert RawPayloadDatasetLoader.can_load(data=bad_data) is False + + def test_can_load_data_dict_without_messages_key_returns_false(self): + assert RawPayloadDatasetLoader.can_load(data={"not_messages": []}) is False + + def test_can_load_data_dict_messages_not_a_list_returns_false(self): + assert ( + RawPayloadDatasetLoader.can_load(data={"messages": "not-a-list"}) is False + ) + + def test_can_load_data_dict_with_conversation_id_returns_false(self): + """Agentic trajectory records must be rejected even with messages.""" + assert ( + RawPayloadDatasetLoader.can_load( + data={"messages": [], "conversation_id": "x"} + ) + is False + ) + + def test_can_load_data_dict_with_top_level_data_list_returns_false(self): + """InputsFile shape (top-level data=list) must not match raw-payload.""" + assert ( + RawPayloadDatasetLoader.can_load(data={"messages": [], "data": []}) is False + ) + + +class TestCanLoadDirectoryPeek: + def test_can_load_file_with_zero_byte_jsonl_returns_false(self, tmp_path): + d = tmp_path / "empty_line" + d.mkdir() + (d / "empty.jsonl").write_bytes(b"") + assert RawPayloadDatasetLoader.can_load(filename=d) is False + + def test_can_load_file_with_null_first_line_returns_false(self, tmp_path): + d = tmp_path / "null_line" + d.mkdir() + (d / "null.jsonl").write_bytes(b"null\n") + assert RawPayloadDatasetLoader.can_load(filename=d) is False + + def test_can_load_file_with_json_array_first_line_returns_false(self, tmp_path): + d = tmp_path / "array_line" + d.mkdir() + (d / "arr.jsonl").write_bytes(b"[1,2,3]\n") + assert RawPayloadDatasetLoader.can_load(filename=d) is False + + def test_can_load_directory_with_non_jsonl_extension_returns_false(self, tmp_path): + """Directory with only .json (not .jsonl) files must not match.""" + d = tmp_path / "wrong_ext" + d.mkdir() + (d / "payload.json").write_bytes( + orjson.dumps({"messages": [{"role": "user", "content": "x"}]}) + ) + assert RawPayloadDatasetLoader.can_load(filename=d) is False + + def test_can_load_directory_with_first_jsonl_malformed_returns_false_currently( + self, tmp_path + ): + """Documents _dir_has_raw_payload_jsonl silent-swallow behavior. + + The helper catches bare Exception on orjson parse errors and continues + to the next file. BUT: the happy-path `return` inside the try-block + only fires when parsing succeeds. On a malformed first file, control + falls through to the next file. Here the malformed file is the only + file, so the final `return False` fires. + + Candidate for Wave 2: narrow the except to orjson.JSONDecodeError so + downstream fall-through is explicit rather than swallowing everything. + """ + d = tmp_path / "malformed_only" + d.mkdir() + (d / "bad.jsonl").write_bytes(b"{not valid json\n") + assert RawPayloadDatasetLoader.can_load(filename=d) is False + + +class TestLoadDatasetLineEdgeCases: + def test_load_dataset_directory_unsorted_multiple_files_all_emitted(self, tmp_path): + """Three single-turn sessions in a dir must all be emitted (sorted).""" + d = tmp_path / "sessions" + d.mkdir() + # Write in non-alphabetical creation order to exercise sorted(glob). + for name in ("zulu", "alpha", "mike"): + (d / f"{name}.jsonl").write_bytes( + orjson.dumps({"messages": [{"role": "user", "content": name}]}) + b"\n" + ) + loader = _make_loader(d) + data = loader.load_dataset() + assert len(data) == 3 + for payloads in data.values(): + assert len(payloads) == 1 + assert "messages" in payloads[0].payload + + def test_load_dataset_single_file_with_multiple_lines_emits_one_conversation_per_line( + self, tmp_path + ): + """In single-file mode, each line becomes its own single-turn session.""" + p = tmp_path / "three.jsonl" + lines = [ + orjson.dumps({"messages": [{"role": "user", "content": f"L{i}"}]}) + for i in range(3) + ] + p.write_bytes(b"\n".join(lines) + b"\n") + loader = _make_loader(p) + data = loader.load_dataset() + assert len(data) == 3 + for payloads in data.values(): + assert len(payloads) == 1 + + def test_load_dataset_file_with_blank_line_in_middle_skips_blank(self, tmp_path): + """Blank (whitespace-only) lines in the middle of a JSONL file are skipped.""" + p = tmp_path / "blanks.jsonl" + first = orjson.dumps({"messages": [{"role": "user", "content": "a"}]}) + second = orjson.dumps({"messages": [{"role": "user", "content": "b"}]}) + p.write_bytes(first + b"\n\n" + second + b"\n") + loader = _make_loader(p) + data = loader.load_dataset() + # Exactly two sessions: the blank line must not produce a phantom entry. + assert len(data) == 2 + contents = sorted( + payloads[0].payload["messages"][0]["content"] for payloads in data.values() + ) + assert contents == ["a", "b"] + + def test_load_dataset_file_with_trailing_newline_parses_clean(self, tmp_path): + """Trailing \\n must not create a phantom empty conversation.""" + p = tmp_path / "trail.jsonl" + p.write_bytes( + orjson.dumps({"messages": [{"role": "user", "content": "only"}]}) + b"\n" + ) + loader = _make_loader(p) + data = loader.load_dataset() + assert len(data) == 1 + + def test_convert_to_conversations_turn_carries_raw_payload_verbatim(self, tmp_path): + """Turn.raw_payload must preserve the entire source dict, including + caller-supplied extra keys beyond the standard chat schema. + """ + p = tmp_path / "extras.jsonl" + payload = { + "messages": [{"role": "user", "content": "hello"}], + "model": "test-model", + "custom_field": "verbatim", + "nested": {"temperature": 0.7, "seed": 42}, + } + p.write_bytes(orjson.dumps(payload) + b"\n") + loader = _make_loader(p) + data = loader.load_dataset() + conversations = loader.convert_to_conversations(data) + assert len(conversations) == 1 + assert len(conversations[0].turns) == 1 + turn = conversations[0].turns[0] + assert turn.raw_payload == payload + assert turn.raw_payload["custom_field"] == "verbatim" + assert turn.raw_payload["nested"]["seed"] == 42 + + +class TestWave2BugCandidates: + def test_can_load_directory_with_unreadable_jsonl_raises_post_fix(self, tmp_path): + """Unreadable .jsonl file in a directory. + + Post Wave-2 fix: _dir_has_raw_payload_jsonl narrows its except to + orjson.JSONDecodeError, letting PermissionError/OSError propagate so + misconfigured inputs fail loudly rather than silently miscategorize. + + Today: the broad `except Exception: continue` swallows the OSError + and can_load returns False. The xfail-strict marker flips as soon as + Wave 2 ships, alerting us to remove the xfail. + """ + import os + + d = tmp_path / "locked" + d.mkdir() + p = d / "blocked.jsonl" + p.write_bytes( + orjson.dumps({"messages": [{"role": "user", "content": "x"}]}) + b"\n" + ) + os.chmod(p, 0o000) + try: + with pytest.raises((PermissionError, OSError)): + RawPayloadDatasetLoader.can_load(filename=d) + finally: + os.chmod(p, 0o644)