diff --git a/examples/10_Agentic_Inference/README.md b/examples/10_Agentic_Inference/README.md index 7f0f84b6c..883d24974 100644 --- a/examples/10_Agentic_Inference/README.md +++ b/examples/10_Agentic_Inference/README.md @@ -15,9 +15,9 @@ Use flat JSONL with one row per message. Rows for each `conversation_id` must be Required fields are `conversation_id`, `turn`, and `role`. User rows normally include `content`; agentic rows can also include `system`, `tools`, `tool_calls`, `tool_results`, `reasoning_content`, and `delay_seconds`. -The official MLPerf dataset is available from [MLCommons storage](https://endpoints.mlcommons-storage.org/index.html#mlperf-agentic-inference). The dataset SHA-256 is `1beb24c882122df96571cf11b390acbea388944038bc55c78b891475459014ae`. Submitters must use this dataset unchanged for official submissions. +The official MLPerf dataset is available from [MLCommons storage](https://endpoints.mlcommons-storage.org/index.html#mlperf-agentic-inference). The predefined `agentic_inference_conversations` dataset downloads and caches the official JSONL automatically under `dataset_cache/agentic_inference_conversations/` on first use. It verifies the dataset SHA-256, `1beb24c882122df96571cf11b390acbea388944038bc55c78b891475459014ae`, before loading it. Submitters must use this dataset unchanged for official submissions. -Place the dataset under `examples/10_Agentic_Inference/datasets/` or point the YAML at another accessible JSONL path. +The runnable configs use this automatic download and do not need a dataset path. To use another accessible JSONL file, set `path` with a non-predefined dataset name while retaining the `agentic_inference` configuration. ## Supported Models @@ -120,7 +120,7 @@ uv run --project src/inference_endpoint/evaluation/swebench_service \ ## Run The Client -Update the first `datasets` entry (`name` and `path`), `model_params.name`, and `endpoint_config.endpoints` as needed. Then select the matching model config and run it from the repo root: +Update `model_params.name` and `endpoint_config.endpoints` as needed. The first `datasets` entry uses the official dataset download by default. Then select the matching model config and run it from the repo root: ```bash CONFIG=examples/10_Agentic_Inference/qwen_agentic_benchmark.yaml diff --git a/examples/10_Agentic_Inference/kimi_agentic_benchmark.yaml b/examples/10_Agentic_Inference/kimi_agentic_benchmark.yaml index 67d7d69d3..ea932caec 100644 --- a/examples/10_Agentic_Inference/kimi_agentic_benchmark.yaml +++ b/examples/10_Agentic_Inference/kimi_agentic_benchmark.yaml @@ -11,9 +11,8 @@ model_params: streaming: "on" datasets: - - name: agentic_combined + - name: agentic_inference_conversations type: performance - path: /path/to/agentic_combined.jsonl accuracy_config: eval_method: agentic_inference_inline # required benchmark default. num_repeats: 1 diff --git a/examples/10_Agentic_Inference/qwen_agentic_benchmark.yaml b/examples/10_Agentic_Inference/qwen_agentic_benchmark.yaml index 256a9217b..e7c64cf26 100644 --- a/examples/10_Agentic_Inference/qwen_agentic_benchmark.yaml +++ b/examples/10_Agentic_Inference/qwen_agentic_benchmark.yaml @@ -14,9 +14,8 @@ model_params: preserve_thinking: true datasets: - - name: agentic_coding + - name: agentic_inference_conversations type: performance - path: ./agentic_combined.jsonl accuracy_config: eval_method: agentic_inference_inline # required benchmark default. agentic_inference: diff --git a/src/inference_endpoint/dataset_manager/agentic_inference_dataset.py b/src/inference_endpoint/dataset_manager/agentic_inference_dataset.py index aa054d15d..ac6cef791 100644 --- a/src/inference_endpoint/dataset_manager/agentic_inference_dataset.py +++ b/src/inference_endpoint/dataset_manager/agentic_inference_dataset.py @@ -17,6 +17,7 @@ import logging from dataclasses import dataclass, field, replace +from pathlib import Path from typing import Any import pandas as pd @@ -24,6 +25,7 @@ from ..config.schema import APIType, ModelParams from ..exceptions import InputValidationError from .dataset import Dataset +from .download import download_r2_artifact, verify_sha256 from .transforms import ( AddStaticColumns, apply_transforms, @@ -241,6 +243,48 @@ class AgenticInferenceDataset(Dataset, dataset_id="agentic_inference_conversatio """ COLUMN_NAMES = ["conversation_id", "turn", "role", "content"] + CACHE_FILENAME = "mlperf_agentic_inference_dataset.jsonl" + R2_DATASET_URI = ( + "https://endpoints.mlcommons-storage.org/metadata/" + "mlperf_agentic_inference_dataset.uri" + ) + DATASET_SHA256 = "1beb24c882122df96571cf11b390acbea388944038bc55c78b891475459014ae" + + @classmethod + def _download_dataset(cls, cache_dir: Path, jsonl_path: Path) -> pd.DataFrame: + """Download, verify, and load the official MLPerf Agentic dataset.""" + download_r2_artifact( + uri=cls.R2_DATASET_URI, + destination_dir=cache_dir, + artifact_name=jsonl_path.name, + expected_sha256=cls.DATASET_SHA256, + ) + logger.info("Loaded Agentic Inference dataset from %s", jsonl_path) + return pd.read_json(jsonl_path, lines=True) + + @classmethod + def generate(cls, datasets_dir: Path, force: bool = False) -> pd.DataFrame: + """Load the official Agentic Inference dataset from the local cache.""" + cache_dir = datasets_dir / cls.DATASET_ID + jsonl_path = cache_dir / cls.CACHE_FILENAME + if jsonl_path.exists() and not force: + try: + verify_sha256(jsonl_path, cls.DATASET_SHA256) + except ValueError as exc: + logger.warning( + "Cached Agentic Inference dataset failed verification (%s); " + "re-downloading.", + exc, + ) + jsonl_path.unlink() + else: + logger.info( + "Loading cached Agentic Inference dataset from %s", jsonl_path + ) + return pd.read_json(jsonl_path, lines=True) + + jsonl_path.unlink(missing_ok=True) + return cls._download_dataset(cache_dir, jsonl_path) def __init__(self, dataframe: pd.DataFrame, **kwargs): """Initialize agentic inference dataset. diff --git a/src/inference_endpoint/dataset_manager/download.py b/src/inference_endpoint/dataset_manager/download.py new file mode 100644 index 000000000..9745dc5ec --- /dev/null +++ b/src/inference_endpoint/dataset_manager/download.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small, guarded helpers for downloading MLCommons R2 artifacts.""" + +from __future__ import annotations + +import hashlib +import os +import subprocess +import tempfile +from pathlib import Path +from urllib.parse import urlparse + +import requests + +_DOWNLOADER_URL_TEMPLATE = ( + "https://raw.githubusercontent.com/mlcommons/r2-downloader/" + "{commit}/mlc-r2-downloader.sh" +) +_DOWNLOADER_HOST = "raw.githubusercontent.com" +_REQUEST_TIMEOUT_S = 30 +_DEFAULT_CHUNK_SIZE = 1 << 20 +DEFAULT_R2_DOWNLOADER_COMMIT = "27da4421877f2831eeb615b43ee5098c4b70be7e" +DEFAULT_R2_ALLOWED_HOST = "mlcommons-storage.org" + + +def _validate_https_host(uri: str, allowed_host: str) -> None: + parsed = urlparse(uri) + host = (parsed.hostname or "").lower() + allowed = allowed_host.lower().rstrip(".") + if ( + parsed.scheme != "https" + or not host + or parsed.username is not None + or parsed.password is not None + or not (host == allowed or host.endswith(f".{allowed}")) + ): + raise ValueError( + f"Refusing to download dataset from untrusted URI {uri!r}: " + f"expected https on {allowed_host}" + ) + + +def verify_sha256( + path: Path, expected: str, chunk_size: int = _DEFAULT_CHUNK_SIZE +) -> None: + """Verify a file's SHA-256 digest without loading it all into memory.""" + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(chunk_size), b""): + digest.update(chunk) + actual = digest.hexdigest() + if actual != expected.lower(): + raise ValueError( + f"SHA-256 mismatch for {path.name}: expected {expected}, got {actual}" + ) + + +def _artifact_target(destination_dir: Path, artifact_name: str) -> Path: + name = Path(artifact_name) + if not artifact_name or name.name != artifact_name or artifact_name in {".", ".."}: + raise ValueError(f"artifact_name must be a plain filename: {artifact_name!r}") + return destination_dir / artifact_name + + +def download_r2_artifact( + uri: str, + destination_dir: Path, + artifact_name: str, + downloader_commit: str = DEFAULT_R2_DOWNLOADER_COMMIT, + expected_sha256: str | None = None, + timeout_s: float = 1800, + allowed_host: str = DEFAULT_R2_ALLOWED_HOST, +) -> Path: + """Download one exact artifact through a pinned MLCommons R2 script.""" + _validate_https_host(uri, allowed_host) + if timeout_s <= 0: + raise ValueError("timeout_s must be positive") + + destination_dir = Path(destination_dir) + destination_dir.mkdir(parents=True, exist_ok=True) + target = _artifact_target(destination_dir, artifact_name) + target.unlink(missing_ok=True) + downloader_url = _DOWNLOADER_URL_TEMPLATE.format(commit=downloader_commit) + _validate_https_host(downloader_url, _DOWNLOADER_HOST) + + script_fd, script_name = tempfile.mkstemp( + prefix=".mlc-r2-downloader-", suffix=".sh", dir=destination_dir + ) + os.close(script_fd) + script_path = Path(script_name) + try: + response = requests.get(downloader_url, timeout=_REQUEST_TIMEOUT_S) + response.raise_for_status() + script_path.write_bytes(response.content) + script_path.chmod(0o755) + + try: + result = subprocess.run( + [ + "bash", + str(script_path.resolve()), + "-d", + str(destination_dir.resolve()), + uri, + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + check=False, + timeout=timeout_s, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError( + f"R2 downloader timed out after {timeout_s} seconds" + ) from exc + if result.returncode != 0: + raise RuntimeError( + f"R2 downloader failed with code {result.returncode}: " + f"{result.stderr}" + ) + + candidates = [ + path + for path in destination_dir.rglob("*") + if path.is_file() and path.name == artifact_name and path != script_path + ] + if len(candidates) != 1: + if not candidates: + raise FileNotFoundError( + f"R2 download completed but exact artifact {artifact_name!r} " + f"is missing under {destination_dir}" + ) + raise RuntimeError( + f"R2 download produced ambiguous artifact {artifact_name!r}: " + f"{len(candidates)} files found under {destination_dir}" + ) + + found = candidates[0] + if found != target: + target.unlink(missing_ok=True) + found.replace(target) + if expected_sha256 is not None: + verify_sha256(target, expected_sha256) + return target + finally: + script_path.unlink(missing_ok=True) diff --git a/src/inference_endpoint/dataset_manager/predefined/bfcl_v4/__init__.py b/src/inference_endpoint/dataset_manager/predefined/bfcl_v4/__init__.py index f91c190ea..d204f03a6 100644 --- a/src/inference_endpoint/dataset_manager/predefined/bfcl_v4/__init__.py +++ b/src/inference_endpoint/dataset_manager/predefined/bfcl_v4/__init__.py @@ -19,19 +19,16 @@ Reference: https://gorilla.cs.berkeley.edu/leaderboard.html """ -import hashlib import json import os -import subprocess from logging import getLogger from pathlib import Path from typing import Any -from urllib.parse import urlparse import pandas as pd -import requests from ...dataset import Dataset +from ...download import download_r2_artifact, verify_sha256 from . import presets logger = getLogger(__name__) @@ -154,7 +151,6 @@ class BFCLv4( PRESETS = presets - # --- MLCommons R2 hosting (mirrors the OpenOrca dataset pattern) --- # SHA-256 of the hosted full single-turn parquet, pinning the exact bytes so # every submitter scores identical data. Computed from the artifact handed to # MLCommons for upload. Verified only on the R2 download path; a parquet built @@ -168,13 +164,6 @@ class BFCLv4( # staging/testing via the BFCL_V4_DATASET_URI env var. R2_DATASET_URI: str | None = None R2_DATASET_URI_ENV = "BFCL_V4_DATASET_URI" - # Pinned mlcommons/r2-downloader commit (same convention as open_orca). - R2_DOWNLOADER_COMMIT = "27da4421877f2831eeb615b43ee5098c4b70be7e" - # The resolved dataset URI is handed to a downloaded shell script, so the - # host is restricted to MLCommons R2 storage. This stops a mis-set - # BFCL_V4_DATASET_URI env var from pointing the downloader at an - # attacker-controlled location. - R2_ALLOWED_HOST = "mlcommons-storage.org" @classmethod def generate( @@ -283,7 +272,7 @@ def generate( # verified R2 download, so a stale or locally-built parquet can't # silently win and make the run score different data. try: - cls._verify_sha256(dst_path, cls.SINGLE_TURN_SHA256) + verify_sha256(dst_path, cls.SINGLE_TURN_SHA256) except ValueError as exc: logger.warning( "Cached %s failed SHA-256 verification (%s); re-downloading " @@ -400,22 +389,10 @@ def _deserialize_complex_columns(df: pd.DataFrame) -> pd.DataFrame: df[col] = df[col].apply(json.loads) return df - @classmethod - def _verify_sha256(cls, path: Path, expected: str) -> None: - """Raise ValueError if path's SHA-256 digest does not match expected.""" - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != expected: - raise ValueError( - f"SHA-256 mismatch for {path.name}: expected {expected}, got {digest}" - ) - @classmethod def _download_full_parquet_from_r2(cls, dst_path: Path) -> bool: """Download the full single-turn parquet from MLCommons R2 storage. - Mirrors the OpenOrca pattern: fetch the pinned ``mlc-r2-downloader.sh``, - run it against the dataset ``.uri``, then SHA-256 verify the result. - Returns True when a verified parquet is in place at ``dst_path``; returns False when no R2 URI is configured, in which case the caller builds the parquet from ``bfcl-eval``. Raises if a URI is configured but the download @@ -426,61 +403,17 @@ def _download_full_parquet_from_r2(cls, dst_path: Path) -> bool: if not uri: return False - # The uri is passed to a downloaded shell script; restrict it to https - # on MLCommons R2 storage before invoking anything. - parsed = urlparse(uri) - host = parsed.hostname or "" - if parsed.scheme != "https" or not ( - host == cls.R2_ALLOWED_HOST or host.endswith("." + cls.R2_ALLOWED_HOST) - ): - raise ValueError( - f"Refusing to download dataset from untrusted URI '{uri}': " - f"expected https on {cls.R2_ALLOWED_HOST}" - ) - download_dir = dst_path.parent download_dir.mkdir(parents=True, exist_ok=True) # Start from a clean target so a stale parquet can't shadow the fresh - # download (the relocate below only fires when dst_path is absent). + # download. dst_path.unlink(missing_ok=True) - downloader_url = ( - "https://raw.githubusercontent.com/mlcommons/r2-downloader/" - f"{cls.R2_DOWNLOADER_COMMIT}/mlc-r2-downloader.sh" + download_r2_artifact( + uri=uri, + destination_dir=download_dir, + artifact_name=dst_path.name, + expected_sha256=cls.SINGLE_TURN_SHA256, ) - script_path = download_dir / "mlc-r2-downloader.sh" - resp = requests.get(downloader_url, timeout=30) - resp.raise_for_status() - script_path.write_bytes(resp.content) - script_path.chmod(0o755) - - try: - result = subprocess.run( - ["bash", str(script_path.resolve()), "-d", str(download_dir), uri], - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, - text=True, - check=False, - ) - if result.returncode != 0: - raise RuntimeError( - f"R2 downloader failed with code {result.returncode}: " - f"{result.stderr}" - ) - finally: - script_path.unlink(missing_ok=True) - - # The downloader may nest the file under a subdir per the .uri manifest; - # relocate it to the expected cache path if so. - if not dst_path.exists(): - found = next(download_dir.rglob(dst_path.name), None) - if found is not None and found != dst_path: - found.replace(dst_path) - if not dst_path.exists(): - raise FileNotFoundError( - f"R2 download completed but {dst_path.name} is missing under " - f"{download_dir}; ensure the hosted artifact is named {dst_path.name}" - ) - cls._verify_sha256(dst_path, cls.SINGLE_TURN_SHA256) return True @classmethod diff --git a/src/inference_endpoint/dataset_manager/predefined/open_orca/__init__.py b/src/inference_endpoint/dataset_manager/predefined/open_orca/__init__.py index ed092c986..deb04e6eb 100644 --- a/src/inference_endpoint/dataset_manager/predefined/open_orca/__init__.py +++ b/src/inference_endpoint/dataset_manager/predefined/open_orca/__init__.py @@ -14,16 +14,14 @@ # limitations under the License. import gzip -import hashlib import shutil -import subprocess from logging import getLogger from pathlib import Path import pandas as pd -import requests from ...dataset import Dataset +from ...download import download_r2_artifact, verify_sha256 from . import presets logger = getLogger(__name__) @@ -39,16 +37,10 @@ class OpenOrca( SOURCE_FILENAME = "open_orca_gpt4_tokenized_llama.sampled_24576.pkl" CACHE_FILENAME = "open_orca_gpt4_tokenized_llama.sampled_24576.jsonl" SOURCE_SHA256 = "b64e66e54b6267f79eb4f9ccec52d466bab3ac94747ed258c3b0f337ed166fab" - - @classmethod - def _verify_sha256(cls, path: Path) -> None: - """Raise ValueError if the file's SHA-256 digest does not match SOURCE_SHA256.""" - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != cls.SOURCE_SHA256: - raise ValueError( - f"SHA-256 mismatch for {path.name}: " - f"expected {cls.SOURCE_SHA256}, got {digest}" - ) + DATASET_URI = ( + "https://inference.mlcommons-storage.org/metadata/" + "llama-2-70b-open-orca-dataset.uri" + ) @classmethod def _extract_gz_files(cls, download_dir: Path, gzip_dir: Path) -> None: @@ -63,7 +55,7 @@ def _extract_gz_files(cls, download_dir: Path, gzip_dir: Path) -> None: @classmethod def _convert_pickle_cache(cls, pickle_path: Path, jsonl_path: Path) -> pd.DataFrame: """Convert the upstream pickle artifact into the local JSONL cache.""" - cls._verify_sha256(pickle_path) + verify_sha256(pickle_path, cls.SOURCE_SHA256) dataframe = pd.read_pickle(pickle_path) tmp_path = jsonl_path.with_suffix(".jsonl.tmp") try: @@ -124,55 +116,21 @@ def generate( ) return cls._convert_pickle_cache(pickle_path, jsonl_path) - # Dataset URL from README - dataset_url = "https://inference.mlcommons-storage.org/metadata/llama-2-70b-open-orca-dataset.uri" - - # Download the r2-downloader script into a temp file in the target dir - COMMIT_HASH = "27da4421877f2831eeb615b43ee5098c4b70be7e" - downloader_url = f"https://raw.githubusercontent.com/mlcommons/r2-downloader/{COMMIT_HASH}/mlc-r2-downloader.sh" - download_dir = cache_dir - script_path = cache_dir / "mlc-r2-downloader.sh" - r = requests.get(downloader_url, timeout=30) - r.raise_for_status() - script_path.write_bytes(r.content) - script_path.chmod(0o755) - - # Run the script with the dataset URL. + downloaded_gzip = download_r2_artifact( + uri=cls.DATASET_URI, + destination_dir=cache_dir, + artifact_name=f"{cls.SOURCE_FILENAME}.gz", + ) try: - # Use absolute path for the script to avoid path doubling when cwd is set - script_abs = str(script_path.resolve()) - result = subprocess.run( - ["bash", script_abs, "-d", str(download_dir), dataset_url], - stdout=subprocess.DEVNULL, # Suppress normal output - stderr=subprocess.PIPE, # Capture errors - text=True, - check=False, - ) - if result.returncode != 0: - raise RuntimeError( - f"R2 downloader failed with code {result.returncode}: {result.stderr}" + cls._extract_gz_files(cache_dir, downloaded_gzip.parent) + if not pickle_path.exists(): + raise FileNotFoundError( + f"OpenOrca was downloaded, but {pickle_path} does not exist" ) - except subprocess.CalledProcessError as e: - raise RuntimeError(f"R2 downloader failed: {e}") from e - - # Script will generate a new 'open_orca' subdirectory with gzip'd pickle files - gzip_dir = download_dir - if (gzip_dir / "open_orca").exists(): - gzip_dir = gzip_dir / "open_orca" - - cls._extract_gz_files(download_dir, gzip_dir) - - if not pickle_path.exists(): - raise FileNotFoundError( - f"OpenOrca was downloaded, but {pickle_path} does not exist" - ) - - try: return cls._convert_pickle_cache(pickle_path, jsonl_path) finally: - for gz_path in gzip_dir.glob("*.pkl.gz"): + for gz_path in downloaded_gzip.parent.glob("*.pkl.gz"): gz_path.unlink() - script_path.unlink(missing_ok=True) __all__ = ["OpenOrca"] diff --git a/tests/unit/dataset_manager/test_agentic_inference_dataset.py b/tests/unit/dataset_manager/test_agentic_inference_dataset.py index 9d0df6a01..497da3492 100644 --- a/tests/unit/dataset_manager/test_agentic_inference_dataset.py +++ b/tests/unit/dataset_manager/test_agentic_inference_dataset.py @@ -13,11 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import json import tempfile from collections.abc import Generator from pathlib import Path +from unittest.mock import MagicMock +import inference_endpoint.dataset_manager.agentic_inference_dataset as agentic_module import pandas as pd import pytest from inference_endpoint.dataset_manager.agentic_inference_dataset import ( @@ -235,6 +238,82 @@ def test_agentic_inference_dataset_validation_missing_fields(missing_fields_json ) +@pytest.mark.unit +def test_agentic_inference_dataset_downloads_dataset(tmp_path, monkeypatch): + payload = b'{"conversation_id":"c1","turn":1,"role":"user","content":"hello"}\n' + expected_sha256 = hashlib.sha256(payload).hexdigest() + monkeypatch.setattr(AgenticInferenceDataset, "DATASET_SHA256", expected_sha256) + + def fake_download_r2_artifact(**kwargs): + downloaded_path = kwargs["destination_dir"] / kwargs["artifact_name"] + downloaded_path.parent.mkdir(parents=True, exist_ok=True) + downloaded_path.write_bytes(payload) + return downloaded_path + + download = MagicMock(side_effect=fake_download_r2_artifact) + monkeypatch.setattr(agentic_module, "download_r2_artifact", download) + + dataframe = AgenticInferenceDataset.generate(tmp_path) + cache_path = ( + tmp_path + / AgenticInferenceDataset.DATASET_ID + / AgenticInferenceDataset.CACHE_FILENAME + ) + assert cache_path.read_bytes() == payload + assert dataframe.iloc[0]["conversation_id"] == "c1" + download.assert_called_once_with( + uri=AgenticInferenceDataset.R2_DATASET_URI, + destination_dir=tmp_path / AgenticInferenceDataset.DATASET_ID, + artifact_name=AgenticInferenceDataset.CACHE_FILENAME, + expected_sha256=expected_sha256, + ) + + +@pytest.mark.unit +def test_agentic_inference_dataset_recovers_from_corrupt_cache(tmp_path, monkeypatch): + payload = b'{"conversation_id":"c1","turn":1,"role":"user","content":"hello"}\n' + monkeypatch.setattr( + AgenticInferenceDataset, "DATASET_SHA256", hashlib.sha256(payload).hexdigest() + ) + cache_dir = tmp_path / AgenticInferenceDataset.DATASET_ID + cache_dir.mkdir() + cache_path = cache_dir / AgenticInferenceDataset.CACHE_FILENAME + cache_path.write_bytes(b"corrupt") + + def fake_download_r2_artifact(**kwargs): + downloaded_path = kwargs["destination_dir"] / kwargs["artifact_name"] + downloaded_path.write_bytes(payload) + return downloaded_path + + download = MagicMock(side_effect=fake_download_r2_artifact) + monkeypatch.setattr(agentic_module, "download_r2_artifact", download) + + dataframe = AgenticInferenceDataset.generate(tmp_path) + + assert cache_path.read_bytes() == payload + assert dataframe.iloc[0]["conversation_id"] == "c1" + download.assert_called_once() + + +@pytest.mark.unit +def test_agentic_inference_dataset_uses_valid_cache(tmp_path, monkeypatch): + payload = b'{"conversation_id":"c1","turn":1,"role":"user","content":"hello"}\n' + monkeypatch.setattr( + AgenticInferenceDataset, "DATASET_SHA256", hashlib.sha256(payload).hexdigest() + ) + cache_dir = tmp_path / AgenticInferenceDataset.DATASET_ID + cache_dir.mkdir() + (cache_dir / AgenticInferenceDataset.CACHE_FILENAME).write_bytes(payload) + + download = MagicMock() + monkeypatch.setattr(agentic_module, "download_r2_artifact", download) + + dataframe = AgenticInferenceDataset.generate(tmp_path) + + download.assert_not_called() + assert dataframe.iloc[0]["conversation_id"] == "c1" + + @pytest.mark.unit def test_agentic_inference_dataset_multiple_conversations(): """Test dataset with multiple conversations of varying lengths.""" diff --git a/tests/unit/dataset_manager/test_bfcl_v4_dataset.py b/tests/unit/dataset_manager/test_bfcl_v4_dataset.py index e9d9a36e5..d1d0b9b5f 100644 --- a/tests/unit/dataset_manager/test_bfcl_v4_dataset.py +++ b/tests/unit/dataset_manager/test_bfcl_v4_dataset.py @@ -13,27 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the BFCL v4 dataset loader's MLCommons R2 download path. - -These exercise the OpenAI-equivalent of OpenOrca's hosting flow without any -network: requests.get (downloader script fetch) and subprocess.run (the -mlc-r2-downloader invocation) are mocked, so only the loader's own logic -(URI resolution, relocate-from-subdir, SHA-256 verification, fallback) runs. -""" - -import hashlib import json from pathlib import Path from unittest.mock import MagicMock, patch +import inference_endpoint.dataset_manager.predefined.bfcl_v4 as bfcl_module import pandas as pd import pytest from inference_endpoint.dataset_manager.predefined.bfcl_v4 import BFCLv4 pytestmark = pytest.mark.unit -_MODULE = "inference_endpoint.dataset_manager.predefined.bfcl_v4" - def _dst(tmp_path: Path) -> Path: """Cache path the loader writes the full single-turn parquet to.""" @@ -71,195 +61,59 @@ def _write_cached_parquet(tmp_path: Path, subset_counts: dict[str, int]) -> Path return dst -def _mock_requests_get() -> MagicMock: - """A requests.get returning a trivial downloader script body.""" - resp = MagicMock() - resp.content = b"#!/usr/bin/env bash\nexit 0\n" - resp.raise_for_status.return_value = None - return resp - - -def _fake_run_writing(target: Path, payload: bytes): - """Build a subprocess.run side effect that drops `payload` at `target`.""" - - def _run(cmd, **kwargs): - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(payload) - return MagicMock(returncode=0, stderr="") - - return _run - - class TestR2DownloadFallback: def test_returns_false_when_no_uri(self, tmp_path, monkeypatch): """No constant and no env var -> caller falls back to bfcl-eval.""" monkeypatch.setattr(BFCLv4, "R2_DATASET_URI", None) monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - # Must short-circuit before touching the network. - with patch(f"{_MODULE}.requests.get") as get: + with patch.object(bfcl_module, "download_r2_artifact") as download: assert BFCLv4._download_full_parquet_from_r2(_dst(tmp_path)) is False - get.assert_not_called() + download.assert_not_called() class TestR2DownloadSuccess: def test_success_via_constant_uri(self, tmp_path, monkeypatch): dst = _dst(tmp_path) payload = b"PARQUET-BYTES" - monkeypatch.setattr( - BFCLv4, "R2_DATASET_URI", "https://inference.mlcommons-storage.org/x.uri" - ) + uri = "https://inference.mlcommons-storage.org/x.uri" + monkeypatch.setattr(BFCLv4, "R2_DATASET_URI", uri) monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - monkeypatch.setattr( - BFCLv4, "SINGLE_TURN_SHA256", hashlib.sha256(payload).hexdigest() - ) - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch( - f"{_MODULE}.subprocess.run", - side_effect=_fake_run_writing(dst, payload), - ), - ): + def fake_download_r2_artifact(**kwargs): + downloaded_path = kwargs["destination_dir"] / kwargs["artifact_name"] + downloaded_path.parent.mkdir(parents=True, exist_ok=True) + downloaded_path.write_bytes(payload) + return downloaded_path + + download = MagicMock(side_effect=fake_download_r2_artifact) + with patch.object(bfcl_module, "download_r2_artifact", download): assert BFCLv4._download_full_parquet_from_r2(dst) is True assert dst.read_bytes() == payload + download.assert_called_once_with( + uri=uri, + destination_dir=dst.parent, + artifact_name=dst.name, + expected_sha256=BFCLv4.SINGLE_TURN_SHA256, + ) def test_env_var_overrides_constant(self, tmp_path, monkeypatch): """BFCL_V4_DATASET_URI takes precedence over the class constant.""" dst = _dst(tmp_path) - payload = b"ENV-WINS" - monkeypatch.setattr( - BFCLv4, - "R2_DATASET_URI", - "https://inference.mlcommons-storage.org/constant.uri", - ) - monkeypatch.setenv( - BFCLv4.R2_DATASET_URI_ENV, "https://inference.mlcommons-storage.org/env.uri" - ) - monkeypatch.setattr( - BFCLv4, "SINGLE_TURN_SHA256", hashlib.sha256(payload).hexdigest() - ) + constant_uri = "https://inference.mlcommons-storage.org/constant.uri" + env_uri = "https://inference.mlcommons-storage.org/env.uri" + monkeypatch.setattr(BFCLv4, "R2_DATASET_URI", constant_uri) + monkeypatch.setenv(BFCLv4.R2_DATASET_URI_ENV, env_uri) - run = MagicMock(side_effect=_fake_run_writing(dst, payload)) - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch(f"{_MODULE}.subprocess.run", run), - ): + download = MagicMock(return_value=dst) + with patch.object(bfcl_module, "download_r2_artifact", download): assert BFCLv4._download_full_parquet_from_r2(dst) is True - # The env URI, not the constant, is the last positional arg passed to bash. - passed_cmd = run.call_args.args[0] - assert passed_cmd[-1] == "https://inference.mlcommons-storage.org/env.uri" - - def test_relocates_file_from_subdir(self, tmp_path, monkeypatch): - """Downloader nests the file in a subdir; loader relocates to dst_path.""" - dst = _dst(tmp_path) - payload = b"NESTED" - nested = dst.parent / "edge-agentic" / dst.name - monkeypatch.setattr( - BFCLv4, "R2_DATASET_URI", "https://inference.mlcommons-storage.org/x.uri" + download.assert_called_once_with( + uri=env_uri, + destination_dir=dst.parent, + artifact_name=dst.name, + expected_sha256=BFCLv4.SINGLE_TURN_SHA256, ) - monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - monkeypatch.setattr( - BFCLv4, "SINGLE_TURN_SHA256", hashlib.sha256(payload).hexdigest() - ) - - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch( - f"{_MODULE}.subprocess.run", - side_effect=_fake_run_writing(nested, payload), - ), - ): - assert BFCLv4._download_full_parquet_from_r2(dst) is True - assert dst.exists() and dst.read_bytes() == payload - - -class TestR2DownloadFailures: - def test_sha256_mismatch_raises(self, tmp_path, monkeypatch): - dst = _dst(tmp_path) - monkeypatch.setattr( - BFCLv4, "R2_DATASET_URI", "https://inference.mlcommons-storage.org/x.uri" - ) - monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - # Real (unmocked) SHA pin; downloaded bytes won't match it. - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch( - f"{_MODULE}.subprocess.run", - side_effect=_fake_run_writing(dst, b"WRONG-BYTES"), - ), - pytest.raises(ValueError, match="SHA-256 mismatch"), - ): - BFCLv4._download_full_parquet_from_r2(dst) - - def test_missing_file_after_download_raises(self, tmp_path, monkeypatch): - dst = _dst(tmp_path) - monkeypatch.setattr( - BFCLv4, "R2_DATASET_URI", "https://inference.mlcommons-storage.org/x.uri" - ) - monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - - def _run_noop(cmd, **kwargs): - return MagicMock(returncode=0, stderr="") - - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch(f"{_MODULE}.subprocess.run", side_effect=_run_noop), - pytest.raises(FileNotFoundError), - ): - BFCLv4._download_full_parquet_from_r2(dst) - - def test_downloader_nonzero_exit_raises(self, tmp_path, monkeypatch): - dst = _dst(tmp_path) - monkeypatch.setattr( - BFCLv4, "R2_DATASET_URI", "https://inference.mlcommons-storage.org/x.uri" - ) - monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - - def _run_fail(cmd, **kwargs): - return MagicMock(returncode=1, stderr="boom") - - with ( - patch(f"{_MODULE}.requests.get", return_value=_mock_requests_get()), - patch(f"{_MODULE}.subprocess.run", side_effect=_run_fail), - pytest.raises(RuntimeError, match="R2 downloader failed"), - ): - BFCLv4._download_full_parquet_from_r2(dst) - - @pytest.mark.parametrize( - "uri", - [ - "https://evil.example.com/x.uri", # wrong host - "http://inference.mlcommons-storage.org/x.uri", # not https - "file:///etc/passwd", # non-http scheme - ], - ) - def test_untrusted_uri_rejected_before_network(self, tmp_path, monkeypatch, uri): - """A mis-set URI is refused before any script fetch / subprocess runs.""" - dst = _dst(tmp_path) - monkeypatch.setattr(BFCLv4, "R2_DATASET_URI", uri) - monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) - with ( - patch(f"{_MODULE}.requests.get") as get, - patch(f"{_MODULE}.subprocess.run") as run, - pytest.raises(ValueError, match="untrusted URI"), - ): - BFCLv4._download_full_parquet_from_r2(dst) - get.assert_not_called() - run.assert_not_called() - - -class TestVerifySha256: - def test_match_passes(self, tmp_path): - f = tmp_path / "a.bin" - f.write_bytes(b"hello") - BFCLv4._verify_sha256(f, hashlib.sha256(b"hello").hexdigest()) - - def test_mismatch_raises(self, tmp_path): - f = tmp_path / "a.bin" - f.write_bytes(b"hello") - with pytest.raises(ValueError, match="SHA-256 mismatch"): - BFCLv4._verify_sha256(f, "0" * 64) class TestGenerateSelection: @@ -333,14 +187,15 @@ def _download(dst_path): return True with ( - patch.object(BFCLv4, "_verify_sha256", side_effect=ValueError("mismatch")), + patch.object( + bfcl_module, "verify_sha256", side_effect=ValueError("mismatch") + ), patch.object( BFCLv4, "_download_full_parquet_from_r2", side_effect=_download ) as dl, ): df = BFCLv4.generate(tmp_path, subsets=["simple_python"]) - # Stale cache failed re-verification and was replaced by the R2 download. dl.assert_called_once() assert len(df) == 1 @@ -352,7 +207,7 @@ def test_cached_parquet_trusted_when_sha_matches(self, tmp_path, monkeypatch): monkeypatch.delenv(BFCLv4.R2_DATASET_URI_ENV, raising=False) with ( - patch.object(BFCLv4, "_verify_sha256", return_value=None), + patch.object(bfcl_module, "verify_sha256", return_value=None), patch.object(BFCLv4, "_download_full_parquet_from_r2") as dl, ): df = BFCLv4.generate(tmp_path, subsets=["simple_python"]) diff --git a/tests/unit/dataset_manager/test_download.py b/tests/unit/dataset_manager/test_download.py new file mode 100644 index 000000000..96d251b43 --- /dev/null +++ b/tests/unit/dataset_manager/test_download.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused tests for the guarded MLCommons R2 download utility.""" + +import hashlib +import subprocess +from pathlib import Path +from unittest.mock import MagicMock + +import inference_endpoint.dataset_manager.download as download_module +import pytest +from inference_endpoint.dataset_manager.download import ( + DEFAULT_R2_DOWNLOADER_COMMIT, + download_r2_artifact, + verify_sha256, +) + +pytestmark = pytest.mark.unit +_MODULE = "inference_endpoint.dataset_manager.download" +_URI = "https://inference.mlcommons-storage.org/data.uri" +_COMMIT = "a" * 40 + + +def _response() -> MagicMock: + response = MagicMock() + response.content = b"#!/bin/sh\nexit 0\n" + return response + + +def _run_result(returncode: int = 0, stderr: str = "") -> MagicMock: + result = MagicMock() + result.returncode = returncode + result.stderr = stderr + return result + + +def test_verify_sha256_incrementally(tmp_path: Path) -> None: + path = tmp_path / "artifact.bin" + payload = b"0123456789" + path.write_bytes(payload) + verify_sha256(path, hashlib.sha256(payload).hexdigest(), chunk_size=3) + + +def test_verify_sha256_mismatch(tmp_path: Path) -> None: + path = tmp_path / "artifact.bin" + path.write_bytes(b"wrong") + with pytest.raises(ValueError, match="SHA-256 mismatch"): + verify_sha256(path, "0" * 64) + + +@pytest.mark.parametrize( + "uri", + [ + "http://inference.mlcommons-storage.org/data.uri", + "https://evil.example/data.uri", + "file:///tmp/data.uri", + "https://mlcommons-storage.org.evil.example/data.uri", + ], +) +def test_rejects_untrusted_uri_before_network(tmp_path: Path, monkeypatch, uri: str): + get = MagicMock() + monkeypatch.setattr(f"{_MODULE}.requests.get", get) + with pytest.raises(ValueError, match="untrusted URI"): + download_r2_artifact(uri, tmp_path, "artifact.bin", _COMMIT) + get.assert_not_called() + + +def test_rejects_untrusted_downloader_url_before_network( + tmp_path: Path, monkeypatch +) -> None: + get = MagicMock() + monkeypatch.setattr( + download_module, "_DOWNLOADER_URL_TEMPLATE", "http://evil/{commit}" + ) + monkeypatch.setattr(download_module.requests, "get", get) + + with pytest.raises(ValueError, match="untrusted URI"): + download_r2_artifact(_URI, tmp_path, "artifact.bin", _COMMIT) + + get.assert_not_called() + + +def test_timeout_cleans_temporary_script(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + + def timeout(*args, **kwargs): + assert any( + path.name.startswith(".mlc-r2-downloader-") for path in tmp_path.iterdir() + ) + raise subprocess.TimeoutExpired(args[0], kwargs["timeout"]) + + monkeypatch.setattr(f"{_MODULE}.subprocess.run", timeout) + with pytest.raises(RuntimeError, match="timed out"): + download_r2_artifact(_URI, tmp_path, "artifact.bin", _COMMIT, timeout_s=2) + assert not list(tmp_path.glob(".mlc-r2-downloader-*.sh")) + + +def test_nonzero_exit_cleans_script(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + monkeypatch.setattr( + f"{_MODULE}.subprocess.run", + lambda *args, **kwargs: _run_result(3, "boom"), + ) + with pytest.raises(RuntimeError, match="code 3"): + download_r2_artifact(_URI, tmp_path, "artifact.bin", _COMMIT) + assert not list(tmp_path.glob(".mlc-r2-downloader-*.sh")) + + +def test_missing_artifact(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + monkeypatch.setattr( + f"{_MODULE}.subprocess.run", lambda *args, **kwargs: _run_result() + ) + with pytest.raises(FileNotFoundError, match="artifact.bin"): + download_r2_artifact(_URI, tmp_path, "artifact.bin", _COMMIT) + + +@pytest.mark.parametrize( + ("artifact_name", "decoy_name"), + [ + ("artifact.bin", None), + ("artifact[1].bin", "artifact1.bin"), + ], +) +def test_relocates_exact_nested_artifact( + tmp_path: Path, + monkeypatch, + artifact_name: str, + decoy_name: str | None, +) -> None: + payload = b"nested artifact" + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + + def run(*args, **kwargs): + nested = tmp_path / "manifest-directory" / artifact_name + nested.parent.mkdir() + nested.write_bytes(payload) + if decoy_name is not None: + decoy = tmp_path / "decoy-directory" / decoy_name + decoy.parent.mkdir() + decoy.write_bytes(b"decoy") + return _run_result() + + monkeypatch.setattr(f"{_MODULE}.subprocess.run", run) + result = download_r2_artifact( + _URI, + tmp_path, + artifact_name, + _COMMIT, + expected_sha256=hashlib.sha256(payload).hexdigest(), + ) + assert result == tmp_path / artifact_name + assert result.read_bytes() == payload + + +def test_ambiguous_artifact_rejected(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + + def run(*args, **kwargs): + for directory in ("one", "two"): + path = tmp_path / directory / "artifact.bin" + path.parent.mkdir() + path.write_bytes(directory.encode()) + return _run_result() + + monkeypatch.setattr(f"{_MODULE}.subprocess.run", run) + with pytest.raises(RuntimeError, match="ambiguous"): + download_r2_artifact(_URI, tmp_path, "artifact.bin", _COMMIT) + + +def test_checksum_mismatch_cleans_script(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(f"{_MODULE}.requests.get", lambda *args, **kwargs: _response()) + + def run(*args, **kwargs): + (tmp_path / "artifact.bin").write_bytes(b"wrong") + return _run_result() + + monkeypatch.setattr(f"{_MODULE}.subprocess.run", run) + with pytest.raises(ValueError, match="SHA-256 mismatch"): + download_r2_artifact( + _URI, tmp_path, "artifact.bin", _COMMIT, expected_sha256="0" * 64 + ) + assert not list(tmp_path.glob(".mlc-r2-downloader-*.sh")) + + +def test_success_uses_pinned_script_and_timeout(tmp_path: Path, monkeypatch) -> None: + payload = b"success" + response = _response() + get = MagicMock(return_value=response) + monkeypatch.setattr(f"{_MODULE}.requests.get", get) + run = MagicMock( + side_effect=lambda *args, **kwargs: ( + (tmp_path / "artifact.bin").write_bytes(payload), + _run_result(), + )[1] + ) + monkeypatch.setattr(f"{_MODULE}.subprocess.run", run) + + result = download_r2_artifact( + _URI, + tmp_path, + "artifact.bin", + expected_sha256=hashlib.sha256(payload).hexdigest(), + timeout_s=17, + ) + assert result == tmp_path / "artifact.bin" + get.assert_called_once_with( + "https://raw.githubusercontent.com/mlcommons/r2-downloader/" + f"{DEFAULT_R2_DOWNLOADER_COMMIT}/" + "mlc-r2-downloader.sh", + timeout=30, + ) + assert run.call_args.kwargs["timeout"] == 17 + assert not list(tmp_path.glob(".mlc-r2-downloader-*.sh"))