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..3a02c9de4c3 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,36 @@ 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]] = {} + unsharded_seen = False + for path in data_dir.rglob("metadata*.parquet"): + match = _SHARDED_PARQUET_RE.fullmatch(path.name) + if match is None: + 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)))) + 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]: """Load GAIA tasks from local repository directory.""" tasks: list[Task] = [] @@ -405,7 +435,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 _has_complete_parquet_snapshot(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 2af57b5f9a0..6fbfacc985e 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,76 @@ 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 + + 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" + 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 + + 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()