Wandb bug/wandb multi (#1360)

* Allow reinits in sub procs

* Dont create an experiment on pickle, name, or project

* Comments consistency

* Fix test

* Apply suggestions from code review

Co-authored-by: Chris Van Pelt <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

Author: 3 people

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

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)