Allowing callers to pass a DatasetSource to dataset(...), and updating docs

CONFIG.md

@@ -116,7 +116,7 @@ For `custom_code` benchmarks, `input` is an arbitrary TOML table that becomes a
 type = "custom_code"
 command = ["python", "eval.py"]
 input = {
-    dataset = "my_data.jsonl",
+    dataset_path = "my_data.jsonl",
     max_samples = 100,
     nested = { foo = "bar" }
 }
@@ -126,14 +126,16 @@ This produces:
 
 ```json
 {
-  "dataset": "my_data.jsonl",
+  "dataset_path": "my_data.jsonl",
   "max_samples": 100,
   "nested": {
     "foo": "bar"
   }
 }
 ```
 
+Quantiles passes these values through to your custom program. It does not assign special meaning to `dataset_path` or load the file automatically; your custom code is responsible for opening the file directly or using the value to construct a dataset source.
+
 #### CLI `--input` overrides
 
 You can override or extend config input at runtime:

README.md

@@ -138,9 +138,11 @@ type = "custom_code"
 command = ["python", "my_eval.py"]
 
 [benchmarks.my-eval.input]
-dataset = "my_dataset.jsonl"
+dataset_path = "my_dataset.jsonl"
 ```
 
+The `input` table is passed to your custom program. Quantiles does not automatically load `dataset_path`; your code decides whether to open the file directly or pass it into a custom Python `DatasetSource`.
+
 Run the evaluation with `qt run my-eval`. If it fails, resume it later with `qt resume <run_id>` — the CLI re-reads the command and stored input automatically.
 
 Use custom evaluations when you need to measure behavior that is specific to your product, workflow, prompt, dataset, rubric, or release process.

python/README.md

@@ -33,6 +33,24 @@ if __name__ == "__main__":
 
 In local development, the SDK executes user code locally. The `qt` server deduplicates steps, triggers workflows, owns durable state, stored outputs, observability records, and metrics.
 
+## Datasets
+
+Use `quantiles.toml` input for configuration, then construct dataset loading behavior in code. For example, pass `dataset_path = "my_data.jsonl"` in config, read `input_value["dataset_path"]` in the workflow handler, and either open that file directly or pass it into a custom `DatasetSource`.
+
+Hugging Face datasets can be loaded by passing a URI string to `dataset(...)`:
+
+```python
+ds = await dataset(
+    ctx,
+    source="huggingface://quantiles/PubMedQA",
+    row_type=Row,
+    config="pqa_labeled",
+    split="train",
+)
+```
+
+For non-Hugging Face public or private sources, implement `DatasetSource` and pass an instance as `source`. Custom sources run inside the Python workflow process, while batch loading is still recorded through Quantiles steps.
+
 ## Development
 
 Run tests:

python/src/quantiles/datasets.py

@@ -81,9 +81,24 @@ async def initialize(self) -> JsonValue:
         raise QuantilesError(f"dataset init failed ({resp.status}): {text}")
       data = await resp.json()
 
-    self._resolved_config = data.get("config")
-    self._resolved_split = data.get("selected_split")
-    return cast(JsonValue, data)
+    metadata = cast(JsonValue, data)
+    self._apply_init_metadata(metadata)
+    return metadata
+
+  def _apply_init_metadata(self, metadata: JsonValue) -> None:
+    if not isinstance(metadata, dict):
+      raise QuantilesError("dataset init returned invalid metadata")
+
+    config = metadata.get("config")
+    if config is not None and not isinstance(config, str):
+      raise QuantilesError("dataset init returned invalid config metadata")
+
+    selected_split = metadata.get("selected_split")
+    if selected_split is not None and not isinstance(selected_split, str):
+      raise QuantilesError("dataset init returned invalid split metadata")
+
+    self._resolved_config = config
+    self._resolved_split = selected_split
 
   async def load_batch(self, offset: int, batch_size: int) -> list[dict[str, JsonValue]]:
     config = self._resolved_config or self.config
@@ -209,7 +224,7 @@ async def _execute() -> JsonValue:
 
 async def dataset[RowT: BaseModel](
   ctx: WorkflowContext,
-  source: str,
+  source: str | DatasetSource,
   row_type: type[RowT],
   *,
   batch_size: int = 100,
@@ -221,15 +236,49 @@ async def dataset[RowT: BaseModel](
   max_rows: int | None = None,
   # **kwargs: JsonValue,
 ) -> Dataset[RowT]:
-  ds_source: DatasetSource = _HttpCliSource(
-    base_url=ctx.client.base_url,
-    source=source,
-    config=config,
-    split=split,
-    revision=revision,
-  )
+  if isinstance(source, str):
+    ds_source: DatasetSource = _HttpCliSource(
+      base_url=ctx.client.base_url,
+      source=source,
+      config=config,
+      split=split,
+      revision=revision,
+    )
+    init_input: dict[str, JsonValue] = {
+      "source": source,
+      "config": config,
+      "split": split,
+      "revision": revision,
+    }
+    init_metadata = await step(
+      ctx,
+      step_key="dataset-init",
+      input_value=cast(JsonValue, init_input),
+      execute=ds_source.initialize,
+    )
+    ds_source._apply_init_metadata(init_metadata)
+  else:
+    if config is not None or split is not None or revision is not None:
+      raise QuantilesError(
+        "config, split, and revision are only supported for Hugging Face dataset URI sources"
+      )
+    ds_source = source
+    init_metadata = await ds_source.initialize()
+    init_input = {
+      "source": ds_source.source_id,
+    }
+
+    async def _record_init_metadata() -> JsonValue:
+      return init_metadata
+
+    await step(
+      ctx,
+      step_key="dataset-init",
+      input_value=cast(JsonValue, init_input),
+      execute=_record_init_metadata,
+    )
 
-  ds = Dataset(
+  return Dataset(
     ctx,
     ds_source,
     row_type,
@@ -238,16 +287,3 @@ async def dataset[RowT: BaseModel](
     transform,
     max_rows,
   )
-
-  init_input: dict[str, JsonValue] = {
-    "source": source,
-    "batch_size": batch_size,
-  }
-  await step(
-    ctx,
-    step_key="dataset-init",
-    input_value=cast(JsonValue, init_input),
-    execute=ds_source.initialize,
-  )
-
-  return ds

python/tests/test_datasets.py

@@ -7,7 +7,7 @@
 import pytest
 from pydantic import BaseModel
 
-from quantiles.datasets import Dataset, _HttpCliSource
+from quantiles.datasets import Dataset, _HttpCliSource, dataset
 from quantiles.types import JsonValue, QuantilesError
 from quantiles.workflow_context import WorkflowContext
 
@@ -22,12 +22,14 @@ class _FakeSource:
 
   def __init__(self, rows: Sequence[Mapping[str, object]]) -> None:
     self._rows: list[dict[str, JsonValue]] = [cast(dict[str, JsonValue], dict(row)) for row in rows]
+    self.initialize_calls = 0
 
   @property
   def source_id(self) -> str:
     return "fake:source"
 
   async def initialize(self) -> JsonValue:
+    self.initialize_calls += 1
     return cast(JsonValue, {"total_rows": len(self._rows)})
 
   async def load_batch(self, offset: int, batch_size: int) -> list[dict[str, JsonValue]]:
@@ -54,6 +56,32 @@ async def fake_run_step(
   return WorkflowContext(run_id=1, workflow_name="test", client=mock_client)
 
 
+def _make_cached_init_ctx(
+  init_metadata: JsonValue,
+) -> tuple[WorkflowContext, list[tuple[str, JsonValue | None]]]:
+  """Create a mock context that reuses dataset-init and executes other steps."""
+  mock_client = AsyncMock()
+  mock_client.base_url = "http://test:8765"
+  step_inputs: list[tuple[str, JsonValue | None]] = []
+
+  async def fake_run_step(
+    *,
+    run_id: int,
+    step_key: str,
+    input_value: JsonValue | None = None,
+    execute: Callable[[], Awaitable[JsonValue]] | None = None,
+  ) -> JsonValue:
+    step_inputs.append((step_key, input_value))
+    if step_key == "dataset-init":
+      return init_metadata
+    if execute is not None:
+      return await execute()
+    return cast(JsonValue, [])
+
+  mock_client.run_step = fake_run_step
+  return WorkflowContext(run_id=1, workflow_name="test", client=mock_client), step_inputs
+
+
 def _make_mock_aiohttp(resp: AsyncMock) -> MagicMock:
   """Build a mock aiohttp session with proper async context manager support."""
   post_cm = MagicMock()
@@ -99,6 +127,31 @@ async def test_initialize_parses_response(self) -> None:
     assert src._resolved_split == "test"
     assert src.source_id == "hf:huggingface://quantiles/PubMedQA:pqa_labeled:test"
 
+  def test_apply_init_metadata_parses_cached_response(self) -> None:
+    src = _HttpCliSource("http://test:8765", "huggingface://quantiles/PubMedQA")
+
+    src._apply_init_metadata(
+      cast(
+        JsonValue,
+        {
+          "total_rows": 1000,
+          "available_splits": ["train", "test"],
+          "selected_split": "test",
+          "config": "pqa_labeled",
+        },
+      )
+    )
+
+    assert src._resolved_config == "pqa_labeled"
+    assert src._resolved_split == "test"
+    assert src.source_id == "hf:huggingface://quantiles/PubMedQA:pqa_labeled:test"
+
+  def test_apply_init_metadata_rejects_invalid_response(self) -> None:
+    src = _HttpCliSource("http://test:8765", "huggingface://quantiles/PubMedQA")
+
+    with pytest.raises(QuantilesError, match="invalid config metadata"):
+      src._apply_init_metadata(cast(JsonValue, {"config": 123, "selected_split": "test"}))
+
   @pytest.mark.asyncio
   async def test_initialize_raises_on_error(self) -> None:
     src = _HttpCliSource("http://test:8765", "huggingface://bad")
@@ -137,6 +190,126 @@ async def test_load_batch_returns_rows(self) -> None:
     assert rows[0]["name"] == "a"
 
 
+class TestDatasetHelper:
+  @pytest.mark.asyncio
+  async def test_hydrates_huggingface_source_from_cached_init(self) -> None:
+    ctx, _step_inputs = _make_cached_init_ctx(
+      cast(
+        JsonValue,
+        {
+          "total_rows": 1000,
+          "available_splits": ["train", "test"],
+          "selected_split": "test",
+          "config": "pqa_labeled",
+        },
+      )
+    )
+    ds = await dataset(
+      ctx,
+      source="huggingface://quantiles/PubMedQA",
+      row_type=_SampleRow,
+      batch_size=10,
+    )
+
+    assert isinstance(ds._source, _HttpCliSource)
+    assert ds._source._resolved_config == "pqa_labeled"
+    assert ds._source._resolved_split == "test"
+
+  @pytest.mark.asyncio
+  async def test_huggingface_init_input_includes_source_options(self) -> None:
+    ctx, step_inputs = _make_cached_init_ctx(
+      cast(
+        JsonValue,
+        {
+          "total_rows": 1000,
+          "available_splits": ["train", "test"],
+          "selected_split": "test",
+          "config": "pqa_labeled",
+        },
+      )
+    )
+
+    _ds = await dataset(
+      ctx,
+      source="huggingface://quantiles/PubMedQA",
+      row_type=_SampleRow,
+      batch_size=10,
+      config="pqa_labeled",
+      split="test",
+      revision="abc123",
+      max_rows=20,
+    )
+
+    assert step_inputs[0] == (
+      "dataset-init",
+      cast(
+        JsonValue,
+        {
+          "source": "huggingface://quantiles/PubMedQA",
+          "config": "pqa_labeled",
+          "split": "test",
+          "revision": "abc123",
+        },
+      ),
+    )
+
+  @pytest.mark.asyncio
+  async def test_custom_source_initialize_runs_when_init_step_is_cached(self) -> None:
+    source = _FakeSource([{"id": 1, "name": "alice"}])
+    ctx, step_inputs = _make_cached_init_ctx(cast(JsonValue, {"total_rows": 1}))
+
+    ds = await dataset(
+      ctx,
+      source=source,
+      row_type=_SampleRow,
+      batch_size=10,
+    )
+
+    assert source.initialize_calls == 1
+    assert step_inputs[0] == (
+      "dataset-init",
+      cast(JsonValue, {"source": "fake:source"}),
+    )
+
+    collected = []
+    async for row in ds.iter_rows():
+      collected.append(row)
+
+    assert len(collected) == 1
+    assert collected[0].name == "alice"
+
+  @pytest.mark.asyncio
+  async def test_accepts_custom_source(self) -> None:
+    rows = [{"id": 1, "name": "alice"}]
+    ctx = _make_mock_ctx()
+    ds = await dataset(
+      ctx,
+      source=_FakeSource(rows),
+      row_type=_SampleRow,
+      batch_size=10,
+    )
+
+    collected = []
+    async for row in ds.iter_rows():
+      collected.append(row)
+
+    assert len(collected) == 1
+    assert collected[0].id == 1
+    assert collected[0].name == "alice"
+
+  @pytest.mark.asyncio
+  async def test_rejects_huggingface_options_for_custom_source(self) -> None:
+    ctx = _make_mock_ctx()
+
+    with pytest.raises(QuantilesError, match="only supported for Hugging Face"):
+      await dataset(
+        ctx,
+        source=_FakeSource([]),
+        row_type=_SampleRow,
+        config="cfg",
+      )
+
+
 class TestDatasetIterator:
   @pytest.mark.asyncio
   async def test_iter_rows_basic(self) -> None: