[WIP] LanceDB dataloader

.gitignore

@@ -218,3 +218,7 @@ cython_debug/
 #  refer to https://docs.cursor.com/context/ignore-files
 .cursorignore
 .cursorindexingignore
+
+# Lance dataloader benchmarks: generated run artifacts
+benchmarks/lance/train_out/
+benchmarks/lance/matrix_results.txt

benchmarks/lance/base_standins.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: OpenMDW-1.1
+"""Benchmark standins over base Cosmos loaders.
+
+Subclasses genuine Cosmos loaders to measure performance in storage regimes
+not natively supported by the base classes.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import boto3
+from transformers import AutoTokenizer
+
+from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+from cosmos_framework.data.generator.local_datasets.sft_dataset import (
+    SFTDataset,
+    _load_sft_metadata_from_s3,
+)
+
+_QWEN_TOKENIZER = "Qwen/Qwen2.5-7B"
+
+
+@contextmanager
+def hf_online_preserved():
+    """Constructing the action base flips HF Hub offline process-wide (env + constant);
+    restore both so HF-dependent loaders (tokenizers, streaming) keep working."""
+    import huggingface_hub.constants as hfc
+
+    prev_const, prev_env = hfc.HF_HUB_OFFLINE, os.environ.get("HF_HUB_OFFLINE")
+    try:
+        yield
+    finally:
+        hfc.HF_HUB_OFFLINE = prev_const
+        if prev_env is None:
+            os.environ.pop("HF_HUB_OFFLINE", None)
+        else:
+            os.environ["HF_HUB_OFFLINE"] = prev_env
+
+
+class S3DROIDLeRobotDataset(DROIDLeRobotDataset):
+    """DROIDLeRobotDataset that materializes the S3-hosted videos, then runs the genuine base.
+
+    Builds a shadow root (same versioned dir name, so the base's version registry
+    resolves): metadata/labels are symlinked from the local tree, the mega-mp4s
+    under ``videos/`` are downloaded from ``s3://{bucket}/{prefix}/videos/...``.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        s3_bucket: str,
+        s3_prefix: str,
+        *,
+        region: str | None = None,
+        cache_dir: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        src = Path(root)
+        key = s3_prefix.strip("/").replace("/", "_")
+        cache = Path(cache_dir or os.path.join(tempfile.gettempdir(), "_s3base_droid", key)) / src.name
+        success = cache / "success"
+        success.mkdir(parents=True, exist_ok=True)
+        for sub in ("meta", "data"):
+            link = success / sub
+            if not (link.exists() or link.is_symlink()):
+                link.symlink_to((src / "success" / sub).resolve())
+
+        s3 = boto3.client("s3", region_name=region) if region else boto3.client("s3")
+        pref = s3_prefix.strip("/") + "/videos/"
+        paginator = s3.get_paginator("list_objects_v2")
+        for page in paginator.paginate(Bucket=s3_bucket, Prefix=pref):
+            for obj in page.get("Contents", []):
+                rel = obj["Key"][len(pref) - len("videos/") :]  # keep the videos/ prefix
+                dst = success / rel
+                if dst.exists():
+                    continue
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                tmp = str(dst) + f".part{os.getpid()}"
+                s3.download_file(s3_bucket, obj["Key"], tmp)
+                os.replace(tmp, dst)
+
+        super().__init__(root=str(cache), **kwargs)
+
+
+def _qwen_tokenizer_config():
+    return SimpleNamespace(tokenizer=AutoTokenizer.from_pretrained(_QWEN_TOKENIZER))
+
+
+def load_sft_metadata(
+    jsonl_path: str, *, s3_bucket: str | None = None, s3_prefix: str | None = None, min_frames: int = 61
+) -> list[dict]:
+    meta = _load_sft_metadata_from_s3(None, jsonl_path, min_frames=min_frames)
+    if s3_bucket and s3_prefix:
+        base_dir = os.path.dirname(os.path.abspath(jsonl_path))
+        pref = s3_prefix.strip("/")
+        for m in meta:
+            vp = m["vision_path"]
+            if os.path.isabs(vp) or os.path.exists(vp):
+                rel = os.path.relpath(vp, base_dir)
+            else:
+                rel = vp
+            m["vision_path"] = f"s3://{s3_bucket}/{pref}/{rel}"
+    return meta
+
+
+class BenchSFTDataset(SFTDataset):
+    """SFTDataset driver for throughput benchmarks."""
+
+    def __init__(
+        self,
+        metadata: list[dict],
+        *,
+        num_video_frames: int = 16,
+        resolution: str = "256",
+        temporal_interval_mode: str = "entire_chunk",
+        frame_selection_mode: str = "first",
+        temporal_compression_factor: int = 4,
+        skip_tokenize: bool = False,
+    ) -> None:
+        super().__init__(
+            metadata=metadata,
+            num_video_frames=num_video_frames,
+            resolution=resolution,
+            s3_credentials={},
+            temporal_interval_mode=temporal_interval_mode,
+            frame_selection_mode=frame_selection_mode,
+            tokenizer_config=_qwen_tokenizer_config(),
+            cfg_dropout_rate=0.0,
+            temporal_compression_factor=temporal_compression_factor,
+        )
+        self.skip_tokenize = bool(skip_tokenize)
+        self.shard_world_size = 1
+        self.shard_rank = 0
+        self.shard_id = 0
+
+    def _tokenize_caption(self, caption: str):
+        if self.skip_tokenize:
+            return ([], caption)
+        return super()._tokenize_caption(caption)
+
+    def __iter__(self):
+        if not hasattr(self, "_meta0"):
+            self._meta0 = list(self.metadata)
+        self.metadata = list(self._meta0)
+        self.is_initialized = False
+        return super().__iter__()
+
+    @classmethod
+    def from_jsonl(
+        cls, jsonl_path: str, *, s3_bucket: str | None = None, s3_prefix: str | None = None, **kw
+    ) -> "BenchSFTDataset":
+        return cls(load_sft_metadata(jsonl_path, s3_bucket=s3_bucket, s3_prefix=s3_prefix), **kw)
+
+
+__all__ = ["S3DROIDLeRobotDataset", "BenchSFTDataset", "load_sft_metadata", "hf_online_preserved"]

benchmarks/lance/bench_action_faithful.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: OpenMDW-1.1
+"""Faithful action-loader dataloader-throughput benchmark.
+
+The production base loader uses EPISODE-SHUFFLE (`ActionIterableShuffleDataset`,
+`iterable_shuffle=True`), not RandomSampler. So the apples-to-apples comparison is
+episode-shuffle on BOTH sides. We also include lance-random to show that batched
+takes + concurrency (LANCE_IO_THREADS) make random S3 reads competitive too.
+
+Pure dataloader throughput (no model). Stressful config: many episodes (decoder cache
+<< episodes), 8+ workers, batch 16, long steady-state. Set LANCE_IO_THREADS=256 for S3.
+
+  modes: base-episode | lance-episode | lance-random
+"""
+
+from __future__ import annotations
+
+import argparse
+import multiprocessing as mp
+import os
+import time
+
+import torch
+from base_standins import S3DROIDLeRobotDataset
+
+from cosmos_framework.data.generator.action.datasets.action_sft_dataset import ActionIterableShuffleDataset
+from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+from cosmos_framework.data.lance import LanceDROIDComposedDataset
+
+_KW = dict(action_space="joint_pos", use_state=True, mode="policy", chunk_length=16)
+
+
+def _collate(items):
+    return torch.stack([s["video"] for s in items])
+
+
+def _build(mode, root, uri, region, cache, s3_bucket=None, s3_prefix=None):
+    so = {"region": region} if region else None
+
+    def _base():
+        # genuine DROIDLeRobotDataset; for S3 the standin materializes the mega-mp4s first.
+        if s3_bucket and s3_prefix:
+            return S3DROIDLeRobotDataset(
+                root=root, s3_bucket=s3_bucket, s3_prefix=s3_prefix, region=region, use_success_only=True, **_KW
+            )
+        return DROIDLeRobotDataset(root=root, use_success_only=True, **_KW)
+
+    if mode == "base-random":
+        return _base(), "random"
+    if mode == "base-episode":
+        return ActionIterableShuffleDataset(_base()), None  # the genuine production shuffle
+    comp = LanceDROIDComposedDataset(uri, decode_device="cpu", decoder_cache_size=cache, storage_options=so, **_KW)
+    if mode == "lance-episode":
+        return ActionIterableShuffleDataset(comp), None
+    return comp, "random"  # lance-random -> RandomSampler
+
+
+def _measure(ds, sampler_kind, *, batch_size, num_workers, num_batches, warmup):
+    kw = dict(
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=_collate,
+        persistent_workers=num_workers > 0,
+        prefetch_factor=4 if num_workers > 0 else None,
+        multiprocessing_context="spawn" if num_workers > 0 else None,
+    )
+    if sampler_kind == "random":
+        g = torch.Generator()
+        g.manual_seed(0)
+        loader = torch.utils.data.DataLoader(ds, sampler=torch.utils.data.RandomSampler(ds, generator=g), **kw)
+    else:
+        loader = torch.utils.data.DataLoader(ds, **kw)  # IterableDataset (episode-shuffle)
+    seen, t0 = 0, None
+    for i, _ in enumerate(loader):
+        if i == warmup:
+            t0 = time.perf_counter()
+        if i >= warmup:
+            seen += 1
+        if seen >= num_batches:
+            break
+    return seen * batch_size / (time.perf_counter() - t0)
+
+
+def _mode_entry(mode, a, q):
+    """Subprocess entrypoint: build+measure one mode, return its samples/s. Each mode runs
+    in its own process so the torchcodec/lance C++ teardown can't SIGABRT a later mode."""
+    ds, sk = _build(
+        mode, a["root"], a["uri"], a["region"], a["cache_size"], s3_bucket=a["s3_bucket"], s3_prefix=a["s3_prefix"]
+    )
+    sps = _measure(
+        ds,
+        sk,
+        batch_size=a["batch_size"],
+        num_workers=a["num_workers"],
+        num_batches=a["num_batches"],
+        warmup=a["warmup"],
+    )
+    q.put(sps)
+    q.close()
+    q.join_thread()
+    os._exit(0)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", required=True)
+    ap.add_argument("--uri", required=True)
+    ap.add_argument("--region", default=None)
+    ap.add_argument(
+        "--s3-bucket", default=None, help="if set, base materializes mega-mp4s from this bucket (S3 regime)"
+    )
+    ap.add_argument("--s3-prefix", default=None, help="key prefix the DROID videos/ tree lives under")
+    ap.add_argument("--cache-size", type=int, default=16)
+    ap.add_argument("--batch-size", type=int, default=16)
+    ap.add_argument("--num-workers", type=int, default=8)
+    ap.add_argument("--num-batches", type=int, default=60)
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--modes", nargs="+", default=["base-episode", "lance-episode", "lance-random"])
+    args = ap.parse_args()
+
+    a = vars(args)
+    print(
+        f"batch={args.batch_size} workers={args.num_workers} cache={args.cache_size} "
+        f"num_batches={args.num_batches} LANCE_IO_THREADS={os.environ.get('LANCE_IO_THREADS', 'default')}\n"
+    )
+    print(f"{'mode':<16}{'samples/s':>12}{'vs base':>10}")
+    ctx = mp.get_context("spawn")
+    base = None
+    for mode in args.modes:
+        q = ctx.Queue()
+        p = ctx.Process(target=_mode_entry, args=(mode, a, q))
+        p.start()
+        sps = q.get()
+        p.join()
+        if mode == "base-episode":
+            base = sps
+        spd = f"{sps / base:.2f}x" if base else "-"
+        print(f"{mode:<16}{sps:>12.1f}{spd:>10}", flush=True)
+
+
+if __name__ == "__main__":
+    main()