feat(connectors): split ARFBench into a series and a plots dataset

packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/connector.py

@@ -1,129 +1,78 @@
-"""Convert ARFBench questions, metric series, and charts to TimeF.
+"""Convert the ARFBench questions and metric series to TimeF.
 
 The release holds one question table, one series file per metric and rollup interval, and one
-chart per metric or cited metric pair. Each metric becomes one Record, each chart one Record, and
-each question a series task and a chart task. The loop that walks the release is in ``convert``.
-:mod:`~timenet_connectors.datasets.datadog.arfbench.series` reads the series files, and
-:mod:`~timenet_connectors.datasets.datadog.arfbench.questions` reads the question table.
+chart per metric or cited metric pair. This connector reads the table and the series. Each metric
+becomes one Record, and each question becomes one series task. The loop over the release is in
+``convert``. :mod:`~timenet_connectors.datasets.datadog.arfbench.series` reads the series files.
+:mod:`~timenet_connectors.datasets.datadog.arfbench.questions` reads the question table. The
+``datadog/arfbench-plots`` dataset uses this dataset as its parent and adds the charts.
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
-from timenet.connectors import BaseConnector
-from timenet.dataset import Record, Source, TimeFDataset
-from timenet.errors import TimeNetConnectorError, TimeNetDownloadError
-from timenet_connectors.datasets.datadog.arfbench.questions import ARFBenchQuestion, read_questions
+from timenet.dataset import TimeFDataset
+from timenet.errors import TimeNetConnectorError
+from timenet_connectors.bases.huggingface import BaseHuggingFaceConnector
+from timenet_connectors.datasets.datadog.arfbench.questions import read_questions
 from timenet_connectors.datasets.datadog.arfbench.series import MetricSeriesFiles
-from timenet_connectors.sources.huggingface_hub import hub_snapshot
-from timenet_connectors.sources.images import image_signal
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from pathlib import Path
+
+    from timenet_connectors.bases.huggingface import HubFileRef
 
 REPO = "Datadog/ARFBench"
 REVISION = "1cc8ae54e9633f596b755f7c4bce54ccb0cb9f5a"
 QA_CSV = "arfbench-qa.csv"
 SERIES_DIR = "arfbench-ts-data"  # {metric}_{interval}.parquet: one metric at one rollup interval, in seconds
-CHART_DIR = "arfbench-images"  # {metric}.png, and {metric}-{metric}.png for the pairs the questions cite
 
 
-def _create_chart_record(path: Path) -> Record:
-    """Create the Record of a chart, with the image as its one signal.
+def release_paths(raw_refs: list[HubFileRef], *parts: str) -> dict[str, Path]:
+    """Map each named part of the release to its local path.
 
     Returns:
-        The Record, not yet added to a dataset.
-    """
-    return Record(
-        metadata={"source_record_id": f"arfbench-chart-{path.stem}"},
-        sources=(Source(name="ARFBench chart", signals=(image_signal(path, name="chart"),)),),
-    )
-
-
-@dataclass(frozen=True)
-class ARFBenchSource:
-    """Paths to the ARFBench question table, series files, and charts."""
-
-    qa_csv: Path
-    series_dir: Path
-    chart_dir: Path
-
-    def read_questions(self) -> Iterator[ARFBenchQuestion]:
-        """Read the question table, in table order.
-
-        Returns:
-            Each parsed question.
-        """
-        return read_questions(self.qa_csv)
-
-    def metric_series_files(self) -> list[MetricSeriesFiles]:
-        """List the series files, grouped by metric.
-
-        Returns:
-            The metrics in ID order.
-        """
-        return MetricSeriesFiles.from_series_dir(self.series_dir)
+        The path of each part: a file, or the directory of a pattern part.
 
-    def chart_paths(self) -> list[Path]:
-        """List the chart files.
-
-        Returns:
-            The PNG files in name order.
-        """
-        return sorted(self.chart_dir.glob("*.png"))
-
-
-class ARFBenchConnector(BaseConnector[ARFBenchSource]):
-    """Connector for ARFBench (Hub repo ``Datadog/ARFBench``)."""
+    Raises:
+        TimeNetConnectorError: If a part is missing.
+    """
+    paths = {ref.part: ref.path for ref in raw_refs}
+    missing = [part for part in parts if part not in paths]
+    if missing:
+        raise TimeNetConnectorError(f"the ARFBench release lacks {', '.join(missing)}")
+    return paths
 
-    values_backend = "zarr"  # the chart tensors need it
 
-    async def download_async(self, cache_dir: Path) -> list[ARFBenchSource]:  # ruff: ignore[no-self-use] - BaseConnector override
-        """Download the question table, series files, and charts from the pinned release.
+class ARFBenchConnector(BaseHuggingFaceConnector):
+    """Connector for the ARFBench questions about the metric series (Hub repo ``Datadog/ARFBench``).
 
-        Args:
-            cache_dir: Directory for cached Hub files.
-
-        Returns:
-            A list with the paths to the table and the two directories.
+    The values are plain series, so the dataset uses the default Parquet backend.
+    """
 
-        Raises:
-            TimeNetDownloadError: If the download fails or the question table is missing.
-        """
-        root = await hub_snapshot(REPO, REVISION, cache_dir, (QA_CSV, f"{SERIES_DIR}/*.parquet", f"{CHART_DIR}/*.png"))
-        if not (root / QA_CSV).is_file():
-            raise TimeNetDownloadError(f"{REPO!r} at {REVISION!r} holds no {QA_CSV} under {root}")
-        return [ARFBenchSource(qa_csv=root / QA_CSV, series_dir=root / SERIES_DIR, chart_dir=root / CHART_DIR)]
+    HF_REPO = REPO
+    SOURCE_REVISION = REVISION
+    SOURCE_FILES: ClassVar = {"questions": QA_CSV, "series": f"{SERIES_DIR}/*.parquet"}
 
-    def convert(self, raw_refs: list[ARFBenchSource]) -> TimeFDataset:
-        """Create a Record for each metric and chart, and two tasks for each question.
+    def convert(self, raw_refs: list[HubFileRef]) -> TimeFDataset:
+        """Create a Record for each metric and a series task for each question.
 
         Args:
-            raw_refs: The one source from :meth:`download_async`.
+            raw_refs: The question table and the series directory from :meth:`download_async`.
 
         Returns:
             The populated dataset.
-
-        Raises:
-            TimeNetConnectorError: If there is not exactly one source.
         """
-        if len(raw_refs) != 1:
-            raise TimeNetConnectorError(f"expected one ARFBench source, got {len(raw_refs)}")
-        source = raw_refs[0]
+        paths = release_paths(raw_refs, "questions", "series")
         dataset = TimeFDataset(metadata=self.metadata())
         metric_records = {
-            files.metric: dataset.add_record(record=files.create_record()) for files in source.metric_series_files()
-        }
-        chart_records = {
-            path.stem: dataset.add_record(record=_create_chart_record(path)) for path in source.chart_paths()
+            files.metric: dataset.add_record(record=files.create_record())
+            for files in MetricSeriesFiles.from_series_dir(paths["series"])
         }
-        for question in source.read_questions():
+        for question in read_questions(paths["questions"]):
             dataset.add_task(task=question.create_series_task(metric_records))
-            dataset.add_task(task=question.create_chart_task(chart_records))
         return dataset
 
 

packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/dataset.yaml

@@ -1,38 +1,43 @@
-# yaml-language-server: $schema=https://docs.timenet.ai/schemas/dataset-card.schema.json
+# yaml-language-server: $schema=https://docs.timenet.ai/schemas/dataset-card-v1.schema.json
 #
 # ARFBench asks 750 multiple-choice questions about 142 metrics from 63 Datadog incidents. A metric
 # is one monitored quantity, such as the request count of a service. Datadog splits a metric by
-# tags, so a metric holds one series per tag group, for example one per host.
+# tags. Thus a metric holds one series per tag group, for example one series per host.
 #
 # Datadog stores a metric at several rollups. A rollup is the metric on a fixed time grid, with one
 # value per interval of 10, 60, 300, 1800, 3600, or 86400 seconds. The rollups of a metric are
-# separate exports, not resamplings of one series: their time ranges and tag groups differ.
+# separate exports, not resamplings of one series. Their time ranges and tag groups differ.
 #
 # Each metric is one record, with one source per tag group and one signal per rollup. A consumer
-# chooses the rollup. The paper and the charts use the finest one, which the annotation
-# `finest_interval_s` names.
+# chooses the rollup. The paper and the charts use the finest rollup. The annotation
+# `finest_interval_s` names this rollup.
 #
 # A rollup file lists the tag groups at the steps of its grid. Where the metric reported nothing
-# for a group, the step holds a null or is left out. Each series is cut to the window from its
-# group's first value to its last, so a group with no value at a rollup has no series there. Inside
-# the window, each such step is a missing timestep on the regular grid, with the spec's reason
-# `not_recorded`. Values are float64, and the release names no units. The release stores naive
-# epochs, which the dataset reads as UTC. The charts and the answer options use the same clock.
+# for a group, the step holds a null or is not in the file. The dataset cuts each series to the
+# window from the first value of its group to the last value. Thus a group with no value at a
+# rollup has no series at that rollup. Inside the window, each such step is a missing timestep on
+# the regular grid, with the reason `not_recorded` from the spec. Values are float64. The release
+# names no units. The release stores naive epochs. The dataset reads them as UTC. The answer
+# options use the same clock.
 #
-# Each question becomes two test tasks: one over the cited metric records, one over the released
-# chart images. Both carry the question's row as `question_index` in their metadata, so the test
-# split holds 1,500 tasks.
+# Each question becomes one test task over the cited metric records. Thus the test split holds 750
+# tasks. A task carries the row of the question as `question_index` in its metadata. The released
+# charts are in `datadog/arfbench-plots`. The chart tasks of that dataset carry the same
+# `question_index`.
 #
-# Kept as released: `query_name`, an identifier that differs from the metric ID for 122 of the 142
-# metrics and that the release does not explain; one end-time answer (row 191, metric 36464_0)
-# 140 seconds after its series ends; and 29 time options that lie outside every rollup.
+# The dataset keeps three items as released:
+# - `query_name`. This identifier differs from the metric ID for 122 of the 142 metrics. The
+#   release does not explain it.
+# - One end-time answer (row 191, metric 36464_0). This answer is 140 seconds after the end of its
+#   series.
+# - 29 time options that are outside every rollup.
 yaml_schema_version: 1
 dataset_id: datadog/arfbench
 dataset_version: 1.0.0
 name: ARFBench
 description: >-
   750 multiple-choice questions about anomalies in 142 metrics from 63 Datadog incidents.
-  The dataset includes every published metric rollup and the released charts.
+  The dataset includes every published metric rollup. The charts are in datadog/arfbench-plots.
 license: Apache-2.0
 source_url: https://huggingface.co/datasets/Datadog/ARFBench
 citation: >-

packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/questions.py

@@ -1,4 +1,4 @@
-"""The ARFBench question table: 750 multiple-choice questions, each a series task and a chart task."""
+"""The ARFBench question table: 750 multiple-choice questions, each a series task or a chart task."""
 
 from __future__ import annotations
 
@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING, Self, TypeVar
 
 from timenet.errors import TimeNetConnectorError
-from timenet.types import Annotation, AnswerTask, InputModality, Split
+from timenet.types import Annotation, ClassificationTask, InputModality, Label, LabelSet, Split
 
 
 if TYPE_CHECKING:
@@ -167,7 +167,7 @@ def chart_stems(self) -> tuple[str, ...]:
             return self.cited_metrics
         return ("-".join(self.cited_metrics), *self.cited_metrics)
 
-    def create_series_task(self, metric_records: Mapping[str, Record]) -> AnswerTask:
+    def create_series_task(self, metric_records: Mapping[str, Record]) -> ClassificationTask:
         """Create the task that asks the question about the cited metric records.
 
         Args:
@@ -185,7 +185,7 @@ def create_series_task(self, metric_records: Mapping[str, Record]) -> AnswerTask
         inputs = tuple(metric_records[metric] for metric in self.cited_metrics)
         return self._create_task(inputs, InputModality.TIME_SERIES)
 
-    def create_chart_task(self, chart_records: Mapping[str, Record]) -> AnswerTask:
+    def create_chart_task(self, chart_records: Mapping[str, Record]) -> ClassificationTask:
         """Create the task that asks the question about the released charts.
 
         Args:
@@ -203,25 +203,28 @@ def create_chart_task(self, chart_records: Mapping[str, Record]) -> AnswerTask:
         inputs = tuple(chart_records[stem] for stem in self.chart_stems())
         return self._create_task(inputs, InputModality.IMAGE)
 
-    def _create_task(self, inputs: tuple[Record, ...], modality: InputModality) -> AnswerTask:
-        """Create a test task with the question as prompt and the correct answer as target.
+    def _create_task(self, inputs: tuple[Record, ...], modality: InputModality) -> ClassificationTask:
+        """Create a test task that asks for the correct answer from the options of the question.
 
-        Both tasks of a question carry its row index in their metadata, so a consumer can pair
-        them. The question's annotations describe it, and the interpolation flags describe what
-        an engineer saw. Nothing here interpolates.
+        The options are the label set of the task. The label set has the name of the question row,
+        so both tasks of a question share it. Both tasks also have the row index in their metadata,
+        so a consumer can pair them. The annotations describe the question. The interpolation flags
+        describe what an engineer saw. This method does not interpolate.
 
         Returns:
             The task, not yet added to a dataset.
         """
-        task = AnswerTask(
+        task = ClassificationTask(
             inputs=inputs,
             prompt=self.question,
             targets=(self.correct_answer,),
             input_modalities=frozenset({InputModality.TEXT, modality}),
             split=Split.TEST,
             metadata={"question_index": self.index},
+            label_set=LabelSet(
+                name=f"question_{self.index}", labels=tuple(Label(name=option) for option in self.answer_options)
+            ),
         )
-        task.annotate(Annotation(key="answer_options", value=list(self.answer_options)))
         task.annotate(Annotation(key="task_category", value=self.task_category.value))
         task.annotate(Annotation(key="difficulty", value=self.difficulty.value))
         task.annotate(Annotation(key="interpolate_1", value=self.interpolate_1))

packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/requirements.txt

@@ -1,7 +1,2 @@
-# huggingface_hub downloads the release from the HuggingFace Hub.
-huggingface_hub>=0.24
-# pyarrow reads the series files.
-pyarrow>=21
-# Pillow decodes chart PNGs; Zarr stores their tensor values.
-pillow>=11
-zarr>=3
+# huggingface_hub downloads the release from the HuggingFace Hub. Version 1.1 adds tqdm_class.
+huggingface_hub>=1.1