cosmos_framework/data/generator/dataflow/collators.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Built-in BatchCollator implementations."""

from __future__ import annotations

import torch
import torch.utils.data
from torch.utils.data.dataloader import default_collate

from cosmos_framework.data.generator.dataflow.base import BatchCollator


class DefaultBatchCollator(BatchCollator):
    """Stacks samples with torch's default_collate — stock DataLoader behavior."""

    def collate(self, samples: list[dict]) -> dict:
        return torch.utils.data.default_collate(samples)


# ---------------------------------------------------------------------------
# VFMListCollator — reproduces the legacy PackingDataLoader packed-batch structure:
#   for each raw sample, inner-collate at batch_size=1 (custom_collate_fn copy),
#   split per _get_next_sample i=0 rules, then _update_output_batch-accumulate.
#
# This produces list[list[Tensor]] for _MULTI_ITEM_KEYS (not flat list[Tensor]).
# ---------------------------------------------------------------------------

_TIMING_KEYS = {"_sample_time", "_aug_time", "_pre_aug_time", "_aug_step_times"}
_BATCH_TIMING_KEYS = {
    "_worker_batch_time",
    "_worker_aug_time",
    "_worker_io_time",
    "_worker_aug_step_times",
    "_worker_id",
}

# Verbatim copy of JointDataLoader._MULTI_ITEM_KEYS
_MULTI_ITEM_KEYS = {"text_token_ids", "images", "video", "action", "sound"}

# Verbatim copy of JointDataLoader._FLATTEN_LIST_KEYS
_FLATTEN_LIST_KEYS = {"image_size"}


def _vfm_inner_collate(batch):
    """
    Verbatim copy of custom_collate_fn from joint_dataloader.py.

    Collate function that works like default_collate for all keys other than "text_token_ids", "images", and "video".
    For "text_token_ids", "images", and "video" it simply returns them in a list, instead of stacking them as a tensor.
    """
    list_collate_keys = {
        "text_token_ids",
        "images",
        "video",
        "action",
        "domain_id",
        "sequence_plan",
        "sound",
        "raw_action_dim",
        "image_size",
    }

    # Data keys where a per-sample value of ``None`` is a meaningful signal
    # (e.g. audio extraction failed for that sample → ``sound=None`` paired
    # with ``plan.has_sound=False``).  These keys must be kept as a list with
    # ``None`` placeholders so the model can align per-sample data 1:1 with
    # per-sample plans.  Dropping the entire key on any None would leave the
    # remaining sound tensors mis-aligned with the plans whose ``has_sound``
    # flag was set BEFORE collation, causing ``sequence_packing`` to index
    # past the end of ``x0_tokens_sound``.
    sparse_data_keys = {"sound"}

    # Handle the case where the batch is already a dictionary (e.g. column-wise batching)
    if isinstance(batch, dict):
        return {key: (value if key in list_collate_keys else default_collate(value)) for key, value in batch.items()}

    # Handle standard list of samples
    elem = batch[0]
    if isinstance(elem, dict):

        # Some Action datasets add optional metadata keys (for example
        # ``additional_view_description`` for concat-view captions) only for a
        # subset of samples.  PyTorch can batch such samples together when
        # DataLoader batch_size > 1; collating only elem's keys and indexing
        # every sample by that key turns the optional field into a fatal
        # KeyError.  Use the union of keys and skip optional keys that are not
        # present in every sample.  Required training keys still fail loudly via
        # downstream assertions if actually missing.
        result = {}
        keys = set().union(*(d.keys() for d in batch))
        for key in keys:
            if key in _TIMING_KEYS:
                continue
            values = [d.get(key) for d in batch]
            if any(value is None for value in values):
                # Sparse data keys keep their None placeholders to preserve
                # 1:1 alignment with sequence_plan.  Other (optional metadata)
                # keys not present in every sample are dropped.
                if key in sparse_data_keys:
                    result[key] = values
                continue
            if key in list_collate_keys:
                result[key] = values
            else:
                result[key] = default_collate(values)
        result.update(_aggregate_worker_timing(batch))
        return result
    else:
        return default_collate(batch)


def _aggregate_worker_timing(samples: list[dict]) -> dict:
    """Extract per-sample timing keys, aggregate into per-batch scalars."""
    info: dict[str, float | int] = {}
    if "_sample_time" in samples[0]:
        info["_worker_batch_time"] = sum(s.get("_sample_time", 0.0) for s in samples)
    if "_aug_time" in samples[0]:
        aug_total = sum(s.get("_aug_time", 0.0) for s in samples)
        info["_worker_aug_time"] = aug_total
        if "_worker_batch_time" in info:
            info["_worker_io_time"] = info["_worker_batch_time"] - aug_total
    if "_aug_step_times" in samples[0]:
        agg: dict[str, float] = {}
        for s in samples:
            for step_name, t in s.get("_aug_step_times", {}).items():
                agg[step_name] = agg.get(step_name, 0.0) + t
        info["_worker_aug_step_times"] = agg
    worker_info = torch.utils.data.get_worker_info()
    info["_worker_id"] = worker_info.id if worker_info is not None else 0
    return info


def _split_one(batch: dict) -> dict:
    """Port of _get_next_sample split rules for i=0 (verbatim from joint_dataloader.py lines 470-490).

    Splitting rules:
        - _BATCH_TIMING_KEYS: passed through as-is.
        - _MULTI_ITEM_KEYS with list value: elem = v[0]; if elem is a list → sample[k]=elem,
          else → sample[k]=v[0:1] (single-element list wrapping the tensor).
        - Other list values: sample[k] = v[0] (bare element, direct-indexed).
        - Non-list (tensor/scalar) values: sample[k] = v[0:1] (preserve batch dim).
    """
    sample = {}
    for k, v in batch.items():
        if k in _BATCH_TIMING_KEYS:
            sample[k] = v
        elif isinstance(v, list) and k in _MULTI_ITEM_KEYS:
            elem = v[0]
            sample[k] = elem if isinstance(elem, list) else v[0:1]
        elif isinstance(v, list):
            sample[k] = v[0]
        else:
            sample[k] = v[0:1]
    return sample


def _accumulate(output_batch: dict, output: dict) -> None:
    """Port of _update_output_batch from joint_dataloader.py lines 405-418."""
    for key, value in output.items():
        if key in _BATCH_TIMING_KEYS:
            if key not in output_batch:
                output_batch[key] = value
        elif key in _FLATTEN_LIST_KEYS and isinstance(value, list):
            if key not in output_batch:
                output_batch[key] = value
            else:
                output_batch[key].extend(value)
        elif key not in output_batch:
            output_batch[key] = [value]
        else:
            output_batch[key].append(value)


# Keep _vfm_collate as an alias for the inner collate (used by legacy callers).
_vfm_collate = _vfm_inner_collate


class VFMListCollator(BatchCollator):
    """Reproduces the legacy PackingDataLoader packed-batch structure.

    For a group of N raw SFTDataset samples, the packed output has:
      - ``_MULTI_ITEM_KEYS`` (``text_token_ids``, ``video``, ``images``,
        ``action``, ``sound``): ``list[list[Tensor]]`` — each inner list
        is a single-element list ``[tensor]``, matching the
        ``v[i:i+1]`` slice from ``_get_next_sample``.
      - Metadata list keys (``sequence_plan``, ``domain_id``,
        ``raw_action_dim``): flat ``list[element]``.
      - ``image_size`` (``_FLATTEN_LIST_KEYS``): flat ``list[Tensor]``
        (extended, not appended).
      - Non-list tensor keys: ``list[Tensor(1,...)]``.
      - ``_BATCH_TIMING_KEYS``: set once from the first sample.

    Implementation: for each sample, inner-collate at batch_size=1 via
    ``_vfm_inner_collate`` (verbatim ``custom_collate_fn``), split
    sample 0 per ``_split_one`` (verbatim ``_get_next_sample`` i=0
    rules), then accumulate via ``_accumulate`` (verbatim
    ``_update_output_batch``).  Byte-identical to the legacy packer.
    """

    def collate(self, samples: list[dict]) -> dict:
        # Reproduce the legacy PackingDataLoader packed batch: for each sample,
        # inner-collate at batch_size=1, split sample 0 per _MULTI_ITEM_KEYS /
        # list / tensor rules, then _update_output_batch-accumulate across the group.
        output_batch: dict = {}
        for s in samples:
            collated = _vfm_inner_collate([s])      # verbatim custom_collate_fn copy
            split = _split_one(collated)            # i=0 split (rules from _get_next_sample)
            _accumulate(output_batch, split)        # _update_output_batch copy
        return output_batch