Repository: Lightning-AI/pytorch-lightning

Issue:
WandbLogger cannot be used with 'ddp'
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

wandb modifies `init` such that a child process calling init returns None if the master process has called init. This seems to cause a bug with ddp, and results in rank zero having experiment = None, which crashes the program.

### To Reproduce

Can be reproduced with the basic MNIST gpu template, simply add a WandbLogger and pass 'ddp' as the distributed backend.
```
-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/rmrao/anaconda3/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/home/rmrao/anaconda3/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 331, in ddp_train
    self.run_pretrain_routine(model)
  File "/home/rmrao/anaconda3/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 757, in run_pretrain_routine
    self.logger.log_hyperparams(ref_model.hparams)
  File "/home/rmrao/anaconda3/lib/python3.6/site-packages/pytorch_lightning/logging/base.py", line 14, in wrapped_fn
    fn(self, *args, **kwargs)
  File "/home/rmrao/anaconda3/lib/python3.6/site-packages/pytorch_lightning/logging/wandb.py", line 79, in log_hyperparams
    self.experiment.config.update(params)
AttributeError: 'NoneType' object has no attribute 'config'
```

This occurs with the latest wandb version and with pytorch-lightning 0.6.



Discussion:
Hi! thanks for your contribution!, great first issue!
Some hacky solutions: calling `reinit=True` for wandb, adding or this terrible hack:

```python
def init_ddp_connection(self, *args, **kwargs):
    super().init_ddp_connection(*args, **kwargs)

    if torch.distributed.get_rank() == 0:
        import wandb
        wandb.run = None
```

These both seem to only kind-of work and result in multiple independent calls to wandb.init. I think the ideal solution is that the experiment is only ever initialized on rank zero. *However* doing this means that wandb *cannot* be initialized in the master thread at all.

Better than this probably requires some changes to the wandb API.

Following up slightly - my hacky solution doesn't really work. It's easy enough though to get the rank zero only solution to work. If this seems like a reasonable solution, let me know and I'll submit a PR.
well, have observed some issues with `wandb` earlier #906 could you check it?
Hmm, I think this is a slightly different issue (I'm running on Ubuntu so I don't think that's the issue). Pickling also works correctly.
This particular problem I think stems from this part of the `wandb.init(...)` code:
```python
def init(...):
    ...
    # If a thread calls wandb.init() it will get the same Run object as
    # the parent. If a child process with distinct memory space calls
    # wandb.init(), it won't get an error, but it will get a result of
    # None.
    # This check ensures that a child process can safely call wandb.init()
    # after a parent has (only the parent will create the Run object).
    # This doesn't protect against the case where the parent doesn't call
    # wandb.init but two children do.
    if run or os.getenv(env.INITED):
        return run
```

Child processes end up getting `None` for the wandb run object, which causes logging to fail. There are probably two reasonable and complementary solutions:

1. The main thread should avoid creating a wandb experiment unless absolutely necessary. 

Right now, [this](https://github.com/PyTorchLightning/pytorch-lightning/blob/e586ed47674fd78b158322bb7b14d00aeb912abd/pytorch_lightning/loggers/wandb.py#L63-L69) is the only part of the logging code that the parent thread calls (I assume it's called when pickling):

```python
def __getstate__(self):
    state = self.__dict__.copy()
    # cannot be pickled
    state['_experiment'] = None
    # args needed to reload correct experiment
    state['_id'] = self.experiment.id
    return state
```

If this is changed to:

```python
def __getstate__(self):
    state = self.__dict__.copy()
    # args needed to reload correct experiment
    if self._experiment is not None:
        state['_id'] = self._experiment.id
    else:
        state['_id'] = None

    # cannot be pickled
    state['_experiment'] = None
    return state
```
That will ensure that unless the user explicitly logs something or creates the wandb experiment first, then the main thread will not try to create an experiment. Since subsequent logging / saving code is wrapped by the `@rank_zero_only` decorator, this will generally solve the issue in the base case.

It's also possible that [these properties](https://github.com/PyTorchLightning/pytorch-lightning/blob/e586ed47674fd78b158322bb7b14d00aeb912abd/pytorch_lightning/loggers/wandb.py#L112-L118) are also called by master. Ideally they would be wrapped to not create the experiment unless it had been already created (i.e. experiment should only be created by a function that is wrapped with the `@rank_zero_only` decorator).

2. If the main thread *has* created an experiment, rank zero should be passed the re-init argument.

`wandb` does allow you to reinitialize the experiment. I tried to play around with this a little bit and got some errors, but in theory adding this:

```python
wandb.init(..., reinit=dist.is_available() and dist.is_initialized() and dist.get_rank() == 0)
```
should force a re-initialization when wandb is already initialzed for rank zero.




Resolving patch:
diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py
--- a/pytorch_lightning/loggers/wandb.py
+++ b/pytorch_lightning/loggers/wandb.py
@@ -65,10 +65,11 @@ def __init__(self, name: Optional[str] = None, save_dir: Optional[str] = None,
 
     def __getstate__(self):
         state = self.__dict__.copy()
+        # args needed to reload correct experiment
+        state['_id'] = self._experiment.id if self._experiment is not None else None
+
         # cannot be pickled
         state['_experiment'] = None
-        # args needed to reload correct experiment
-        state['_id'] = self.experiment.id
         return state
 
     @property
@@ -87,7 +88,7 @@ def experiment(self) -> Run:
                 os.environ['WANDB_MODE'] = 'dryrun'
             self._experiment = wandb.init(
                 name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,
-                id=self._id, resume='allow', tags=self._tags, entity=self._entity)
+                reinit=True, id=self._id, resume='allow', tags=self._tags, entity=self._entity)
             # save checkpoints in wandb dir to upload on W&B servers
             if self._log_model:
                 self.save_dir = self._experiment.dir
@@ -109,8 +110,11 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
 
     @property
     def name(self) -> str:
-        return self.experiment.project_name()
+        # don't create an experiment if we don't have one
+        name = self._experiment.project_name() if self._experiment else None
+        return name
 
     @property
     def version(self) -> str:
-        return self.experiment.id
+        # don't create an experiment if we don't have one
+        return self._experiment.id if self._experiment else None


Test patch:
diff --git a/tests/loggers/test_wandb.py b/tests/loggers/test_wandb.py
--- a/tests/loggers/test_wandb.py
+++ b/tests/loggers/test_wandb.py
@@ -50,6 +50,8 @@ class Experiment:
     trainer_options = dict(max_epochs=1, logger=logger)
 
     trainer = Trainer(**trainer_options)
+    # Access the experiment to ensure it's created
+    trainer.logger.experiment
     pkl_bytes = pickle.dumps(trainer)
     trainer2 = pickle.loads(pkl_bytes)