florence_forge/data/loader.py

"""FlorenceForge任务数据加载器模块

提供高效的数据加载、批处理和采样功能
"""

import torch
import os
from torch.utils.data import DataLoader, Sampler
import random
import logging
import numpy as np
from typing import Optional, Iterator, Dict, Any
from collections import defaultdict, deque

from .dataset import MultiTaskDataset
from .collate import Florence2Collator
from ..core.config import DataConfig

logger = logging.getLogger(__name__)


def _get_distributed_info() -> tuple[int, int, int]:
    """获取分布式训练信息

    自动检测当前是否处于分布式环境，返回 (world_size, rank, local_rank)。
    如果未初始化分布式，则返回 (1, 0, 0)。

    Returns:
        (world_size, rank, local_rank) 元组
    """
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        # local_rank 无法直接从 torch.distributed 获取，从环境变量推断
        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        return world_size, rank, local_rank
    return 1, 0, 0

class TaskBalancedSampler(Sampler):
    """任务平衡采样器
    
    确保每个批次中包含来自不同任务的样本
    """
    
    def __init__(
        self,
        dataset: MultiTaskDataset,
        batch_size: int,
        drop_last: bool = True,
        shuffle: bool = True,
        seed: Optional[int] = None
    ):
        """初始化采样器
        
        Args:
            dataset: 多任务数据集
            batch_size: 批次大小
            drop_last: 是否丢弃最后不完整的批次
            shuffle: 是否打乱顺序
            seed: 随机种子
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.shuffle = shuffle
        self.seed = seed
        
        self.task_indices = dataset.task_indices
        self.task_weights = dataset.task_weights
        
        # 计算总批次数
        self.num_samples = len(dataset)
        if self.drop_last:
            self.num_batches = self.num_samples // self.batch_size
        else:
            self.num_batches = (self.num_samples + self.batch_size - 1) // self.batch_size
    
    def __iter__(self) -> Iterator[int]:
        """生成采样索引（预计算优化版本）"""
        if self.seed is not None:
            random.seed(self.seed)
            np.random.seed(self.seed)

        # 为每个任务创建索引池（使用 deque 支持 O(1) pop）
        task_pools = {}
        for task_type, indices in self.task_indices.items():
            pool = deque(indices)
            if self.shuffle:
                # 转为 list 打乱后重新构建 deque
                shuffled = list(pool)
                random.shuffle(shuffled)
                pool = deque(shuffled)
            task_pools[task_type] = pool

        # 计算任务权重
        task_types = list(self.task_indices.keys())
        if self.task_weights:
            weights = [self.task_weights.get(task, 1.0) for task in task_types]
        else:
            weights = [1.0] * len(task_types)

        # 归一化权重
        total_weight = sum(weights)
        weights = [w / total_weight for w in weights]

        # 预计算所有任务选择（比逐样本 np.random.choice 快 10-50 倍）
        total_indices_needed = self.num_batches * self.batch_size
        task_choices = np.random.choice(
            len(task_types), size=total_indices_needed, p=weights
        )

        # 按预计算的选择分配索引
        idx = 0
        for batch_idx in range(self.num_batches):
            batch_indices = []

            for sample_idx in range(self.batch_size):
                if idx >= len(task_choices):
                    break
                task_type = task_types[task_choices[idx]]
                idx += 1

                # 从任务池中获取样本 (deque.popleft() O(1))
                if task_pools[task_type]:
                    batch_indices.append(task_pools[task_type].popleft())
                else:
                    # 如果任务池为空，重新填充
                    pool = deque(self.task_indices[task_type])
                    if self.shuffle:
                        shuffled = list(pool)
                        random.shuffle(shuffled)
                        pool = deque(shuffled)
                    task_pools[task_type] = pool

                    if pool:
                        batch_indices.append(task_pools[task_type].popleft())

            # 如果批次不完整且需要丢弃，则跳过
            if len(batch_indices) < self.batch_size and self.drop_last:
                continue

            yield from batch_indices
    
    def __len__(self) -> int:
        """返回采样器长度"""
        return self.num_batches * self.batch_size if self.drop_last else self.num_samples

class TaskRoundRobinSampler(Sampler):
    """任务轮询采样器
    
    按轮询方式从不同任务中采样
    """
    
    def __init__(
        self,
        dataset: MultiTaskDataset,
        shuffle: bool = True,
        seed: Optional[int] = None
    ):
        """初始化采样器
        
        Args:
            dataset: 多任务数据集
            shuffle: 是否打乱任务内顺序
            seed: 随机种子
        """
        self.dataset = dataset
        self.shuffle = shuffle
        self.seed = seed
        
        self.task_indices = dataset.task_indices
        self.num_samples = len(dataset)
    
    def __iter__(self) -> Iterator[int]:
        """生成采样索引"""
        if self.seed is not None:
            random.seed(self.seed)
        
        # 为每个任务创建索引迭代器
        task_iterators = {}
        for task_type, indices in self.task_indices.items():
            task_list = indices.copy()
            if self.shuffle:
                random.shuffle(task_list)
            task_iterators[task_type] = iter(task_list)
        
        task_types = list(self.task_indices.keys())
        if self.shuffle:
            random.shuffle(task_types)
        
        # 轮询生成样本
        samples_generated = 0
        task_idx = 0
        
        while samples_generated < self.num_samples:
            current_task = task_types[task_idx % len(task_types)]
            
            try:
                idx = next(task_iterators[current_task])
                yield idx
                samples_generated += 1
            except StopIteration:
                # 当前任务已耗尽，重新创建迭代器
                task_list = self.task_indices[current_task].copy()
                if self.shuffle:
                    random.shuffle(task_list)
                task_iterators[current_task] = iter(task_list)
                
                # 尝试再次获取
                try:
                    idx = next(task_iterators[current_task])
                    yield idx
                    samples_generated += 1
                except StopIteration:
                    # 如果任务为空，记录警告并跳过
                    logger.warning(f"任务 {current_task} 的数据已完全耗尽，跳过此轮采样")
                    # 从任务列表中移除空任务，避免无限循环
                    if current_task in task_types:
                        task_types.remove(current_task)
                    
                    # 如果所有任务都为空，提前结束
                    if not task_types:
                        logger.error("所有任务数据都已耗尽，提前结束采样")
                        break
            
            task_idx += 1
    
    def __len__(self) -> int:
        """返回采样器长度"""
        return self.num_samples


class DistributedTaskSampler(Sampler):
    """分布式任务采样器包装器

    将任意底层采样器（如 TaskBalancedSampler、TaskRoundRobinSampler）
    包装为分布式版本，确保多进程训练中每个 rank 只处理不同的数据子集。

    实现方式：先由底层采样器生成完整的 epoch 索引序列，
    再按 rank 和 world_size 切片分配，避免重复采样。

    用法示例::

        base_sampler = TaskBalancedSampler(dataset, batch_size=4)
        dist_sampler = DistributedTaskSampler(
            base_sampler, world_size=4, rank=0, shuffle=True
        )
        loader = DataLoader(dataset, sampler=dist_sampler, ...)
    """

    def __init__(
        self,
        sampler: Sampler,
        world_size: int,
        rank: int,
        shuffle: bool = True,
        seed: int = 0,
        drop_last: bool = False,
    ):
        """初始化分布式采样器

        Args:
            sampler: 底层采样器（如 TaskBalancedSampler）
            world_size: 分布式世界大小（进程总数）
            rank: 当前进程 rank
            shuffle: 是否在每个 epoch 打乱顺序
            seed: 随机种子（配合 epoch 使用）
            drop_last: 是否丢弃不完整的尾部数据
        """
        self.sampler = sampler
        self.world_size = world_size
        self.rank = rank
        self.shuffle = shuffle
        self.seed = seed
        self.drop_last = drop_last
        self.epoch = 0

        # 计算当前 rank 应处理的样本数
        full_length = len(sampler)
        if drop_last:
            self.num_samples = full_length // world_size
        else:
            self.num_samples = (full_length + world_size - 1) // world_size

        self.total_size = self.num_samples * world_size

    def __iter__(self) -> Iterator[int]:
        """生成分布式索引序列"""
        # 1. 获取底层采样器的完整索引序列
        indices = list(self.sampler)

        # 2. 如果需要，填充到 world_size 的整数倍（保证每个 rank 样本数一致）
        if not self.drop_last:
            # 循环填充
            padding_size = self.total_size - len(indices)
            if padding_size > 0:
                indices += indices[:padding_size]
        else:
            # 截断到 world_size 整数倍
            indices = indices[: self.total_size]

        # 3. 根据 epoch 设置随机种子进行打乱（保证不同 epoch 顺序不同）
        if self.shuffle:
            # 使用确定性随机，保证所有 rank 看到相同的打乱顺序
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            # 将 indices 转为 tensor 进行随机打乱，再转回 list
            indices_tensor = torch.tensor(indices, dtype=torch.int64)
            perm = torch.randperm(len(indices_tensor), generator=g)
            indices = indices_tensor[perm].tolist()

        # 4. 按 rank 切片
        start_idx = self.rank * self.num_samples
        end_idx = start_idx + self.num_samples
        rank_indices = indices[start_idx:end_idx]

        return iter(rank_indices)

    def __len__(self) -> int:
        """返回当前 rank 的样本数"""
        return self.num_samples

    def set_epoch(self, epoch: int) -> None:
        """设置当前 epoch（用于同步打乱顺序）

        应在每个 epoch 开始时调用，确保不同 epoch 的采样顺序不同。

        Args:
            epoch: 当前训练轮次
        """
        self.epoch = epoch
        # 同时传播给底层采样器（如果支持）
        if hasattr(self.sampler, 'set_epoch'):
            self.sampler.set_epoch(epoch)


class TaskDataLoader:
    """任务数据加载器
    
    提供多种采样策略和批处理选项，支持分布式训练。
    """
    
    def __init__(
        self,
        dataset: MultiTaskDataset,
        config: Optional[DataConfig] = None,
        sampling_strategy: str = "balanced",
        **kwargs
    ):
        """初始化数据加载器
        
        Args:
            dataset: 多任务数据集
            config: 数据配置
            sampling_strategy: 采样策略 ("balanced", "round_robin", "random")
            **kwargs: 其他DataLoader参数
        """
        self.dataset = dataset
        self.config = config or DataConfig()
        self.sampling_strategy = sampling_strategy
        
        # 检测分布式环境
        self._distributed, self._world_size, self._rank, self._local_rank = self._detect_distributed()
        
        # 设置默认参数
        # pin_memory仅在CUDA/MPS设备可用时有效
        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
        pin_memory = self.config.pin_memory and (torch.cuda.is_available() or mps_available)
        
        # 使用 Florence2Collator 替代旧 collate_fn，支持动态 padding
        collator = Florence2Collator()

        self.dataloader_kwargs = {
            "batch_size": self.config.batch_size,
            "num_workers": self.config.num_workers,
            "pin_memory": pin_memory,
            "drop_last": self.config.drop_last,
            "collate_fn": collator  # 使用专业的 Florence2Collator
        }

        # DataLoader 性能调优参数
        if self.config.num_workers > 0:
            if self.config.prefetch_factor is not None:
                self.dataloader_kwargs["prefetch_factor"] = self.config.prefetch_factor
            if self.config.persistent_workers:
                self.dataloader_kwargs["persistent_workers"] = True
        
        # 更新用户提供的参数
        self.dataloader_kwargs.update(kwargs)

        if (
            getattr(self.dataset, "processor", None) is not None
            and self.dataloader_kwargs.get("num_workers", 0) > 0
        ):
            logger.warning(
                "检测到数据集依赖 processor 进行在线编码，已将 num_workers 设为 0，"
                "避免 spawn worker 丢失 processor 后返回未编码样本。"
            )
            self.dataloader_kwargs["num_workers"] = 0
            self.dataloader_kwargs.pop("prefetch_factor", None)
            self.dataloader_kwargs.pop("persistent_workers", None)
        
        # 设置采样器
        self._setup_sampler()
    
    def _detect_distributed(self) -> tuple[bool, int, int, int]:
        """检测分布式训练环境

        优先使用配置中的显式设置，否则自动检测 torch.distributed 环境。

        Returns:
            (is_distributed, world_size, rank, local_rank)
        """
        # 1. 如果配置显式启用分布式，使用配置值
        if self.config.distributed:
            ws = self.config.world_size or 1
            r = self.config.rank or 0
            lr = self.config.local_rank or 0
            logger.info(f"分布式训练已显式启用: world_size={ws}, rank={r}, local_rank={lr}")
            return True, ws, r, lr

        # 2. 自动检测 torch.distributed 环境
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            ws = torch.distributed.get_world_size()
            r = torch.distributed.get_rank()
            lr = int(os.environ.get("LOCAL_RANK", 0))
            logger.info(f"自动检测到分布式环境: world_size={ws}, rank={r}, local_rank={lr}")
            return True, ws, r, lr

        # 3. 检测环境变量（torchrun / launch 方式）
        if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
            ws = int(os.environ["WORLD_SIZE"])
            r = int(os.environ["RANK"])
            lr = int(os.environ.get("LOCAL_RANK", 0))
            if ws > 1:
                logger.info(f"通过环境变量检测到分布式: world_size={ws}, rank={r}, local_rank={lr}")
                return True, ws, r, lr

        # 4. 非分布式环境
        return False, 1, 0, 0
    
    def _setup_sampler(self) -> None:
        """设置采样器（含分布式包装）"""
        # 1. 创建底层采样器
        if self.sampling_strategy == "balanced":
            base_sampler = TaskBalancedSampler(
                dataset=self.dataset,
                batch_size=self.dataloader_kwargs["batch_size"],
                drop_last=self.dataloader_kwargs["drop_last"],
                shuffle=self.config.shuffle
            )
            # 使用自定义采样器时，不能同时设置shuffle
            self.dataloader_kwargs["shuffle"] = False
        
        elif self.sampling_strategy == "round_robin":
            base_sampler = TaskRoundRobinSampler(
                dataset=self.dataset,
                shuffle=self.config.shuffle
            )
            self.dataloader_kwargs["shuffle"] = False
        
        elif self.sampling_strategy == "random":
            # 使用默认的随机采样
            base_sampler = None
            self.dataloader_kwargs["shuffle"] = self.config.shuffle
        
        else:
            raise ValueError(f"未知采样策略: {self.sampling_strategy}")

        # 2. 如果处于分布式环境，包装为 DistributedTaskSampler
        if self._distributed and base_sampler is not None:
            seed = self.config.distributed_seed or 42
            self.sampler = DistributedTaskSampler(
                sampler=base_sampler,
                world_size=self._world_size,
                rank=self._rank,
                shuffle=self.config.shuffle,
                seed=seed,
                drop_last=self.config.drop_last,
            )
            self.dataloader_kwargs["sampler"] = self.sampler
            logger.info(
                f"已启用分布式采样: rank={self._rank}/{self._world_size}, "
                f"本地样本数={len(self.sampler)}"
            )
        elif base_sampler is not None:
            self.sampler = base_sampler
            self.dataloader_kwargs["sampler"] = self.sampler
        else:
            self.sampler = None
    
    def get_dataloader(self) -> DataLoader:
        """获取PyTorch DataLoader
        
        Returns:
            配置好的DataLoader
        """
        return DataLoader(self.dataset, **self.dataloader_kwargs)
    
    def __iter__(self):
        """迭代器接口"""
        return iter(self.get_dataloader())
    
    def __len__(self) -> int:
        """返回批次数量"""
        batch_size = self.dataloader_kwargs["batch_size"]
        if self.sampler:
            if self.dataloader_kwargs["drop_last"]:
                return len(self.sampler) // batch_size
            return (len(self.sampler) + batch_size - 1) // batch_size
        else:
            dataset_size = len(self.dataset)
            if self.dataloader_kwargs["drop_last"]:
                return dataset_size // batch_size
            else:
                return (dataset_size + batch_size - 1) // batch_size
    
    def get_task_distribution(self, num_batches: int = 10) -> Dict[str, float]:
        """分析任务分布
        
        Args:
            num_batches: 分析的批次数量
            
        Returns:
            任务分布统计
        """
        task_counts = defaultdict(int)
        total_samples = 0
        
        dataloader = self.get_dataloader()
        
        for i, batch in enumerate(dataloader):
            if i >= num_batches:
                break
            
            # 兼容不同的批次结构
            if "task_types" in batch:
                # 多任务类型批次
                for task_type in batch["task_types"]:
                    task_counts[task_type] += 1
                    total_samples += 1
            elif "task_type" in batch:
                # 单任务类型批次
                task_type = batch["task_type"]
                batch_size = batch["input_ids"].size(0) if "input_ids" in batch else 1
                task_counts[task_type] += batch_size
                total_samples += batch_size
            else:
                logger.warning(f"批次 {i} 中未找到任务类型信息")
        
        # 计算比例
        if total_samples > 0:
            task_distribution = {
                task: count / total_samples
                for task, count in task_counts.items()
            }
        else:
            task_distribution = {}
        
        return task_distribution
    
    def set_epoch(self, epoch: int) -> None:
        """设置训练轮次（用于分布式训练）
        
        Args:
            epoch: 当前轮次
        """
        if hasattr(self.sampler, 'set_epoch'):
            self.sampler.set_epoch(epoch)
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取数据加载器统计信息
        
        Returns:
            统计信息字典
        """
        stats = {
            "dataset_size": len(self.dataset),
            "batch_size": self.dataloader_kwargs["batch_size"],
            "num_batches": len(self),
            "sampling_strategy": self.sampling_strategy,
            "num_workers": self.dataloader_kwargs["num_workers"],
            "pin_memory": self.dataloader_kwargs["pin_memory"],
            "drop_last": self.dataloader_kwargs["drop_last"],
            "dataset_statistics": self.dataset.get_task_statistics()
        }
        # 分布式信息
        stats["distributed"] = self._distributed
        if self._distributed:
            stats["world_size"] = self._world_size
            stats["rank"] = self._rank
            stats["local_rank"] = self._local_rank
            if self.sampler is not None:
                stats["local_samples"] = len(self.sampler)
        return stats