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/3] 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/3] 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/3] 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()