feat(agentic): auto-download official dataset

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

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

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:

src/inference_endpoint/dataset_manager/agentic_inference_dataset.py

@@ -17,13 +17,15 @@
 
 import logging
 from dataclasses import dataclass, field, replace
+from pathlib import Path
 from typing import Any
 
 import pandas as pd
 
 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.

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)