From ddd62b7f05d2733ca912c7d671d756e251a3c46f Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:50:58 +0530 Subject: [PATCH 1/4] Python: fix(lab-gaia): accept local parquet data in _ensure_data _ensure_data only recognized the old metadata.jsonl layout, so a data dir holding only the parquet files the loader itself prefers was ignored and the runner demanded HF_TOKEN for a fresh download. Accept metadata*.parquet in the same existence check. Fixes #9337 Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../lab/gaia/agent_framework_lab_gaia/gaia.py | 4 ++- python/packages/lab/gaia/tests/test_gaia.py | 31 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py index 28c9b8a07eb..34919b52f6d 100644 --- a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py +++ b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py @@ -405,7 +405,9 @@ async def _default_evaluator(self, task: Task, prediction: Prediction) -> Evalua def _ensure_data(self) -> Path: """Ensure GAIA data is available locally.""" - if self.data_dir.exists() and any(self.data_dir.rglob("metadata.jsonl")): + if self.data_dir.exists() and ( + any(self.data_dir.rglob("metadata.jsonl")) or any(self.data_dir.rglob("metadata*.parquet")) + ): return self.data_dir # Download data if not available diff --git a/python/packages/lab/gaia/tests/test_gaia.py b/python/packages/lab/gaia/tests/test_gaia.py index 2af57b5f9a0..1deee4b99db 100644 --- a/python/packages/lab/gaia/tests/test_gaia.py +++ b/python/packages/lab/gaia/tests/test_gaia.py @@ -2,6 +2,8 @@ """Tests for GAIA benchmark implementation.""" +import pytest + from agent_framework_lab_gaia import gaia_scorer @@ -34,3 +36,32 @@ def test_none_handling(self): """Test handling of None values.""" assert gaia_scorer("None", "test") is False assert gaia_scorer("", "test") is False + + +class TestEnsureData: + """Test local data detection in GAIA._ensure_data.""" + + def test_parquet_only_data_dir_is_accepted(self, tmp_path): + """A data dir in the parquet layout _load_gaia_local supports must not force a download (#9337).""" + pytest.importorskip("pyarrow") + from agent_framework_lab_gaia.gaia import GAIA + + validation = tmp_path / "2023" / "validation" + validation.mkdir(parents=True) + rows = [ + { + "task_id": "v1", + "Question": "What is 2+2?", + "Final answer": "4", + "Level": 1, + "file_name": "", + "Annotator Metadata": "{}", + } + ] + import pyarrow as pa + import pyarrow.parquet as pq + + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00001.parquet") + + runner = GAIA(data_dir=str(tmp_path)) + assert runner._ensure_data() == tmp_path From 47891ef94bfe6024b34ef6c7df153682faad97e6 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Wed, 14 Oct 2026 13:39:45 +0530 Subject: [PATCH 2/4] Python: fix(lab-gaia): require every declared parquet shard before using local data A partial metadata-XXXXX-of-NNNNN family from an interrupted copy would have counted as ready and run the benchmark over a truncated task set. Group shards by directory and prefix and only treat the snapshot as complete when each group holds all its declared shards. Follow-up to the review on #9340. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../lab/gaia/agent_framework_lab_gaia/gaia.py | 26 ++++++++++++++++++- python/packages/lab/gaia/tests/test_gaia.py | 23 ++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py index 34919b52f6d..a8ac09def69 100644 --- a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py +++ b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py @@ -264,6 +264,30 @@ def _read_jsonl(path: Path) -> Iterable[dict[str, Any]]: yield record +_SHARDED_PARQUET_RE = re.compile(r"(metadata(?:-[^-]+)?)-(\d{5})-of-(\d{5})\.parquet") + + +def _has_complete_parquet_snapshot(data_dir: Path) -> bool: + """Whether every sharded metadata parquet group holds all of its declared shards. + + An interrupted copy can leave only part of a ``metadata-XXXXX-of-NNNNN.parquet`` + family in the directory. Treating any matching filename as ready would run the + benchmark over a silently truncated task set, so an incomplete family forces the + normal download path instead. Plain unsharded metadata parquet files count as + complete on their own. + """ + sharded: dict[tuple[Path, str], tuple[set[int], int]] = {} + for path in data_dir.rglob("metadata*.parquet"): + match = _SHARDED_PARQUET_RE.fullmatch(path.name) + if match is None: + return True + key = (path.parent, match.group(1)) + indexes, total = sharded.get(key, (set(), int(match.group(3)))) + indexes.add(int(match.group(2))) + sharded[key] = (indexes, max(total, int(match.group(3)))) + return bool(sharded) and all(len(indexes) == total for indexes, total in sharded.values()) + + def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max_n: int | None = None) -> list[Task]: """Load GAIA tasks from local repository directory.""" tasks: list[Task] = [] @@ -406,7 +430,7 @@ async def _default_evaluator(self, task: Task, prediction: Prediction) -> Evalua def _ensure_data(self) -> Path: """Ensure GAIA data is available locally.""" if self.data_dir.exists() and ( - any(self.data_dir.rglob("metadata.jsonl")) or any(self.data_dir.rglob("metadata*.parquet")) + any(self.data_dir.rglob("metadata.jsonl")) or _has_complete_parquet_snapshot(self.data_dir) ): return self.data_dir diff --git a/python/packages/lab/gaia/tests/test_gaia.py b/python/packages/lab/gaia/tests/test_gaia.py index 1deee4b99db..12c834adeb5 100644 --- a/python/packages/lab/gaia/tests/test_gaia.py +++ b/python/packages/lab/gaia/tests/test_gaia.py @@ -65,3 +65,26 @@ def test_parquet_only_data_dir_is_accepted(self, tmp_path): runner = GAIA(data_dir=str(tmp_path)) assert runner._ensure_data() == tmp_path + + def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path): + """A partial metadata-...-of-N family must not count as ready local data (#9337).""" + pytest.importorskip("pyarrow") + from agent_framework_lab_gaia.gaia import GAIA + + validation = tmp_path / "2023" / "validation" + validation.mkdir(parents=True) + rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}] + + import pyarrow as pa + import pyarrow.parquet as pq + + # interrupted copy: only shard 0 of 2 landed + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet") + + runner = GAIA(data_dir=str(tmp_path)) + with pytest.raises(RuntimeError, match="HF_TOKEN"): + runner._ensure_data() + + # completing the family makes the directory ready again + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet") + assert runner._ensure_data() == tmp_path From a16b33d153789de122b34828157ef37a9085e5ef Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:55:11 +0530 Subject: [PATCH 3/4] Python: require every sharded family even beside unsharded files An unsharded metadata parquet made the completeness check return early, so a directory mixing a plain file with an incomplete sharded family was treated as ready. Keep scanning after unsharded files and let any incomplete family disqualify the directory. Also clear HF_TOKEN in the incomplete snapshot test so it cannot attempt a real download. Follow-up to the review on #9340. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../lab/gaia/agent_framework_lab_gaia/gaia.py | 10 ++++++-- python/packages/lab/gaia/tests/test_gaia.py | 23 ++++++++++++++++++- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py index a8ac09def69..3a02c9de4c3 100644 --- a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py +++ b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py @@ -277,15 +277,21 @@ def _has_complete_parquet_snapshot(data_dir: Path) -> bool: complete on their own. """ sharded: dict[tuple[Path, str], tuple[set[int], int]] = {} + unsharded_seen = False for path in data_dir.rglob("metadata*.parquet"): match = _SHARDED_PARQUET_RE.fullmatch(path.name) if match is None: - return True + unsharded_seen = True + continue key = (path.parent, match.group(1)) indexes, total = sharded.get(key, (set(), int(match.group(3)))) indexes.add(int(match.group(2))) sharded[key] = (indexes, max(total, int(match.group(3)))) - return bool(sharded) and all(len(indexes) == total for indexes, total in sharded.values()) + if not sharded: + return unsharded_seen + # The loader reads every matching parquet file, so a single incomplete sharded + # family disqualifies the directory even when unsharded files also exist. + return all(len(indexes) == total for indexes, total in sharded.values()) def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max_n: int | None = None) -> list[Task]: diff --git a/python/packages/lab/gaia/tests/test_gaia.py b/python/packages/lab/gaia/tests/test_gaia.py index 12c834adeb5..6fbfacc985e 100644 --- a/python/packages/lab/gaia/tests/test_gaia.py +++ b/python/packages/lab/gaia/tests/test_gaia.py @@ -66,9 +66,10 @@ def test_parquet_only_data_dir_is_accepted(self, tmp_path): runner = GAIA(data_dir=str(tmp_path)) assert runner._ensure_data() == tmp_path - def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path): + def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path, monkeypatch): """A partial metadata-...-of-N family must not count as ready local data (#9337).""" pytest.importorskip("pyarrow") + monkeypatch.delenv("HF_TOKEN", raising=False) from agent_framework_lab_gaia.gaia import GAIA validation = tmp_path / "2023" / "validation" @@ -88,3 +89,23 @@ def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path): # completing the family makes the directory ready again pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet") assert runner._ensure_data() == tmp_path + + def test_unsharded_file_does_not_rescue_an_incomplete_family(self, tmp_path, monkeypatch): + """An unsharded file must not mask a partial sharded family (#9337 review).""" + pytest.importorskip("pyarrow") + monkeypatch.delenv("HF_TOKEN", raising=False) + from agent_framework_lab_gaia.gaia import GAIA + + validation = tmp_path / "2023" / "validation" + validation.mkdir(parents=True) + rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}] + + import pyarrow as pa + import pyarrow.parquet as pq + + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet") + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-extra.parquet") + + runner = GAIA(data_dir=str(tmp_path)) + with pytest.raises(RuntimeError, match="HF_TOKEN"): + runner._ensure_data() From 60b4e0aa97524802e013af08c2df501d66bf2907 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:54:09 +0000 Subject: [PATCH 4/4] Python: require a complete parquet snapshot whenever metadata parquet is present Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../lab/gaia/agent_framework_lab_gaia/gaia.py | 20 +++++++-- python/packages/lab/gaia/tests/test_gaia.py | 43 ++++++++++++++++++- 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py index 3a02c9de4c3..ac69b7d9feb 100644 --- a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py +++ b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py @@ -294,6 +294,22 @@ def _has_complete_parquet_snapshot(data_dir: Path) -> bool: return all(len(indexes) == total for indexes, total in sharded.values()) +def _has_usable_local_data(data_dir: Path) -> bool: + """Whether the directory holds data that is safe to load without a download. + + A JSONL-only layout is usable when ``metadata.jsonl`` is present, and a parquet + layout is usable when every sharded family is complete. A mixed layout must + not pass merely because JSONL exists: ``_load_gaia_local`` reads parquet files + first and only falls back to JSONL when no parquet file is present, so an + incomplete parquet family would silently truncate the task set even though a + complete ``metadata.jsonl`` sits beside it. When any metadata parquet file is + present, the parquet snapshot must therefore be complete. + """ + if any(data_dir.rglob("metadata*.parquet")): + return _has_complete_parquet_snapshot(data_dir) + return any(data_dir.rglob("metadata.jsonl")) + + def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max_n: int | None = None) -> list[Task]: """Load GAIA tasks from local repository directory.""" tasks: list[Task] = [] @@ -435,9 +451,7 @@ async def _default_evaluator(self, task: Task, prediction: Prediction) -> Evalua def _ensure_data(self) -> Path: """Ensure GAIA data is available locally.""" - if self.data_dir.exists() and ( - any(self.data_dir.rglob("metadata.jsonl")) or _has_complete_parquet_snapshot(self.data_dir) - ): + if self.data_dir.exists() and _has_usable_local_data(self.data_dir): return self.data_dir # Download data if not available diff --git a/python/packages/lab/gaia/tests/test_gaia.py b/python/packages/lab/gaia/tests/test_gaia.py index 6fbfacc985e..a3b09cc560a 100644 --- a/python/packages/lab/gaia/tests/test_gaia.py +++ b/python/packages/lab/gaia/tests/test_gaia.py @@ -3,7 +3,6 @@ """Tests for GAIA benchmark implementation.""" import pytest - from agent_framework_lab_gaia import gaia_scorer @@ -90,6 +89,48 @@ def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path, mon pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet") assert runner._ensure_data() == tmp_path + def test_jsonl_does_not_rescue_an_incomplete_parquet_family(self, tmp_path, monkeypatch): + """A mixed layout must not pass on the JSONL shortcut (#9340 review). + + _load_gaia_local reads parquet first and only falls back to JSONL when + no parquet file is present, so metadata.jsonl beside an incomplete + sharded family would silently truncate the task set. + """ + pytest.importorskip("pyarrow") + monkeypatch.delenv("HF_TOKEN", raising=False) + from agent_framework_lab_gaia.gaia import GAIA + + validation = tmp_path / "2023" / "validation" + validation.mkdir(parents=True) + rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}] + + import pyarrow as pa + import pyarrow.parquet as pq + + (validation / "metadata.jsonl").write_text("{}\n", encoding="utf-8") + # interrupted copy: only shard 0 of 2 landed beside the JSONL file + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet") + + runner = GAIA(data_dir=str(tmp_path)) + with pytest.raises(RuntimeError, match="HF_TOKEN"): + runner._ensure_data() + + # completing the family makes the mixed directory ready again + pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet") + assert runner._ensure_data() == tmp_path + + def test_jsonl_only_layout_is_still_accepted(self, tmp_path, monkeypatch): + """A JSONL-only directory keeps working without any parquet files (#9340 review).""" + monkeypatch.delenv("HF_TOKEN", raising=False) + from agent_framework_lab_gaia.gaia import GAIA + + validation = tmp_path / "2023" / "validation" + validation.mkdir(parents=True) + (validation / "metadata.jsonl").write_text("{}\n", encoding="utf-8") + + runner = GAIA(data_dir=str(tmp_path)) + assert runner._ensure_data() == tmp_path + def test_unsharded_file_does_not_rescue_an_incomplete_family(self, tmp_path, monkeypatch): """An unsharded file must not mask a partial sharded family (#9337 review).""" pytest.importorskip("pyarrow")