Repository: Lightning-AI/pytorch-lightning

Issue:
Passing dataloader to trainer.fit() doesn't work with tpu (and maybe ddp)
## 🐛 Bug

Receive a 
```
AttributeError: Can't pickle local object 'Trainer.__set_fit_dataloaders.<locals>.patch_train_dataloader'
``` 
error when passing the dataloader directly to `trainer.fit(model, train_loader)`

### To Reproduce

Steps to reproduce the behavior:

Try to call `trainer.fit(model, train_loader)` in TPU mode.

(I suspect that anything that calls `mp.spawn` will cause this problem, so ddp probably will face this issue too.)

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
```python
import os

import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

import torch_xla.core.xla_model as xm


class CoolSystem(pl.LightningModule):
    def __init__(self, use_tpu=False):
        super(CoolSystem, self).__init__()
        # not the best model...
        self.use_tpu = use_tpu
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'val_loss': F.cross_entropy(y_hat, y)}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        tensorboard_logs = {'val_loss': avg_loss}
        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

    def test_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'test_loss': F.cross_entropy(y_hat, y)}

    def test_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        tensorboard_logs = {'test_loss': avg_loss}
        return {'avg_test_loss': avg_loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        # REQUIRED
        # can return multiple optimizers and learning_rate schedulers
        # (LBFGS it is automatically supported, no need for closure function)
        return torch.optim.Adam(self.parameters(), lr=0.0004)


if __name__ == '__main__':
    from pytorch_lightning import Trainer

    model = CoolSystem(use_tpu=True)

    dataset = MNIST(os.getcwd(),
                        train=True,
                        download=True,
                        transform=transforms.ToTensor())

    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True)

    loader = DataLoader(dataset, sampler=sampler, batch_size=32)

    # most basic trainer, uses good defaults
    trainer = Trainer(num_tpu_cores=8)
    trainer.fit(model, loader)
```

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
Ideally, specifying the dataloaders as part of the LightningModule should work just the same as passing the dataloaders into `trainer.fit()`

### Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
 - Docker image: gcr.io/tpu-pytorch/xla:nightly
 - build steps:
     - pip install git+git://github.com/williamFalcon/pytorch-lightning.git@master --upgrade 

(I don't have access to the machine right now, so please forgive me on the specific version info temporarily)

### Proposed solution

The issue is [here](https://github.com/PyTorchLightning/pytorch-lightning/blob/d856989120b078581f3f694fd7a1c036703f67a9/pytorch_lightning/trainer/trainer.py#L1001), trying to assign a local function to the model
```python
def __set_fit_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
        # when dataloader is passed via fit, patch the train_dataloader
        # functions to overwrite with these implementations
        if train_dataloader is not None:
            if not self.is_overriden('training_step', model):
                m = 'You called .fit() with a train_dataloader but did not define training_step()'
                raise MisconfigurationException(m)

            def patch_train_dataloader():
                return train_dataloader

            model.train_dataloader = patch_train_dataloader
```

Instead of using a closure or a local function, you could [use a callable](https://stackoverflow.com/questions/573569/python-serialize-lexical-closures) defined at the top-level.  This will be pickleable.

```python

class DataLoaderPatcher
    def __init__(self, loader):
        self.loader = loader
    def __call__():
        return self.loader

def __set_fit_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
        # when dataloader is passed via fit, patch the train_dataloader
        # functions to overwrite with these implementations
        if train_dataloader is not None:
            if not self.is_overriden('training_step', model):
                m = 'You called .fit() with a train_dataloader but did not define training_step()'
                raise MisconfigurationException(m)

            model.train_dataloader = DataLoaderPatcher(train_dataloader)
```

<!-- Add any other context about the problem here. -->



Discussion:
I don't mind tackling the issue myself, but I'd like some input first.  Thanks everyone!
lightning automates the sampler. pass just the dataloader
that's super cool.  I'm giving that a try, both in the GCP XLA docker image and in colab.  

It seems to run fine in a colab modified from the TPU-MNIST example, but is still failing in the GCP XLA docker image.

I suppose that's due to the different `start_method`s?
```python
#  COLAB_GPU is an env var available by default in Colab environments.
start_method = 'fork' if os.getenv('COLAB_GPU') else 'spawn'
xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
```

Reading the flow of `trainer.fit()`, I don't see anything that would change the local functions set by `__set_fit_dataloaders` before hitting `mp.spawn` or `xmp.spawn(start_method='spawn')`.  

```python
### excerpt from trainer.fit()

        # set up the passed in dataloaders (if needed)
        ### !! this will set the unpickleable local functions
        self.__set_fit_dataloaders(model, train_dataloader, val_dataloaders, test_dataloaders)

        ...

        elif self.use_ddp:
            if self.is_slurm_managing_tasks:
                task = int(os.environ['SLURM_LOCALID'])
                self.ddp_train(task, model)
            else:
                ### !! I expect this will fail.  need testing.
                mp.spawn(self.ddp_train, nprocs=self.num_gpus, args=(model,))

        ...

        elif self.use_tpu:
            log.info(f'training on {self.num_tpu_cores} TPU cores')

            #  COLAB_GPU is an env var available by default in Colab environments.
            start_method = 'fork' if os.getenv('COLAB_GPU') else 'spawn'

            ### !! using start_method=spawn here requires pickling
            xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
```

EDIT:
[here](https://colab.research.google.com/drive/1ke-69jV0xm-PsToqeN5kGGdH4KBv7wAj) is the modified colab I was running.  I exported the same thing and tried to run it in non-colab setting and got the failure
so you think the spawn method is causing issues?
> so you think the spawn method is causing issues?

Yes.  `spawn` pickles the function and arguments to hand them to a new process, and `patch_train_dataloader` is not pickleable
![image](https://user-images.githubusercontent.com/16643856/75493796-95577480-596f-11ea-832f-f46fdaef15f4.png)

A quick run of the same example in `ddp` mode (colab, so non-slurm) results in the same error
From the [pickle docs](https://docs.python.org/3/library/pickle.html):

The following types can be pickled:
 - None, True, and False
 - integers, floating point numbers, complex numbers
 - strings, bytes, bytearrays
 - tuples, lists, sets, and dictionaries containing only picklable objects
 - **functions defined at the top level of a module (using def, not lambda)**
 - built-in functions defined at the top level of a module
 - classes that are defined at the top level of a module

The "functions defined at the top level" requirement is what we're failing currently, by defining the `patch_train_dataloader()` function inside another function.
oh yeah. that’s just one way to do it. we could do the same another way. Basically just wanted to plug it into the framework instead of coming up with new functionality. 

want to submit a PR?
or @ethanwharris @Borda 
I can take a stab at it with my proposed solution if that's cool
looks great! give it a shot

Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -1002,30 +1002,21 @@ def __set_fit_dataloaders(self, model, train_dataloader, val_dataloaders, test_d
                 m = 'You called .fit() with a train_dataloader but did not define training_step()'
                 raise MisconfigurationException(m)
 
-            def patch_train_dataloader():
-                return train_dataloader
-
-            model.train_dataloader = patch_train_dataloader
+            model.train_dataloader = _PatchDataLoader(train_dataloader)
 
         if val_dataloaders is not None:
             if not self.is_overriden('validation_step', model):
                 m = 'You called .fit() with a val_dataloaders but did not define validation_step()'
                 raise MisconfigurationException(m)
 
-            def patch_val_dataloader():
-                return val_dataloaders
-
-            model.val_dataloader = patch_val_dataloader
+            model.val_dataloader = _PatchDataLoader(val_dataloaders)
 
         if test_dataloaders is not None:
             if not self.is_overriden('test_step', model):
                 m = 'You called .fit() with a test_dataloaders but did not define test_step()'
                 raise MisconfigurationException(m)
 
-            def patch_test_dataloader():
-                return test_dataloaders
-
-            model.test_dataloader = patch_test_dataloader
+            model.test_dataloader = _PatchDataLoader(test_dataloaders)
 
     def init_optimizers(
             self,
@@ -1189,6 +1180,21 @@ def test(self, model: Optional[LightningModule] = None):
         self.run_evaluation(test_mode=True)
 
 
+class _PatchDataLoader(object):
+    r'''
+    Callable object for patching dataloaders passed into trainer.fit().
+    Use this class to override model.*_dataloader() and be pickle-compatible.
+
+    Args:
+        dataloader: Dataloader object to return when called.
+    '''
+    def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):
+        self.dataloader = dataloader
+
+    def __call__(self) -> Union[List[DataLoader], DataLoader]:
+        return self.dataloader
+
+
 def _set_dataloader(model, dataloader, attribute):
     r'''
     Check dataloaders passed to .fit() method if they are pytorch DataLoader


Test patch:
diff --git a/tests/test_gpu_models.py b/tests/test_gpu_models.py
--- a/tests/test_gpu_models.py
+++ b/tests/test_gpu_models.py
@@ -66,6 +66,31 @@ def test_multi_gpu_model_ddp(tmpdir):
     tutils.run_model_test(trainer_options, model)
 
 
+def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
+    """Make sure DDP works with dataloaders passed to fit()"""
+    if not tutils.can_run_gpu_test():
+        return
+
+    tutils.reset_seed()
+    tutils.set_random_master_port()
+
+    model, hparams = tutils.get_model()
+    trainer_options = dict(default_save_path=tmpdir,
+                           show_progress_bar=False,
+                           max_epochs=1,
+                           train_percent_check=0.4,
+                           val_percent_check=0.2,
+                           gpus=[0, 1],
+                           distributed_backend='ddp')
+
+    fit_options = dict(train_dataloader=model.train_dataloader(),
+                       val_dataloaders=model.val_dataloader())
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model, **fit_options)
+    assert result == 1, "DDP doesn't work with dataloaders passed to fit()."
+
+
 def test_optimizer_return_options():
     tutils.reset_seed()