MMaDA/tools/run_dataloaders.py

#!/usr/bin/env python3
"""
Utility script to sanity-check data loaders defined in train_omada_inst.py
without constructing the full training stack.

Example:
    python MMaDA/tools/run_dataloaders.py config=MMaDA/configs/omada_instruction_tuning.yaml \
        --flows v2t --num-workers 0 --max-batches 10
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from typing import Any, Dict, Iterable, List, Optional, Tuple

import torch
from omegaconf import DictConfig, OmegaConf
from torch.utils.data import DataLoader
from transformers import AutoTokenizer

# Ensure repository root is importable when executing from arbitrary cwd.
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

from training.data import VideoCaptionDataset  # noqa: E402
from training.utils import image_transform  # noqa: E402

LOGGER = logging.getLogger("run_dataloaders")


def _parse_args() -> Tuple[argparse.Namespace, DictConfig]:
    parser = argparse.ArgumentParser(description="Run Omada dataloaders without the trainer.")
    parser.add_argument(
        "--flows",
        default="v2t",
        help="Comma separated list of dataloaders to exercise (currently supports: v2t). "
             "Use 'all' to run every available flow.",
    )
    parser.add_argument(
        "--max-batches",
        type=int,
        default=0,
        help="Stop after this many batches per loader (0 means iterate the entire epoch).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=None,
        help="Override DataLoader num_workers (falls back to config.dataset.params.num_workers).",
    )
    parser.add_argument(
        "--persistent-workers",
        dest="persistent_workers",
        action="store_true",
        help="Force persistent_workers=True regardless of config.",
    )
    parser.add_argument(
        "--no-persistent-workers",
        dest="persistent_workers",
        action="store_false",
        help="Force persistent_workers=False regardless of config.",
    )
    parser.set_defaults(persistent_workers=None)
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Torch manual seed for reproducibility.",
    )

    args, unknown = parser.parse_known_args()

    cli_conf = OmegaConf.from_cli(unknown)
    if "config" not in cli_conf:
        parser.error("Please provide the training config via 'config=/path/to/config.yaml'.")

    yaml_conf = OmegaConf.load(cli_conf.config)
    merged = OmegaConf.merge(yaml_conf, cli_conf)

    return args, merged


def _collate_v2t(batch: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    """Minimal collate fn mirroring train_omada_inst.collate_fn_v2t."""
    filtered: List[Dict[str, Any]] = [sample for sample in batch if sample is not None]
    if not filtered:
        return None

    videos: List[torch.Tensor] = []
    captions: List[Any] = []
    for sample in filtered:
        frames = sample.get("video")
        caption = sample.get("caption")
        if frames is None:
            continue
        try:
            tensor = torch.stack(frames, dim=0)
        except Exception as exc:
            LOGGER.exception("Failed to stack frames for sample %s", sample)
            raise exc
        videos.append(tensor)
        captions.append(caption)

    if not videos:
        return None

    return {
        "video": torch.stack(videos, dim=0),
        "captions": captions,
    }


def _build_v2t_loader(
    cfg: DictConfig,
    tokenizer,
    *,
    num_workers: int,
    persistent_workers: bool,
    pin_memory: bool,
) -> DataLoader:
    speech_cfg = getattr(cfg.dataset.params, "video_speech_dataset", {})
    if not isinstance(speech_cfg, dict):
        speech_cfg = OmegaConf.to_container(speech_cfg, resolve=True)
    speech_cfg = speech_cfg or {}

    dataset = VideoCaptionDataset(
        transform=image_transform,
        tokenizer=tokenizer,
        max_seq_length=int(cfg.dataset.preprocessing.max_seq_length),
        resolution=int(cfg.dataset.preprocessing.resolution),
        sample_method=speech_cfg.get("sample_method", "uniform"),
        dataset_name=speech_cfg.get("llavavid_dataset_name", "llavavid"),
        num_frames=int(speech_cfg.get("num_frames", 8)),
    )

    batch_size = int(max(1, cfg.training.batch_size_v2t))
    LOGGER.info(
        "Instantiated VideoCaptionDataset with %d samples; batch_size=%d num_workers=%d",
        len(dataset),
        batch_size,
        num_workers,
    )

    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory,
        persistent_workers=persistent_workers if num_workers > 0 else False,
        collate_fn=_collate_v2t,
        drop_last=False,
    )


def _iterate_loader(name: str, loader: DataLoader, max_batches: int) -> None:
    LOGGER.info("Starting iteration over '%s' (max_batches=%s)", name, max_batches or "full epoch")
    start = time.time()
    failures = 0
    processed = 0

    try:
        for step, batch in enumerate(loader, start=1):
            if batch is None:
                failures += 1
                LOGGER.warning("[%s] Received empty batch at step %d", name, step)
                continue

            processed += batch["video"].size(0)

            if max_batches and step >= max_batches:
                break
    except Exception as exc:
        LOGGER.exception("Loader '%s' raised an exception at batch %d", name, step)
        raise exc
    finally:
        duration = time.time() - start
        LOGGER.info(
            "Finished '%s': steps=%d samples=%d failures=%d elapsed=%.2fs",
            name,
            step if 'step' in locals() else 0,
            processed,
            failures,
            duration,
        )


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
        datefmt="%H:%M:%S",
    )

    args, cfg = _parse_args()
    torch.manual_seed(args.seed)

    pin_memory = bool(getattr(cfg.dataset.params, "pin_memory", False))

    if args.num_workers is None:
        num_workers = int(getattr(cfg.dataset.params, "num_workers", 0))
    else:
        num_workers = max(0, args.num_workers)

    if args.persistent_workers is None:
        persistent_workers = bool(getattr(cfg.dataset.params, "persistent_workers", False))
    else:
        persistent_workers = bool(args.persistent_workers)

    flows_arg = [item.strip().lower() for item in args.flows.split(",") if item.strip()]
    if "all" in flows_arg:
        flows = {"v2t"}
    else:
        flows = set(flows_arg)

    tokenizer = AutoTokenizer.from_pretrained(cfg.model.omada.tokenizer_path, padding_side="left")

    loaders: Dict[str, DataLoader] = {}
    if "v2t" in flows:
        loaders["v2t"] = _build_v2t_loader(
            cfg,
            tokenizer,
            num_workers=num_workers,
            persistent_workers=persistent_workers,
            pin_memory=pin_memory,
        )

    if not loaders:
        LOGGER.error("No loaders selected. Supported flows: v2t")
        sys.exit(1)

    for name, loader in loaders.items():
        _iterate_loader(name, loader, args.max_batches)


if __name__ == "__main__":
    main()