source/src/speculators/train/dataloader.py

from __future__ import annotations

import copy
import logging
import os
from typing import TYPE_CHECKING, Any, Literal

if TYPE_CHECKING:
    from collections.abc import Callable

import torch
from torch.utils.data import DataLoader

from hs_connectors import HiddenStatesTransfer
from speculators.train.data import (
    ArrowDataset,
    BaseDataset,
    CollateFn,
)
from speculators.train.distributed import get_dp_rank, get_dp_size
from speculators.train.distributed_batch_sampler import (
    MultipackDistributedBatchSamplerV2,
)
from speculators.train.noise_transforms import AddUniformNoise

logger = logging.getLogger(__name__)

BatchType = dict[str, Any]


def _limit_worker_threads() -> None:
    """Limit per-worker thread pools to avoid thread exhaustion.

    With ``multiprocessing_context='spawn'``, each worker is a full process
    that re-imports numpy (OpenBLAS) and torch, each creating thread pools
    sized to the core count.  DataLoader workers only do I/O and tensor
    slicing — they don't benefit from intra-op parallelism.

    The env vars must be set before numpy/torch are imported to take effect
    on OpenBLAS/OMP.  Call this at the top of the training entry point,
    before DataLoader construction.
    """
    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")


def _worker_init_fn(worker_id: int) -> None:  # noqa: ARG001
    torch.set_num_threads(1)


def _make_worker_dataset(dataset: BaseDataset) -> BaseDataset:
    """Drop parent-only sampler state from the dataset sent to spawned workers.

    ``approx_lengths`` can contain hundreds of thousands of Python integers. The
    batch sampler copies it into its own NumPy array before worker startup, and
    ``Dataset.__getitem__`` never reads it. Avoiding that field in every spawn pickle
    keeps the parent from blocking while each child imports the training stack.
    """

    worker_dataset = copy.copy(dataset)
    worker_dataset.approx_lengths = []
    return worker_dataset


def _setup_dataloader(
    dataset: BaseDataset,
    total_seq_len: int,
    hidden_size: int,
    num_workers: int = 12,
    num_target_layers: int = 3,
    prefetch_factor: int | None = 4,
    preprocess: Callable[[BatchType], BatchType] | None = None,
    in_order: bool = True,
    max_batches: int | None = None,
) -> DataLoader:
    batch_sampler = MultipackDistributedBatchSamplerV2(
        batch_max_length=total_seq_len,
        lengths=dataset.approx_lengths,
        num_replicas=get_dp_size(),
        rank=get_dp_rank(),
        max_batches=max_batches,
    )
    use_workers = num_workers > 0
    if not in_order and not use_workers:
        # Ordering is a property of the multi-worker result queue; with a single
        # process there is nothing to reorder, so silently ignoring the flag would
        # hide a misconfigured run behind unchanged (slow) behaviour.
        raise ValueError("in_order=False requires num_workers > 0")
    worker_dataset = _make_worker_dataset(dataset) if use_workers else dataset
    return DataLoader(
        worker_dataset,
        batch_sampler=batch_sampler,
        num_workers=num_workers,
        prefetch_factor=prefetch_factor if use_workers else None,
        pin_memory=True,
        collate_fn=CollateFn(
            total_seq_len,
            hidden_size,
            num_target_layers=num_target_layers,
            dtype=dataset.hidden_states_dtype,
            preprocess=preprocess,
        ),
        persistent_workers=use_workers,
        multiprocessing_context="spawn" if use_workers else None,
        worker_init_fn=_worker_init_fn if use_workers else None,
        # Online hidden-state fetches are heavy-tailed, so the strict round-robin
        # of in-order delivery makes the rank wait on worker i even when workers
        # i+1.. already hold finished batches -- and every peer rank then waits on
        # this one at the all-reduce. Out-of-order delivery lets the prefetch queue
        # absorb the tail instead of serialising on it. The epoch still covers the
        # same batches; only mid-epoch resume degrades, since its fast-skip slices
        # the sampler by count and no longer names the exact batches consumed.
        in_order=in_order or not use_workers,
    )


def create_train_val_loaders(
    *,
    data_path: str,
    total_seq_len: int,
    hidden_states_dtype: torch.dtype,
    noise_std: float,
    transfer: HiddenStatesTransfer | None = None,
    vllm_endpoint: str,
    on_missing: Literal["generate", "skip", "warn", "raise"],
    on_generate: Literal["cache", "delete"],
    verifier_name_or_path: str,
    request_timeout: float | None,
    max_retries: int,
    generation_validation_retries: int,
    max_consecutive_generation_failures: int,
    hidden_size: int,
    num_target_layers: int,
    num_workers: int,
    prefetch_factor: int,
    preprocess: Callable[[BatchType], BatchType] | None,
    train_data_ratio: float = 0.9,
    fail_on_hidden_state_error: bool = False,
    fetch_threads: int = 1,
    dataloader_in_order: bool = True,
    vllm_http_keepalive: bool = True,
    max_train_batches: int | None = None,
) -> tuple[DataLoader, DataLoader]:
    """Create training and validation DataLoaders.

    Non-data SP ranks get lightweight loaders with no workers (they receive
    batches via scatter).  Reads DP/SP topology from
    :mod:`speculators.train.distributed`.
    """
    _limit_worker_threads()
    noise_transform = AddUniformNoise(std=noise_std)

    if not (0.0 < train_data_ratio < 1.0):
        raise ValueError(f"train_data_ratio must be in (0, 1), got {train_data_ratio}")

    train_dataset: BaseDataset = ArrowDataset(
        datapath=data_path,
        max_len=total_seq_len,
        transfer=transfer,
        vllm_endpoint=vllm_endpoint,
        on_missing=on_missing,
        on_generate=on_generate,
        transform=noise_transform,
        train_ratio=train_data_ratio,
        split="train",
        model=verifier_name_or_path,
        hidden_states_dtype=hidden_states_dtype,
        request_timeout=request_timeout,
        max_retries=max_retries,
        fail_on_hidden_state_error=fail_on_hidden_state_error,
        fetch_threads=fetch_threads,
        http_keepalive=vllm_http_keepalive,
        generation_validation_retries=generation_validation_retries,
        max_consecutive_generation_failures=max_consecutive_generation_failures,
    )
    val_dataset: BaseDataset = ArrowDataset(
        datapath=data_path,
        max_len=total_seq_len,
        transfer=transfer,
        vllm_endpoint=vllm_endpoint,
        on_missing=on_missing,
        on_generate=on_generate,
        train_ratio=train_data_ratio,
        split="val",
        model=verifier_name_or_path,
        hidden_states_dtype=hidden_states_dtype,
        request_timeout=request_timeout,
        max_retries=max_retries,
        fail_on_hidden_state_error=fail_on_hidden_state_error,
        fetch_threads=fetch_threads,
        http_keepalive=vllm_http_keepalive,
        generation_validation_retries=generation_validation_retries,
        max_consecutive_generation_failures=max_consecutive_generation_failures,
    )

    train_loader = _setup_dataloader(
        train_dataset,
        total_seq_len,
        hidden_size,
        num_target_layers=num_target_layers,
        num_workers=num_workers,
        prefetch_factor=prefetch_factor,
        preprocess=preprocess,
        in_order=dataloader_in_order,
        max_batches=max_train_batches,
    )
    # Validation reduces a mean over the whole split, so its batch order is
    # irrelevant too and it sees the same heavy-tailed fetch latency.
    val_loader = _setup_dataloader(
        val_dataset,
        total_seq_len,
        hidden_size,
        num_target_layers=num_target_layers,
        num_workers=num_workers,
        prefetch_factor=prefetch_factor,
        preprocess=preprocess,
        in_order=dataloader_in_order,
    )

    return train_loader, val_loader