diff --git a/.gitignore b/.gitignore index 2592e48c..ebbabb74 100644 --- a/.gitignore +++ b/.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 diff --git a/benchmarks/lance/base_standins.py b/benchmarks/lance/base_standins.py new file mode 100644 index 00000000..f67f1bc3 --- /dev/null +++ b/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"] diff --git a/benchmarks/lance/bench_action_faithful.py b/benchmarks/lance/bench_action_faithful.py new file mode 100644 index 00000000..fbdcf028 --- /dev/null +++ b/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() diff --git a/benchmarks/lance/bench_combined_faithful.py b/benchmarks/lance/bench_combined_faithful.py new file mode 100644 index 00000000..fd72eaf2 --- /dev/null +++ b/benchmarks/lance/bench_combined_faithful.py @@ -0,0 +1,363 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Faithful combined 3-dataloader throughput benchmark: base-trio vs lance-trio. + +Every base side is a GENUINE shipped Cosmos loader (no reconstructions): + + * ACTION — DROIDLeRobotDataset (LOCAL) / S3DROIDLeRobotDataset (S3 standin, which + just materializes the mega-mp4s from S3 then runs the identical base + decode). EPISODE-SHUFFLE on both sides (the production shuffle). + * VLM — get_llava_ov_streaming: the shipped HF-Hub streaming factory, imported + and called directly. Cosmos has no local/S3 VLM base, so this is the base + in every regime (sequential shards + shuffle buffer, no random access). + * VISION-SFT — the shipped SFTDataset (via BenchSFTDataset). LOCAL reads local mp4s; + S3 rewrites vision_path to s3:// so SFTDataset downloads each sample's mp4 + via boto3 — the genuine per-sample remote path. + +Two regimes, reported separately: + LOCAL — all loaders local (cosmos's pre-download-then-train workflow). + S3 — action via the S3 standin, vision-SFT via genuine boto3 per-sample, VLM via HF + streaming (its only mode); Lance reads s3:// natively. + +RAW mode (no Qwen image-processor — that is model work, not the dataloader's job). The +1:1:1 mixer aggregate is gated by the SLOWEST loader, so report the combined number WITH +the per-loader breakdown, never as a bare multiple. Run ``--trios base`` and +``--trios lance`` in SEPARATE processes (one process hits the torchcodec/lance teardown +SIGABRT between trios). +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) + +import bench_vision_sft # noqa: E402 (kept loader benches) +import bench_vlm # noqa: E402 +from base_standins import S3DROIDLeRobotDataset, hf_online_preserved # noqa: E402 + +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import ( # noqa: E402 + ActionIterableShuffleDataset, +) +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset # noqa: E402 +from cosmos_framework.data.lance import ( # noqa: E402 + LanceDROIDComposedDataset, + LanceVisionSFTDataset, +) +from cosmos_framework.data.lance.vlm_dataset import LanceVLMShuffleScan # noqa: E402 + +_ACTION_KW = dict(action_space="joint_pos", use_state=True, mode="policy", chunk_length=16) +_VSFT_KW = dict(num_video_frames=16, frame_selection_mode="first", temporal_interval_mode="entire_chunk") + + +# ── runner helpers ── +def _action_collate(samples): + out = {} + for k in samples[0]: + v = samples[0][k] + out[k] = torch.stack([s[k] for s in samples]) if torch.is_tensor(v) else [s[k] for s in samples] + return out + + +def _batch_count(batch, batch_size): + if isinstance(batch, (list, tuple)): + return len(batch) + if isinstance(batch, dict): + for v in batch.values(): + try: + return len(v) + except TypeError: + continue + return batch_size + if torch.is_tensor(batch): + return batch.shape[0] + try: + return len(batch) + except TypeError: + return batch_size + + +class _InfiniteLoader: + def __init__(self, loader, name): + self.loader, self.name, self.it = loader, name, iter(loader) + + def next_batch(self): + try: + return next(self.it) + except StopIteration: + self.it = iter(self.loader) + return next(self.it) + + +def _standalone_sps(loader, *, batch_size, rounds, warmup): + inf = _InfiniteLoader(loader, "standalone") + seen, t0 = 0, None + for i in range(rounds + warmup): + b = inf.next_batch() + n = _batch_count(b, batch_size) + if i == warmup: + t0 = time.perf_counter() + if i >= warmup: + seen += n + return seen / (time.perf_counter() - t0) + + +def _combined_sps(loaders, names, *, batch_size, rounds, warmup): + infs = [_InfiniteLoader(ld, nm) for ld, nm in zip(loaders, names)] + seen, t0 = 0, None + for r in range(rounds + warmup): + if r == warmup: + t0 = time.perf_counter() + for inf in infs: + n = _batch_count(inf.next_batch(), batch_size) + if r >= warmup: + seen += n + return seen / (time.perf_counter() - t0) + + +def _so(region, uri): + """storage_options only for s3:// uris — lets one run mix local + S3 loaders.""" + return {"region": region} if (region and str(uri).startswith("s3://")) else None + + +# ── per-loader builders (genuine bases) ── +def build_action_loader(which, root, uri, region, cache, batch_size, num_workers, s3_bucket=None, s3_prefix=None): + if which == "base": + with hf_online_preserved(): + if s3_bucket and s3_prefix: # genuine base + S3 materialization standin + base = S3DROIDLeRobotDataset( + root=root, + s3_bucket=s3_bucket, + s3_prefix=s3_prefix, + region=region, + use_success_only=True, + **_ACTION_KW, + ) + else: + base = DROIDLeRobotDataset(root=root, use_success_only=True, **_ACTION_KW) + ds = ActionIterableShuffleDataset(base) # the genuine production shuffle + else: + comp = LanceDROIDComposedDataset( + uri, + decode_device="cpu", + decoder_cache_size=cache, + storage_options=_so(region, uri), + **_ACTION_KW, + ) + ds = ActionIterableShuffleDataset(comp) + return torch.utils.data.DataLoader( + ds, + batch_size=batch_size, + num_workers=num_workers, + collate_fn=_action_collate, + drop_last=True, + persistent_workers=num_workers > 0, + prefetch_factor=4 if num_workers > 0 else None, + multiprocessing_context="spawn" if num_workers > 0 else None, + ) + + +def build_vlm_loader(which, uri, region, batch_size, num_workers, hf_subset): + collate = bench_vlm.Collate("raw") + if which == "base": + return torch.utils.data.DataLoader( + bench_vlm.GenuineVLMBase(hf_subset), + 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, + ) + ds = LanceVLMShuffleScan(uri, "llava", buffer_size=1000, storage_options=_so(region, uri)) + return torch.utils.data.DataLoader( + ds, + 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, + ) # lance not fork-safe + + +def build_vsft_loader(which, jsonl, uri, region, batch_size, num_workers, n_total, s3_bucket, s3_prefix): + if which == "base": + # genuine SFTDataset (iterable): local mp4s, or boto3 per-sample for s3:// + ds = bench_vision_sft.build_base(jsonl, tokenize=False, s3_bucket=s3_bucket, s3_prefix=s3_prefix) + return torch.utils.data.DataLoader( + ds, + batch_size=batch_size, + num_workers=num_workers, + collate_fn=bench_vision_sft._collate, + drop_last=True, + persistent_workers=num_workers > 0, + prefetch_factor=4 if num_workers > 0 else None, + multiprocessing_context="spawn" if num_workers > 0 else None, + ) + ds = LanceVisionSFTDataset( + uri, table="vision_sft", decode_device="cpu", storage_options=_so(region, uri), **_VSFT_KW + ) + ds.skip_tokenize = True + g = torch.Generator().manual_seed(42) + sampler = torch.utils.data.RandomSampler(ds, replacement=True, num_samples=n_total, generator=g) + return torch.utils.data.DataLoader( + ds, + batch_size=batch_size, + sampler=sampler, + num_workers=num_workers, + collate_fn=bench_vision_sft._collate, + drop_last=True, + persistent_workers=num_workers > 0, + prefetch_factor=4 if num_workers > 0 else None, + multiprocessing_context="spawn" if num_workers > 0 else None, + ) + + +def run_trio( + which, + paths, + *, + region, + cache, + batch_size, + workers, + rounds, + warmup, + vsft_n_total, + action_s3_bucket, + action_s3_prefix, + vsft_s3_bucket, + vsft_s3_prefix, + vlm_hf_subset, +): + aw, vw, sw = workers["action"], workers["vlm"], workers["vision-sft"] + print(f"\n========== {which.upper()}-TRIO (faithful) workers a={aw}/v={vw}/s={sw} ==========", flush=True) + a = build_action_loader( + which, + paths["action_root"], + paths["action_uri"], + region, + cache, + batch_size, + aw, + s3_bucket=action_s3_bucket if which == "base" else None, + s3_prefix=action_s3_prefix if which == "base" else None, + ) + v = build_vlm_loader(which, paths["vlm_uri"], region, batch_size, vw, vlm_hf_subset) + s = build_vsft_loader( + which, + paths["vsft_jsonl"], + paths["vsft_uri"], + region, + batch_size, + sw, + vsft_n_total, + vsft_s3_bucket if which == "base" else None, + vsft_s3_prefix if which == "base" else None, + ) + loaders, names = [a, v, s], ["action", "vlm", "vision-sft"] + standalone = {} + for ld, nm in zip(loaders, names): + standalone[nm] = _standalone_sps(ld, batch_size=batch_size, rounds=rounds, warmup=warmup) + print(f" [{which}] standalone {nm:<12} {standalone[nm]:10.1f} samples/s", flush=True) + agg = _combined_sps(loaders, names, batch_size=batch_size, rounds=rounds, warmup=warmup) + print(f" [{which}] combined mixer (1:1:1) {agg:10.1f} samples/s", flush=True) + return standalone, agg + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--action-root", required=True, help="local DROID root (parquet/meta index; videos local or via S3 standin)" + ) + ap.add_argument("--action-uri", required=True) + ap.add_argument( + "--action-s3-bucket", + default=None, + help="if set, base action materializes mega-mp4s from this bucket (S3 regime)", + ) + ap.add_argument("--action-s3-prefix", default=None, help="key prefix the DROID videos/ tree lives under") + ap.add_argument("--vlm-uri", required=True) + ap.add_argument( + "--vlm-hf-subset", + default="figureqa(cauldron,llava_format)", + help="lmms-lab/LLaVA-OneVision-Data subset the base streams from HF (cosmos default)", + ) + ap.add_argument("--vsft-jsonl", required=True) + ap.add_argument("--vsft-uri", required=True) + ap.add_argument( + "--vsft-s3-bucket", default=None, help="if set, base vsft downloads each mp4 via boto3 (genuine S3 path)" + ) + ap.add_argument("--vsft-s3-prefix", default=None, help="key prefix the jsonl-relative vision_path lives under") + ap.add_argument("--region", default=None) + 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=6, help="default per-loader worker count") + ap.add_argument("--action-workers", type=int, default=None) + ap.add_argument("--vlm-workers", type=int, default=None) + ap.add_argument("--vsft-workers", type=int, default=None) + ap.add_argument("--rounds", type=int, default=30) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--trios", nargs="+", default=["base", "lance"]) + args = ap.parse_args() + + paths = dict( + action_root=args.action_root, + action_uri=args.action_uri, + vlm_uri=args.vlm_uri, + vsft_jsonl=args.vsft_jsonl, + vsft_uri=args.vsft_uri, + ) + vsft_n_total = (args.rounds + args.warmup + 8) * args.batch_size + regime = "S3" if args.region else "LOCAL" + print( + f"FAITHFUL COMBINED RAW [{regime}] — genuine bases; action=EPISODE-SHUFFLE both sides\n" + f"batch={args.batch_size} workers={args.num_workers}/loader rounds={args.rounds} " + f"LANCE_IO_THREADS={os.environ.get('LANCE_IO_THREADS', 'default')}", + flush=True, + ) + + workers = { + "action": args.action_workers or args.num_workers, + "vlm": args.vlm_workers or args.num_workers, + "vision-sft": args.vsft_workers or args.num_workers, + } + results = {} + for which in args.trios: + results[which] = run_trio( + which, + paths, + region=args.region, + cache=args.cache_size, + batch_size=args.batch_size, + workers=workers, + rounds=args.rounds, + warmup=args.warmup, + vsft_n_total=vsft_n_total, + action_s3_bucket=args.action_s3_bucket, + action_s3_prefix=args.action_s3_prefix, + vsft_s3_bucket=args.vsft_s3_bucket, + vsft_s3_prefix=args.vsft_s3_prefix, + vlm_hf_subset=args.vlm_hf_subset, + ) + + if "base" in results and "lance" in results: + print("\n--- per-loader RAW samples/s ---") + print(f"{'loader':<14}{'base':>12}{'lance':>12}{'speedup':>10}") + for nm in ["action", "vlm", "vision-sft"]: + b, l = results["base"][0].get(nm), results["lance"][0].get(nm) + print(f"{nm:<14}{b:>12.1f}{l:>12.1f}{l / b:>9.2f}x") + ba, la = results["base"][1], results["lance"][1] + print(f"\ncombined (1:1:1) base={ba:.1f} lance={la:.1f} speedup={la / ba:.2f}x") + + +if __name__ == "__main__": + main() + os._exit(0) diff --git a/benchmarks/lance/bench_memory.py b/benchmarks/lance/bench_memory.py new file mode 100644 index 00000000..8e1bd03b --- /dev/null +++ b/benchmarks/lance/bench_memory.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Memory-footprint benchmark for the action / DROID loader: base vs LanceDB. + +Measures two axes, one side per process (``--side base|lance``) so RSS is clean: + + 1. INDEX memory — RSS after constructing the dataset (before any iteration), plus the + spawn payload (the pickled bytes each spawn worker receives). + 2. RUNTIME memory — peak total RSS (main + all DataLoader workers) during steady-state + iteration, and the per-worker average — that is what multiplies by ``num_workers`` + and decides how many workers fit in RAM. + +PSS (proportional set size) is reported alongside RSS for fork/COW fairness. +""" + +from __future__ import annotations + +import argparse +import gc +import os +import pickle + +import psutil +import torch +from base_standins import S3DROIDLeRobotDataset + +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) +_MB = 1024 * 1024 + + +def _collate(items): + return torch.stack([s["video"] for s in items]) + + +def _build(side, root, uri, cache, s3_bucket=None, s3_prefix=None, region=None): + if side == "base": + 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) + so = {"region": region} if (region and str(uri).startswith("s3://")) else None + return LanceDROIDComposedDataset(uri, decode_device="cpu", decoder_cache_size=cache, storage_options=so, **_KW) + + +def _mem_tree(proc): + """Return (total_rss, total_pss, per_worker_rss, per_worker_pss) for proc + children. + + PSS (proportional set size) splits each shared page across the procs mapping it, so it + is the fair physical-RAM metric when fork shares pages copy-on-write; RSS double-counts + those shared pages.""" + + def _pss(p): + try: + return p.memory_full_info().pss + except (psutil.Error, AttributeError): + return p.memory_info().rss # fallback if PSS unavailable + + total_rss = proc.memory_info().rss + total_pss = _pss(proc) + per_rss, per_pss = [], [] + for c in proc.children(recursive=True): + try: + per_rss.append(c.memory_info().rss) + per_pss.append(_pss(c)) + total_rss += per_rss[-1] + total_pss += per_pss[-1] + except psutil.Error: + pass + return total_rss, total_pss, per_rss, per_pss + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--side", choices=["base", "lance"], required=True) + 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) + ap.add_argument("--s3-prefix", default=None) + ap.add_argument("--cache-size", type=int, default=16) + ap.add_argument( + "--mp-context", + choices=["spawn", "fork"], + default="spawn", + help="DataLoader worker start method. fork shares the parent's index via copy-on-write " + "(measure with PSS); lance fork support is experimental.", + ) + 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=40) + ap.add_argument("--warmup", type=int, default=10) + args = ap.parse_args() + + proc = psutil.Process() + gc.collect() + rss_before = proc.memory_info().rss + + ds = _build( + args.side, + args.root, + args.uri, + args.cache_size, + s3_bucket=args.s3_bucket, + s3_prefix=args.s3_prefix, + region=args.region, + ) + gc.collect() + rss_after_init = proc.memory_info().rss + n_frames = int(getattr(ds, "_num_valid_indices", 0)) or len(ds) # valid samples + + # spawn per-worker payload: the bytes each spawn worker receives (pickle applies the + # loader's __getstate__, so this is exactly what is shipped). With spawn this duplicates + # into every worker; with fork the parent's pages are COW-shared instead. + spawn_payload_mb = len(pickle.dumps(ds, protocol=pickle.HIGHEST_PROTOCOL)) / _MB + + # steady-state runtime RSS (main + workers) + loader = torch.utils.data.DataLoader( + ds, + batch_size=args.batch_size, + num_workers=args.num_workers, + collate_fn=_collate, + persistent_workers=args.num_workers > 0, + prefetch_factor=4 if args.num_workers > 0 else None, + multiprocessing_context=args.mp_context if args.num_workers > 0 else None, + ) + peak_rss, peak_pss, rss_s, pss_s = 0, 0, [], [] + for i, _ in enumerate(loader): + if i >= args.warmup: + t_rss, t_pss, per_rss, per_pss = _mem_tree(proc) + peak_rss = max(peak_rss, t_rss) + peak_pss = max(peak_pss, t_pss) + if per_rss: + rss_s.append(sum(per_rss) / len(per_rss)) + pss_s.append(sum(per_pss) / len(per_pss)) + if i >= args.warmup + args.num_batches: + break + per_worker_mb = (sum(rss_s) / len(rss_s) / _MB) if rss_s else float("nan") + per_worker_pss_mb = (sum(pss_s) / len(pss_s) / _MB) if pss_s else float("nan") + + print( + f"MEM_RESULT side={args.side} ctx={args.mp_context} workers={args.num_workers} frames={n_frames} " + f"init_index_mb={(rss_after_init - rss_before) / _MB:.0f} " + f"peak_rss_mb={peak_rss / _MB:.0f} peak_pss_mb={peak_pss / _MB:.0f} " + f"per_worker_rss_mb={per_worker_mb:.0f} per_worker_pss_mb={per_worker_pss_mb:.0f} " + f"spawn_payload_mb={spawn_payload_mb:.1f}", + flush=True, + ) + + +if __name__ == "__main__": + main() + os._exit(0) # skip torchcodec/lance teardown SIGABRT diff --git a/benchmarks/lance/bench_vision_sft.py b/benchmarks/lance/bench_vision_sft.py new file mode 100644 index 00000000..87e1ba0a --- /dev/null +++ b/benchmarks/lance/bench_vision_sft.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Throughput benchmark: the GENUINE vision-SFT base vs the LanceDB loader. + +Base is the shipped ``SFTDataset`` (driven by :class:`BenchSFTDataset`, which only adds +single-shard setup + a direct Qwen tokenizer + a raw-mode flag — the per-sample hot path +``process_one_sample`` is unchanged). The SAME class is the base for both regimes: + + LOCAL — vision_path points at local mp4s; SFTDataset reads them (download_from_s3 + falls back to Path.read_bytes), spawns ffmpeg to decode+resize per sample. + S3 — vision_path is rewritten to s3://; SFTDataset downloads each sample's mp4 via + boto3 (genuine per-sample remote GET, no amortization) then decodes. + + lance — LanceVisionSFTDataset: decode a pre-resized, short-GOP per-clip blob; on S3 a + columnar take of plain-binary clips (parallel IO). + +``--mode raw`` skips tokenization on both sides to isolate the video I/O (the win is in +video I/O, not the storage-independent tokenize compute). Token-ids are otherwise exact. +""" + +from __future__ import annotations + +import argparse +import multiprocessing as mp +import os +import time + +import torch +from base_standins import BenchSFTDataset + +from cosmos_framework.data.lance import LanceVisionSFTDataset + +_KW = dict(num_video_frames=16, frame_selection_mode="first", temporal_interval_mode="entire_chunk") + + +def _collate(samples): + out = {} + for k in samples[0]: + v = samples[0][k] + if torch.is_tensor(v): + try: + out[k] = torch.stack([s[k] for s in samples]) + except Exception: + out[k] = [s[k] for s in samples] # ragged (e.g. text_token_ids) + else: + out[k] = [s[k] for s in samples] + return out + + +def build_base(jsonl, tokenize, *, s3_bucket=None, s3_prefix=None): + """Genuine SFTDataset (iterable) over local or s3:// vision paths.""" + return BenchSFTDataset.from_jsonl( + jsonl, s3_bucket=s3_bucket, s3_prefix=s3_prefix, skip_tokenize=not tokenize, **_KW + ) + + +def build_lance(uri, tokenize, *, region=None, table="vision_sft"): + so = {"region": region} if (region and str(uri).startswith("s3://")) else None + ds = LanceVisionSFTDataset(uri, table=table, decode_device="cpu", storage_options=so, **_KW) + ds.skip_tokenize = not tokenize + return ds + + +def _measure_iter(ds, *, batch_size, num_workers, num_batches, warmup): + """Steady-state samples/s for an IterableDataset (genuine SFT base).""" + loader = torch.utils.data.DataLoader( + ds, + batch_size=batch_size, + num_workers=num_workers, + collate_fn=_collate, + drop_last=True, + persistent_workers=num_workers > 0, + prefetch_factor=4 if num_workers > 0 else None, + multiprocessing_context="spawn" if num_workers > 0 else None, + ) + 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 _measure_map(ds, *, batch_size, num_workers, num_batches, warmup): + """Steady-state samples/s for the map-style Lance loader (global-shuffle RandomSampler).""" + n_total = (num_batches + warmup + 4) * batch_size + g = torch.Generator().manual_seed(42) + sampler = torch.utils.data.RandomSampler(ds, replacement=True, num_samples=n_total, generator=g) + loader = torch.utils.data.DataLoader( + ds, + batch_size=batch_size, + sampler=sampler, + num_workers=num_workers, + collate_fn=_collate, + drop_last=True, + persistent_workers=num_workers > 0, + prefetch_factor=4 if num_workers > 0 else None, + multiprocessing_context="spawn" if num_workers > 0 else None, + ) + 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 _side_entry(side, workers, a, q): + """Subprocess entrypoint: build+measure one (side, workers) cell. Isolated per process + so the ffmpeg/torchcodec/lance teardown can't SIGABRT a later cell.""" + tokenize = a["mode"] == "e2e" + if side == "base": + ds = build_base(a["jsonl"], tokenize, s3_bucket=a["s3_bucket"], s3_prefix=a["s3_prefix"]) + sps = _measure_iter( + ds, batch_size=a["batch_size"], num_workers=workers, num_batches=a["num_batches"], warmup=a["warmup"] + ) + else: + ds = build_lance(a["uri"], tokenize, region=a["region"], table=a["table"]) + sps = _measure_map( + ds, batch_size=a["batch_size"], num_workers=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("--jsonl", required=True) + ap.add_argument("--uri", required=True) + ap.add_argument("--table", default="vision_sft") + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument("--num-workers", nargs="+", type=int, default=[4, 8]) + ap.add_argument("--num-batches", type=int, default=25) + ap.add_argument("--warmup", type=int, default=6) + ap.add_argument( + "--mode", choices=["raw", "e2e"], default="e2e", help="raw = video only (no tokenize); e2e = video + tokenize" + ) + ap.add_argument("--modes", nargs="+", default=["base", "lance"]) + ap.add_argument("--region", default=None, help="storage_options region for an s3:// --uri") + ap.add_argument( + "--s3-bucket", default=None, help="if set, base reads each sample's mp4 from s3://bucket//" + ) + ap.add_argument("--s3-prefix", default=None, help="key prefix the jsonl-relative vision_path lives under") + args = ap.parse_args() + + a = vars(args) + regime = "S3" if (args.s3_bucket and args.s3_prefix) else "LOCAL" + print( + f"mode={args.mode} regime={regime} batch_size={args.batch_size} " + f"num_batches={args.num_batches} warmup={args.warmup}\n" + ) + print(f"{'workers':>8}{'base sps':>12}{'lance sps':>12}{'speedup':>10}") + ctx = mp.get_context("spawn") + for workers in args.num_workers: + sps = {} + for side in ("base", "lance"): + if side not in args.modes: + continue + q = ctx.Queue() + p = ctx.Process(target=_side_entry, args=(side, workers, a, q)) + p.start() + sps[side] = q.get() + p.join() + spd = sps["lance"] / sps["base"] if sps.get("base") else float("nan") + print( + f"{workers:>8}{sps.get('base', float('nan')):>12.1f}{sps.get('lance', float('nan')):>12.1f}{spd:>9.2f}x", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/lance/bench_vlm.py b/benchmarks/lance/bench_vlm.py new file mode 100644 index 00000000..da3cc73a --- /dev/null +++ b/benchmarks/lance/bench_vlm.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""VLM dataloader throughput: the GENUINE cosmos VLM base vs the LanceDB loader. + +The base is the shipped source factory ``get_llava_ov_streaming`` (imported, not +reconstructed) — ``lmms-lab/LLaVA-OneVision-Data`` streamed from the HuggingFace Hub +(``streaming=True`` + the same image/conversation filter), which is cosmos's actual +default VLM read pattern (sequential shard reads + a bounded shuffle buffer, no random +access). Cosmos has no local/S3 VLM base, so this is the base in every regime. + + base — get_llava_ov_streaming(subset): HF-Hub streaming IterableDataset + lance — LanceVLMDataset (Permutation API: O(1) random access + true global shuffle) + or LanceVLMShuffleScan (chunked-shuffle columnar scan — the S3 pattern) + +Both paths feed the SAME tokenize+image-process step (a faithful stand-in for cosmos +``VLMProcessor``), so only the data-access layer differs. Two measurements: raw access +(no processing — isolates the access bottleneck) and end-to-end (with processing). +""" + +from __future__ import annotations + +import argparse +import io +import os +import time + +import torch +from PIL import Image +from transformers import AutoProcessor + +from cosmos_framework.configs.base.reasoner.experiment.llava_ov_vlm import get_llava_ov_streaming +from cosmos_framework.data.lance.vlm_dataset import LanceVLMDataset, LanceVLMShuffleScan + + +def _decode_image(image): + if isinstance(image, dict): + raw = image.get("bytes") + return Image.open(io.BytesIO(raw)).convert("RGB") if raw else None + return image.convert("RGB") if image is not None else None + + +def _sharegpt_to_messages(conversations, image): + msgs, inserted = [], False + for turn in conversations: + role = "user" if turn["from"] == "human" else "assistant" + text = turn["value"].replace("", "").strip() + if role == "user" and not inserted and image is not None: + content = [{"type": "image", "image": image}, {"type": "text", "text": text}] + inserted = True + else: + content = [{"type": "text", "text": text}] + msgs.append({"role": role, "content": content}) + return msgs + + +def make_processor(): + return AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + +def process(item, processor): + """Faithful stand-in for VLMProcessor.process: ShareGPT image+convo -> tensors.""" + image = _decode_image(item.get("image")) + messages = _sharegpt_to_messages(item.get("conversations", []), image) + inputs = processor.apply_chat_template( + messages, tokenize=True, add_generation_prompt=False, return_dict=True, return_tensors="pt" + ) + return inputs["input_ids"] + + +def _measure(loader, *, num_batches, warmup, batch_size): + 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 + dt = time.perf_counter() - t0 + return seen * batch_size / dt + + +# ── base: the genuine cosmos VLM source (HF-Hub streaming) ────────────── +class GenuineVLMBase(torch.utils.data.IterableDataset): + """Cosmos's actual default VLM base: ``get_llava_ov_streaming`` from the shipped + config module. Built fresh in __iter__ (the HF filter lambda isn't picklable for + spawn workers), yielding the raw ``{id, image(PIL), conversations}`` dict.""" + + def __init__(self, subset: str): + self.subset = subset + + def __iter__(self): + yield from get_llava_ov_streaming(subset=self.subset) + + +_PROC = None + + +def _get_proc(): + global _PROC + if _PROC is None: + _PROC = make_processor() + return _PROC + + +class Collate: + """Module-level (picklable for spawn workers). raw -> ids; e2e -> tokenize.""" + + def __init__(self, mode): + self.mode = mode + + def __call__(self, items): + if self.mode == "raw": + return [it.get("id") for it in items] + proc = _get_proc() + return [process(it, proc) for it in items] + + +def _build_loader(side, a): + """Build the (loader, label) for one side from a plain args-dict ``a``.""" + collate = Collate(a["mode"]) + kw = dict( + batch_size=a["batch_size"], + num_workers=a["num_workers"], + collate_fn=collate, + persistent_workers=a["num_workers"] > 0, + prefetch_factor=4 if a["num_workers"] > 0 else None, + ) + so = {"region": a["region"]} if a["region"] else None + if side == "base": + return torch.utils.data.DataLoader( + GenuineVLMBase(a["subset"]), multiprocessing_context="spawn" if a["num_workers"] > 0 else None, **kw + ), "hf-stream" + if a["lance_scan"]: + ds = LanceVLMShuffleScan(a["lance_uri"], a["lance_table"], storage_options=so, buffer_size=1000) + return torch.utils.data.DataLoader(ds, **kw), "lance-scan" + ds = LanceVLMDataset(a["lance_uri"], a["lance_table"], storage_options=so) + g = torch.Generator().manual_seed(42) + sampler = torch.utils.data.RandomSampler(ds, generator=g) + return torch.utils.data.DataLoader( + ds, sampler=sampler, multiprocessing_context="spawn" if a["num_workers"] > 0 else None, **kw + ), "lance-random" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--subset", + default="figureqa(cauldron,llava_format)", + help="lmms-lab/LLaVA-OneVision-Data subset for the genuine HF-streaming base", + ) + ap.add_argument("--lance-uri", required=True) + ap.add_argument("--lance-table", default="llava") + ap.add_argument("--region", default=None, help="storage_options region for an s3:// lance-uri") + ap.add_argument( + "--lance-scan", + action="store_true", + help="use chunked-shuffle sequential scan (right for S3) instead of random point-lookups", + ) + ap.add_argument( + "--side", + choices=["base", "lance"], + required=True, + help="measure ONE side per process (run twice + divide) — each backend torn down in " + "its own process avoids the HF/lance C++ finalization crashes of an in-process compare", + ) + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument("--num-workers", type=int, default=4) + ap.add_argument("--num-batches", type=int, default=40) + ap.add_argument("--warmup", type=int, default=8) + ap.add_argument("--mode", choices=["raw", "e2e"], default="raw") + args = ap.parse_args() + + a = vars(args) + loader, label = _build_loader(args.side, a) + sps = _measure(loader, num_batches=args.num_batches, warmup=args.warmup, batch_size=args.batch_size) + print( + f"VLM_RESULT side={args.side} label={label} mode={args.mode} workers={args.num_workers} samples_per_s={sps:.1f}", + flush=True, + ) + + +if __name__ == "__main__": + main() + os._exit(0) # skip the HF/lance C++ teardown SIGABRT (result already printed) diff --git a/benchmarks/lance/forward_equivalence.py b/benchmarks/lance/forward_equivalence.py new file mode 100644 index 00000000..6943a57b --- /dev/null +++ b/benchmarks/lance/forward_equivalence.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Vision-SFT forward-equivalence: base vs Lance loader through the REAL Cosmos3-Nano (16B MoT). + +For a handful of indices, take the SAME clip from the base ``SFTDataset`` and the +``LanceVisionSFTDataset`` (aligned by index -> no shuffle), pack each into a model +batch, and run it through ``model.training_step`` with FIXED weights + seed. The +only thing that differs is the loader's decoded video (an offline H.264 re-encode +on the Lance side), so the loss delta measures exactly that -- it sits inside the +re-encode tolerance (~2%), well below the loss spread across different clips. + +The real-model counterpart to ``tests/data/lance/`` (which proves the per-sample +batches are byte/token-equal at the data level). Single GPU, forward only -- no +FSDP, no optimizer. + + python benchmarks/lance/forward_equivalence.py # vision-SFT + +ACTION and VLM run through the GENUINE training recipe with ``optimizer.lr=0`` (weights +never change, so each step is a forward loss on the same sample) + ``shuffle off`` (so +base step-k and Lance step-k are the SAME sample). The Lance swap is one line: point the +recipe's dataset ``_target_`` at our ``get_lance_*`` factory + add the uri/table. Nothing +else in the recipe changes -- that one-line swap IS the drop-in. The trainer is used here +(not this standalone batcher) because action's ``ActionProcessingRecord`` / VLM's records +need the recipe's own collation. + + * ACTION -- per-step base-vs-Lance loss within ~1.4% (measured against the pre-2026-07 + base; labels re-verified bit-exact against the rewritten lazy-LeRobot base — see + tests/data/lance/test_action.py): + torchrun --nproc_per_node=4 -m cosmos_framework.scripts.train \ + --sft-toml=examples/toml/sft_config/action_policy_droid_repro.toml --deterministic -- \ + optimizer.lr=0.0 trainer.max_iter=5 model.parallelism.data_parallel_shard_degree=4 \ + model.compile.enabled=false model.ema.enabled=false \ + dataloader_train.max_samples_per_batch=null dataloader_train.max_sequence_length=2048 \ + dataloader_train.dataloader.datasets.droid.dataset.iterable_shuffle=false \ + dataloader_train.dataloader.datasets.droid.dataset.resolution=256 \ + dataloader_train.dataloader.datasets.droid.dataset._target_=cosmos_framework.data.lance.action_dataset.get_lance_action_droid_sft_dataset \ + ~dataloader_train.dataloader.datasets.droid.dataset.root \ + ~dataloader_train.dataloader.datasets.droid.dataset.use_success_only \ + +dataloader_train.dataloader.datasets.droid.dataset.lance_uri= \ + +dataloader_train.dataloader.datasets.droid.dataset.table=droid_composed \ + +dataloader_train.dataloader.datasets.droid.dataset.decode_device=cpu + (drop the '~'/'+' lines for the base arm — the Lance loader reads labels + video + from LanceDB, so the base 'root'/'use_success_only' args are removed rather than + passed through.) + + * VLM -- byte-identical records, so the loss matches EXACTLY (measured: base 0.8149 == + Lance 0.8149, 0.00%): + torchrun --nproc_per_node=4 -m cosmos_framework.scripts.train \ + --sft-toml=examples/toml/sft_config/llava_ov_mapstyle_dataloader.toml --deterministic -- \ + optimizer.lr=0.0 trainer.max_iter=1 model.parallelism.data_parallel_shard_degree=4 \ + dataloader_train.distributor.shuffle=false \ + dataloader_train.distributor.dataset.subset="'figureqa(cauldron,llava_format)'" \ + dataloader_train.distributor.dataset._target_=cosmos_framework.data.lance.vlm_dataset.get_lance_vlm_dataset \ + +dataloader_train.distributor.dataset.uri= \ + +dataloader_train.distributor.dataset.table_name=llava + (drop the last two '+' lines for the base arm.) + +Env: HF_TOKEN, and LD_LIBRARY_PATH must include the venv's nvidia/*/lib (for +torchcodec). Requires the converted Cosmos3-Nano + Wan VAE (see docs/training.md). +""" + +from __future__ import annotations + +import argparse +import os +from types import SimpleNamespace + +from cosmos_framework.inference.common.init import init_script + +init_script(env={"COSMOS_DEVICE": "cuda"}) + +import torch +from transformers import AutoTokenizer + +from cosmos_framework.data.generator.dataflow.batchers import SequentialPackingBatcher +from cosmos_framework.data.generator.dataflow.collators import VFMListCollator +from cosmos_framework.inference.args import OmniSetupOverrides +from cosmos_framework.inference.common.args import CheckpointOverrides +from cosmos_framework.inference.common.public_model_config import build_public_model_config +from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel + +_D = "/home/ubuntu/work/data" +_VAE = "/home/ubuntu/work/cosmos-framework/examples/checkpoints/wan22_vae/Wan2.2_VAE.pth" +_TOKENIZER = "Qwen/Qwen2.5-7B" # both arms share one tokenizer; only the video differs + + +def build_model(): + """Build the real Cosmos3-Nano on one GPU with weights loaded, forward-only.""" + ckpt = CheckpointOverrides(checkpoint_path="Cosmos3-Nano").build_checkpoint( + checkpoints=OmniSetupOverrides.CHECKPOINTS + ) + hf_path = ckpt.download_checkpoint() + from cosmos_framework.scripts.convert_model_to_dcp import _redirect_avae_to_local + + _redirect_avae_to_local(hf_path) + pub = build_public_model_config(ckpt.load_model_config_dict()) + tk = pub["config"]["tokenizer"] + tk["vae_path"] = _VAE # local Wan VAE + tk["bucket_name"] = "" + tk["object_store_credential_path_pretrained"] = "" # don't auth to GCS + pub["config"]["sound_gen"] = False # no audio -> skip the AVAE + pub["config"]["sound_tokenizer"] = None + model = Cosmos3OmniModel.from_pretrained_dcp(hf_path, config=Cosmos3OmniConfig(model=pub)).model + model = model.cuda().eval() # config precision handles dtype; don't cast fp32 buffers (inv_freq) + for p in model.parameters(): + p.requires_grad_(False) + return model + + +def _pack(sample: dict) -> dict: + batcher = SequentialPackingBatcher( + max_sequence_length=8192, + tokenizer_spatial_compression_factor=16, + tokenizer_temporal_compression_factor=4, + patch_spatial=2, + max_samples_per_batch=None, + sound_latent_fps=0, + audio_sample_rate=48000, + ) + group = next(batcher.batches(iter([sample]))) + return VFMListCollator().collate(group) + + +def _loss(model, sample: dict) -> float: + batch = {k: (v.cuda() if torch.is_tensor(v) else v) for k, v in _pack(sample).items()} + torch.manual_seed(0) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = model.training_step(batch, 0) + loss = out[1] if isinstance(out, (tuple, list)) else out + return float(loss.item() if torch.is_tensor(loss) else loss) + + +def _vision_pair(tok): + from cosmos_framework.data.generator.local_datasets.sft_dataset import ( + SFTDataset, + _flatten_metadata_by_window, + _load_sft_metadata_from_s3, + ) + from cosmos_framework.data.lance import LanceVisionSFTDataset + + jsonl = f"{_D}/bridge_src/sft_dataset_bridge/train/video_dataset_file.jsonl" + base_dir = os.path.dirname(jsonl) + metas = _flatten_metadata_by_window(_load_sft_metadata_from_s3(None, jsonl, min_frames=61)) + for m in metas: + vp = m["vision_path"] + m["vision_path"] = vp if ("://" in vp or vp.startswith("/")) else os.path.join(base_dir, vp) + vkw = dict(num_video_frames=16, frame_selection_mode="first", temporal_interval_mode="entire_chunk") + base = SFTDataset( + metadata=metas, resolution="256", s3_credentials={}, tokenizer_config=tok, cfg_dropout_rate=0.0, **vkw + ) + base.s3_client = None + lance = LanceVisionSFTDataset(f"{_D}/lance/vision_sft_plain", table="vision_sft", decode_device="cpu", **vkw) + + def get_base(i): + s = base.process_one_sample(metas[i]) + s["conditioning_fps"] = 24.0 + return s + + def get_lance(i): + s = lance[i] + s["conditioning_fps"] = 24.0 + return s + + return get_base, get_lance + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--indices", type=int, nargs="+", default=[0, 1, 2, 3]) + args = ap.parse_args() + + tok = SimpleNamespace(tokenizer=AutoTokenizer.from_pretrained(_TOKENIZER)) + get_base, get_lance = _vision_pair(tok) + + print("building real Cosmos3-Nano (this loads ~16B params)...", flush=True) + model = build_model() + print("[vision] forward-equivalence: base vs Lance through the real model\n", flush=True) + print(f"{'idx':>4} {'base':>10} {'lance':>10} {'%diff':>7}") + diffs = [] + for i in args.indices: + b, l = _loss(model, get_base(i)), _loss(model, get_lance(i)) + d = abs(b - l) / b * 100 + diffs.append(d) + print(f"{i:>4} {b:>10.4f} {l:>10.4f} {d:>6.2f}%", flush=True) + print(f"\nmax %diff = {max(diffs):.2f}% (within the H.264 re-encode tolerance)") + + +if __name__ == "__main__": + main() + os._exit(0) diff --git a/benchmarks/lance/run_matrix.sh b/benchmarks/lance/run_matrix.sh new file mode 100755 index 00000000..3bf563b7 --- /dev/null +++ b/benchmarks/lance/run_matrix.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Full combined-dataloader benchmark matrix: 3 storage regimes (LOCAL / full-S3 / MIXED) +# x 2 worker allocations x {base, lance}. Each cell is an isolated process. Results -> $RES. +# +# Env (defaults match the dev box; override for another machine): +# REPO repo root (default: this script's ../../..) +# DATA local dataset root (default: /home/ubuntu/work/data) +# S3 s3:// uri of the cosmos/ prefix (default: s3://lancedb-datasets-dev-us-east-2-devrel/cosmos) +# BUCKET bucket name (for the boto3 vsft base) (default: lancedb-datasets-dev-us-east-2-devrel) +# REGION AWS region (default: us-east-2) +# ALLOCS worker allocations to sweep, "a v s" per entry, ';'-separated +# (default: "4 4 4;18 4 18" — RE-TUNE the 2nd for this machine's core count) +# RES output file (default: ./matrix_results.txt) +# Requires: .venv (uv sync --extra train --group cu130-train) + AWS creds (profile "cosmosbench" or the default chain). +set +u +REPO="${REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +cd "$REPO" +export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}" +source .venv/bin/activate +export LD_LIBRARY_PATH="$(python -c "import glob;print(':'.join(sorted(glob.glob('$REPO/.venv/lib/python3.13/site-packages/nvidia/*/lib'))))"):${LD_LIBRARY_PATH}" +export PYTHONPATH="$REPO" AWS_PROFILE="${AWS_PROFILE:-cosmosbench}" LANCE_IO_THREADS="${LANCE_IO_THREADS:-256}" + +DATA="${DATA:-/home/ubuntu/work/data}" +S="${S:-s3://lancedb-datasets-dev-us-east-2-devrel/cosmos}" +BUCKET="${BUCKET:-lancedb-datasets-dev-us-east-2-devrel}" +REGION="${REGION:-us-east-2}" +JSONL="$DATA/bridge_src/sft_dataset_bridge/train/video_dataset_file.jsonl" +RES="${RES:-./matrix_results.txt}" +: > "$RES" +R="--rounds 20 --warmup 6 --batch-size 16" + +run() { # label trio aw vw sw + local label="$1" trio="$2" aw="$3" vw="$4" sw="$5"; shift 5 + echo ">>> $label | $trio | $aw/$vw/$sw" | tee -a "$RES" + python benchmarks/lance/bench_combined_faithful.py "$@" $R --trios "$trio" \ + --action-workers "$aw" --vlm-workers "$vw" --vsft-workers "$sw" 2>&1 \ + | grep -iE "standalone (action|vlm|vision)|combined mixer" | grep -v warn \ + | sed "s/^/ [$label|$trio|$aw\/$vw\/$sw] /" | tee -a "$RES" +} + +# Bases are the GENUINE shipped loaders. action_root is always the LOCAL DROID root +# (parquet/meta index); for S3 the base materializes the mega-mp4s from --action-s3-*. +# The VLM base is HF-Hub streaming (--vlm-hf-subset) in every regime — cosmos has no +# local/S3 VLM base. +HF_SUBSET="figureqa(cauldron,llava_format)" +LOCAL_ARGS=(--action-root $DATA/droid_plus_lerobot_320x180_20260406 --action-uri $DATA/lance/droid_composed327_plain + --vlm-uri $DATA/lance/llava_figureqa --vlm-hf-subset "$HF_SUBSET" + --vsft-jsonl $JSONL --vsft-uri $DATA/lance/vision_sft_plain) +S3_ARGS=(--action-root $DATA/droid_plus_lerobot_320x180_20260406 --action-uri $S/droid327/lance/droid_composed327_plain + --action-s3-bucket $BUCKET --action-s3-prefix cosmos/droid327/base/success + --vlm-uri $S/llava/lance/llava_figureqa --vlm-hf-subset "$HF_SUBSET" + --vsft-jsonl $JSONL --vsft-uri $S/vision_sft/lance/vision_sft_plain + --vsft-s3-bucket $BUCKET --vsft-s3-prefix cosmos/vision_sft/base/sft_dataset_bridge/train --region $REGION) +MIXED_ARGS=(--action-root $DATA/droid_plus_lerobot_320x180_20260406 --action-uri $DATA/lance/droid_composed327_plain + --vlm-uri $S/llava/lance/llava_figureqa --vlm-hf-subset "$HF_SUBSET" + --vsft-jsonl $JSONL --vsft-uri $S/vision_sft/lance/vision_sft_plain + --vsft-s3-bucket $BUCKET --vsft-s3-prefix cosmos/vision_sft/base/sft_dataset_bridge/train + --region $REGION) + +IFS=';' read -ra ALLOC_LIST <<< "${ALLOCS:-4 4 4;18 4 18}" +for alloc in "${ALLOC_LIST[@]}"; do + set -- $alloc; A=$1 V=$2 Sw=$3 + for trio in base lance; do + run LOCAL "$trio" $A $V $Sw "${LOCAL_ARGS[@]}" + run S3 "$trio" $A $V $Sw "${S3_ARGS[@]}" + run MIXED "$trio" $A $V $Sw "${MIXED_ARGS[@]}" + done +done +echo "=== MATRIX DONE ($RES) ===" | tee -a "$RES" diff --git a/benchmarks/lance/table_sizes.py b/benchmarks/lance/table_sizes.py new file mode 100644 index 00000000..f7b40833 --- /dev/null +++ b/benchmarks/lance/table_sizes.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Source-vs-Lance storage comparison for the video-carrying tables. + +Answers "does the re-encoded table blow up disk?": for each modality, the size of +the original video files the base loader reads vs the converted Lance table +(video + label tables — everything the Lance loader needs). + + python table_sizes.py --droid-root /success --droid-uri \ + --vsft-jsonl --vsft-uri + +Either pair may be omitted. Local paths only (S3 tables are byte-identical copies). +""" + +from __future__ import annotations + +import argparse +import json +import os + + +def _dir_bytes(path: str) -> int: + total = 0 + for root, _, files in os.walk(path, followlinks=True): + for f in files: + p = os.path.join(root, f) + if os.path.exists(p): + total += os.path.getsize(p) + return total + + +def _row(label: str, src: int, lance: int) -> None: + print(f"{label:<12} source={src / 1e9:6.2f} GB lance={lance / 1e9:6.2f} GB ratio={lance / src:.2f}x") + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--droid-root", help="DROID success/ root (source videos under videos/)") + ap.add_argument("--droid-uri", help="composed-DROID Lance dir") + ap.add_argument("--vsft-jsonl", help="vision-SFT video_dataset_file.jsonl") + ap.add_argument("--vsft-uri", help="vision-SFT Lance dir") + args = ap.parse_args() + + if args.droid_root and args.droid_uri: + _row("action", _dir_bytes(os.path.join(args.droid_root, "videos")), _dir_bytes(args.droid_uri)) + if args.vsft_jsonl and args.vsft_uri: + base = os.path.dirname(os.path.abspath(args.vsft_jsonl)) + clips: dict[str, int] = {} + with open(args.vsft_jsonl) as fh: + for line in fh: + vp = json.loads(line)["vision_path"] + p = vp if vp.startswith("/") else os.path.join(base, vp) + clips[p] = os.path.getsize(p) + _row("vision-sft", sum(clips.values()), _dir_bytes(args.vsft_uri)) + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/data/lance/README.md b/cosmos_framework/data/lance/README.md new file mode 100644 index 00000000..2b3a605e --- /dev/null +++ b/cosmos_framework/data/lance/README.md @@ -0,0 +1,211 @@ +# LanceDB-powered Cosmos Dataloaders + +This directory contains LanceDB-backed implementations of the three main dataloaders used in Cosmos training: +- **Action (DROID/LeRobot)**: `LanceDROIDComposedDataset` +- **Vision-SFT (Local clips)**: `LanceVisionSFTDataset` +- **VLM (LLaVA-OneVision)**: `LanceVLMDataset` + +Each is a drop-in for the corresponding base loader, reading from a converted LanceDB table +instead of the original source (LeRobot tree / local clips / HuggingFace stream). Output is +equivalent to the base — VLM records byte-identical, vision-SFT token-ids exact, action labels +(action/pose/caption) bit-exact, and video within one offline H.264 re-encode (< 2.5% pixel MAD) — and the +table can be read directly from object storage (S3) without FUSE or full downloads. + +## Performance Summary + +Numbers below use a **327-episode subset of the public [`lerobot/droid_1.0.1`](https://huggingface.co/datasets/lerobot/droid_1.0.1)** +dataset (LeRobot v3.0; materialized via `tools/lance_datagen/prepare_droid_subset.py` into a +version-named root the base loader's registry resolves), whose camera views are 320×180 → +270×320 composed. Production DROID uses 640×360 views → 540×640 (see Dataset Size). + +### Combined Throughput (samples/s) +Combined 3-loader throughput, 327 DROID episodes, batch 16: + +| Workers (Action/VLM/VSFT) | Base (Local) | Lance (Local) | Base (S3) | Lance (S3) | +| ------------------------- | ------------ | ------------- | --------- | ---------- | +| 4/4/4 (Default) | 87.4 | 225.7 (2.6x) | 69.1 | 228.2 (3.3x)| +| 18/4/18 (Tuned) | 251.2 | 947.5 (3.8x) | 232.4 | 970.5 (4.2x)| + +### Per-Loader Throughput (samples/s) +Each loader standalone, tuned workers (Action/VSFT 18, VLM 4): + +| Loader | Base (Local) | Lance (Local) | Base (S3) | Lance (S3) | +| -------------- | ------------ | ------------- | ---------- | ----------- | +| Action (DROID) | 143.2 | 269.1 (1.9x) | 131.9 | 268.7 (2.0x)| +| Vision-SFT | 120.7 | 1016.4 (8.4x) | 102.4 | 875.8 (8.6x)| +| VLM (LLaVA) | 118.1 (hf) | 392.3 (3.3x) | 118.1 (hf) | 328.6 (2.8x)| + +Action decodes one composed clip instead of three runtime views; Vision-SFT decodes a pre-resized +short-GOP clip in-process instead of the base's per-sample ffmpeg resize. The VLM base has no +local/S3 form — it streams from the HuggingFace Hub (marked `hf`, so the same number appears in +both columns) — and the VLM row is measured end-to-end (image decode + tokenize) to be comparable +to the video-decoding loaders. + +### Dataset Size +Source video files the base loader reads vs the full converted Lance table (video + +label tables). Reproduce with `benchmarks/lance/table_sizes.py`: + +| modality | Source | Lance table | ratio | +| ---------- | ------------------------------- | ---------------------------------- | ----- | +| Action | 1.54 GB (3 views, AV1 long-GOP) | 0.59 GB (1 composed view, `gop=1`) | 0.38× | +| Vision-SFT | 0.10 GB (200 clips, native res) | 0.11 GB (pre-resized 256, `gop=1`) | 1.13× | + +Action — 327 DROID episodes, original three views vs the composed table: + +| metric | Original (3 views) | Composed (Lance) | +| -------- | ---------------------- | ------------------------------ | +| encoding | AV1, long-GOP | H.264, all-intra (`gop=1`) | +| streams | 3 views @ 320×180 RGB | 1 composed view @ 270×320 RGB | +| size | 1.54 GB | 0.59 GB (0.38×) | + +The `3`/`1` are the number of video **streams** (three camera views vs one composed view), not +channels — every frame is RGB. The composed 270×320 frame is the wrist view on top of the two +half-size exterior views. + +The composed table is ~2.7× smaller even though all-intra `gop=1` H.264 is *less* space-efficient +per pixel than the source's AV1 long-GOP: it stores one stream at reduced resolution (the two +exterior views are downscaled to half) rather than three full views, which outweighs the codec/GOP +cost. `gop=1` is a deliberate trade — exact, cheap random-window seeks in exchange for size (a +larger GOP would shrink the table further at some seek cost). + +The composed resolution is derived from the source (`1.5×h × w`), not fixed: this public subset has +320×180 views → 270×320. + +Vision-SFT comes out near parity (~1.1×): the pre-resize to the training resolution roughly +offsets the all-intra cost at this source resolution. A larger `--gop` shrinks either table +below source size at some seek cost. + +## Memory + +Memory is not a differentiator in either direction. The current base loader is index-light — +lazy per-shard LeRobot readers behind an LRU, metadata-only init, near-zero spawn payload — and +the Lance loader is comparable: at 1× (87.6k samples, 8 spawn workers) per-worker PSS is +~0.60 GB (base) vs ~0.71 GB (Lance; it holds the compact label arrays in memory plus the +torchcodec decoder cache, trading a little RSS for not touching the LeRobot tree at all). +The Lance wins are **throughput and native S3**, not memory. + +(Historical note: the pre-2026-07 base materialized a per-frame dict index that scaled to a +~12 GB resident index at DROID scale; the upstream rewrite to lazy LeRobot readers fixed that +wholesale, so earlier memory-scaling comparisons against it are obsolete.) + +## How it works + +There are two phases: an offline conversion (`tools/lance_datagen/`) writes one LanceDB table per +modality, and the training-time loader reads that table and decodes clips in-process. Tables are +read through the lancedb Permutation API, and media is stored one clip/image per row in a plain +`large_binary` column. Loaders null their DB/decoder handles in `__getstate__`, so each spawn +worker reopens them lazily (lancedb is not fork-safe). + +> Note: video is currently stored as plain `large_binary`. It will move to blob encoding +> (blob-v2) once the lancedb-level blob API is available. + +### Action — `LanceDROIDComposedDataset` + +Fully Lance-backed: labels **and** video come from LanceDB, so the loader takes only a +`lance_uri` — no LeRobot tree at train time. The base loader queries lazy per-shard LeRobot +readers for each sample's label windows and three camera-view windows, then resizes + concatenates +the views into one composed frame at runtime; the converter stores that composed frame once per +episode, plus the per-frame labels dumped verbatim from the base's LeRobot table. +`LanceDROIDComposedDataset` subclasses `DROIDLeRobotDataset`: the train/val split and episode-span +index are built with the base's own helpers (`split_episode_ids` / `build_episode_spans`), and +`_fetch_sample` assembles the same windowed sample dict the LeRobot readers would return — so the +inherited `__getitem__` (pose math, gripper handling, action assembly for every action space) runs +unchanged. Labels are bit-exact for the same split parameters; video is within one offline H.264 +re-encode plus the base's own decoder-backend difference (< 2.5% pixel MAD). A `version` parameter +selects the same per-dataset feature config the base resolves from its root name. + +Clips are encoded all-intra (`gop=1`), so torchcodec's `seek_mode="approximate"` lands on each +window exactly; a per-worker LRU cache keeps recently used episode decoders open. `take` returns +rows sorted by offset, so the byte read keys results by row rather than the requested order. + +Four tables (one video + three label tables, named `{table}` / `{table}_*`): + +`droid_composed` — one row per episode: + +| column | type | description | +| --------------- | -------------- | ---------------------------------------- | +| `episode_index` | int64 | episode id (used to locate the clip) | +| `ep_start` | int64 | first global frame index (build metadata) | +| `length` | int64 | number of frames (build metadata) | +| `video_bytes` | large_binary | composed 270×320 mp4 for the episode | + +`droid_composed_frames` — one row per frame (feature names store `.` as `__`): + +| column | type | description | +| --- | --- | --- | +| `episode_index`, `task_index` | int64 | frame → episode/task | +| `timestamp` | float64 | frame timestamp | +| `action__joint_position` | fixed_list[7] | commanded joints | +| `action__gripper_position` | float32 | commanded gripper | +| `observation__state__joint_positions` | fixed_list[7] | observed joints | +| `observation__state__gripper_position` | float32 | observed gripper | +| `observation__state__cartesian_position` | fixed_list[6] | EE pose (ee_pose space) | + +`droid_composed_tasks` (`task_index` int64, `task` string) and +`droid_composed_episodes` (`episode_index` int64, `episode_id` string — for keep-ranges +filtering) complete the label set. + +### Vision-SFT — `LanceVisionSFTDataset` + +The base `SFTDataset` fetches each source clip, decodes it at native size, and resizes it per +sample every epoch through an ffmpeg subprocess. The Lance table stores each clip already resized +to the training resolution with a short GOP, so the loader decodes fewer pixels in-process and +seeks windows cheaply. Caption selection, post-processing (CFG dropout, duration/resolution +conditioning suffixes), and tokenization reuse the base code, so `text_token_ids` are token-exact; +the converter uses the base's own metadata load (same duration/min-frames filters). + +| column | type | description | +| --------------------- | ------------ | ------------------------------------ | +| `clip_id` | string | `{uuid}_w{window}` | +| `width`, `height` | int64 | original resolution | +| `start_frame`, `end_frame` | int64 | window bounds | +| `temporal_interval` | int64 | frame stride | +| `enc_h`, `enc_w` | int64 | stored (resized) resolution | +| `fps` | float64 | source fps | +| `caption_json` | string | structured caption (JSON) or `""` | +| `caption` | string | dense caption fallback | +| `video_bytes` | large_binary | pre-resized clip mp4 | + +### VLM — `LanceVLMDataset` + +The base streams LLaVA-OneVision from the HuggingFace Hub (sequential shards + a bounded shuffle +buffer). The Lance table stores each sample's image bytes and conversation; the Permutation API +reads them by row, so a global shuffle is just a shuffled list of row indices. `LanceVLMShuffleScan` +instead reads contiguous row-chunks in shuffled order for S3-friendly sequential access. Records are +byte-identical to the base, so downstream image decoding and tokenization are unchanged. + +| column | type | description | +| --------------- | ------------ | ------------------------------- | +| `sample_id` | string | sample id | +| `image_bytes` | large_binary | raw image (PNG/JPEG) | +| `conversations` | string | conversation turns (JSON) | + +## Usage + +### 1. Build Tables +The conversion scripts live in [`tools/lance_datagen/`](../../../../tools/lance_datagen) (VLM uses +`convert_llava_to_lance` in [`vlm_dataset.py`](./vlm_dataset.py)): +```bash +# Action — tools/lance_datagen/build_composed_droid.py +python tools/lance_datagen/build_composed_droid.py --root --uri --gop 1 + +# Vision-SFT — tools/lance_datagen/build_vision_sft.py +python tools/lance_datagen/build_vision_sft.py --jsonl --uri + +# VLM — convert_llava_to_lance() in cosmos_framework/data/lance/vlm_dataset.py +python -c "from datasets import load_dataset; from cosmos_framework.data.lance.vlm_dataset import convert_llava_to_lance; \ +convert_llava_to_lance(load_dataset('lmms-lab/LLaVA-OneVision-Data', name='', split='train', streaming=True), '')" +``` +`tools/lance_datagen/prepare_droid_subset.py` materializes a Cosmos-canonical DROID subset from the public LeRobot release. + +### 2. Integration +Replace the standard datasets with their Lance counterparts in your configuration. +```python +from cosmos_framework.data.lance import LanceDROIDComposedDataset, LanceVisionSFTDataset, LanceVLMDataset +``` + +## Testing +Run equivalence tests to verify parity with base loaders: +```bash +pytest tests/data/lance/ +``` diff --git a/cosmos_framework/data/lance/__init__.py b/cosmos_framework/data/lance/__init__.py new file mode 100644 index 00000000..d8109d50 --- /dev/null +++ b/cosmos_framework/data/lance/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""LanceDB-powered Cosmos dataloaders (Permutation API reads, plain-binary media).""" + +from cosmos_framework.data.lance.action_dataset import LanceDROIDComposedDataset +from cosmos_framework.data.lance.vision_sft_dataset import ( + LanceVisionSFTDataset, + LanceVisionSFTIterable, +) +from cosmos_framework.data.lance.vlm_dataset import ( + LanceVLMDataset, + LanceVLMShuffleScan, +) + +__all__ = [ + "LanceDROIDComposedDataset", + "LanceVLMDataset", + "LanceVLMShuffleScan", + "LanceVisionSFTDataset", + "LanceVisionSFTIterable", +] diff --git a/cosmos_framework/data/lance/action_dataset.py b/cosmos_framework/data/lance/action_dataset.py new file mode 100644 index 00000000..d50e340d --- /dev/null +++ b/cosmos_framework/data/lance/action_dataset.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""LanceDB-backed DROID action dataset. + +Drop-in for DROIDLeRobotDataset that reads everything from LanceDB — per-frame +labels from ``{table}_frames`` / ``{table}_tasks`` / ``{table}_episodes`` and the +pre-composed video from ``{table}`` (see tools/lance_datagen/build_composed_droid.py). +The split/span index is built with the base's own helpers and the sample dict is +assembled to match what the lazy LeRobot readers return, so the inherited +``__getitem__`` (pose math, gripper handling, action assembly) runs unchanged — +labels stay bit-exact; only the H.264 re-encode of the video is lossy. +""" + +from __future__ import annotations + +import json +from typing import Any + +import lancedb +import numpy as np +import pyarrow as pa +import torch +from lancedb.permutation import Permutation +from torchcodec.decoders import VideoDecoder + +from cosmos_framework.data.generator.action.action_processing import resolve_action_normalization +from cosmos_framework.data.generator.action.datasets import droid_lerobot_dataset_config as _cfg +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import ( + ActionIterableShuffleDataset, + ActionSFTDataset, +) +from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import ( + _normalize_split, + build_episode_spans, + split_episode_ids, +) +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import ( + _DROID_TO_OPENCV, + DROIDLeRobotDataset, +) +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline + +_DEFAULT_VERSION = "droid_plus_lerobot_320x180_20260406" + + +def _resolve_device(device: str | None) -> torch.device | None: + if device == "auto": + return torch.device("cuda") if torch.cuda.is_available() else None + if device in (None, "cpu"): + return None + return torch.device(device) + + +def _read_all(db, name: str, columns: list[str]) -> pa.RecordBatch: + tbl = db.open_table(name) + perm = Permutation.identity(tbl).select_columns(columns).with_format("arrow") + return perm.__getitems__(list(range(tbl.count_rows()))) + + +class LanceDROIDComposedDataset(DROIDLeRobotDataset): + """Action loader reading labels + pre-composed episodes from LanceDB (no LeRobot tree). + + Decodes a single composed video stream per episode instead of 3 views. + """ + + def __init__( + self, + lance_uri: str, + *, + table: str = "droid_composed", + version: str = _DEFAULT_VERSION, + fps: float = 15.0, + chunk_length: int = 16, + split_seed: int = 42, + split_val_ratio: float = 0.03, + split: str = "train", + mode: str = "policy", + viewpoint: str = "concat_view", + action_space: str = "midtrain", + use_state: bool = False, + action_normalization: str | None = None, + use_filter_dict: bool = False, + filter_dict_path: str | None = None, + sample_stride: int = 1, + decode_device: str | None = "cpu", + decoder_cache_size: int = 32, + storage_options: dict | None = None, + ) -> None: + # Same argument surface as the base loader, minus what only applies to + # raw-LeRobot reading (root/video_mode/history/augmentation/temp-seg). + if viewpoint != "concat_view": + raise NotImplementedError("LanceDROIDComposedDataset only supports concat_view.") + if use_filter_dict and not filter_dict_path: + raise ValueError("use_filter_dict=True requires filter_dict_path") + if split.lower() == "val_temp_seg": + raise NotImplementedError("val_temp_seg is not supported by the Lance loader.") + + # -- config attributes the inherited code reads (base + DROID init, sans LeRobot IO) -- + self._fps = float(fps) + self._dt = 1.0 / self._fps + self._chunk_length = int(chunk_length) + self._split_seed = split_seed + self._split_val_ratio = split_val_ratio + self._split = _normalize_split(split) + self._mode = mode + self._embodiment_type = "droid_lerobot" + self._viewpoint = viewpoint + self._pose_convention = "backward_framewise" + self._rotation_format = "rot6d" + self._action_normalizer = None + if action_normalization is not None: + self._action_normalizer = resolve_action_normalization( + action_normalization, self._load_norm_stats(action_normalization) + ) + self._skip_video_loading = False + self._sample_stride = int(sample_stride) + self._min_episode_length_frames = None + self._domain_id = get_domain_id(self._embodiment_type) + self._to_opencv = _DROID_TO_OPENCV + + self._use_success_only = True # subset selection happened at convert time + self._video_mode = None + self._action_space = action_space + self._use_state = use_state + self._use_filter_dict = use_filter_dict + self._filter_dict_path = filter_dict_path + self._max_num_history_actions = 0 + self._use_image_augmentation = False + self._image_augmentor = None + self._is_val_temp_seg = False + + self._image_features = _cfg.IMAGE_FEATURES[version] + self._state_features = _cfg.STATE_FEATURES[version] + self._action_features = _cfg.ACTION_FEATURES[version] + self._is_flat_action = _cfg.IS_FLAT_ACTION[version] + self._has_multi_language_annotations = _cfg.HAS_MULTI_LANGUAGE_ANNOTATIONS[version] + self._is_gripper_action_flipped = _cfg.IS_GRIPPER_ACTION_FLIPPED[version] + + # Label-window plan: feature -> window length, mirroring the base's + # delta_timestamps ([0..k]*dt lists; our frames are contiguous per episode, + # so a window is a plain row slice of that length). + obs_len, act_len = self._chunk_length + 1, self._chunk_length + self._label_windows: dict[str, int] = { + self._state_features: obs_len, + self._action_features: act_len, + } + if action_space == "joint_pos": + self._label_windows[_cfg._JOINT_ACTION_FEATURE] = act_len + if use_state: + self._label_windows[_cfg._JOINT_STATE_FEATURE] = obs_len + self._label_windows[_cfg._GRIPPER_STATE_FEATURE] = obs_len + elif use_state: + self._label_windows[_cfg._GRIPPER_STATE_FEATURE] = obs_len + + # -- labels from Lance: same per-frame arrays the LeRobot parquets hold -- + db = lancedb.connect(lance_uri, storage_options=storage_options) + frames = _read_all( + db, + f"{table}_frames", + ["episode_index", "task_index", *[c.replace(".", "__") for c in self._label_windows]], + ) + self._row_task = frames.column("task_index").to_numpy(zero_copy_only=False).astype(np.int64) + row_episode = frames.column("episode_index").to_numpy(zero_copy_only=False).astype(np.int64) + self._feat: dict[str, np.ndarray] = {} + for c in self._label_windows: + arr = frames.column(c.replace(".", "__")) + if pa.types.is_fixed_size_list(arr.type): + self._feat[c] = np.asarray(arr.values).reshape(len(arr), arr.type.list_size) + else: + self._feat[c] = arr.to_numpy(zero_copy_only=False) + + assert np.all(np.diff(row_episode) >= 0), "episode_index is not contiguous in the frames table" + ep_vals, ep_starts, ep_counts = np.unique(row_episode, return_index=True, return_counts=True) + self._ep_row_start = {int(v): int(s) for v, s in zip(ep_vals, ep_starts)} + # episode_id in span records is positional (0..N-1) — map to the table's episode_index. + self._ep_index_of = {i: int(v) for i, v in enumerate(ep_vals)} + + tasks = _read_all(db, f"{table}_tasks", ["task_index", "task"]) + self._tasks = dict(zip(tasks.column("task_index").to_pylist(), tasks.column("task").to_pylist())) + eps = _read_all(db, f"{table}_episodes", ["episode_index", "episode_id"]) + ep_id_str = dict(zip(eps.column("episode_index").to_pylist(), eps.column("episode_id").to_pylist())) + + # -- split + span index via the base's own helpers (identical semantics) -- + episodes_meta = { + "dataset_from_index": [int(s) for s in ep_starts], + "dataset_to_index": [int(s + c) for s, c in zip(ep_starts, ep_counts)], + "length": [int(c) for c in ep_counts], + } + episode_ids = split_episode_ids( + total_episodes=len(ep_vals), seed=self._split_seed, val_ratio=self._split_val_ratio, split=self._split + ) + episode_spans, _, _ = build_episode_spans( + episodes=episodes_meta, + episode_ids=episode_ids, + chunk_length=self._chunk_length, + sample_stride=self._sample_stride, + ) + self._episode_records: list[tuple[int, int, int, int]] = [] + self._episode_cum_ends: list[int] = [] + self._num_valid_indices = 0 + if not use_filter_dict: + for episode_id, sample_start, valid_len in episode_spans: + self._episode_records.append((0, sample_start, valid_len, episode_id)) + self._num_valid_indices += valid_len + self._episode_cum_ends.append(self._num_valid_indices) + else: + # Keep-ranges filter — same construction as the base loader. + with open(filter_dict_path) as f: + filter_dict = json.load(f) + for episode_id, sample_start, valid_len in episode_spans: + eid = str(ep_id_str.get(self._ep_index_of[episode_id], "")) + key = ( + f"gs://xembodiment_data/r2d2/r2d2-data-full/{eid}/recordings/" + f"MP4--gs://xembodiment_data/r2d2/r2d2-data-full/{eid}/trajectory.h5" + ) + ranges = filter_dict.get(key) + if ranges is None: + continue + for s, e in ranges: + sub_start = max(s, 0) + sub_end = min(e - self._chunk_length, valid_len) + sub_valid_len = max(0, sub_end - sub_start) + if sub_valid_len > 0: + self._episode_records.append((0, sample_start + sub_start, sub_valid_len, episode_id)) + self._num_valid_indices += sub_valid_len + self._episode_cum_ends.append(self._num_valid_indices) + if self._num_valid_indices == 0: + # Fail here (like the base) rather than as an opaque empty-sampler error: + # the usual cause is a table whose episodes have no episode_id strings. + raise ValueError( + "filter_dict matched no episodes — the composed table's episode_id " + "strings are empty or do not correspond to the filter keys." + ) + + # -- video: lazy per-worker handles into the composed table -- + self._lance_uri = lance_uri + self._table = table + self._decode_device = _resolve_device(decode_device) + self._cache_size = decoder_cache_size + self._storage_options = storage_options + self._perm = None + self._ep_row: dict[int, int] | None = None + self._decoders: dict[int, VideoDecoder] | None = None + self._pending_window: tuple[int, int] | None = None + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + for k in ("_perm", "_ep_row", "_decoders"): + state[k] = None + return state + + # -- labels: assemble the LeRobot-shaped sample dict from the Lance arrays -- + + def _fetch_sample(self, idx: int) -> tuple[str, int, int, dict[str, Any]]: + mode = self._choose_mode() + dataset_idx, row_idx, episode_id, _ = self._resolve_index(idx) + sample: dict[str, Any] = { + c: torch.from_numpy(self._feat[c][row_idx : row_idx + n].copy()).float() + for c, n in self._label_windows.items() + } + sample["task"] = self._tasks[int(self._row_task[row_idx])] + ep_index = self._ep_index_of[episode_id] + self._pending_window = (ep_index, row_idx - self._ep_row_start[ep_index]) + return mode, dataset_idx, row_idx, sample + + # -- video: decode the composed clip window instead of composing 3 views -- + + def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor: + self._ensure_open() + ep_index, offset = self._pending_window + self._ensure_decoders([ep_index]) + idxs = [offset + k for k in range(self._chunk_length + 1)] + frames = self._decoders[ep_index].get_frames_at(indices=idxs).data # (T,C,H,W) uint8 + return frames.to(torch.float32) / 255.0 # [0,1] float, as the base expects + + def _ensure_open(self) -> None: + if self._decoders is not None: + return + tbl = lancedb.connect(self._lance_uri, storage_options=self._storage_options).open_table(self._table) + ep = Permutation.identity(tbl).select_columns(["episode_index"]).with_format("arrow") + rows = ep.__getitems__(list(range(tbl.count_rows()))) + self._ep_row = {int(rows.column("episode_index")[i].as_py()): i for i in range(rows.num_rows)} + self._perm = Permutation.identity(tbl).select_columns(["video_bytes"]).with_format("arrow") + self._decoders = {} + + def _read_clip_bytes(self, rows: list[int]) -> dict[int, bytes]: + # Plain large_binary via the Permutation API. TODO: move to blob-v2 after optimizations. + # take returns rows sorted by offset, so key by row instead of relying on order. + rows = sorted({int(r) for r in rows}) + col = self._perm.__getitems__(rows).column("video_bytes") + return {r: col[i].as_py() for i, r in enumerate(rows)} + + def _build_decoder(self, data: bytes) -> VideoDecoder: + device = str(self._decode_device) if self._decode_device else None + return VideoDecoder(data, seek_mode="approximate", device=device) + + def _ensure_decoders(self, ep_indices: list[int]) -> None: + needed = list(dict.fromkeys(ep_indices)) + needed_set = set(needed) + missing = [e for e in needed if e not in self._decoders] + if not missing: + return + clips = self._read_clip_bytes([self._ep_row[e] for e in missing]) + for e in missing: + while len(self._decoders) >= self._cache_size: + victim = next((k for k in self._decoders if k not in needed_set), None) + if victim is None: + break + self._decoders.pop(victim) + self._decoders[e] = self._build_decoder(clips[self._ep_row[e]]) + + def __getitems__(self, indices: list[int]) -> list[dict[str, Any]]: + # Warm the decoder cache for the whole batch (one batched byte read), then + # let the inherited per-sample path assemble each result. + self._ensure_open() + eps = {self._ep_index_of[self._resolve_index(int(i))[2]] for i in indices} + self._ensure_decoders(list(eps)) + return [self[int(i)] for i in indices] + + +def get_lance_action_droid_sft_dataset( + *, + lance_uri: str, + table: str = "droid_composed", + version: str = _DEFAULT_VERSION, + decode_device: str | None = "cpu", + storage_options: dict | None = None, + fps: float = 15.0, + chunk_length: int = 32, + action_space: str = "joint_pos", + mode: str = "policy", + use_state: bool = True, + action_normalization: str | None = None, + viewpoint: str = "concat_view", + use_filter_dict: bool = False, + filter_dict_path: str | None = None, + resolution: str | int = "256", + max_action_dim: int = 64, + tokenizer_config: Any = None, + cfg_dropout_rate: float = 0.1, + append_viewpoint_info: bool = True, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + append_idle_frames: bool = False, + format_prompt_as_json: bool = False, + iterable_shuffle: bool = False, + episode_shuffle_seed: int = 42, +): + """Lance drop-in for ``get_action_droid_sft_dataset``: same DROID action SFT + stack (``ActionTransformPipeline`` + ``ActionSFTDataset``), reading labels and + pre-composed episodes from LanceDB instead of the raw LeRobot tree.""" + dataset = LanceDROIDComposedDataset( + lance_uri, + table=table, + version=version, + decode_device=decode_device, + storage_options=storage_options, + fps=fps, + chunk_length=chunk_length, + viewpoint=viewpoint, + action_space=action_space, + mode=mode, + use_state=use_state, + action_normalization=action_normalization, + use_filter_dict=use_filter_dict, + filter_dict_path=filter_dict_path, + ) + transform = ActionTransformPipeline( + tokenizer_config=tokenizer_config, + cfg_dropout_rate=cfg_dropout_rate, + max_action_dim=max_action_dim, + append_viewpoint_info=append_viewpoint_info, + append_duration_fps_timestamps=append_duration_fps_timestamps, + append_resolution_info=append_resolution_info, + append_idle_frames=append_idle_frames, + format_prompt_as_json=format_prompt_as_json, + ) + sft = ActionSFTDataset(dataset, transform, resolution) + return ActionIterableShuffleDataset(sft, seed=episode_shuffle_seed) if iterable_shuffle else sft + + +__all__ = [ + "LanceDROIDComposedDataset", + "get_lance_action_droid_sft_dataset", +] diff --git a/cosmos_framework/data/lance/design.md b/cosmos_framework/data/lance/design.md new file mode 100644 index 00000000..ddd43e22 --- /dev/null +++ b/cosmos_framework/data/lance/design.md @@ -0,0 +1,223 @@ +# Design: LanceDB-backed Cosmos Dataloaders + +How the three Lance loaders work and what keeps their output equivalent to the base +loaders. Usage and benchmark results are in the [README](./README.md). + +## Constraints + +1. **Drop-in equivalence** — same samples as the base loader: labels/tokens exact, video + within one offline H.264 re-encode. Wherever possible the loaders reuse the base code + rather than reimplement it, so equivalence is structural. +2. **Per-epoch work moved offline** — the converters do the repeated work (multi-view + compose, per-sample resize, subprocess decode) once at build time; the hot path is a + columnar read + one in-process decode. +3. **Object-store native** — tables readable straight from S3 (selective, parallel reads); + no FUSE, no full downloads. +4. **lancedb-level APIs only** — all reads via the `Permutation` API, no pylance. Video is + plain `large_binary` until the lancedb-level blob API (blob-v2) lands. + +## Shared mechanics + +- **Permutation reads.** One `Permutation.identity(tbl).select_columns([...]).with_format("arrow")` + handle per table: full-column scans at init, point lookups in the hot path. `take` + returns rows **sorted by offset and deduplicated**, so batched reads key results by row + (`_read_clip_bytes` → `{row: bytes}`), never positionally. The equivalence tests use + unsorted/duplicate index lists to pin this contract. +- **Worker safety.** lancedb connections and decoders are not fork/pickle-safe. Every + loader nulls its handles in `__getstate__`; spawn workers reopen lazily + (`_ensure_open`). This also keeps the spawn payload small. +- **In-process decode.** Clips are short mp4s decoded with torchcodec + (`VideoDecoder(bytes, seek_mode="approximate")`). Converters encode all-intra + (`gop=1`), which makes approximate seeking exact and random window reads cheap. A + per-worker LRU (`decoder_cache_size`, default 32) keeps recent clip decoders open; + eviction skips decoders the current batch still needs. `gop` is a build-time size/seek + knob. +- **Batched `__getitems__`.** Two-pass: plan the batch (group requested windows by clip, + record which output slot owns which slice), then decode each clip once and scatter. +- **Lossiness.** The offline re-encode is the only lossy step (~1–2% pixel MAD; the + action gate is 2.5% because the base's decoder backend also differs). Verified + training-irrelevant by the real-model forward-equivalence runs. + +--- + +## Action — `LanceDROIDComposedDataset` + +Base behavior: `DROIDLeRobotDataset` registers LeRobot sources metadata-only, splits +episodes deterministically (`split_episode_ids`), builds a span index +(`build_episode_spans`), and keeps per-shard `LeRobotDataset` readers lazy behind an LRU. +Per sample, `_fetch_sample` returns windowed label features + three camera-view windows +(`delta_timestamps`); `__getitem__` composes the views into one `1.5·h × w` frame +(`_compose_multi_view`) and assembles the action for the chosen action space. Feature +names/flags resolve from a version registry keyed by the root's directory name. + +### Tables (`tools/lance_datagen/build_composed_droid.py`) + +Four tables. `{table}` — one row per **episode** (the video): + +| column | type | description | +| --- | --- | --- | +| `episode_index` | int64 | episode id | +| `ep_start` | int64 | first global frame index (build metadata) | +| `length` | int64 | number of frames (build metadata) | +| `video_bytes` | large_binary | the base's exact composition over the full episode, re-encoded `gop=1` | + +`{table}_frames` — one row per **frame**, dumped verbatim from the base's LeRobot table +(bit-exact roundtrip; feature names store `.` as `__`): + +| column | type | description | +| --- | --- | --- | +| `episode_index`, `task_index` | int64 | frame → episode / task | +| `timestamp` | float64 | frame timestamp | +| `action__joint_position` | fixed_size_list\[7] | commanded joints | +| `action__gripper_position` | float32 | commanded gripper | +| `observation__state__joint_positions` | fixed_size_list\[7] | observed joints | +| `observation__state__gripper_position` | float32 | observed gripper | +| `observation__state__cartesian_position` | fixed_size_list\[6] | EE pose | + +`{table}_tasks` — one row per **task**: + +| column | type | description | +| --- | --- | --- | +| `task_index` | int64 | task id | +| `task` | string | task/caption text | + +`{table}_episodes` — one row per **episode** (keep-ranges filter only): + +| column | type | description | +| --- | --- | --- | +| `episode_index` | int64 | episode id | +| `episode_id` | string | source episode-id string the filter dict keys on | + +`--labels-only` rewrites the label tables against an existing video table. The converter +raises on multi-shard roots (one frames table = one shard). + +### Loader + +Subclasses `DROIDLeRobotDataset`, **bypassing its LeRobot-reading `__init__`**: + +- Takes `lance_uri` (+ `storage_options`) and a `version` (same registry); sets the same + config attributes the base would; loads `{table}_frames` label columns in one + full-column read (~10 MB / 96k frames). +- Split/span index built with the base's **own helpers** (`split_episode_ids` + + `build_episode_spans`) → `split`, `split_seed`, `split_val_ratio`, `sample_stride`, + keep-ranges filter behave identically. An empty filter match raises. +- `_fetch_sample` override returns the same windowed sample dict the LeRobot readers + would (contiguous row slices per the base's `delta_timestamps` plan + task string) — + the inherited `__getitem__` then does actions, captions, gripper flips, idle frames, + and normalization unchanged. +- `_compose_multi_view` override decodes the requested window from the stored composed + clip (uint8 → the `[0,1]` float layout the base expects). +- `get_shuffle_blocks` + the base `ActionIterableShuffleDataset` give the production + episode-shuffle stream; `__getitems__` pre-warms the decoder cache with one batched + byte read. + +Labels are bit-exact for `joint_pos` and `midtrain`, with and without `use_state`; video +< 2.5% pixel MAD. Not supported: image augmentation (applies to raw views before +composition), `max_num_history_actions` (needs pre-window history rows), `val_temp_seg`, +multi-shard roots. `get_lance_action_droid_sft_dataset` mirrors the base factory +(`ActionSFTDataset` + `ActionTransformPipeline`) — the recipe swap is one `_target_` +change. + +--- + +## Vision-SFT — `LanceVisionSFTDataset` + +Base behavior: `SFTDataset` fetches each source clip (S3/local), decodes at native +resolution through an ffmpeg subprocess with a `scale_hw` resize, selects a frame window, +center-crops, picks a caption, post-processes it, and tokenizes. + +### Table (`tools/lance_datagen/build_vision_sft.py`) + +| column | type | description | +| --- | --- | --- | +| `clip_id` | string | `{uuid}_w{window}` | +| `width`, `height` | int64 | original resolution | +| `start_frame`, `end_frame`, `temporal_interval` | int64 | window bounds / stride | +| `enc_h`, `enc_w` | int64 | stored (resized) resolution | +| `fps` | float64 | source fps | +| `caption_json` | string | structured caption (verbatim JSON) or `""` | +| `caption` | string | dense caption fallback | +| `video_bytes` | large_binary | clip resized to the training resolution, `gop=1` | + +One row per clip-window; the resize is the base's exact op (same `scale_hw` ratio; the +center-crop is left to decode time). The metadata pass is the base's own loader +(`_load_sft_metadata_from_s3` + `_flatten_metadata_by_window`), so the duration and +`min_frames` filters match the base population exactly. Windows carrying caption keys +the schema does not persist (`CAPTION_TYPES` styles, `qwen3_32b_rewrite-dense`) are +rejected at build time. + +### Loader + +- Metadata columns read once into `_rows` at init; `video_bytes` fetched per clip via + batched take + the LRU decoder cache. +- `_window_plan` recomputes the base's window arithmetic (`temporal_interval_mode` / + `frame_selection_mode` / `num_video_frames`); frame indices are clamped to the stored + clip the way the base's sequential decode clamps. +- Center crop from `enc_h/enc_w` to the `VIDEO_RES_SIZE_INFO` target; a table built at a + smaller `--resolution` than requested raises. +- Caption pipeline is the base's: `_select_caption`, then the same post-processing order + (`caption_suffix`, CFG dropout, duration/FPS + resolution suffixes for non-structured + captions), then `tokenize_caption` — `text_token_ids` are token-exact for structured + and dense captions. +- Samples the base would skip (short window, no usable caption) return `None` — the + `process_one_sample` contract; the iterable filters them. +- `LanceVisionSFTIterable` + `get_lance_vision_sft_dataset` provide the packing stack's + iterable contract: per-(rank, worker) shard of a seeded shuffle with a + `torch.distributed` fallback, `conditioning_fps` added to match the base sample dict. + +--- + +## VLM — `LanceVLMDataset` / `LanceVLMShuffleScan` + +Base behavior: streams LLaVA-OneVision from the HuggingFace Hub (`streaming=True`) — +sequential shard reads, bounded shuffle buffer, image+conversation filter. Image decode +and tokenization happen downstream in the processor. + +### Table (`convert_llava_to_lance` in `vlm_dataset.py`) + +| column | type | description | +| --- | --- | --- | +| `sample_id` | string | sample id | +| `image_bytes` | large_binary | raw PNG/JPEG bytes, byte-identical to the source | +| `conversations` | string | ShareGPT turns as JSON | + +Records are byte-identical, so everything downstream is unchanged by construction. + +### Loaders + +- **`LanceVLMDataset`** (map-style): row point-lookups; a global shuffle is just a + shuffled index list from the sampler. `__getitems__` keys the sorted-take result by + row and maps back to the requested order (duplicate-safe). +- **`LanceVLMShuffleScan`** (iterable): for S3 — contiguous row-chunks in shuffled chunk + order (reshuffled each pass), sharded per (rank, worker) with a `torch.distributed` + fallback, through a local shuffle buffer. Sequential I/O with decorrelated output — + the same access pattern the HF base gets from shard streaming. +- `get_lance_vlm_dataset` mirrors `get_llava_ov_map`; `n` caps the dataset like the + base's `.select(range(n))`. + +--- + +## Equivalence methodology + +1. **Data-level tests** (`tests/data/lance/`): per-sample comparison against the genuine + base loaders — action labels/captions bit-exact (both action spaces, ± `use_state`), + vision-SFT token-ids exact (structured + dense captions), VLM records byte-identical; + video within the re-encode tolerance. Index lists are unsorted (with duplicates for + VLM) to pin the sorted-take contract. +2. **Real-model forward equivalence** (`benchmarks/lance/forward_equivalence.py`): the + same samples through the real Cosmos3-Nano with fixed weights and seed (`lr=0`), + comparing per-step loss. Base-vs-Lance within the re-encode tolerance (action ≤1.4%, + vision ≤2.1%, VLM exact 0.00%), while different samples differ ~10× more; a + base-vs-base control gives 0.000%. (The action run predates the upstream loader + rewrite; label bit-exactness was re-verified against the rewritten base.) + +## Benchmark methodology + +- The base side is always the **genuine shipped loader**, never a reconstruction. +- S3 regimes give the base a materialization standin (download-then-run) — it has no + native S3 path. +- Access patterns match production: episode-shuffle for action on both sides (the + shipped `ActionIterableShuffleDataset`), not a random sampler. +- VLM throughput is end-to-end (image decode + tokenize) since the loaders emit raw + records. +- Memory is per-worker under spawn (PSS for fork/COW fairness), one side per process. diff --git a/cosmos_framework/data/lance/vision_sft_dataset.py b/cosmos_framework/data/lance/vision_sft_dataset.py new file mode 100644 index 00000000..8c148342 --- /dev/null +++ b/cosmos_framework/data/lance/vision_sft_dataset.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""LanceDB-backed local vision-SFT (video+caption) dataset. + +Alternative to SFTDataset that decodes pre-resized, short-GOP per-clip mp4s from LanceDB. +Reuses the base's caption selection and tokenization logic. +""" + +from __future__ import annotations + +import json +import random +from typing import Any, Optional + +import lancedb +import numpy as np +import torch +from lancedb.permutation import Permutation +from torchcodec.decoders import VideoDecoder +from transformers import AutoTokenizer + +from cosmos_framework.data.generator.local_datasets.helper import get_aspect_ratio +from cosmos_framework.data.generator.local_datasets.sft_dataset import ( + _DURATION_TEMPLATE, + _RESOLUTION_TEMPLATE, + _select_caption, +) +from cosmos_framework.data.generator.sequence_packing.modalities import add_special_tokens +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.utils import log + +_MAX_CAPTION_TOKENS = 1024 +_META_COLS = [ + "clip_id", + "width", + "height", + "start_frame", + "end_frame", + "temporal_interval", + "enc_h", + "enc_w", + "fps", + "caption_json", + "caption", +] + + +def _resolve_device(device: str | None) -> torch.device | None: + if device == "auto": + return torch.device("cuda") if torch.cuda.is_available() else None + if device in (None, "cpu"): + return None + return torch.device(device) + + +class LanceVisionSFTDataset(torch.utils.data.Dataset): + """Map-style local vision-SFT loader backed by LanceDB. + + Decodes pre-resized clips in-process, avoiding ffmpeg subprocess overhead. + """ + + def __init__( + self, + lance_uri: str, + *, + table: str = "vision_sft", + resolution: str = "256", + num_video_frames: int = 16, + temporal_interval_mode: str = "entire_chunk", + frame_selection_mode: str = "first", + temporal_compression_factor: int = 4, + tokenizer: Optional[Any] = None, + tokenizer_name: str = "Qwen/Qwen2.5-7B", + use_system_prompt: bool = False, + max_caption_tokens: int = _MAX_CAPTION_TOKENS, + cfg_dropout_rate: float = 0.0, + cfg_dropout_keep_metadata: bool = False, + caption_suffix: str = "", + conditioning_fps: float = 24, + conditioning_fps_noise_std: float = 0.0, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + decode_device: str | None = "cpu", + decoder_cache_size: int = 32, + storage_options: dict | None = None, + ) -> None: + assert temporal_interval_mode in ("force_one", "max_30fps", "entire_chunk") + assert frame_selection_mode in ("center", "first", "random") + self._lance_uri = lance_uri + self._table = table + self._resolution_str = resolution + self.num_video_frames = num_video_frames + self.temporal_interval_mode = temporal_interval_mode + self.frame_selection_mode = frame_selection_mode + self.temporal_compression_factor = temporal_compression_factor + self.use_system_prompt = use_system_prompt + self.max_caption_tokens = max_caption_tokens + self.cfg_dropout_rate = cfg_dropout_rate + self.cfg_dropout_keep_metadata = cfg_dropout_keep_metadata + self.caption_suffix = caption_suffix.strip() + self.conditioning_fps = conditioning_fps + self.conditioning_fps_noise_std = conditioning_fps_noise_std + self.append_duration_fps_timestamps = append_duration_fps_timestamps + self.append_resolution_info = append_resolution_info + self.tokenizer_name = tokenizer_name + self._decode_device = _resolve_device(decode_device) + self._cache_size = decoder_cache_size + self._storage_options = storage_options + self._tokenizer = tokenizer + + self._perm = None + self._rows: list[dict] | None = None + self._decoders: dict[int, VideoDecoder] | None = None + + self._length = lancedb.connect(lance_uri, storage_options=storage_options).open_table(table).count_rows() + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + for k in ("_perm", "_rows", "_decoders", "_tokenizer"): + state[k] = None + return state + + def _ensure_open(self) -> None: + if self._decoders is not None: + return + tbl = lancedb.connect(self._lance_uri, storage_options=self._storage_options).open_table(self._table) + meta_perm = Permutation.identity(tbl).select_columns(_META_COLS).with_format("arrow") + self._rows = meta_perm.__getitems__(list(range(tbl.count_rows()))).to_pylist() + self._perm = Permutation.identity(tbl).select_columns(["video_bytes"]).with_format("arrow") + self._decoders = {} + + def _read_clip_bytes(self, rows: list[int]) -> dict[int, bytes]: + # Plain large_binary via the Permutation API. TODO: move to blob-v2 after optimizations. + # take returns rows sorted by offset, so key by row instead of relying on order. + rows = sorted({int(r) for r in rows}) + col = self._perm.__getitems__(rows).column("video_bytes") + return {r: col[i].as_py() for i, r in enumerate(rows)} + + def _build_decoder(self, data: bytes) -> VideoDecoder: + device = str(self._decode_device) if self._decode_device else None + return VideoDecoder(data, seek_mode="approximate", device=device) + + def _ensure_decoders(self, rows: list[int]) -> None: + needed = list(dict.fromkeys(rows)) + needed_set = set(needed) + missing = [r for r in needed if r not in self._decoders] + if not missing: + return + clips = self._read_clip_bytes(missing) + for r in missing: + while len(self._decoders) >= self._cache_size: + victim = next((k for k in self._decoders if k not in needed_set), None) + if victim is None: + break + self._decoders.pop(victim) + self._decoders[r] = self._build_decoder(clips[r]) + + def _ensure_tokenizer(self): + if self._tokenizer is None: + tok = AutoTokenizer.from_pretrained(self.tokenizer_name) + tok, _ = add_special_tokens(tok) + self._tokenizer = tok + return self._tokenizer + + def _decoder(self, row: int) -> VideoDecoder: + d = self._decoders.get(row) + if d is None: + d = self._build_decoder(self._read_clip_bytes([row])[row]) + if len(self._decoders) >= self._cache_size: + self._decoders.pop(next(iter(self._decoders))) + self._decoders[row] = d + return d + + def __len__(self) -> int: + return self._length + + skip_tokenize: bool = False + + def _tokenize(self, caption: str) -> list[int]: + if self.skip_tokenize: + return [] + ids = tokenize_caption( + caption, self._ensure_tokenizer(), is_video=True, use_system_prompt=self.use_system_prompt + ) + return ids[: self.max_caption_tokens] + + def _window_plan(self, meta: dict) -> tuple[int, int, int] | None: + window_start, window_end = meta["start_frame"], meta["end_frame"] + clip_total = meta["_clip_total"] + actual_end = min(window_end, clip_total - 1) + frames_in_window = actual_end - window_start + 1 + if self.num_video_frames == -1: + return window_start, actual_end, meta["temporal_interval"] + if frames_in_window < self.num_video_frames: + # base behavior: warn and skip the sample, never crash the worker + log.warning(f"Not enough frames in window for {meta['clip_id']}. Skipping sample.") + return None + + if self.temporal_interval_mode == "force_one": + temporal_interval = 1 + elif self.temporal_interval_mode == "max_30fps": + temporal_interval = max(1, int(meta["fps"] / 30.0)) + else: + temporal_interval = max(1, frames_in_window // self.num_video_frames) + + num_before = (self.num_video_frames - 1) * temporal_interval + 1 + if self.frame_selection_mode == "first": + start_frame = window_start + elif self.frame_selection_mode == "center": + start_frame = window_start + (frames_in_window - num_before) // 2 + else: + start_frame = window_start + random.randint(0, max(0, frames_in_window - num_before)) + return start_frame, start_frame + num_before - 1, temporal_interval + + def _finalize_caption( + self, + caption: str, + used_structured_json: bool, + *, + num_decoded_frames: int, + fps: float, + target_h: int, + target_w: int, + ) -> str: + """The base's caption post-processing (suffix, CFG dropout, duration/resolution + conditioning text), applied in the same order as SFTDataset.process_one_sample.""" + cond_fps = fps if self.conditioning_fps < 0 else self.conditioning_fps + if self.conditioning_fps_noise_std > 0: + cond_fps = cond_fps * float(np.exp(np.random.randn() * self.conditioning_fps_noise_std)) + if self.caption_suffix and not used_structured_json: + caption = (caption + " " + self.caption_suffix).strip() + if self.cfg_dropout_keep_metadata and self.cfg_dropout_rate > 0 and random.random() < self.cfg_dropout_rate: + caption = "" + if self.append_duration_fps_timestamps and not used_structured_json: + caption = caption + " " + _DURATION_TEMPLATE.format(duration=num_decoded_frames / cond_fps, fps=cond_fps) + if self.append_resolution_info and not used_structured_json: + caption = caption + " " + _RESOLUTION_TEMPLATE.format(height=target_h, width=target_w) + caption = caption.strip() + if not self.cfg_dropout_keep_metadata and self.cfg_dropout_rate > 0 and random.random() < self.cfg_dropout_rate: + caption = "" + return caption + + def __getitem__(self, idx: int) -> dict[str, Any] | None: + return self.__getitems__([int(idx)])[0] + + def __getitems__(self, indices: list[int]) -> list[dict[str, Any] | None]: + """Batched fetch. A slot is ``None`` when the base loader would skip the sample + (short window / no usable caption) — the same contract as ``process_one_sample``.""" + self._ensure_open() + n = len(indices) + self._ensure_decoders([int(i) for i in indices]) + + specs: list[dict | None] = [] + plan: dict[int, dict] = {} + for sp, idx in enumerate(indices): + row = int(idx) + r = self._rows[row] + dec = self._decoder(row) + clip_total = dec.metadata.num_frames + r = {**r, "_clip_total": clip_total} + wp = self._window_plan(r) + sel = _select_caption(self._window_dict(r)) + if sel is None: + log.warning(f"No known caption key found for sample {r['clip_id']}. Skipping sample.") + if wp is None or sel is None: + specs.append(None) + continue + start_frame, end_frame, ti = wp + # Clamp to the stored clip exactly as the base's sequential decode loop does + # (out-of-range indices simply don't yield frames there). + frame_idx = [i for i in range(start_frame, end_frame + 1, ti) if 0 <= i < clip_total] + + target_w, target_h = self._target_size(r) # (w, h) per VIDEO_RES_SIZE_INFO + if r["enc_h"] < target_h or r["enc_w"] < target_w: + raise ValueError( + f"stored clip {r['clip_id']} is {r['enc_h']}x{r['enc_w']} but resolution=" + f"{self._resolution_str!r} needs {target_h}x{target_w}: the table was built " + f"at a smaller --resolution than requested" + ) + crop_y = round((r["enc_h"] - target_h) / 2) + crop_x = round((r["enc_w"] - target_w) / 2) + caption_key, caption, used_structured_json = sel + cid = r["clip_id"] + win_idx = int(cid.rsplit("_w", 1)[1]) if "_w" in cid and cid.rsplit("_w", 1)[1].isdigit() else 0 + specs.append( + { + "row": row, + "clip_id": cid, + "fps": r["fps"], + "clip_total": clip_total, + "win_idx": win_idx, + "temporal_interval": ti, + "start_frame": start_frame, + "end_frame": end_frame, + "crop": (crop_y, crop_x, target_h, target_w), + "caption": caption, + "caption_key": caption_key, + "used_structured_json": used_structured_json, + } + ) + e = plan.setdefault(row, {"frames": [], "owners": []}) + lo = len(e["frames"]) + e["frames"].extend(frame_idx) + e["owners"].append((sp, lo, lo + len(frame_idx))) + + decoded: list[torch.Tensor | None] = [None] * n + for row, e in plan.items(): + dec = self._decoder(row) + frames = dec.get_frames_at(indices=e["frames"]).data + for sp, lo, hi in e["owners"]: + decoded[sp] = frames[lo:hi] + + results: list[dict[str, Any] | None] = [] + for sp in range(n): + s = specs[sp] + if s is None: + results.append(None) + continue + vid = decoded[sp] + cy, cx, th, tw = s["crop"] + t = vid.shape[0] + target_t = (t - 1) // self.temporal_compression_factor * self.temporal_compression_factor + 1 + vid = vid[:target_t, :, cy : cy + th, cx : cx + tw] + video = vid.permute(1, 0, 2, 3).contiguous().to(torch.uint8) + + caption = self._finalize_caption( + s["caption"], + s["used_structured_json"], + num_decoded_frames=video.shape[1], + fps=s["fps"], + target_h=th, + target_w=tw, + ) + text_ids = self._tokenize(caption) + image_size = torch.tensor([th, tw, th, tw], dtype=torch.float32) + padding_mask = torch.zeros((1, th, tw), dtype=torch.float32) + results.append( + dict( + __key__=s["clip_id"], + __url__=s["clip_id"], + fps=s["fps"], + n_orig_video_frames=s["clip_total"], + chunk_index=s["win_idx"], + frame_start=s["start_frame"], + frame_end=s["end_frame"], + num_frames=video.shape[1], + video=video, + num_multiplier=s["temporal_interval"], + padding_mask=padding_mask, + image_size=image_size, + ai_caption=caption, + sampled_caption_style=s["caption_key"], + text_token_ids=torch.tensor(text_ids, dtype=torch.long), + ) + ) + return results + + def _target_size(self, r: dict) -> tuple[int, int]: + ar = get_aspect_ratio(r["width"], r["height"]) + return VIDEO_RES_SIZE_INFO[self._resolution_str][ar] + + def _window_dict(self, r: dict) -> dict: + w: dict[str, Any] = {} + if r.get("caption_json"): + w["caption_json"] = json.loads(r["caption_json"]) + if r.get("caption"): + w["caption"] = r["caption"] + return w + + +class LanceVisionSFTIterable(torch.utils.data.IterableDataset): + """Streams clip-windows from LanceVisionSFTDataset with per-(rank, worker) shuffle. + + Mirrors SFTDataset's iterable/self-sharding contract so it drops into the + training packing stack; adds conditioning_fps to match the SFTDataset sample. + """ + + def __init__(self, dataset: LanceVisionSFTDataset, conditioning_fps: float = 24.0, seed: int = 42): + super().__init__() + self._ds = dataset + self._cond_fps = float(conditioning_fps) + self._seed = int(seed) + # Set by RankPartitionedDataLoader; None falls back to torch.distributed + # (the same contract as the base SFTDataset.__iter__). + self.shard_world_size = None + self.shard_rank = None + + def __len__(self) -> int: + return len(self._ds) + + def _shard(self) -> tuple[int, int]: + ws, rk = self.shard_world_size, self.shard_rank + if ws is None or rk is None: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + ws, rk = torch.distributed.get_world_size(), torch.distributed.get_rank() + else: + ws, rk = 1, 0 + return int(ws), int(rk) + + def __iter__(self): + info = torch.utils.data.get_worker_info() + wid = info.id if info is not None else 0 + nw = info.num_workers if info is not None else 1 + ws, rk = self._shard() + shard = rk * nw + wid + total = max(1, ws * nw) + n = len(self._ds) + epoch = 0 + while True: + g = torch.Generator().manual_seed(self._seed + epoch) + for i in torch.randperm(n, generator=g).tolist()[shard::total]: + s = self._ds[i] + if s is None: # base contract: skipped sample + continue + s["conditioning_fps"] = self._cond_fps + yield s + epoch += 1 + + +def get_lance_vision_sft_dataset( + *, + lance_uri: str, + table: str = "vision_sft", + resolution: str = "256", + num_video_frames: int = 16, + frame_selection_mode: str = "first", + temporal_interval_mode: str = "entire_chunk", + tokenizer_config: Any = None, + cfg_dropout_rate: float = 0.1, + cfg_dropout_keep_metadata: bool = False, + caption_suffix: str = "", + conditioning_fps: float = 24.0, + conditioning_fps_noise_std: float = 0.0, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + decode_device: str | None = "cpu", + seed: int = 42, +) -> LanceVisionSFTIterable: + """Build the iterable Lance vision-SFT dataset for the training packing stack. + + Caption knobs default to the base ``get_sft_dataset`` factory's values.""" + tok = getattr(tokenizer_config, "tokenizer", None) if tokenizer_config is not None else None + ds = LanceVisionSFTDataset( + lance_uri, + table=table, + resolution=resolution, + num_video_frames=num_video_frames, + frame_selection_mode=frame_selection_mode, + temporal_interval_mode=temporal_interval_mode, + tokenizer=tok, + cfg_dropout_rate=cfg_dropout_rate, + cfg_dropout_keep_metadata=cfg_dropout_keep_metadata, + caption_suffix=caption_suffix, + conditioning_fps=conditioning_fps, + conditioning_fps_noise_std=conditioning_fps_noise_std, + append_duration_fps_timestamps=append_duration_fps_timestamps, + append_resolution_info=append_resolution_info, + decode_device=decode_device, + ) + return LanceVisionSFTIterable(ds, conditioning_fps=conditioning_fps, seed=seed) + + +__all__ = ["LanceVisionSFTDataset", "LanceVisionSFTIterable", "get_lance_vision_sft_dataset"] diff --git a/cosmos_framework/data/lance/vlm_dataset.py b/cosmos_framework/data/lance/vlm_dataset.py new file mode 100644 index 00000000..bd4c430f --- /dev/null +++ b/cosmos_framework/data/lance/vlm_dataset.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""LanceDB-backed VLM (LLaVA-OneVision) dataset. + +Row-level random access and global shuffle for VLM datasets. +Drop-in replacement for HF streaming or WebDataset sources. +""" + +from __future__ import annotations + +import io +import json +import random +from typing import Any + +import lancedb +import pyarrow as pa +import torch +from lancedb.permutation import Permutation + +_COLS = ["sample_id", "image_bytes", "conversations"] +_SCHEMA = pa.schema( + [ + pa.field("sample_id", pa.string()), + pa.field("image_bytes", pa.large_binary()), + pa.field("conversations", pa.string()), + ] +) + + +def _record_batches(hf_dataset, batch_rows: int = 512): + schema = _SCHEMA + ids, imgs, convs = [], [], [] + for i, rec in enumerate(hf_dataset): + img = rec.get("image") + if isinstance(img, dict): + raw = img.get("bytes") or b"" + elif img is not None: + buf = io.BytesIO() + img.save(buf, format=img.format or "PNG") + raw = buf.getvalue() + else: + raw = b"" + ids.append(str(rec.get("id", i))) + imgs.append(raw) + convs.append(json.dumps(rec.get("conversations") or [])) + if len(ids) >= batch_rows: + yield pa.RecordBatch.from_arrays( + [pa.array(ids, pa.string()), pa.array(imgs, pa.large_binary()), pa.array(convs, pa.string())], + schema=schema, + ) + ids, imgs, convs = [], [], [] + if ids: + yield pa.RecordBatch.from_arrays( + [pa.array(ids, pa.string()), pa.array(imgs, pa.large_binary()), pa.array(convs, pa.string())], schema=schema + ) + + +def convert_llava_to_lance(hf_dataset, uri: str, table_name: str = "llava") -> str: + reader = pa.RecordBatchReader.from_batches(_SCHEMA, _record_batches(hf_dataset)) + db = lancedb.connect(uri) + if table_name in db.table_names(): + db.drop_table(table_name) + db.create_table(table_name, data=reader, schema=_SCHEMA) + return table_name + + +class LanceVLMDataset(torch.utils.data.Dataset): + """Map-style LLaVA-OneVision source backed by LanceDB.""" + + def __init__(self, uri: str, table_name: str = "llava", storage_options: dict | None = None): + self.uri = uri + self.table_name = table_name + self.storage_options = storage_options + self._perm = None + self.length = self._connect().open_table(table_name).count_rows() + + def _connect(self): + return lancedb.connect(self.uri, storage_options=self.storage_options) + + def __len__(self) -> int: + return self.length + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state["_perm"] = None + return state + + def _ensure_open(self) -> None: + if self._perm is None: + db = self._connect() + table = db.open_table(self.table_name) + self._perm = Permutation.identity(table).select_columns(_COLS).with_format("arrow") + + def _row_to_item(self, batch: pa.RecordBatch, i: int) -> dict[str, Any]: + return { + "id": batch.column("sample_id")[i].as_py(), + "image": {"bytes": batch.column("image_bytes")[i].as_py()}, + "conversations": json.loads(batch.column("conversations")[i].as_py()), + } + + def __getitem__(self, idx: int) -> dict[str, Any]: + self._ensure_open() + return self._row_to_item(self._perm.__getitems__([int(idx)]), 0) + + def __getitems__(self, indices: list[int]) -> list[dict[str, Any]]: + self._ensure_open() + # take returns rows sorted by offset (and deduplicated), so key results by + # row and map back to the requested order — never zip positionally. + rows = sorted({int(i) for i in indices}) + batch = self._perm.__getitems__(rows) + by_row = {r: self._row_to_item(batch, i) for i, r in enumerate(rows)} + return [dict(by_row[int(i)]) for i in indices] + + +class LanceVLMShuffleScan(torch.utils.data.IterableDataset): + """Chunked-shuffle scan over a Lance table for efficient S3 training. + + Permutation API only (no pylance): shuffle the order of contiguous row-chunks, + read each chunk as a columnar range (sequential -> S3-friendly), and emit + through a local shuffle buffer. + """ + + def __init__( + self, + uri: str, + table_name: str = "llava", + storage_options: dict | None = None, + buffer_size: int = 1000, + batch_size: int = 256, + seed: int = 42, + ): + self.uri = uri + self.table_name = table_name + self.storage_options = storage_options + self.buffer_size = buffer_size + self.batch_size = batch_size + self.seed = seed + self._perm = None + self._epoch = 0 + # Set by RankPartitionedDataLoader; None falls back to torch.distributed. + self.shard_world_size = None + self.shard_rank = None + self.length = self._open_table().count_rows() + + def _open_table(self): + return lancedb.connect(self.uri, storage_options=self.storage_options).open_table(self.table_name) + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state["_perm"] = None + return state + + def _ensure_perm(self): + if self._perm is None: + self._perm = Permutation.identity(self._open_table()).select_columns(_COLS).with_format("arrow") + return self._perm + + def __len__(self) -> int: + return self.length + + def __iter__(self): + info = torch.utils.data.get_worker_info() + wid, nw = (info.id, info.num_workers) if info else (0, 1) + ws, rk = self.shard_world_size, self.shard_rank + if ws is None or rk is None: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + ws, rk = torch.distributed.get_world_size(), torch.distributed.get_rank() + else: + ws, rk = 1, 0 + shard = int(rk) * nw + wid + total = max(1, int(ws) * nw) + perm = self._ensure_perm() + chunks = [(s, min(s + self.batch_size, self.length)) for s in range(0, self.length, self.batch_size)] + epoch, self._epoch = self._epoch, self._epoch + 1 # reshuffle each pass + rng = random.Random(self.seed + epoch) + rng.shuffle(chunks) + buf = [] + for start, end in chunks[shard::total]: + batch = perm.__getitems__(list(range(start, end))) + ids = batch.column("sample_id").to_pylist() + imgs = batch.column("image_bytes").to_pylist() + convs = batch.column("conversations").to_pylist() + for sid, raw, cv in zip(ids, imgs, convs): + buf.append({"id": sid, "image": {"bytes": raw}, "conversations": json.loads(cv)}) + if len(buf) >= self.buffer_size: + yield buf.pop(rng.randrange(len(buf))) + rng.shuffle(buf) + yield from buf + + +def get_lance_vlm_dataset( + *, + uri: str, + table_name: str = "llava", + storage_options: dict | None = None, + subset: str | None = None, + split: str | None = None, + n: int | None = None, +): + """Lance drop-in for ``get_llava_ov_map``: the same map-style image+conversation + records, read from LanceDB. ``n`` caps the dataset to the first ``n`` rows like + the base's ``.select(range(n))``; ``subset``/``split`` are accepted for + signature-compatibility but must match what the table was built from (the + conversion bakes them in).""" + ds = LanceVLMDataset(uri, table_name=table_name, storage_options=storage_options) + if n is not None: + ds.length = min(ds.length, int(n)) + return ds + + +__all__ = [ + "LanceVLMDataset", + "LanceVLMShuffleScan", + "convert_llava_to_lance", + "get_lance_vlm_dataset", +] diff --git a/tests/data/lance/test_action.py b/tests/data/lance/test_action.py new file mode 100644 index 00000000..95c83076 --- /dev/null +++ b/tests/data/lance/test_action.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Equivalence test for the composed Action (DROID) loader vs the base DROIDLeRobotDataset. + +Labels (action/pose/caption/idle) are bit-exact; video is within one offline H.264 re-encode. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.lance import LanceDROIDComposedDataset + +AROOT = os.environ.get("DROID_LEROBOT_ROOT") # versioned dir, e.g. .../droid_plus_lerobot_320x180_20260406 +ACOMP = os.environ.get("DROID_COMPOSED_LANCE_URI") + +pytestmark = pytest.mark.skipif( + not (AROOT and ACOMP and os.path.isdir(AROOT)), reason="set DROID_LEROBOT_ROOT and DROID_COMPOSED_LANCE_URI" +) + +_IDXS = [17000, 1, 26000, 0, 123, 5000] # unsorted: batched take must map back to the right episode + + +@pytest.fixture(autouse=True) +def _hf_offline(monkeypatch): + # the base loader sets HF_HUB_OFFLINE itself; pre-set it so the env guard stays clean + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + + +@pytest.mark.parametrize("action_space,use_state", [("joint_pos", True), ("midtrain", False), ("midtrain", True)]) +def test_action_composed(action_space, use_state): + kw = dict(action_space=action_space, use_state=use_state, mode="policy", chunk_length=16) + base = DROIDLeRobotDataset(root=AROOT, use_success_only=True, **kw) + lance = LanceDROIDComposedDataset(ACOMP, decode_device="cpu", **kw) + assert len(base) == len(lance) + idxs = [i for i in _IDXS if i < len(base)] + batch = lance.__getitems__(idxs) + for j, i in enumerate(idxs): + b, l = base[i], batch[j] + assert torch.equal(b["action"], l["action"]) # labels bit-exact + assert b["ai_caption"] == l["ai_caption"] + assert torch.equal(b["idle_frames"], l["idle_frames"]) + if "initial_pose" in b: + assert torch.equal(b["initial_pose"], l["initial_pose"]) + mad = (b["video"].float() - l["video"].float()).abs().mean().item() / 255.0 + # One offline H.264 re-encode + the base's own decoder backend difference. + assert mad < 0.025 diff --git a/tests/data/lance/test_vision_sft.py b/tests/data/lance/test_vision_sft.py new file mode 100644 index 00000000..deaa5526 --- /dev/null +++ b/tests/data/lance/test_vision_sft.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Equivalence test for the Vision-SFT loader vs the genuine SFTDataset.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +from transformers import AutoTokenizer + +from cosmos_framework.data.generator.local_datasets.sft_dataset import ( + SFTDataset, + _flatten_metadata_by_window, + _load_sft_metadata_from_s3, +) +from cosmos_framework.data.lance import LanceVisionSFTDataset + +JSONL = os.environ.get("BRIDGE_JSONL") +URI = os.environ.get("VISION_SFT_LANCE_URI") + +pytestmark = pytest.mark.skipif( + not (JSONL and URI and os.path.isfile(JSONL)), reason="set BRIDGE_JSONL and VISION_SFT_LANCE_URI" +) + +_VKW = dict(num_video_frames=16, frame_selection_mode="first", temporal_interval_mode="entire_chunk") + + +@pytest.fixture(scope="module", autouse=True) +def _hf_online(): + # the action base flips HF Hub offline process-wide; this module loads a tokenizer from the hub cache + import huggingface_hub.constants as hfc + + mp = pytest.MonkeyPatch() + mp.setattr(hfc, "HF_HUB_OFFLINE", False) + mp.delenv("HF_HUB_OFFLINE", raising=False) + yield + mp.undo() + + +@pytest.fixture(scope="module") +def base_and_metas(): + # the same metadata load + per-window flattening the converter uses (min_frames=61) + metas = _flatten_metadata_by_window(_load_sft_metadata_from_s3(None, JSONL, min_frames=61)) + base_dir = os.path.dirname(os.path.abspath(JSONL)) + for m in metas: + vp = m["vision_path"] + m["vision_path"] = vp if ("://" in vp or vp.startswith("/")) else os.path.join(base_dir, vp) + tok_cfg = SimpleNamespace(tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")) + ds = SFTDataset( + metadata=metas, + num_video_frames=16, + resolution="256", + s3_credentials={}, + frame_selection_mode="first", + temporal_interval_mode="entire_chunk", + tokenizer_config=tok_cfg, + cfg_dropout_rate=0.0, + ) + ds.s3_client = None + return ds, metas + + +def test_vision_sft(base_and_metas): + base, metas = base_and_metas + lance = LanceVisionSFTDataset(URI, table="vision_sft", decode_device="cpu", **_VKW) + assert len(metas) == len(lance) + idxs = [i for i in [50, 1, 123, 0, 17] if i < len(lance)] # unsorted: batched take must map back to the right clip + batch = lance.__getitems__(idxs) + for j, i in enumerate(idxs): + ref, l = base.process_one_sample(metas[i]), batch[j] + assert torch.equal(ref["text_token_ids"], l["text_token_ids"]) + assert ref["ai_caption"] == l["ai_caption"] + mad = (ref["video"].float() - l["video"].float()).abs().mean().item() / 255.0 + assert mad < 0.02 + + +def test_vision_sft_dense_caption(base_and_metas): + """Dense (non-structured) captions get the base's duration/resolution suffixes.""" + base, metas = base_and_metas + lance = LanceVisionSFTDataset(URI, table="vision_sft", decode_device="cpu", **_VKW) + lance._ensure_open() + for i in [0, 7]: + meta = { + **metas[i], + "t2w_windows": [{k: v for k, v in metas[i]["t2w_windows"][0].items() if k != "caption_json"}], + } + lance._rows[i] = {**lance._rows[i], "caption_json": ""} + ref, l = base.process_one_sample(meta), lance[i] + assert "seconds long" in ref["ai_caption"] # the dense path really appends the suffixes + assert ref["ai_caption"] == l["ai_caption"] + assert torch.equal(ref["text_token_ids"], l["text_token_ids"]) diff --git a/tests/data/lance/test_vlm.py b/tests/data/lance/test_vlm.py new file mode 100644 index 00000000..d8295fb6 --- /dev/null +++ b/tests/data/lance/test_vlm.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Equivalence test for the VLM (LLaVA-OneVision) loader vs the base HF stream.""" + +from __future__ import annotations + +import io +import os +import tempfile + +import pytest +from datasets import load_dataset + +from cosmos_framework.data.lance.vlm_dataset import LanceVLMDataset, convert_llava_to_lance + +pytestmark = pytest.mark.skipif( + not (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")), reason="set HF_TOKEN" +) + + +@pytest.fixture(scope="module", autouse=True) +def _hf_online(): + # the action base flips HF Hub offline process-wide; this module streams from the hub + import huggingface_hub.constants as hfc + + mp = pytest.MonkeyPatch() + mp.setattr(hfc, "HF_HUB_OFFLINE", False) + mp.delenv("HF_HUB_OFFLINE", raising=False) + yield + mp.undo() + + +def _norm_image_bytes(rec): + img = rec.get("image") + if isinstance(img, dict): + return img.get("bytes") or b"" + if img is not None: + buf = io.BytesIO() + img.save(buf, format=img.format or "PNG") + return buf.getvalue() + return b"" + + +def test_vlm(): + subset = os.environ.get("LLAVA_SUBSET", "figureqa(cauldron,llava_format)") + stream = load_dataset("lmms-lab/LLaVA-OneVision-Data", name=subset, split="train", streaming=True) + stream = stream.filter(lambda x: x.get("image") is not None and len(x.get("conversations") or []) >= 2) + base = [] + for rec in stream: + base.append(rec) + if len(base) >= 8: + break + with tempfile.TemporaryDirectory() as tmp: + convert_llava_to_lance(iter(base), tmp, table_name="llava") + lance = LanceVLMDataset(tmp, table_name="llava") + assert len(lance) == len(base) + # unsorted + duplicate indices: batched take must map back to the requested order + idxs = [5, 1, 5, 0, 7, 3] + batch = lance.__getitems__(idxs) + assert len(batch) == len(idxs) + for i, l in zip(idxs, batch): + assert l["conversations"] == (base[i].get("conversations") or []) + assert l["image"]["bytes"] == _norm_image_bytes(base[i]) diff --git a/tools/lance_datagen/build_composed_droid.py b/tools/lance_datagen/build_composed_droid.py new file mode 100644 index 00000000..8d52fa83 --- /dev/null +++ b/tools/lance_datagen/build_composed_droid.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Build a training-optimized DROID representation for LanceDB (video + labels). + +For each episode, compose the 3 camera views EXACTLY as the base loader does +(wrist on top; the two exteriors resized to half and concatenated on the +bottom), then re-encode that single composed clip with a tiny GOP (all-intra by +default) and store it as one per-episode large_binary row. + +Alongside the video table, three label tables are written so the Lance loader +needs no LeRobot tree at train time: + {table}_frames — per-frame labels (episode/task/timestamp + action & state + features), dumped verbatim from the base's LeRobot table + {table}_tasks — task_index -> task string + {table}_episodes — episode_index -> episode_id (for keep-ranges filtering) + +Why: the base loader decodes 3 full-resolution views + resizes + concatenates +*per sample*. Decoding one pre-composed, pre-resized, short-GOP clip is far less +work — fewer pixels, one stream, no resize/concat, and short-GOP makes random +window seeks cheap. Still fully video-encoded (no per-frame JPEG / disk blowup). +The composition is byte-for-byte the base's; only the H.264 re-encode is lossy — +labels roundtrip bit-exact. + +``--labels-only`` (re)writes just the label tables against an existing video table. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import tempfile + +import lancedb +import numpy as np +import pyarrow as pa +import torch +from torchcodec.decoders import VideoDecoder + +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset + +# Every per-frame feature column either action space reads ('.' -> '__' in Lance). +FEATURE_COLUMNS = [ + "action.joint_position", + "action.gripper_position", + "observation.state.joint_positions", + "observation.state.gripper_position", + "observation.state.cartesian_position", +] + + +def lance_col(name: str) -> str: + return name.replace(".", "__") + + +def _encode(frames_thwc_u8: np.ndarray, fps: int, gop: int) -> bytes: + """Raw RGB frames -> H.264 mp4 bytes via ffmpeg (short GOP, faststart). + + mp4+faststart needs seekable output, so encode to a temp file then read.""" + t, h, w, _ = frames_thwc_u8.shape + fd, path = tempfile.mkstemp(suffix=".mp4") + os.close(fd) + try: + cmd = [ + "ffmpeg", "-y", "-loglevel", "error", + "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(fps), "-i", "pipe:0", + "-c:v", "libx264", "-preset", "veryfast", "-g", str(gop), "-keyint_min", str(gop), + "-pix_fmt", "yuv420p", "-movflags", "+faststart", path, + ] # fmt: skip + subprocess.run(cmd, input=frames_thwc_u8.tobytes(), stdout=subprocess.DEVNULL, check=True) + with open(path, "rb") as fh: + return fh.read() + finally: + os.unlink(path) + + +def _feature_array(a: np.ndarray) -> pa.Array: + if a.ndim == 2: + return pa.FixedSizeListArray.from_arrays(pa.array(a.ravel().astype(np.float32)), a.shape[1]) + return pa.array(a.astype(np.float32)) + + +def _replace(db: lancedb.DBConnection, name: str, data, schema: pa.Schema) -> None: + if name in db.table_names(): + db.drop_table(name) + db.create_table(name, data=data, schema=schema) + + +def _build_base(root: str) -> DROIDLeRobotDataset: + # split="full" + joint_pos/use_state registers every episode and label column. + base = DROIDLeRobotDataset( + root=root, + split="full", + use_success_only=True, + action_space="joint_pos", + use_state=True, + mode="policy", + chunk_length=16, + ) + if len(base._datasets) != 1: + raise NotImplementedError( + f"root registered {len(base._datasets)} LeRobot shards; this converter (and the " + "Lance loader's single frames table) currently supports single-shard roots only." + ) + return base + + +def write_label_tables(db: lancedb.DBConnection, table: str, base: DROIDLeRobotDataset) -> None: + """Dump the base's LeRobot label table verbatim (bit-exact roundtrip).""" + lr = base._get_dataset(0) + cols = lr.hf_dataset.with_format("numpy")[:] + + arrays = [ + pa.array(np.asarray(cols["episode_index"]).astype(np.int64)), + pa.array(np.asarray(cols["task_index"]).astype(np.int64)), + pa.array(np.asarray(cols["timestamp"]).astype(np.float64)), + ] + names = ["episode_index", "task_index", "timestamp"] + for c in FEATURE_COLUMNS: + arrays.append(_feature_array(np.asarray(cols[c]))) + names.append(lance_col(c)) + frames = pa.table(arrays, names=names) + _replace(db, f"{table}_frames", frames, frames.schema) + + tasks_df = lr.meta.tasks # DataFrame indexed by task string, column task_index + tasks = pa.table( + [pa.array(tasks_df["task_index"].astype("int64").tolist()), pa.array([str(t) for t in tasks_df.index])], + names=["task_index", "task"], + ) + _replace(db, f"{table}_tasks", tasks, tasks.schema) + + eps_meta = lr.meta.episodes + n_eps = len(eps_meta) + if "episode_id" not in eps_meta.column_names: + print("WARNING: source has no episode_id strings; the loader's keep-ranges filter (use_filter_dict) needs them") + ep_ids = eps_meta["episode_id"] if "episode_id" in eps_meta.column_names else [""] * n_eps + episodes = pa.table( + [pa.array(list(range(n_eps)), pa.int64()), pa.array([str(e) for e in ep_ids], pa.string())], + names=["episode_index", "episode_id"], + ) + _replace(db, f"{table}_episodes", episodes, episodes.schema) + print( + f"wrote {table}_frames ({frames.num_rows} frames), {table}_tasks ({tasks.num_rows}), " + f"{table}_episodes ({episodes.num_rows})" + ) + + +def _episode_view_frames(base: DROIDLeRobotDataset, lr, episode_id: int, feature: str) -> torch.Tensor: + """Decode every frame of one episode's view directly from its source mp4.""" + ep = lr.meta.episodes[episode_id] + n = int(ep["length"]) + chunk = int(ep.get(f"videos/{feature}/chunk_index", ep.get("data/chunk_index", 0))) + fil = int(ep.get(f"videos/{feature}/file_index", ep.get("data/file_index", 0))) + from_ts = float(ep.get(f"videos/{feature}/from_timestamp", 0.0)) + rel = lr.meta.info["video_path"].format( + video_key=feature, chunk_index=chunk, file_index=fil, episode_chunk=chunk, episode_file=fil + ) + dec = VideoDecoder(str(lr.root / rel), seek_mode="exact") + ts = [from_ts + i * base._dt for i in range(n)] + return dec.get_frames_played_at(seconds=ts).data.to(torch.float32) / 255.0 # (T,C,H,W) in [0,1] + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--root", required=True, help="versioned DROID LeRobot root (see droid_lerobot_dataset_config)") + ap.add_argument("--uri", required=True, help="output LanceDB dir") + ap.add_argument("--table", default="droid_composed") + ap.add_argument("--gop", type=int, default=1, help="keyframe interval (1=all-intra)") + ap.add_argument("--labels-only", action="store_true", help="(re)write label tables only; keep the video table") + args = ap.parse_args() + + db = lancedb.connect(args.uri) + base = _build_base(args.root) + if args.labels_only: + write_label_tables(db, args.table, base) + return + + lr = base._get_dataset(0) + fps = int(round(base._fps)) + schema = pa.schema( + [ + pa.field("episode_index", pa.int64()), + pa.field("ep_start", pa.int64()), + pa.field("length", pa.int64()), + # video_bytes is plain large_binary. TODO: move to blob-v2 after optimizations. + pa.field("video_bytes", pa.large_binary()), + ] + ) + + def _rows(): + ep_start = 0 + for episode_id in range(len(lr.meta.episodes)): + views = {f: _episode_view_frames(base, lr, episode_id, f) for f in base._image_features.values()} + composed = base._compose_multi_view(views) # (T,C,H,W) in [0,1] + thwc = (composed.permute(0, 2, 3, 1) * 255.0).round().clamp(0, 255).to(torch.uint8).numpy() + vb = _encode(np.ascontiguousarray(thwc), fps, args.gop) + n = int(lr.meta.episodes[episode_id]["length"]) + yield pa.RecordBatch.from_arrays( + [ + pa.array([episode_id], pa.int64()), + pa.array([ep_start], pa.int64()), + pa.array([n], pa.int64()), + pa.array([vb], pa.large_binary()), + ], + schema=schema, + ) + ep_start += n + + reader = pa.RecordBatchReader.from_batches(schema, _rows()) + _replace(db, args.table, reader, schema) + print(f"wrote {args.table}: {db.open_table(args.table).count_rows()} episodes (gop={args.gop}, fps={fps})") + write_label_tables(db, args.table, base) + + +if __name__ == "__main__": + main() diff --git a/tools/lance_datagen/build_vision_sft.py b/tools/lance_datagen/build_vision_sft.py new file mode 100644 index 00000000..64cb74f3 --- /dev/null +++ b/tools/lance_datagen/build_vision_sft.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Build a training-optimized vision-SFT representation for LanceDB. + +One row per SFT clip: decode once, resize to the training resolution exactly as +SFTDataset.process_one_sample does (crop left to decode time), re-encode with a +short GOP (all-intra by default, so window seeks are exact), and store the mp4 +plus caption/sizing metadata. This moves the per-epoch resize offline; only the +H.264 re-encode is lossy, and captions are stored verbatim so tokenization stays +byte-identical to the base loader. + +Schema: clip_id, width/height (orig), start_frame/end_frame/temporal_interval, +enc_h/enc_w (stored size), fps, caption_json, caption, video_bytes (large_binary). +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import tempfile + +import lancedb +import numpy as np +import pyarrow as pa + +from cosmos_framework.data.generator.local_datasets.helper import ( + ffmpeg_decode_video, + get_aspect_ratio, + get_video_metadata, +) +from cosmos_framework.data.generator.local_datasets.sft_dataset import ( + CAPTION_TYPES, + _flatten_metadata_by_window, + _load_sft_metadata_from_s3, +) +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY + +# Caption sources the schema persists; the base's _select_caption also reads these +# other keys — reject at build time rather than silently losing captions. +_UNSUPPORTED_CAPTION_KEYS = {"qwen3_32b_rewrite-dense", *CAPTION_TYPES} + + +def _encode(frames_thwc_u8: np.ndarray, fps: int, gop: int) -> bytes: + """Raw RGB frames -> H.264 mp4 bytes via ffmpeg (short GOP, faststart). + + mp4+faststart needs seekable output, so encode to a temp file then read. + Byte-for-byte the encode path of ``build_composed_droid._encode``.""" + t, h, w, _ = frames_thwc_u8.shape + fd, path = tempfile.mkstemp(suffix=".mp4") + os.close(fd) + try: + cmd = [ + "ffmpeg", + "-y", + "-loglevel", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "-s", + f"{w}x{h}", + "-r", + str(fps), + "-i", + "pipe:0", + "-c:v", + "libx264", + "-preset", + "veryfast", + "-g", + str(gop), + "-keyint_min", + str(gop), + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + path, + ] + subprocess.run(cmd, input=frames_thwc_u8.tobytes(), stdout=subprocess.DEVNULL, check=True) + with open(path, "rb") as fh: + return fh.read() + finally: + os.unlink(path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--jsonl", required=True, help="SFT video_dataset_file.jsonl") + ap.add_argument("--uri", required=True, help="output LanceDB dir") + ap.add_argument("--table", default="vision_sft") + ap.add_argument("--resolution", default="256") + ap.add_argument("--gop", type=int, default=1, help="keyframe interval (1=all-intra)") + ap.add_argument("--min-frames", type=int, default=61, help="window filter, same default as the base metadata load") + args = ap.parse_args() + + base_dir = os.path.dirname(os.path.abspath(args.jsonl)) + output_sizes = VIDEO_RES_SIZE_INFO[args.resolution] + + # video_bytes is plain large_binary. TODO: move to blob-v2 after optimizations. + schema = pa.schema( + [ + pa.field("clip_id", pa.string()), + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + pa.field("start_frame", pa.int64()), + pa.field("end_frame", pa.int64()), + pa.field("temporal_interval", pa.int64()), + pa.field("enc_h", pa.int64()), + pa.field("enc_w", pa.int64()), + pa.field("fps", pa.float64()), + pa.field("caption_json", pa.string()), + pa.field("caption", pa.string()), + pa.field("video_bytes", pa.large_binary()), + ] + ) + + # The base's own metadata load + per-window flattening: applies the same + # duration (> 61 s) and min-frames window filters the base pipeline uses, so + # the table's sample population matches what the base loader would serve. + metas = _flatten_metadata_by_window(_load_sft_metadata_from_s3(None, args.jsonl, min_frames=args.min_frames)) + for m in metas: + bad = _UNSUPPORTED_CAPTION_KEYS & set(m["t2w_windows"][0]) + if bad: + raise NotImplementedError( + f"window {m['uuid']} carries caption keys {sorted(bad)} that this schema does not " + "store (only caption_json + caption); extend the schema before converting." + ) + + def _gen(): + for m in metas: + rec, window = m, m["t2w_windows"][0] + vp = rec["vision_path"] + vp = vp if ("://" in vp or vp.startswith("/")) else os.path.join(base_dir, vp) + input_w, input_h = rec["width"], rec["height"] + aspect_ratio = get_aspect_ratio(input_w, input_h) + target_w, target_h = output_sizes[aspect_ratio] + resize_ratio = max(target_w / input_w, target_h / input_h) + resize_h, resize_w = (round(input_h * resize_ratio), round(input_w * resize_ratio)) + + meta = get_video_metadata(vp) + fps = int(round(meta["fps"])) + # decode the WHOLE clip at training resolution (resize only; the crop is + # done at decode time so the stored clip is a clean rectangle). + frames = list(ffmpeg_decode_video(vp, scale_hw=(resize_h, resize_w), num_threads=2)) + thwc = np.ascontiguousarray(np.stack(frames, axis=0)) # [T, resize_h, resize_w, 3] + vb = _encode(thwc, fps, args.gop) + + cj = window.get(CAPTION_JSON_KEY) + cj_str = json.dumps(cj, ensure_ascii=False) if cj is not None else "" + caption = str(window.get("caption", "")) + clip_id = rec["uuid"] # already "{uuid}_w{win}" from _flatten_metadata_by_window + yield pa.RecordBatch.from_arrays( + [ + pa.array([clip_id], pa.string()), + pa.array([input_w], pa.int64()), + pa.array([input_h], pa.int64()), + pa.array([window["start_frame"]], pa.int64()), + pa.array([window["end_frame"]], pa.int64()), + pa.array([window["temporal_interval"]], pa.int64()), + pa.array([resize_h], pa.int64()), + pa.array([resize_w], pa.int64()), + pa.array([float(meta["fps"])], pa.float64()), + pa.array([cj_str], pa.string()), + pa.array([caption], pa.string()), + pa.array([vb], pa.large_binary()), + ], + schema=schema, + ) + + reader = pa.RecordBatchReader.from_batches(schema, _gen()) + db = lancedb.connect(args.uri) + if args.table in db.table_names(): + db.drop_table(args.table) + db.create_table(args.table, data=reader, schema=schema) + t = db.open_table(args.table) + print(f"wrote {args.table}: {t.count_rows()} clips (gop={args.gop}, res={args.resolution}) at {args.uri}") + + +if __name__ == "__main__": + main() diff --git a/tools/lance_datagen/prepare_droid_subset.py b/tools/lance_datagen/prepare_droid_subset.py new file mode 100644 index 00000000..c7b62677 --- /dev/null +++ b/tools/lance_datagen/prepare_droid_subset.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: OpenMDW-1.1 +"""Materialize a small Cosmos-canonical DROID subset from the public +``lerobot/droid_1.0.1`` LeRobot v3.0 dataset. + +The public release names a few features differently from what +``cosmos_framework.data.generator.action.datasets.DROIDLeRobotDataset`` expects. +This script renames them and writes a self-contained ``/success`` tree +(``meta/``, ``data/``, ``videos/``) that the base Cosmos loader reads as-is, +so the base and the LanceDB loader run on byte-identical inputs. + +The (large, concatenated) source mp4s are symlinked, not copied — episode +``from_timestamp`` offsets index into them unchanged. + +Name ``--out`` after a supported version key (see droid_lerobot_dataset_config +LEROBOT_ROOTS, e.g. ``droid_plus_lerobot_320x180_20260406``) so the base loader's +version registry resolves it; pass ``use_success_only=True`` when loading. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +# public droid_1.0.1 name -> Cosmos-canonical name +VIDEO_KEY_MAP = { + "observation.images.wrist_left": "observation.image.wrist_image_left", + "observation.images.exterior_1_left": "observation.image.exterior_image_1_left", + "observation.images.exterior_2_left": "observation.image.exterior_image_2_left", +} +COLUMN_MAP = {"observation.state.joint_position": "observation.state.joint_positions"} + +# Reserved meta columns + the numeric features the Cosmos DROID loader uses +# (post-rename names). We prune everything else (string metadata, velocities, +# extrinsics, …) so the data parquet and info.json stay consistent and the +# downstream lerobot-lancedb converter only sees numeric + video features. +RESERVED = ["index", "episode_index", "frame_index", "task_index", "timestamp"] +NUMERIC = [ + "observation.state.cartesian_position", + "observation.state.joint_positions", + "observation.state.gripper_position", + "action.joint_position", + "action.gripper_position", +] +DATA_COLS = RESERVED + NUMERIC +# Features kept in info.json (numeric + the 3 renamed video keys). +KEEP_FEATURES = set(DATA_COLS) | set(VIDEO_KEY_MAP.values()) + + +def _rename_info(info: dict) -> dict: + feats = {} + for k, v in info["features"].items(): + nk = VIDEO_KEY_MAP.get(k, COLUMN_MAP.get(k, k)) + if nk in KEEP_FEATURES: + feats[nk] = v + info = dict(info) + info["features"] = feats + return info + + +def _rename_table(tbl: pa.Table, mapping: dict[str, str]) -> pa.Table: + return tbl.rename_columns([mapping.get(n, n) for n in tbl.column_names]) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--src", required=True, help="droid_1.0.1 root (has meta/ data/ videos/)") + ap.add_argument("--out", required=True, help="output root; writes /success/") + ap.add_argument("--num-episodes", type=int, default=100) + args = ap.parse_args() + + src = Path(args.src) + out = Path(args.out) / "success" + n = args.num_episodes + + info = json.loads((src / "meta" / "info.json").read_text()) + + # ---- data: keep episodes [0, n); rename, then prune to DATA_COLS so the + # data parquet and info.json features stay consistent for downstream + # converters (they iterate every declared feature). ---- + data = pq.read_table(src / "data" / "chunk-000" / "file-000.parquet") + data = _rename_table(data, COLUMN_MAP) + data = data.select(DATA_COLS) + data = data.filter(pc.less(data["episode_index"], n)) + n_frames = data.num_rows + + (out / "data" / "chunk-000").mkdir(parents=True, exist_ok=True) + pq.write_table(data, out / "data" / "chunk-000" / "file-000.parquet") + + # ---- episode meta: keep [0, n), drop bulky stats/*, rename video keys ---- + ep = pq.read_table(src / "meta" / "episodes" / "chunk-000" / "file-000.parquet") + ep = ep.filter(pc.less(ep["episode_index"], n)) + col_map = {} + for old, new in VIDEO_KEY_MAP.items(): + for suf in ("chunk_index", "file_index", "from_timestamp", "to_timestamp"): + col_map[f"videos/{old}/{suf}"] = f"videos/{new}/{suf}" + keep_cols = [c for c in ep.column_names if not c.startswith("stats/")] + ep = ep.select(keep_cols) + ep = _rename_table(ep, col_map) + (out / "meta" / "episodes" / "chunk-000").mkdir(parents=True, exist_ok=True) + pq.write_table(ep, out / "meta" / "episodes" / "chunk-000" / "file-000.parquet") + + # ---- tasks: LeRobot v3 format (DataFrame indexed by task string, task_index column) ---- + + tasks = pq.read_table(src / "meta" / "tasks.parquet").to_pandas() + if tasks.index.name != "task": + task_col = "task" if "task" in tasks.columns else "__index_level_0__" + tasks = tasks.rename(columns={task_col: "task"}).set_index("task")[["task_index"]] + tasks.to_parquet(out / "meta" / "tasks.parquet") + info = _rename_info(info) + info["total_episodes"] = n + info["total_frames"] = n_frames + (out / "meta" / "info.json").write_text(json.dumps(info, indent=2)) + + # ---- videos: symlink concatenated source mp4s under Cosmos key dirs ---- + for old, new in VIDEO_KEY_MAP.items(): + dst = out / "videos" / new / "chunk-000" + dst.mkdir(parents=True, exist_ok=True) + srcmp4 = (src / "videos" / old / "chunk-000" / "file-000.mp4").resolve() + link = dst / "file-000.mp4" + if link.exists() or link.is_symlink(): + link.unlink() + link.symlink_to(srcmp4) + + print(f"wrote {out} ({n} episodes, {n_frames} frames)") + + +if __name__ == "__main__": + main()