guidellm.data.loaders

class DataLoaderRegistry(Generic[DataT_co], RegistryMixin[type[DataLoader]]):
    @classmethod
    def create(
        cls,
        config: DataLoaderArgs,
        datasets: list[DatasetType],
        preprocessors: list[DatasetPreprocessor],
        finalizer: DatasetFinalizer[DataT_co],
        random_seed: int,
        **kwargs: Any,
    ) -> DataLoader[DataT_co]:
        """
        Factory method to create a DataLoader instance based on provided configuration.

        :param config: A DataEntrypointArgs object containing the configuration.
        """
        kind = config.kind
        data_loader_cls = cls.get_registered_object(kind)

        if data_loader_cls is None:
            raise ValueError(
                f"DataLoader type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return data_loader_cls(
            config=config,
            datasets=datasets,
            preprocessors=preprocessors,
            finalizer=finalizer,
            random_seed=random_seed,
            **kwargs,
        )

@DataLoaderRegistry.register("pytorch")
class TorchDataLoader(PyTorchDataLoader[DataT], InfoMixin, DataLoader[DataT]):
    def __init__(
        self,
        config: TorchDataLoaderArgs,
        datasets: list[DatasetType],
        preprocessors: list[DatasetPreprocessor | DataDependentPreprocessor],
        finalizer: DatasetFinalizer[DataT],
        random_seed: int = 42,
        **kwargs: Any,
    ):
        self._samples = config.samples
        iterator: DatasetsIterator[DataT] = DatasetsIterator(
            datasets=datasets,
            data_samples=config.samples,
            preprocessors=preprocessors,
            finalizer=finalizer,
        )
        self._info: dict[str, Any] = config.model_dump(mode="json")
        self.epoch = 0

        gen = torch.Generator()
        gen.manual_seed(random_seed)
        super().__init__(
            dataset=iterator,
            batch_size=1,
            shuffle=config.shuffle,
            collate_fn=_collate_first,
            num_workers=config.num_workers,
            prefetch_factor=config.prefetch_factor,
            generator=gen,
            **kwargs,
        )

    def __iter__(self):
        if isinstance(self.dataset, DatasetsIterator):
            self.dataset.set_epoch(self.epoch)
        self.epoch += 1

        return super().__iter__()

    def is_infinite(self) -> bool:
        """
        Whether prefetching every request would run without end.

        A positive sample cap ends generation. Otherwise the loader is infinite
        only when every dataset reports that it does not end.

        :return: ``True`` when iteration does not stop on its own
        """
        if self._samples > 0:
            return False
        datasets = (
            self.dataset.datasets if isinstance(self.dataset, DatasetsIterator) else []
        )
        return datasets_are_infinite(datasets, samples=self._samples)

    @property
    def info(self) -> dict[str, Any]:
        return self._info