Diff Coverage

from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import Any

import torch


def build_dense_attention_masks(
        *,

        )

    # Attention mask (lower triangular).
    att_mask_batch = micro_batch_size if reset_attention_mask else 1
    attention_mask = torch.ones((att_mask_batch, seq_length, seq_length), dtype=torch.bool, device=device).tril()
    attention_mask = attention_mask.view(att_mask_batch, 1, seq_length, seq_length)

    if reset_attention_mask:
        for seq_start in boundaries[:-1]:

            attention_mask[batch_idx, 0, row_seq_start:, :row_seq_start] = False

    swa_mask = None
    if sliding_window is not None:
        positions = torch.arange(seq_length, dtype=torch.int64, device=device)
        query_positions = positions.unsqueeze(1)
        key_positions = positions.unsqueeze(0)
        token_distance = query_positions - key_positions
        outside_window = token_distance > sliding_window

from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any

import torch
from torch.utils.data import default_collate

from hyper_parallel.data.constants import IGNORE_INDEX

        packed_batch = {}
        for field in ("input_ids", "labels"):
            values = [model_sample[field] for model_sample in model_samples]
            packed_batch[field] = torch.cat(values, dim=-1).unsqueeze(0)

        packed_seq_len = packed_batch["input_ids"].shape[-1]
        pad_len = (-packed_seq_len) % self.sequence_parallel_size
        if pad_len:

        pad_len = (-packed_seq_len) % self.sequence_parallel_size
        if pad_len:
            input_padding = packed_batch["input_ids"].new_zeros((1, pad_len))
            label_padding = packed_batch["labels"].new_full((1, pad_len), IGNORE_INDEX)
            packed_batch["input_ids"] = torch.cat((packed_batch["input_ids"], input_padding), dim=-1)
            packed_batch["labels"] = torch.cat((packed_batch["labels"], label_padding), dim=-1)

        seq_lens = model_samples[0]["input_ids"].new_tensor(
            [model_sample["input_ids"].shape[-1] for model_sample in model_samples]
        )

            # Represent the alignment tail as one synthetic packed sequence so
            # attention metadata covers every physical Q/KV token. Its labels
            # remain IGNORE_INDEX and therefore do not contribute to the loss.
            padded_end = seq_ends[-1:] + pad_len
            seq_ends = torch.cat((seq_ends, padded_end))
        cu_seq_lens = torch.cat((zero, seq_ends))
        cu_seq_lens = cu_seq_lens.to(torch.int32)
        packed_batch["cu_seq_lens"] = cu_seq_lens

        return packed_batch

from collections.abc import Mapping
from typing import Any

import torch

from hyper_parallel.data.batching.attention_runtime import (
    AttentionRuntimeAdapter,
    build_dense_attention_masks,

        cp_seq_end = cp_seq_start + local_seq_len

        # This produces the same values as building global [B, S] positions
        # and then taking the CP slice, without allocating the global tensor.
        position_ids = torch.arange(
            cp_seq_start, cp_seq_end, dtype=torch.int64, device=input_ids.device
        )
        local_position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)

        return local_position_ids

    def _build_loss_mask(self, parallel_batch: Mapping[str, Any]) -> Any:
        """Build the local loss mask from labels and input IDs."""
        loss_mask = (parallel_batch["labels"] >= 0).to(dtype=torch.int64)

        if self.eod_mask_loss:
            eod_token_id = getattr(self.tokenizer, "eod", None)
            loss_mask = loss_mask.masked_fill(parallel_batch["input_ids"] == eod_token_id, 0)

from collections.abc import Mapping
from typing import Any

import torch


class OnlineBoundaryResolver:
    """Read global cumulative sequence boundaries emitted by Online packing."""

        Returns:
            Global int32 cumulative sequence boundaries.
        """
        raw_cu_seq_lens = canonical_batch["cu_seq_lens"]
        cu_seq_lens = raw_cu_seq_lens.to(torch.int32)

        return cu_seq_lens

            Global int32 cumulative sequence boundaries.
        """
        input_ids = canonical_batch["input_ids"]
        batch_size, seq_len = input_ids.shape
        token_indices = torch.arange(seq_len, dtype=input_ids.dtype, device=input_ids.device)

        # Collect flattened, nonzero cumulative sequence ends. A leading zero
        # is prepended below to form ``cu_seq_lens``.
        seq_ends = []

            if not seq_ends or seq_ends[-1] != row_end:
                seq_ends.append(row_end)

        # Packed attention uses the standard leading-zero cumulative form.
        cu_seq_lens = torch.tensor([0, *seq_ends], dtype=torch.int32)

        return cu_seq_lens

_DATASET_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f"


def _get_rank() -> int:
    try:
        rank = int(dist.get_rank())
    except (RuntimeError, ValueError):
        rank = 0
    return rank


class _DatasetDebugRankFilter(logging.Filter):
    """Filter Dataset DEBUG records by distributed rank."""

        self.ranks: frozenset[int] | None = frozenset({0})

    def filter(self, record: logging.LogRecord) -> bool:
        """Keep regular logs and Dataset DEBUG logs from selected ranks."""
        rank = _get_rank()
        record.rank_id = rank
        if record.levelno != logging.DEBUG or not record.name.startswith(_DATASET_LOGGER_NAME):
            return True
        rank_enabled = getattr(record, "dataset_rank_enabled", None)
        if rank_enabled is not None:
            return bool(rank_enabled)
        return self.ranks is None or rank in self.ranks


class _DatasetLogFormatter(logging.Formatter):
    """Format Dataset log timestamps with microsecond precision."""

    """Format Dataset log timestamps with microsecond precision."""

    def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str:  # pylint: disable=C0103
        """Format a log record timestamp in local time."""
        record_time = datetime.fromtimestamp(record.created).astimezone()
        return record_time.strftime(datefmt) if datefmt else record_time.isoformat(timespec="microseconds")


_DEBUG_RANK_FILTER = _DatasetDebugRankFilter()

    def debug(self, msg: object, *args: object, **kwargs: Any) -> None:
        """Log on default ranks, or on ranks selected by ``enabled`` when provided."""
        if not self.isEnabledFor(logging.DEBUG):
            return
        enabled: RankCondition | None = kwargs.pop("enabled", None)
        if enabled is not None:
            rank_enabled = enabled() if callable(enabled) else enabled
            extra = dict(kwargs.get("extra", {}))
            extra["dataset_rank_enabled"] = rank_enabled
            kwargs["extra"] = extra
        kwargs.setdefault("stacklevel", 2)
        self.logger.debug(msg, *args, **kwargs)


def get_dataset_logger(name: str) -> DatasetLogger:
    """Return a Dataset logger supporting the ``enabled`` DEBUG argument."""

    Raises:
        ValueError: If ``level`` is unsupported, or ``ranks`` is invalid.
    """
    log_levels = {"debug": logging.DEBUG, "info": logging.INFO, "warn": logging.WARNING}
    if level not in log_levels:
        raise ValueError(f"Unsupported Dataset log level: {level!r}")

    if ranks is None:
        selected_ranks = None
    else:
        selected_ranks = frozenset(ranks)
        if not selected_ranks or any(isinstance(rank, bool) or not isinstance(rank, int) or rank < 0
                                     for rank in selected_ranks):
            raise ValueError("ranks must contain non-negative integers, or be None for all ranks")

    _DEBUG_RANK_FILTER.ranks = selected_ranks
    dataset_logger = logging.getLogger(_DATASET_LOGGER_NAME)
    dataset_logger.setLevel(log_levels[level])
    dataset_logger.propagate = False
    dataset_handlers = [handler for handler in dataset_logger.handlers if _DEBUG_RANK_FILTER in handler.filters]
    if not dataset_handlers:
        dataset_handler = logging.StreamHandler()
        dataset_handler.setFormatter(_DatasetLogFormatter(_DATASET_LOG_FORMAT, datefmt=_DATASET_DATE_FORMAT))
        dataset_handler.addFilter(_DEBUG_RANK_FILTER)
        dataset_logger.addHandler(dataset_handler)
        dataset_handlers.append(dataset_handler)
    for dataset_handler in dataset_handlers:
        dataset_handler.setLevel(log_levels[level])

# limitations under the License.
# ============================================================================
"""Standard deterministic weighted Dataset blending."""

from __future__ import annotations

import hashlib
import json
import os
from collections import OrderedDict
from collections.abc import Mapping, Sequence
from typing import Any

import numpy as np

from hyper_parallel.data.dataset_logging import get_dataset_logger
from hyper_parallel.data.indexed.indexed_data_config import GPTDatasetConfig
from hyper_parallel.data.indexed.indexed_helpers import build_blending_indices

logger = get_dataset_logger(__name__)


class BlendedDataset:
    """Expose several Datasets through deterministic weighted sample indices."""

    def __init__(
            self,
            datasets: Sequence[Any],
            weights: Sequence[float],
            size: int,

            size: int,
            config: GPTDatasetConfig,
    ) -> None:
        """Validate the inputs and create weighted sample indices."""
        if not datasets or len(datasets) != len(weights):
            raise ValueError("datasets and weights must be non-empty and have the same length")
        if len(datasets) >= np.iinfo(np.int16).max:
            raise ValueError("number of datasets must be less than 32767")
        if not all(isinstance(dataset, type(datasets[0])) for dataset in datasets):
            raise ValueError("all datasets must be of the same type")
        self.datasets = list(datasets)
        self.weights = _normalize_weights(weights)
        self.size = size
        self.config = config
        self.unique_identifiers = OrderedDict()
        self.unique_identifiers["class"] = type(self).__name__
        self.unique_identifiers["datasets"] = self._collect_dataset_identifiers()
        self.unique_identifiers["weights"] = self.weights
        self.unique_identifiers["size"] = self.size
        self.unique_description = json.dumps(
            self.unique_identifiers,
            indent=4,
            default=lambda value: value.unique_identifiers,
        )
        self.unique_description_hash = hashlib.md5(
            self.unique_description.encode("utf-8")
        ).hexdigest()
        self.dataset_index, self.dataset_sample_index = self._build_indices()

    def _collect_dataset_identifiers(self) -> list[Any]:
        """Collect component Dataset cache identities."""
        identifiers = [dataset.unique_identifiers for dataset in self.datasets]
        return identifiers

    def __len__(self) -> int:
        """Return the requested blended sample count."""
        return self.size

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        """Read one sample from the Dataset selected by the blend indices."""
        dataset_id = int(self.dataset_index[index])
        sample_id = int(self.dataset_sample_index[index])
        sample = {"dataset_id": dataset_id, **self.datasets[dataset_id][sample_id]}
        return sample

    def _build_indices(self) -> tuple[np.ndarray, np.ndarray]:
        """Load or build the top-level blend index cache."""
        cache_directory = self.config.path_to_cache
        if cache_directory is None:
            logger.debug(
                "Building blended indices in memory: size=%d, sources=%d, weights are %s",
                self.size, len(self.weights), self.weights[:3],
            )
            indices = self._build_indices_in_memory()
            return indices

        cache_prefix = os.path.join(
            cache_directory,
            f"{self.unique_description_hash}-{type(self).__name__}",
        )
        description_path = f"{cache_prefix}-description.txt"
        dataset_index_path = f"{cache_prefix}-dataset_index.npy"
        sample_index_path = f"{cache_prefix}-dataset_sample_index.npy"
        cache_paths = (description_path, dataset_index_path, sample_index_path)
        if all(os.path.isfile(path) for path in cache_paths):
            logger.debug("Loading blended index cache: prefix=%s", cache_prefix)
            dataset_index = np.load(dataset_index_path, allow_pickle=True, mmap_mode="r")
            sample_index = np.load(sample_index_path, allow_pickle=True, mmap_mode="r")
            return dataset_index, sample_index

        dataset_index, sample_index = self._build_indices_in_memory()
        os.makedirs(cache_directory, exist_ok=True)
        with open(description_path, "w", encoding="utf-8") as description_file:
            description_file.write(self.unique_description)
        np.save(dataset_index_path, dataset_index, allow_pickle=True)
        np.save(sample_index_path, sample_index, allow_pickle=True)
        logger.debug("Saved blended index cache: prefix=%s", cache_prefix)
        return dataset_index, sample_index

    def _build_indices_in_memory(self) -> tuple[np.ndarray, np.ndarray]:
        """Build a deterministic weighted source schedule."""
        dataset_index = np.empty(self.size, dtype=np.int16)
        sample_index = np.empty(self.size, dtype=np.int64)
        build_blending_indices(
            dataset_index=dataset_index,
            dataset_sample_index=sample_index,
            weights=self.weights,
        )
        requested_counts = np.bincount(dataset_index, minlength=len(self.datasets))
        for dataset_id, requested_count in enumerate(requested_counts):
            dataset_size = len(self.datasets[dataset_id])
            if requested_count > dataset_size:
                raise ValueError(
                    f"Dataset {dataset_id} has only {dataset_size} samples, "
                    f"but the blend requested {requested_count} samples"
                )
        return dataset_index, sample_index


def _normalize_weights(weights: Sequence[float]) -> list[float]:
    """Validate and normalize positive blend weights."""
    if not weights or any(weight <= 0.0 for weight in weights):
        raise ValueError("Dataset blend weights must be positive")
    weight_array = np.asarray(weights, dtype=np.float64)
    normalized_weights = (weight_array / np.sum(weight_array)).tolist()
    return normalized_weights

from typing import Any

import numpy as np

import torch.distributed as dist

from hyper_parallel.data.dataset_logging import get_dataset_logger

logger = get_dataset_logger(__name__)

                top_files.append(os.path.splitext(path)[0])
        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

        raise ValueError("Rank 0 did not provide the indexed directory layout")
    subdirectories, file_paths = rank_zero_layout
    local_files = _walk_bin_directories(subdirectories[rank::world_size])
    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:
            file_paths.extend(rank_files)
    return file_paths

    # 2. Use distributed walking only after the process group is available.
    # Before initialization, one process walks all configured directories.
    try:
        rank = int(dist.get_rank())
        world_size = int(dist.get_world_size())
    except (RuntimeError, ValueError):
        rank, world_size = 0, 1
    use_distributed_walk = distributed_walk and world_size > 1

# limitations under the License.
# ============================================================================
"""Read token sequences from paired ``.idx/.bin`` files."""

from __future__ import annotations

import os
import struct
import time
from functools import lru_cache
from typing import ClassVar

import numpy as np

from hyper_parallel.data.dataset_logging import get_dataset_logger
logger = get_dataset_logger(__name__)

_INDEX_HEADER = b"MMIDIDX\x00\x00"
_INDEX_VERSION = 1
_DTYPES = {
    1: np.uint8,
    2: np.int8,
    3: np.int16,
    4: np.int32,

    8: np.uint16,
}


class _IndexReader:
    """Memory-map the compact sequence metadata stored in an index file."""

    def __init__(self, idx_path: str) -> None:
        """Parse and memory-map one index metadata file."""
        # Read the fixed-size header, which describes how the following index arrays are decoded.
        with open(idx_path, "rb") as stream:
            header = stream.read(len(_INDEX_HEADER))
            if header != _INDEX_HEADER:
                raise ValueError(f"Invalid indexed Dataset header in {idx_path!r}")

            version = struct.unpack("<Q", stream.read(8))[0]
            if version != _INDEX_VERSION:
                raise ValueError(f"Unsupported indexed Dataset version {version}; expected {_INDEX_VERSION}")

            dtype_code = struct.unpack("<B", stream.read(1))[0]
            self.dtype = np.dtype(_DTYPES[dtype_code])
            self.dtype_size = self.dtype.itemsize

            self.sequence_count = struct.unpack("<Q", stream.read(8))[0]
            self.document_count = struct.unpack("<Q", stream.read(8))[0]
            offset = stream.tell()

        # Map the three arrays stored after the header: token count per sequence, byte offset in .bin per sequence,
        # and sequence boundaries per document. The first two arrays have sequence_count entries, while the last
        # has document_count entries.
        self.bin_buffer_mmap = np.memmap(idx_path, mode="r", order="C")
        self.bin_buffer = memoryview(self.bin_buffer_mmap)
        self.sequence_lengths = np.frombuffer(
            self.bin_buffer, dtype=np.int32, count=self.sequence_count, offset=offset,
        )

        self.sequence_pointers = np.frombuffer(
            self.bin_buffer, dtype=np.int64, count=self.sequence_count, offset=offset + self.sequence_lengths.nbytes,
        )

        self.document_indices = np.frombuffer(
            self.bin_buffer, dtype=np.int64, count=self.document_count,
            offset=offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes,
        )

            offset=offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes,
        )

        # Each document is a half-open sequence range; the final boundary must equal the sequence count.
        if self.document_indices.size == 0 or self.document_indices[-1] != self.sequence_count:
            raise ValueError("Indexed Dataset document boundaries do not match its sequence count")

    def __del__(self) -> None:
        """Close the index metadata mmap when it is no longer referenced."""
        bin_buffer_mmap = getattr(self, "bin_buffer_mmap", None)
        mmap_handle = getattr(bin_buffer_mmap, "_mmap", None)
        if mmap_handle is not None:
            mmap_handle.close()

    def __len__(self) -> int:
        """Return the number of indexed sequences."""
        return self.sequence_count

    @lru_cache(maxsize=8)
    def __getitem__(self, index: int | np.integer) -> tuple[np.int64, np.int32]:
        """Return the byte pointer and token length of one sequence."""
        return self.sequence_pointers[index], self.sequence_lengths[index]


class IndexedDataReader:
    """Read individual token sequences from a paired ``.idx/.bin`` Dataset.

    Args:
        path_prefix: Dataset path without the ``.idx`` or ``.bin`` suffix.

        reuse_index: Whether to reuse the first loaded index metadata. This
            matches corpora whose shards share an identical index layout.
    """

    cached_index_reader: ClassVar[_IndexReader | None] = None

    def __init__(self, path_prefix: str, mmap: bool = True, reuse_index: bool = False) -> None:
        """Open the index metadata and token payload files."""
        self.path_prefix = ""
        self.mmap = False
        self.reuse_index = reuse_index
        self.index: _IndexReader | None = None
        self.bin_buffer_mmap: np.memmap | None = None
        self.bin_buffer: memoryview | None = None
        self.initialize(path_prefix, mmap)

    def initialize(self, path_prefix: str, mmap: bool) -> None:
        """Open the index metadata and token payload.

        Args:
            path_prefix: Dataset path without the ``.idx`` or ``.bin`` suffix.

        Args:
            path_prefix: Dataset path without the ``.idx`` or ``.bin`` suffix.
            mmap: Whether to memory-map the token payload.
        """
        start_time = time.time()
        index_path = path_prefix + ".idx"
        data_path = path_prefix + ".bin"
        if not os.path.isfile(index_path) or not os.path.isfile(data_path):
            raise FileNotFoundError(f"Expected indexed Dataset files {index_path!r} and {data_path!r}")

        self.path_prefix = path_prefix
        self.mmap = mmap
        reused_cached_index = self.reuse_index and IndexedDataReader.cached_index_reader is not None
        if self.index is None:
            if reused_cached_index:
                self.index = IndexedDataReader.cached_index_reader
            else:
                self.index = _IndexReader(index_path)
                if self.reuse_index:
                    IndexedDataReader.cached_index_reader = self.index

        self.bin_buffer_mmap = None
        self.bin_buffer = None
        if mmap:
            self.bin_buffer_mmap = np.memmap(data_path, mode="r", order="C")
            self.bin_buffer = memoryview(self.bin_buffer_mmap)
        logger.debug(
            "Opened indexed Dataset: sequences=%d, documents=%d, dtype=%s, mmap=%s, prefix=%s, "
            "reused_index=%s, elapsed=%.4f seconds",
            len(self.index), self.index.document_indices.size - 1, self.index.dtype, mmap, path_prefix,
            reused_cached_index, time.time() - start_time,

            len(self.index), self.index.document_indices.size - 1, self.index.dtype, mmap, path_prefix,
            reused_cached_index, time.time() - start_time,
        )

    def __getstate__(self) -> tuple[str, bool, bool]:
        """Serialize construction inputs instead of open mmap resources."""
        return self.path_prefix, self.mmap, self.reuse_index

    def __setstate__(self, state: tuple[str, bool, bool]) -> None:
        """Reopen index and data resources inside a DataLoader worker."""
        path_prefix, mmap, reuse_index = state
        self.__init__(path_prefix, mmap=mmap, reuse_index=reuse_index)

    def __del__(self) -> None:
        """Close the token mmap when this reader is released."""
        bin_buffer_mmap = getattr(self, "bin_buffer_mmap", None)
        mmap_handle = getattr(bin_buffer_mmap, "_mmap", None)
        if mmap_handle is not None:
            mmap_handle.close()

    def __len__(self) -> int:
        """Return the number of indexed token sequences."""
        return len(self._require_index())

    def _getitem_mmap(self, index: int | np.integer | slice) -> np.ndarray | list[np.ndarray]:
        """Return one sequence or a contiguous sequence slice from the mapped token payload."""
        index_reader = self._require_index()

        if isinstance(index, (int, np.integer)):
            sequence_pointer, sequence_length = index_reader[index]
            sequence = np.frombuffer(
                self.bin_buffer, dtype=index_reader.dtype, count=sequence_length, offset=sequence_pointer
            )
            return sequence

        if isinstance(index, slice):
            start, stop, step = index.indices(len(self))
            if step != 1:
                raise ValueError("Slices into indexed Dataset must be contiguous")
            if start == stop:
                return []

            sequence_lengths = index_reader.sequence_lengths[index]
            sequence_offsets = np.cumsum(sequence_lengths)
            sequences = np.frombuffer(
                self.bin_buffer, dtype=index_reader.dtype, count=int(sequence_lengths.sum()),
                offset=int(index_reader.sequence_pointers[start]),
            )
            sequences = np.split(sequences, sequence_offsets[:-1])
            return sequences

        raise TypeError(f"Indexed Dataset indices must be integers or slices, got {type(index).__name__}")

    def _getitem_file(self, index: int | np.integer | slice) -> np.ndarray:
        """Return one sequence through direct file reading when mmap is disabled."""
        if isinstance(index, slice):
            raise NotImplementedError("Slicing is not implemented when mmap is disabled")

        if not isinstance(index, (int, np.integer)):
            raise TypeError(f"Indexed Dataset indices must be integers or slices, got {type(index).__name__}")

        return self.get(int(index))

    def __getitem__(self, index: int | np.integer | slice) -> np.ndarray | list[np.ndarray]:
        """Return one sequence, or a contiguous slice when mmap is enabled."""
        if self.bin_buffer is not None:
            return self._getitem_mmap(index)

        return self._getitem_file(index)

    def get(self, index: int, offset: int = 0, length: int | None = None) -> np.ndarray:
        """Read a contiguous token range from one sequence.

        Args:
            index: Sequence index.

        Returns:
            A NumPy view of the requested tokens.
        """
        index_reader = self._require_index()
        sequence_pointer, sequence_length = index_reader[index]
        sequence_length = int(sequence_length)
        read_length = sequence_length - offset if length is None else length
        if offset < 0 or read_length < 0 or offset + read_length > sequence_length:
            raise ValueError(
                f"Token range [{offset}, {offset + read_length}) exceeds sequence length {sequence_length}"
            )

        byte_offset = int(sequence_pointer) + offset * index_reader.dtype.itemsize
        if self.bin_buffer:
            sequence = np.frombuffer(self.bin_buffer, dtype=index_reader.dtype, count=read_length, offset=byte_offset)
        else:
            sequence = np.empty(read_length, dtype=index_reader.dtype)
            with open(self.path_prefix + ".bin", mode="rb", buffering=0) as bin_buffer_file:
                bin_buffer_file.seek(byte_offset)
                bytes_read = bin_buffer_file.readinto(sequence)

            byte_length = read_length * index_reader.dtype.itemsize
            if bytes_read != byte_length:
                raise ValueError(f"Indexed Dataset read returned {bytes_read} bytes; expected {byte_length}")

        return sequence

    @property
    def sequence_lengths(self) -> np.ndarray:
        """Return the token length of every indexed sequence."""
        return self._require_index().sequence_lengths

    @property
    def document_indices(self) -> np.ndarray:
        """Return sequence indices that delimit documents."""
        return self._require_index().document_indices

    def get_document_indices(self) -> np.ndarray:
        """Return sequence indices that delimit documents."""
        return self.document_indices

    def set_document_indices(self, document_indices: np.ndarray) -> None:
        """Replace sequence indices that delimit documents."""
        self._require_index().document_indices = document_indices

    def _require_index(self) -> _IndexReader:
        """Return the initialized index reader."""
        if self.index is None:
            raise ValueError("Indexed Dataset metadata is not initialized")
        return self.index

# limitations under the License.
# ============================================================================
"""Lazy proxy for deferred indexed Dataset construction."""

from __future__ import annotations

from collections.abc import Callable, Mapping
from typing import Any

from hyper_parallel.data.dataset_logging import get_dataset_logger

logger = get_dataset_logger(__name__)


class LazyDatasetProxy:
    """Construct and cache a Dataset when it is first accessed."""

    def __init__(
            self,
            dataset_factory: Callable[[], Any],
            *,
            unique_identifiers: Any,

                sample access.
            unique_identifiers: Stable Dataset cache identity assembled from
                construction inputs.
        """
        self._dataset_factory = dataset_factory
        self._unique_identifiers = unique_identifiers
        self._dataset = None

    def __len__(self) -> int:
        """Initialize the Dataset and return its length."""
        return len(self._get_dataset())

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        """Construct the Dataset if necessary and read one sample."""
        dataset = self._get_dataset()
        sample = dataset[index]
        return sample

    @property
    def unique_identifiers(self) -> Any:
        """Initialize the Dataset and return its cache identity."""
        return self._get_dataset().unique_identifiers

    def _get_dataset(self) -> Any:
        """Construct the Dataset once and reuse it for later accesses."""
        if self._dataset is None:
            identifiers = self._unique_identifiers if isinstance(self._unique_identifiers, Mapping) else {}
            logger.debug(
                "Initializing lazy Dataset: class=%s, path=%s, split=%s, samples=%s",
                identifiers.get("class"), identifiers.get("dataset_path"), identifiers.get("index_split"),
                identifiers.get("num_samples"),
            )
            self._dataset = self._dataset_factory()
        return self._dataset

# limitations under the License.
# ============================================================================
"""Index-free blending for pre-shuffled MR Datasets."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, Literal

SimpleBlendMode = Literal["inter", "intra"]


class SimpleBlendedDataset:
    """Blend equally weighted MR Datasets in interleaved or contiguous order."""

    def __init__(
            self,
            datasets: Sequence[Any],
            size: int,
            mode: SimpleBlendMode,

            size: int,
            mode: SimpleBlendMode,
    ) -> None:
        """Validate the inputs and create the configured simple ordering."""
        if not datasets:
            raise ValueError("simple blend requires at least one Dataset")
        self.datasets = list(datasets)
        self.size = size
        self.mode = mode
        self._locations = self._build_locations()
        if self.size > len(self._locations):
            raise ValueError(
                f"Requested {self.size} samples from a simple blend containing "
                f"{len(self._locations)} samples"
            )

    def __len__(self) -> int:
        """Return the requested blended sample count."""
        return self.size

    def __getitem__(self, index: int) -> Mapping[str, Any]:
        """Read one MR sample in interleaved or Dataset-contiguous order."""
        dataset_id, sample_id = self._locations[index]
        sample = {"dataset_id": dataset_id, **self.datasets[dataset_id][sample_id]}
        return sample

    def _build_locations(self) -> list[tuple[int, int]]:
        """Map blend positions to Dataset and sample indices."""
        if self.mode == "intra":
            locations = [
                (dataset_id, sample_id)
                for dataset_id, dataset in enumerate(self.datasets)
                for sample_id in range(len(dataset))
            ]
            return locations

        # mode = inter
        locations = self._build_interleaved_locations()
        return locations

    def _build_interleaved_locations(self) -> list[tuple[int, int]]:
        """Alternate over non-empty Datasets until every sample is exposed."""
        max_size = max(len(dataset) for dataset in self.datasets)
        locations = [
            (dataset_id, sample_id)
            for sample_id in range(max_size)
            for dataset_id, dataset in enumerate(self.datasets)
            if sample_id < len(dataset)

            for sample_id in range(max_size)
            for dataset_id, dataset in enumerate(self.datasets)
            if sample_id < len(dataset)
        ]
        return locations

            raise ValueError(f"seq_len ({seq_len}) must be divisible by cp_size ({cp_size})")

        cp_rank = self.parallel_context.cp_rank
        for field, value in cp_batch.items():
            local_value = torch.chunk(value, cp_size, dim=1)[cp_rank]
            cp_batch[field] = local_value.contiguous()

        return cp_batch

                for field, value in cp_local_batch.items()
            }
            local_cu_seq_lens = None
            if cu_seq_lens is not None:
                local_cu_seq_lens = cu_seq_lens.to(self.device, dtype=torch.int32, non_blocking=True)
        else:
            parallel_batch = None
            local_cu_seq_lens = None

        # Broadcast compact shape metadata before receivers allocate tensors.
        if tp_rank == 0:
            batch_size, seq_len = parallel_batch["input_ids"].shape
            num_boundaries = 0 if local_cu_seq_lens is None else local_cu_seq_lens.numel()
            batch_meta = torch.tensor(
                [batch_size, seq_len, num_boundaries],
                dtype=torch.int64,
                device=self.device,
            )

                dtype=torch.int64,
                device=self.device,
            )
        else:
            batch_meta = torch.empty((3,), dtype=torch.int64, device=self.device)

        tp_group = self.parallel_context.tp_group
        dist.broadcast(batch_meta, group=tp_group, group_src=0)

        # Non-source TP ranks allocate the same device shapes and dtypes.
        if tp_rank != 0:
            batch_size, seq_len, num_boundaries = [int(value) for value in batch_meta.tolist()]

                "input_ids": torch.empty(shape, dtype=torch.int64, device=self.device),
                "labels": torch.empty(shape, dtype=torch.int64, device=self.device),
            }
            if num_boundaries > 0:
                local_cu_seq_lens = torch.empty((num_boundaries,), dtype=torch.int32, device=self.device)

        # Broadcast only model fields that cannot be regenerated locally.
        for field in ("input_ids", "labels"):
            dist.broadcast(parallel_batch[field], group=tp_group, group_src=0)

        # Packed boundaries remain global across CP ranks and variable in size.
        if local_cu_seq_lens is not None:
            dist.broadcast(local_cu_seq_lens, group=tp_group, group_src=0)

        parallel_batch["cu_seq_lens"] = local_cu_seq_lens

        return parallel_batch

def _validate_positive_integer(value: int, name: str) -> None:
    """Validate an integer boundary shared by sampler options."""
    if value <= 0:
        raise ValueError(f"{name} must be a positive integer")


class _DatasetBatchSampler:
    """Share validation, DP slicing, epoch, and checkpoint state.

            drop_last: bool,
            index_mapping: IndexMapping | None,
    ) -> None:
        """Validate and store DP slicing and checkpoint state."""
        _validate_positive_integer(total_samples, "total_samples")
        _validate_positive_integer(micro_batch_size, "micro_batch_size")
        _validate_positive_integer(dp_world_size, "dp_world_size")
        self._validate_consumed_samples(consumed_samples, total_samples)
        if not 0 <= dp_rank < dp_world_size:
            raise ValueError(f"dp_rank must be in [0, {dp_world_size}), but got {dp_rank!r}")

        self.total_samples = total_samples
        self.consumed_samples = consumed_samples
        self.micro_batch_size = micro_batch_size
        self.dp_rank = dp_rank
        self.dp_world_size = dp_world_size
        self.drop_last = drop_last
        self.index_mapping = index_mapping
        self.global_micro_batch_size = micro_batch_size * dp_world_size
        _validate_positive_integer(global_batch_size, "global_batch_size")
        if global_batch_size % self.global_micro_batch_size != 0:
            raise ValueError("global_batch_size must be divisible by micro_batch_size * dp_world_size")

        self.global_batch_size = global_batch_size
        self.epoch = 0
        self._resume_at_source_batch = False

    def _validate_consumed_samples(self, consumed_samples: int, total_samples: int) -> None:
        """Allow sequential sampling to stop exactly at the epoch boundary."""
        if not 0 <= consumed_samples <= total_samples:
            raise ValueError("consumed_samples must be in [0, total_samples]")

    def __iter__(self) -> Iterator[list[int]]:
        """Yield sequential rank-local indices for ``single`` mode."""
        while self.consumed_samples < self.total_samples:
            block_start = self.consumed_samples
            block_stop = min(block_start + self.global_micro_batch_size, self.total_samples)
            block_size = block_stop - block_start
            if self.drop_last and block_size < self.global_micro_batch_size:
                self.consumed_samples = self.total_samples
                break

            local_start = block_start + self.dp_rank * self.micro_batch_size
            local_stop = min(local_start + self.micro_batch_size, block_stop)
            self.consumed_samples = block_stop
            if local_start >= block_stop:
                continue

            local_indices = [
                self._resolve_index(index)
                for index in range(local_start, local_stop)
            ]
            yield local_indices

    def __len__(self) -> int:
        """Return the remaining number of micro-batches for this DP rank."""
        remaining_samples = self.total_samples - self.consumed_samples
        full_batches, partial_samples = divmod(remaining_samples, self.global_micro_batch_size)
        if self.drop_last or partial_samples <= self.dp_rank * self.micro_batch_size:
            return full_batches
        return full_batches + 1

    def _resolve_index(self, logical_index: int) -> int:
        """Apply an optional logical-to-physical sample index mapping."""
        if self.index_mapping is None:
            return logical_index

        resolved_index = self.index_mapping[logical_index]
        if isinstance(resolved_index, bool):
            raise TypeError(f"index_mapping[{logical_index}] must be an integer")

        physical_index = operator.index(resolved_index)
        return physical_index

    def state_dict(self) -> dict[str, int]:
        """Return the Dataset epoch and position needed to resume iteration."""
        sampler_state = {
            "consumed_samples": self.consumed_samples,
            "epoch": self.epoch,
            "global_batch_size": self.global_batch_size,
            "dp_world_size": self.dp_world_size,