GRID/src/data/loading/components/dataloading.py

import copy
from typing import Any, List, Optional

from torch.utils.data import IterableDataset, get_worker_info

from src.data.loading.components.interfaces import BaseDatasetConfig
from src.utils.pylogger import RankedLogger

command_line_logger = RankedLogger(__name__, rank_zero_only=True)


class BaseDataset:
    def __init__(
        self,
        dataset_config: BaseDatasetConfig,
        data_folder: str,
        should_shuffle_rows: bool = False,
        batch_size: int = 1,
        is_for_training: bool = True,
        assign_all_files_per_worker: bool = False,
    ):
        """
        Base class for all datasets. This class is used to set up the dataset and provide the list of files to be used.
        Args:
            dataset_config (BaseDatasetConfig): Configuration for the dataset.
            data_folder (str): Path to the folder where the data is stored.
            should_shuffle_rows (bool): Whether to shuffle the rows of the dataset.
            batch_size (int): Batch size to be used for the dataset.
            is_for_training (bool): Whether the dataset is for training or not.
            assign_all_files_per_worker (bool): Whether to assign all files to each worker or not.
                This will enable each worker to access all files. Each worker will locally shuffle the files.
                This would be useful for small datasets. In smaller datasets, if each worker only observes a subset of the files,
                it may not be able to learn the distribution of the data.
        """
        self.dataset_config = dataset_config
        self.should_shuffle_rows = should_shuffle_rows
        self.data_folder = data_folder
        self.list_of_file_paths = []
        self.batch_size = batch_size
        self.is_for_training = is_for_training
        self.assign_all_files_per_worker = assign_all_files_per_worker

    def set_list_of_files(self, list_of_files: List[str]):
        self.list_of_file_paths = list_of_files

    def set_distributed_params(self, total_workers: int, global_worker_id: int):
        # TODO (lneves): figure out how to do this in LightningDataModule
        self.total_workers = total_workers
        self.global_worker_id = global_worker_id

    def get_worker_id_and_num_workers(self):
        worker_info = get_worker_info()

        if worker_info is None:
            # Single-worker setup (no multiprocessing)
            worker_id = 0
            num_workers = 1
        else:
            # Multi-worker setup
            worker_id = worker_info.id
            num_workers = worker_info.num_workers

        self.global_dataloader_worker_id = (
            self.global_worker_id * num_workers + worker_id
        )

        return worker_id, num_workers

    def get_list_of_worker_files(self):
        # Get information about worker and then separate only files that belong to this worker
        worker_id, num_workers = self.get_worker_id_and_num_workers()
        if self.assign_all_files_per_worker:
            worker_files = self.list_of_file_paths
        else:
            worker_files = self.list_of_file_paths[worker_id::num_workers]
        command_line_logger.debug(
            f"GPU Worker: {self.global_worker_id}/{self.total_workers} CPU Worker {worker_id} has {len(worker_files)} files"
        )
        return worker_files

    def setup(self):
        pass


class UnboundedSequenceIterable(BaseDataset, IterableDataset):
    """An unbounded dataset is a dataset that we don't know the size of beforehand.
    For training, we will iterate over the dataset infinitely. For evaluation, we will iterate over the dataset once.
    """

    def __init__(
        self,
        dataset_config: BaseDatasetConfig,
        data_folder: str,
        should_shuffle_rows: bool = False,
        batch_size: int = 1,
        is_for_training: bool = True,
        assign_all_files_per_worker: bool = False,
    ):
        super().__init__(
            dataset_config=dataset_config,
            data_folder=data_folder,
            should_shuffle_rows=should_shuffle_rows,
            batch_size=batch_size,
            is_for_training=is_for_training,
            assign_all_files_per_worker=assign_all_files_per_worker,
        )
        self.data_iterator = dataset_config.data_iterator
        self.dataset_to_iterate = None

    def setup(self):
        # We update each worker's data iterator with the files just for that worker.
        self.data_iterator.update_list_of_file_paths(self.get_list_of_worker_files())
        self.data_iterator = (
            # here we use global_dataloader_worker_id as the seed for shuffling
            # this doesn't matter for the case where workers have non-overlapping files
            # but it does matter for the case where workers have all files
            # (e.g. when using assign_all_files_per_worker)
            # the same seed is used for all workers would cause duplicated examples returned by different workers
            self.data_iterator.shuffle(seed=self.global_dataloader_worker_id)
            if self.should_shuffle_rows
            else self.data_iterator
        )
        self.data_iterator.should_shuffle_rows = self.should_shuffle_rows
        # We provide the flexibility to iterate per row, if per row preprocessing is needed, or per batch.
        self.dataset_to_iterate = (
            self.data_iterator.iterrows()
            if self.dataset_config.iterate_per_row
            else self.data_iterator.iter_batches(self.batch_size)
        )

        command_line_logger.debug(
            f"GLOBAL ID {self.global_dataloader_worker_id} GPU Worker: {self.global_worker_id}/{self.total_workers} with {len(self.data_iterator.list_of_file_paths)} files\
                First five files are: {self.data_iterator.list_of_file_paths[:5]}"
        )

    def __iter__(self):
        if self.dataset_to_iterate is None:
            # If it has not been set up, it means it is a forkserver worker. We need to set it up.
            self.setup()
        # If the dataset is for training, we want to keep iterating over the dataset infinitely.
        # On a streaming dataset, we will always be on Epoch 0.
        finished_iteration = False
        while not finished_iteration:

            for row_or_batch in self.dataset_to_iterate:
                for (
                    preprocessing_function
                ) in self.dataset_config.preprocessing_functions:
                    row_or_batch = preprocessing_function(
                        row_or_batch, dataset_config=self.dataset_config
                    )
                    if row_or_batch is None:
                        break
                if row_or_batch:
                    yield row_or_batch
            # if the dataset is not for training, we stop the loop. Otherwise, we continue.
            finished_iteration = not self.is_for_training
            if not finished_iteration:
                self.setup()
        # We reset the dataset to iterate to None, so that it is set up again in the next iteration.
        # This is required for validation when persitent_workers = True.
        self.dataset_to_iterate = None
        return None