Add streaming, bounded, credential-re-vending HF ingest CLI

.gitignore

@@ -3,6 +3,8 @@
 config.yaml
 # local MP4 download cache (ingest-videos)
 video_cache/
+# Hugging Face cache (HF_HOME for ingest-images / ingest-hf-streaming)
+huggingface_cache/
 
 # UDF Studio sample data: keep the directory layout, ignore user-provided media.
 studio_data/images/*

README.md

@@ -90,6 +90,82 @@ dataset, plus `seed-video-clips` for load-testing the frame stages without a ful
 chunk run. Run any CLI with `--help` for its options (e.g. `--chunk-seconds`,
 `--model-name`/`--pretrained`/`--dim` on `frame-embed`).
 
+## Streaming ingest
+
+For a **large** dataset on a `db://` enterprise connection, the naive one-shot
+load dies after ~1h with S3 `ExpiredToken`:
+
+```python
+# Anti-pattern: one Arrow dataset over the whole repo, one long write.
+staged = ds.dataset("~/datacomp_meta", format="parquet")  # ~3TB metadata
+tbl.add(staged)   # single Lance write outlives the vended STS token -> ExpiredToken
+```
+
+The table is opened once, the vended STS credentials are baked into the Lance
+backend, and that single write can't finish inside the token's lifetime.
+`ingest-hf-streaming` is the copy-pasteable mitigation: it **streams the source
+in `--chunk-rows` batches** (so no single write outlives the TTL) and **re-vends
+fresh credentials before every chunk** (so late chunks never use a stale token).
+
+```bash
+# zero-config demo (datasets mode, small public text set)
+uv run ingest-hf-streaming --limit 1000 --chunk-rows 200
+
+# scale path: stream datacomp_xlarge parquet metadata straight from hf://
+# (no local download) into a db:// table, re-vending creds every chunk.
+HF_TOKEN=... uv run ingest-hf-streaming \
+    --source-mode parquet --hf-dataset mlfoundations/datacomp_xlarge \
+    --db-uri db://training-playbook --table-name datacomp_xlarge \
+    --chunk-rows 50000 --revend-mode connect
+
+# restart a failed load: skip rows already written, append the rest
+uv run ingest-hf-streaming --resume --no-overwrite --table-name datacomp_xlarge ...
+```
+
+### Options
+
+Run `uv run ingest-hf-streaming --help` for the live list. All options:
+
+| Option | Default | Description |
+| ------ | ------- | ----------- |
+| `--config` | `./config.yaml` | Path to the config YAML. |
+| `--log-level` | `INFO` | Python logging level (`DEBUG`/`INFO`/`WARNING`/…). |
+| `--db-uri` | config `db_uri` | Override the database URI (e.g. `db://training-playbook` or a local path). |
+| `--table-name` | config `table_name` | Override the target table name. |
+| `--hf-dataset` | `cornell-movie-review-data/rotten_tomatoes` | Hugging Face dataset, as a `namespace/name` repo id. |
+| `--hf-split` | `train` | Dataset split — **datasets mode only**; parquet mode reads every shard. |
+| `--limit` | _none (whole dataset)_ | Cap on total **source** rows ingested. Counts skipped rows too, so resuming a `--limit N` load tops the table up to `N` total. |
+| `--chunk-rows` | `50000` | Bounded sub-write size: rows per `append`. Size it so each append finishes well under the STS TTL. |
+| `--source-mode` | `datasets` | Source reader — `datasets` or `parquet` (see below). |
+| `--revend-mode` | `connect` | Per-chunk credential re-vend lever — `connect`, `reopen`, or `latest` (see below). |
+| `--overwrite` / `--no-overwrite` | `--overwrite` | Drop the table first if it exists. Mutually exclusive with `--resume`. |
+| `--resume` / `--no-resume` | `--no-resume` | Skip rows already in the table and append the rest. Requires `--no-overwrite`. |
+| `--table-write-retries` | `5` | Attempts for each create/add op. |
+| `--table-write-retry-sleep-s` | `2.0` | Base sleep (seconds) between retries; backoff is linear (`sleep × attempt`). |
+
+For gated datasets, set `hf_token` in `config.yaml` (exported to `HF_TOKEN`) or
+pass `HF_TOKEN=...` in the environment.
+
+**`--source-mode`**
+
+| Value | Behavior |
+| ----- | -------- |
+| `datasets` (default) | `load_dataset(..., streaming=True)` buffered into batches. Works for any `datasets`-streamable repo whose columns are Arrow-serializable scalars (text/metadata). Honors `--hf-split`. Decode-heavy feature types (PIL images, audio) are out of scope. |
+| `parquet` | Streams parquet straight from `hf://` via `HfFileSystem` + `pyarrow.dataset` — exact schema, no decode, no full download. Scales to datacomp-style sharded metadata. Reads **all** `*.parquet` in the repo (ignores `--hf-split`). |
+
+**`--revend-mode`** — each chunk logs the vended `aws_session_token` prefix
+(`token=...`) so you can confirm credentials actually rotate.
+
+| Value | Behavior |
+| ----- | -------- |
+| `connect` (default) | Reconstructs the connection (`geneva.connect`) before each chunk → fresh namespace client → fresh `describe_table` vend. The only lever **guaranteed** to re-vend. |
+| `reopen` | Reuses the connection and re-opens the table (`conn.open_table`) per chunk. Lighter, but whether it re-vends happens inside opaque native code and is unconfirmed. |
+| `latest` | Holds one table and refreshes its credentials in place via the underlying `latest_storage_options()` vend primitive (a private, internal lancedb API). |
+
+> **Tuning `--chunk-rows`.** 50k metadata rows finishes far inside a ~1h STS TTL.
+> For wide rows or rows carrying blobs, lower it so each append stays comfortably
+> under the TTL.
+
 ## Inspecting state
 
 ```bash

geneva_examples/core/utils/hf_streaming.py

@@ -0,0 +1,225 @@
+"""Streaming, bounded, credential-re-vending Hugging Face -> Geneva ingest.
+
+A one-shot ``tbl.add(staged)`` of a multi-TB dataset on a ``db://`` enterprise
+connection dies after ~1h with S3 ``ExpiredToken``: the table is opened once, the
+vended STS credentials are baked into the Lance backend, and the single long
+write outlives the token TTL. The client-only mitigation has two independent
+levers, both implemented here:
+
+* **Bounded sub-writes** -- :func:`iter_hf_batches` streams the source into
+  ``chunk_rows``-sized :class:`pyarrow.RecordBatch` es so each ``tbl.add(...)``
+  finishes well inside the STS TTL.
+* **Re-vend per chunk** -- :func:`fresh_table` returns a table handle with
+  freshly vended credentials before each chunk. The default ``connect`` mode
+  reconstructs the connection (the only lever guaranteed to re-vend); ``reopen``
+  and ``latest`` are lighter alternatives (see the docstring).
+
+The network-bound source readers carry ``# pragma: no cover`` -- the bounded
+iterator and the re-vend lever are pure and unit-tested.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+if TYPE_CHECKING:
+    from geneva_examples.core.config import Config
+
+logger = logging.getLogger(__name__)
+
+SOURCE_MODES = ("datasets", "parquet")
+REVEND_MODES = ("connect", "reopen", "latest")
+
+
+def iter_hf_batches(
+    hf_dataset: str,
+    hf_split: str,
+    chunk_rows: int,
+    limit: int | None = None,
+    skip_rows: int = 0,
+    mode: str = "datasets",
+    hf_token: str | None = None,
+) -> Iterator[pa.RecordBatch]:
+    """Yield bounded ``RecordBatch`` es from a Hugging Face dataset.
+
+    ``skip_rows`` source rows are dropped first (resume), then rows are yielded
+    until ``limit`` *source* rows have been consumed (``skip_rows`` counts toward
+    ``limit``, so resuming a ``--limit N`` load tops the table up to ``N`` total).
+    Both bounds slice within the boundary batch, so the row counts are exact
+    regardless of the source's native batch size.
+    """
+    if mode not in SOURCE_MODES:
+        raise ValueError(
+            f"unknown source mode: {mode!r} (expected one of {SOURCE_MODES})"
+        )
+
+    skipped = 0
+    yielded = 0
+    for batch in _raw_batches(mode, hf_dataset, hf_split, chunk_rows, hf_token):
+        # Drop the resume prefix, slicing within the batch that straddles it.
+        if skipped < skip_rows:
+            need = skip_rows - skipped
+            if batch.num_rows <= need:
+                skipped += batch.num_rows
+                continue
+            batch = batch.slice(need)
+            skipped = skip_rows
+
+        # Cap at `limit` source rows (skipped + yielded), slicing the boundary.
+        if limit is not None:
+            consumed = skipped + yielded
+            if consumed >= limit:
+                return
+            remaining = limit - consumed
+            if batch.num_rows > remaining:
+                batch = batch.slice(0, remaining)
+
+        if batch.num_rows == 0:
+            continue
+
+        yield batch
+        yielded += batch.num_rows
+        if limit is not None and skipped + yielded >= limit:
+            return
+
+
+def _raw_batches(
+    mode: str,
+    hf_dataset: str,
+    hf_split: str,
+    chunk_rows: int,
+    hf_token: str | None,
+) -> Iterator[pa.RecordBatch]:  # pragma: no cover - network-bound source readers
+    """Dispatch to the per-mode raw batch reader (no bounding applied)."""
+    if mode == "parquet":
+        return _parquet_batches(hf_dataset, chunk_rows, hf_token)
+    return _datasets_batches(hf_dataset, hf_split, chunk_rows, hf_token)
+
+
+def _parquet_batches(
+    hf_dataset: str,
+    chunk_rows: int,
+    hf_token: str | None,
+) -> Iterator[pa.RecordBatch]:  # pragma: no cover - streams parquet from hf://
+    """Stream parquet metadata straight from ``hf://`` -- no full download.
+
+    Scales to datacomp-style metadata repos (e.g. ``mlfoundations/datacomp_xlarge``):
+    point :mod:`pyarrow.dataset` at the repo's parquet files over
+    :class:`huggingface_hub.HfFileSystem` and read fixed-size batches. Set
+    ``HF_TOKEN`` (``hf_token``) for gated repos.
+
+    Reads *every* ``*.parquet`` under the repo as one dataset (the intended
+    behavior for sharded, single-split metadata repos); ``hf_split`` is not
+    applied here -- use ``datasets`` mode for per-split selection.
+    """
+    import pyarrow.dataset as ds
+    from huggingface_hub import HfFileSystem
+
+    fs = HfFileSystem(token=hf_token)
+    root = f"datasets/{hf_dataset}"
+    files = fs.glob(f"{root}/**/*.parquet")
+    if not files:
+        raise RuntimeError(f"no parquet files under hf://{root}")
+    dataset = ds.dataset(files, filesystem=fs, format="parquet")
+    yield from dataset.to_batches(batch_size=chunk_rows)
+
+
+def _datasets_batches(
+    hf_dataset: str,
+    hf_split: str,
+    chunk_rows: int,
+    hf_token: str | None,
+) -> Iterator[pa.RecordBatch]:  # pragma: no cover - streams via `datasets`
+    """Stream a general dataset via ``datasets`` and buffer rows into batches.
+
+    The first batch's schema is reused for every later batch so all ``add``
+    calls are schema-consistent. Columns must be Arrow-serializable scalars
+    (str/int/float/bool/list/dict) -- e.g. parquet metadata or text; decode-heavy
+    feature types (PIL images, audio) are out of scope for this example.
+    """
+    from datasets import load_dataset
+
+    stream = load_dataset(hf_dataset, split=hf_split, streaming=True, token=hf_token)
+    buffer: list[dict] = []
+    schema: pa.Schema | None = None
+
+    def _flush(rows: list[dict]) -> pa.RecordBatch:
+        nonlocal schema
+        batch = (
+            pa.RecordBatch.from_pylist(rows)
+            if schema is None
+            else pa.RecordBatch.from_pylist(rows, schema=schema)
+        )
+        schema = batch.schema
+        return batch
+
+    for row in stream:
+        buffer.append(dict(row))
+        if len(buffer) >= chunk_rows:
+            yield _flush(buffer)
+            buffer = []
+    if buffer:
+        yield _flush(buffer)
+
+
+def fresh_table(
+    cfg: Config,
+    table_name: str,
+    *,
+    mode: str = "connect",
+    conn: object | None = None,
+    table: object | None = None,
+) -> tuple[object, object]:
+    """Return ``(conn, table)`` with freshly vended credentials for the next add.
+
+    * ``connect`` (default) -- reconstruct the connection via
+      :func:`geneva_examples.core.common.connect`, then ``open_table``. This is
+      the only lever guaranteed to re-vend by construction (fresh namespace
+      client -> fresh ``describe_table`` -> fresh STS creds).
+    * ``reopen`` -- reuse ``conn`` and ``open_table`` again. Lighter, but whether
+      it re-vends happens inside opaque Rust and is unconfirmed; reuse with care.
+    * ``latest`` -- hold one ``table`` and refresh its credentials in place via
+      the underlying ``latest_storage_options()`` vend primitive (a private,
+      internal lancedb API).
+    """
+    if mode == "connect":
+        from geneva_examples.core.common import connect
+
+        conn = connect(cfg)
+        return conn, conn.open_table(table_name)
+
+    if mode == "reopen":
+        if conn is None:
+            raise ValueError("reopen mode requires an existing conn")
+        return conn, conn.open_table(table_name)
+
+    if mode == "latest":
+        if conn is None:
+            raise ValueError("latest mode requires an existing conn")
+        if table is None:
+            table = conn.open_table(table_name)
+        try:
+            table._ltbl.latest_storage_options()
+        except Exception:  # noqa: BLE001 - best-effort refresh; add() still proceeds
+            logger.warning("latest_storage_options_unavailable")
+        return conn, table
+
+    raise ValueError(f"unknown revend mode: {mode!r} (expected one of {REVEND_MODES})")
+
+
+def vended_token_prefix(table: object, length: int = 8) -> str:
+    """Best-effort short prefix of the table's vended ``aws_session_token``.
+
+    Logged per chunk so a ``db://`` run self-verifies that credentials actually
+    rotate across chunks. Returns ``"<none>"`` when no token is available.
+    """
+    try:
+        opts = table._ltbl.latest_storage_options() or {}  # type: ignore[attr-defined]
+        token = opts.get("aws_session_token") or ""
+        return token[:length] if token else "<none>"
+    except Exception:  # noqa: BLE001 - logging helper must never raise
+        return "<none>"