Repository: mosaicml/composer

Issue:
Validate that `dataloader._iterator` is `None`
I am very skeptical that the algorithms that use `add_dataset_transform` in `composer.utils.data` work as intended, as the dataloader workers are already created before the INIT event runs (and the algorithms attempt to monkeypatch the dataset). Instead, `add_dataset_transform` should be replaced by modifying the data on the `AFTER_DATALOADER` event. This will require that `composer.utils.augmentation_primiatves` are reimplemented to operate on batches of images (rather than on individual PIL images), and likely switch to using `torchvision` instead of PIL

It would be good to run a test to confirm that randaugment, augmix, and colout do not work as intentend.


Discussion:
🤔 cc: @coryMosaicML @growlix 

I dont think the proposed solution would work, as the `AFTER_DATALOADER` event now runs on GPU. We've been using `add_dataset_transform` since launch, so unless something broke, it was working previously. 
I think things are working fine on `main`, but going to look into this on `dev` now. @hanlint is correct that this would require doing augmentations on GPU.
Testing this made me aware that `randaugment` and `augmix` don't work with our CIFAR10 dataset, but not for the reason that @ravi-mosaicml mentioned. It's because the the `cifar10` dataset already has `to_tensor()` in the list of transformations, and `add_dataset_transform()` just appends `randaugment`, so it ends up acting on tensor. About to test w/ ImageNet.
Seems to work on ImageNet for `main` and `dev`. Going to update `add_dataset_transform()` to be cifar10-friendly.
Should we support a 
```
add_dataset_transform(.., before_to_tensor=True)
```
option?
@ravi-mosaicml OK with closing this issue then? the CIFAR-10 issue is a different issue.
Good to close, cifar10 compatibility PR incoming!
@growlix can you elaborate how you confirmed this? My mental model for how the pytorch dataloaders works is telling me that these algorithms would be no-ops when using more than 0 dataloader workers.
Ah, since pytorch dataloader lazily forks on the first iteration, this works. However, if the user already spun the dataloader (and persistent workers is True, and num workers > 0), then the dataloader would have already forked, rendering these augmentations as no-ops.

I'll add a check into the trainer to ensure that `dataloader._iterator is None` to ensure that the augmentations will work as intended.
I confirmed by running `run_mosaic_trainer.py -f {resnet50, resnet56} --algorithms randaugment --dataloader.num_workers 8` on a cota-g1-3080 interactive instance for `dev` and `main` and set a breakpoint in [`algorithms/randaugment/randaugment.py`](https://github.com/mosaicml/composer/blob/fa1b992c213094f63ca3b860aaf06e2ea1c82d88/composer/algorithms/randaugment/randaugment.py#L42) and checking that `randaugment()` was executing and the tensor was being modified.
Is there a fix that would accommodate the scenario you're describing? (pre-spun, persistent workers, n workers > 0)
The easiest 'fix' would be to recreate the dataloader e.g. something like `dataloader = Dataloader(**vars(dataloader))`, but I am not a fan of this approach since this wouldn't support a custom dataloader and just feels kinda messy.
Yeah, agree i think the edge case can be not supported.

Resolving patch:
diff --git a/composer/datasets/dataloader.py b/composer/datasets/dataloader.py
--- a/composer/datasets/dataloader.py
+++ b/composer/datasets/dataloader.py
@@ -42,6 +42,20 @@ def __setattr__(self, name: str, value: Any) -> None:
         return super().__setattr__(name, value)
 
 
+def unwrap_data_loader(dataloader: DataLoader) -> DataLoader:
+    """Recursively unwraps a dataloader if it is of type :class:`WrappedDataLoader`.
+
+    Args:
+        dataloader (DataLoader): The dataloader to unwrap
+
+    Returns:
+        DataLoader: The underlying dataloader
+    """
+    if isinstance(dataloader, WrappedDataLoader):
+        return unwrap_data_loader(dataloader.dataloader)
+    return dataloader
+
+
 @dataclass
 class DataloaderHparams(hp.Hparams):
     """Hyperparameters to initialize a :class:`~torch.utils.data.Dataloader`.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -27,6 +27,7 @@
 from composer.core.time import TimeUnit
 from composer.core.types import (Batch, BreakEpochException, DataLoader, Evaluators, Metrics, Optimizers, Precision,
                                  Schedulers)
+from composer.datasets.dataloader import unwrap_data_loader
 from composer.loggers.tqdm_logger import TQDMLoggerBackend
 from composer.models.base import ComposerModel
 from composer.optim import ComposedScheduler
@@ -271,6 +272,19 @@ def __init__(
             train_dataloader = DataSpec(train_dataloader)
 
         self._train_data_spec = train_dataloader
+        unwrapped_data_loader = unwrap_data_loader(self._train_data_spec.dataloader)
+        if isinstance(unwrapped_data_loader, torch.utils.data.DataLoader):
+            if unwrapped_data_loader._iterator is not None:
+                raise ValueError(
+                    textwrap.dedent("""\
+                    The `train_dataloader` has an active iterator. This could occur
+                    if `persistent_workers=True` and the dataloader has already been iterated,
+                    or if the dataloader is mid-epoch. It is required that the training dataloader
+                    does not have an active iterator, so CPU dataset augmentations can be
+                    correctly inserted.
+
+                    To fix, please do not iterate over the dataloader before passing it into
+                    the trainer."""))
 
         if eval_subset_num_batches is not None:
             for evaluator in self.evaluators:


Test patch: