Diff Coverage

def _sft_row_boundaries(boundaries: Any, sequence_length: int) -> list[int]:
    """Remove boundary sentinels and validate one complete SFT row."""
    row = torch.as_tensor(boundaries)
    if row.ndim != 1 or row.dtype not in (torch.int32, torch.int64):
        raise ValueError("Each cu_seq_lens row must be a one-dimensional integer array")
    row = row[row != -1]
    if (row.numel() < 2 or row[0] != 0 or row[-1] != sequence_length
            or torch.any(row[1:] <= row[:-1])):
        raise ValueError("Document boundaries must strictly increase from zero to row length")
    return row.tolist()


def _collate_sft_samples(batch: list[Any], *, alignment_size: int) -> Any:
    """Stack SFT rows and retain per-row boundaries without sentinel padding.

    Records from SFTFromMRDataset have 1D tokens/labels; already grouped
    records may have [rows, S]. Variable-length rows are padded, never truncated.
    """
    if not isinstance(batch[0], Mapping) or "cu_seq_lens" not in batch[0]:
        return default_collate(batch)

    token_key = "tokens" if "tokens" in batch[0] else "input_ids"
    token_rows, label_rows, boundary_rows = [], [], []
    for sample in batch:
        tokens, labels = torch.as_tensor(sample[token_key]), torch.as_tensor(sample["labels"])
        if tokens.ndim not in (1, 2) or tokens.shape != labels.shape or not tokens.numel():
            raise ValueError("SFT tokens and labels must have matching nonempty [S] or [rows, S] shapes")
        boundaries = sample["cu_seq_lens"]
        if tokens.ndim == 1:
            tokens, labels, boundaries = tokens.unsqueeze(0), labels.unsqueeze(0), [boundaries]
        if len(boundaries) != len(tokens):
            raise ValueError("cu_seq_lens must have one boundary array per token row")
        for row in boundaries:
            boundary_rows.append(_sft_row_boundaries(row, tokens.shape[-1]))
        token_rows.extend(tokens.unbind(0))
        label_rows.extend(labels.unbind(0))

    sequence_length = max(row.numel() for row in token_rows)
    sequence_length += (-sequence_length) % alignment_size
    packed_keys = {token_key, "labels", "cu_seq_lens"}
    result = default_collate([
        {key: value for key, value in sample.items() if key not in packed_keys} for sample in batch
    ])
    for key, rows, value in ((token_key, token_rows, 0), ("labels", label_rows, IGNORE_INDEX)):
        result[key] = torch.stack([
            torch.nn.functional.pad(row, (0, sequence_length - row.numel()), value=value) for row in rows
        ])
    for row in boundary_rows:
        if row[-1] < sequence_length:
            # Isolate ignored tail padding from every real document.
            row.append(sequence_length)
    result["cu_seq_lens"] = boundary_rows
    return result


def build_sft_collate_fn(mesh_context: Any | None = None) -> Callable[[list[Any]], Any]:
    """Build SFT collation preserving [B, S] and each row's document boundaries.

    Returns:
        A CPU collator accepting 1D records or already grouped 2D records.
        Ordinary records without cu_seq_lens use default collation.
    """
    return partial(_collate_sft_samples, alignment_size=_get_token_alignment_size(mesh_context))


def build_omni_collate_fn(
        mesh_context: Any | None = None,

        self.device = device
        self.tokenizer = tokenizer
        self.data_config = dict(data_config)
        self.source_type = source_type
        dataset_spec = data_config.get("low_level_dataset")
        self.is_sft_dataset = (source_type == "indexed" and dataset_spec is not None
                               and issubclass(resolve_dataset_type(dataset_spec), SFTFromMRDataset))
        self.causal = causal
        self.sliding_window = sliding_window
        if runtime_input_adapter is not None and not isinstance(

            or bool(self.data_config.get("reset_attention_mask", False))
            or source_type == "online" or self.is_sft_dataset
        )
        self.eod_mask_loss = eod_mask_loss or bool(self.data_config.get("eod_mask_loss", False))
        if self.is_sft_dataset:
            if not self.labels_are_shifted or self.eod_mask_loss or not self.causal:
                raise ValueError("Prepared records require shifted labels, causal attention and eod_mask_loss=False")
            if (attention_mode == "compressed" or not create_attention_mask) and runtime_input_adapter is None:
                raise ValueError("Prepared records without a dense mask require a boundary-aware runtime adapter")
        self.attention_runtime = AttentionRuntime(
            mode=attention_mode,
            create_mask=create_attention_mask,
            reset_mask=resolved_reset_attention_mask,

    def resolve(self, canonical_batch: Mapping[str, Any]) -> Any:
        """Recover leading-zero ``cu_seq_lens`` from Indexed tokens."""
        input_ids = canonical_batch["input_ids"]
        batch_size, seq_len = input_ids.shape
        if "cu_seq_lens" in canonical_batch:
            rows = canonical_batch["cu_seq_lens"]
            if isinstance(rows, torch.Tensor) and rows.ndim == 1:
                return rows.to(torch.int32)
            if len(rows) != batch_size:
                raise ValueError("cu_seq_lens must have one boundary array per batch row")
            # Only metadata is flattened; tokens and labels remain [B, S].
            ends = [end + row_idx * seq_len for row_idx, row in enumerate(rows) for end in row[1:]]
            return torch.tensor([0, *ends], dtype=torch.int32)
        token_indices = torch.arange(seq_len, dtype=input_ids.dtype, device=input_ids.device)

        seq_ends = []
        for batch_idx in range(batch_size):

    def __post_init__(self) -> None:
        """Validate the blended Dataset fields followed by GPT-specific fields."""
        super().__post_init__()

        if self.low_level_dataset is None:
            raise ValueError("low_level_dataset must be explicitly configured")
        if self.simple_blend is None:
            raise ValueError("simple_blend must be explicitly configured")
        if self.sequence_length <= 0:
            raise ValueError("sequence_length must be positive")
        if self.tokenizer is None and not self.mock:
            raise ValueError("Attribute 'tokenizer' must not be None")

    file_suffix: str,
) -> list[str]:
    """Discover indexed prefixes under one directory."""
    if not distributed_walk:
        return _walk_data_directories([current_path], file_suffix)

    root_layout = None
    if rank == 0:
        subdirectories = []

        for name in os.listdir(current_path):
            path = os.path.join(current_path, name)
            if os.path.isdir(path):
                subdirectories.append(path)
            elif name.endswith(file_suffix):
                top_files.append(os.path.splitext(path)[0] if file_suffix == ".bin" else path)
        root_layout = (subdirectories, top_files)

    gathered_layouts: list[tuple[list[str], list[str]] | None] = [None] * world_size
    dist.all_gather_object(gathered_layouts, root_layout)

    rank_zero_layout = gathered_layouts[0]
    if rank_zero_layout is None:
        raise ValueError("Rank 0 did not provide the indexed directory layout")
    subdirectories, file_paths = rank_zero_layout
    local_files = _walk_data_directories(subdirectories[rank::world_size], file_suffix)
    gathered_files: list[list[str] | None] = [None] * world_size
    dist.all_gather_object(gathered_files, local_files)
    for rank_files in gathered_files:
        if rank_files:

    # 3. Accept a direct prefix, or recursively discover prefixes in a directory.
    file_paths: list[str] = []
    for current_path in data_paths:
        if file_suffix == ".arrow" and current_path.endswith(file_suffix) and os.path.isfile(current_path):
            file_paths.append(current_path)
            continue
        if file_suffix == ".bin" and os.path.isfile(current_path + ".bin"):
            file_paths.append(current_path)
            continue
        if not os.path.isdir(current_path):
            # Keep args.data_path unchanged when filesystem discovery

            # ["30", prefix_a, "70", prefix_b] for the blend parser.
            original_paths = list(data_paths)
            return original_paths

        file_paths.extend(_discover_indexed_paths(current_path, rank, world_size, use_distributed_walk, file_suffix))

    if not file_paths:
        raise ValueError(f"data_path {data_paths!r} has no data file ending with {file_suffix}")

    if len(file_paths) == 1:
        return file_paths

    if file_suffix == ".arrow":
        # Every rank must blend files in the same order after directory discovery.
        file_paths.sort()
    else:
        # 4. Prefer the numeric shard order encoded by the corpus filenames.
        # Unknown naming schemes remain usable in their original walk order.
        try:
            def numeric_shard_key(path: str) -> tuple[int, int]:
                """Extract numeric corpus and shard identifiers from a path."""
                path_parts = os.path.basename(path).split("_")
                return int(path_parts[2]), int(path_parts[-1])

            file_paths.sort(key=numeric_shard_key)
        except (IndexError, TypeError, ValueError):
            logger.warning("Cannot sort indexed files with the numeric filename rule; keeping walk order")

    # 5. Automatically discovered sources use equal blend weights.
    weighted_paths = [value for path in file_paths for value in ("1", path)]
    return weighted_paths

        data_config.get("test_data_path"),
    ]
    has_independent_blends = any(blend_per_split)
    try:
        mock_data = bool(data_config.get("mock_data", False))
        tokenizer = data_config.get("tokenizer")
        if tokenizer is None and not mock_data:
            raise ValueError("tokenizer is required when mock_data is false")

    file_paths = []
    for directory in directories:
        for root, _, filenames in os.walk(directory):
            for filename in filenames:
                if filename.endswith(file_suffix):
                    path = os.path.abspath(os.path.join(root, filename))
                    file_paths.append(os.path.splitext(path)[0] if file_suffix == ".bin" else path)
    return file_paths


def resolve_dataset_type(spec: Mapping[str, Any] | type) -> type:

def resolve_dataset_type(spec: Mapping[str, Any] | type) -> type:
    """Resolve an explicitly selected class without constructing a Dataset."""
    if isinstance(spec, type):
        return spec
    if not isinstance(spec, Mapping) or "_target_" not in spec:
        raise ValueError("low_level_dataset requires an explicit Dataset _target_")
    module, _, name = spec["_target_"].rpartition(".")
    dataset_type = getattr(import_module(module), name)
    if not isinstance(dataset_type, type):
        raise ValueError("low_level_dataset._target_ must identify a Dataset class")
    return dataset_type

    """Memory-map one Arrow file, the low-level source of one SFT Dataset."""

    def __init__(self, dataset_path: str) -> None:
        """Open one HF Arrow IPC stream; directory discovery belongs to the builder."""
        path = Path(dataset_path)
        if path.suffix != ".arrow" or not path.is_file():
            raise ValueError("ArrowDataReader requires one .arrow file, not a directory")
        self.dataset = Dataset.from_file(str(path))

    def __len__(self) -> int:
        """Return the number of complete prepared records in this file."""
        return len(self.dataset)

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        """Return source columns for one record, including labels and boundaries."""
        return self.dataset[int(index)]

        """Delegate indexed path discovery to the config build module."""
        if self.data_path is None:
            raise ValueError("data_path is required when mock_data is false")

        dataset_type = resolve_dataset_type(self.data_config["low_level_dataset"])
        data_paths = resolve_data_paths(
            self.data_path,
            distributed_walk=bool(self.data_config["distributed_walk"]),
            file_suffix=".arrow" if issubclass(dataset_type, SFTFromMRDataset) else ".bin",

    @staticmethod
    def _select_dataset_type(config: GPTDatasetConfig) -> type:
        """Select GPT, Mock GPT, or MR GPT Dataset."""
        return config.low_level_dataset

    def _select_split_builder(self, config: GPTDatasetConfig) -> SplitBuilder:
        """
        Select the standard or simple train/valid/test builder.

        self, dataset: ArrowDataReader, dataset_path: str, indices: np.ndarray,
        num_samples: int, index_split: Any, config: GPTDatasetConfig,
    ) -> None:
        """Store whole-record sampling inputs without building GPT packing caches."""
        self.dataset = dataset
        self.indices = indices
        self.num_samples = num_samples
        self.config = config
        self.unique_identifiers = self.build_unique_identifiers(dataset_path, num_samples, index_split, config)
        if self.config.simple_blend not in ("inter", "intra"):
            raise ValueError("Prepared Arrow datasets require simple_blend='inter' or 'intra'")
        if self.num_samples > len(self.indices):
            raise ValueError("Requested samples exceed prepared records in this split")

    @staticmethod
    def numel_low_level_dataset(low_level_dataset: ArrowDataReader) -> int:
        """Split by complete records rather than by inner documents."""
        return len(low_level_dataset)

    @staticmethod
    def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> ArrowDataReader:
        """Open the Arrow reader used by the existing split builder."""
        del config
        return ArrowDataReader(dataset_path)

    @classmethod
    def build_unique_identifiers(
        cls, dataset_path: str, num_samples: int, index_split: Any, config: GPTDatasetConfig,

    def build_unique_identifiers(
        cls, dataset_path: str, num_samples: int, index_split: Any, config: GPTDatasetConfig,
    ) -> dict[str, Any]:
        """Identify the record view without opening Arrow files during lazy setup."""
        return {"class": cls.__name__, "dataset_path": dataset_path, "num_samples": num_samples,
                "index_split": index_split.name.lower(), "split": config.split,
                "dataset_options": config.dataset_options}

    def __len__(self) -> int:

                "dataset_options": config.dataset_options}

    def __len__(self) -> int:
        """Return the requested record count."""
        return self.num_samples

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        """Read labels unchanged except for explicit binary mask suppression."""
        sample = self.dataset[int(self.indices[index])]
        fields = self.config.dataset_options
        tokens = np.asarray(sample[fields.get("input_ids_key", "input_ids")])
        labels = np.asarray(sample[fields.get("labels_key", "labels")])
        boundaries = np.asarray(sample.get(fields.get("boundaries_key", "cu_seq_lens"), [0, tokens.size]))
        boundaries = self._validate_record_layout(tokens, labels, boundaries, int(self.indices[index]))
        if (np.any(tokens < 0) or np.any(tokens >= self.config.tokenizer.vocab_size)
                or np.any((labels < 0) & (labels != IGNORE_INDEX))
                or np.any(labels >= self.config.tokenizer.vocab_size)):
            raise ValueError("Prepared token/label IDs must fit the vocabulary; ignored labels must be -100")
        # Arrow buffers may be shared/read-only; never modify their supervision in place.
        labels = labels.astype(np.int64, copy=True)
        mask = sample.get(fields.get("loss_mask_key", "loss_mask"))
        if mask is not None:
            mask = np.asarray(mask)
            if mask.shape != labels.shape or not np.isin(mask, (0, 1)).all():
                raise ValueError("loss_mask must be binary and match labels")
            labels[mask == 0] = IGNORE_INDEX
        if not np.any(labels != IGNORE_INDEX):
            raise ValueError("Prepared record has no effective targets")
        return {"tokens": tokens.astype(np.int64), "labels": labels,
                "cu_seq_lens": boundaries.astype(np.int32)}

    def _validate_record_layout(
        self, tokens: np.ndarray, labels: np.ndarray, boundaries: np.ndarray, record_index: int,

    def _validate_record_layout(
        self, tokens: np.ndarray, labels: np.ndarray, boundaries: np.ndarray, record_index: int,
    ) -> np.ndarray:
        """Check record shapes, configured length and internal document boundaries."""
        if (tokens.ndim != 1 or not tokens.size or labels.shape != tokens.shape
                or tokens.dtype.kind not in "iu" or labels.dtype.kind not in "iu"):
            raise ValueError("Prepared tokens and labels must be matching nonempty 1D integer arrays")
        if not self.config.skip_data_check and tokens.size != self.config.sequence_length:
            raise ValueError(
                "Prepared record must match seq_length; "
                f"record {record_index}: expected {self.config.sequence_length}, got {tokens.size}"
            )
        boundaries = boundaries[boundaries != -1] if boundaries.ndim == 1 else boundaries
        if (boundaries.ndim != 1 or boundaries.dtype.kind not in "iu" or boundaries.size < 2
                or boundaries[0] != 0 or boundaries[-1] != tokens.size
                or np.any(boundaries[1:] <= boundaries[:-1]) or tokens.size > np.iinfo(np.int32).max):
            raise ValueError("Document boundaries must strictly increase from zero to token count")
        return boundaries

        if any(size < 0 for size in train_valid_test_num_samples):
            raise ValueError("train_valid_test_num_samples must be non-negative")
        if blend_mode not in {"no", "inter", "intra"}:
            raise ValueError(f"Unsupported blend mode: {blend_mode!r}")
        if config.mock != (dataset_type is MockGPTDataset):
            raise ValueError("mock_data must match the selected MockGPTDataset class")
        if not config.mock and (blend_mode == "no") != (dataset_type is GPTDataset):
            raise ValueError("simple_blend='no' is only for GPTDataset; other Datasets require inter/intra")

        logger.debug(
            "Building indexed splits: dataset=%s, num_samples=%s, mock=%s, blend_mode=%s, lazy=%s",
            dataset_type.__name__,

    Raises:
        ValueError: If required path or training information is missing.
    """
    dataset_config = dict(data_config)
    dataset_config.setdefault("distributed_walk", True)
    training_seed = getattr(training_config, "seed", None)
    dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
    if tokenizer is not None:
        dataset_config["tokenizer"] = tokenizer

    training_mode = None if training_mode_node is None else normalize_value(
        training_mode_node, Literal["sft", "pretrain"], path=f"{path}.training_mode"
    )
    # Dataset classes need path/split arguments supplied later by the indexed builder.
    data_options = dict(target_node.get("data_config", {}))
    dataset_class = data_options.pop("low_level_dataset", None)
    if dataset_class is not None:
        target_node["data_config"] = data_options
    target = _resolve_target(target_node, path=path)
    if dataset_class is not None:
        target = target.replace(data_config={**target.data_config, "low_level_dataset": dataset_class})
    if isinstance(model_assets_node, Mapping) and "_target_" in model_assets_node:
        model_assets = _resolve_target(
            model_assets_node,
            path=f"{path}.model_assets",