Repository: Lightning-AI/pytorch-lightning
Pull request: .fit() returns last not best weights in ddp_spawn
Fixes #2547

Problem statement:
Can't use None (anymore) in checkpoint_callback
## 🐛 Bug
using None in checkpoint_callback now errors out

```
-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 562, in ddp_train
    q.put(self.checkpoint_callback.best_model_path)
AttributeError: 'NoneType' object has no attribute 'best_model_path'
```

### To Reproduce
`trainer = Trainer(checkpoint_callback=None)`

Ran into this issue from upgrading to masters, was using masters from a few commits ago before

Edit: `False` casuses the same error as well

Hints from issue discussion:

Commit message: added base tests for tpu

Patch to review:
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
index 5b8d79e47564f..a2eaca6a7166a 100644
--- a/pytorch_lightning/trainer/distrib_data_parallel.py
+++ b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -189,6 +189,7 @@ class TrainerDDPMixin(ABC):
     num_nodes: int
     node_rank: int
     tpu_cores: int
+    testing: bool
 
     @property
     @abstractmethod
@@ -559,9 +560,13 @@ def ddp_train(self, process_idx, q, model, is_master=False, proc_offset=0):
         torch.cuda.empty_cache()
 
         if self.global_rank == 0 and q is not None:
+            rank_zero_warn('cleaning up ddp environment...')
             q.put(self.checkpoint_callback.best_model_path)
             q.put(results)
 
+            if not self.testing:
+                self.save_spawn_weights(model)
+
         if self.global_rank == 0 and self.distributed_backend != 'ddp_spawn':
             return results
 
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index 1b3e053387e96..f7d1e9ebaa79c 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -1043,9 +1043,14 @@ def __run_ddp_spawn(self, model, nprocs):
         # restore main state with best weights
         best_path = q.get()
         results = q.get()
+
+        # transfer back the best path to the trainer
         if best_path is not None and len(best_path) > 0:
             self.checkpoint_callback.best_model_path = best_path
-            model.load_from_checkpoint(best_path)
+
+        # load last model weights
+        if self.testing:
+            self.load_spawn_weights(model)
 
         self.model = model
         return results


Code review comments:
diff_hunk: @@ -559,9 +560,13 @@ def ddp_train(self, process_idx, q, model, is_master=False, proc_offset=0):
         torch.cuda.empty_cache()
 
         if self.global_rank == 0 and q is not None:
+            rank_zero_warn('cleaning up ddp environment...')
             q.put(self.checkpoint_callback.best_model_path)
path: pytorch_lightning/trainer/distrib_data_parallel.py
text: @user1:
could we add a None check here for checkpoint_callback? because the user can set it to None, if they want. 
See #2547

@author:
fixed!