# Buggy code
def transfer_distrib_spawn_state_on_fit_end(self, results):
    # TODO: is there a better way than accessing callback through model -> trainer -> callback?
    best_model_path = self.lightning_module.trainer.checkpoint_callback.best_model_path

    if self.mp_queue is not None:
        rank_zero_warn("cleaning up ddp environment...")

        # save the last weights
        last_path = None
        # TODO: is there a better way than accessing trainer through model -> trainer?
        if (
            not self.lightning_module.trainer.testing
            and best_model_path is not None
            and len(best_model_path) > 0
        ):
            last_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
            xm.save(self.lightning_module.state_dict(), last_path)

        if self.global_rank == 0:
            # todo, pass complete checkpoint as state dictionary
            self.mp_queue.put(best_model_path)
            self.mp_queue.put(last_path)
            self.mp_queue.put(results)


# Traceback (RuntimeError)
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-587e2a9e3858> in <module>
----> 1 trainer.test(datamodule=dm)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in test(self, model, test_dataloaders, ckpt_path, verbose, datamodule)
922             results = self.__test_given_model(model, test_dataloaders)
923         else:
--> 924             results = self.__test_using_best_weights(ckpt_path, test_dataloaders)
925
926         self.teardown('test')

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in __test_using_best_weights(self, ckpt_path, test_dataloaders)
950                 return {}
951             if not self._device_type == DeviceType.TPU:
--> 952                 self.accelerator.barrier()
953
954             ckpt = pl_load(ckpt_path, map_location=lambda storage, loc: storage)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py in barrier(self, name)
375
376     def barrier(self, name: Optional[str] = None) -> None:
--> 377         self.training_type_plugin.barrier(name=name)
378
379     def broadcast(self, obj: object, src: int = 0) -> object:

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/tpu_spawn.py in barrier(self, name)
113
114     def barrier(self, name: Optional[str] = None) -> None:
--> 115         rendezvous(f"pl.Trainer.{name}")
116
117     def transfer_distrib_spawn_state_on_fit_end(self, results):

/opt/conda/lib/python3.7/site-packages/torch_xla/core/xla_model.py in rendezvous(tag, payload, replicas)
859     ordinal `i` at position `i` in the returned tuple.
860   """
--> 861   return torch_xla._XLAC._xla_rendezvous(get_ordinal(), tag, payload, replicas)
862
863

RuntimeError: tensorflow/compiler/xla/xla_client/mesh_service.cc:316 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds))

# Fixed code
def transfer_distrib_spawn_state_on_fit_end(self, results):
    # TODO: is there a better way than accessing callback through model -> trainer -> callback?
    best_model_path = self.lightning_module.trainer.checkpoint_callback.best_model_path

    if self.mp_queue is not None:
        rank_zero_warn("cleaning up ddp environment...")

        # save the last weights
        last_path = None
        # TODO: is there a better way than accessing trainer through model -> trainer?
        if (
            not self.lightning_module.trainer.testing
            and best_model_path is not None
            and len(best_model_path) > 0
        ):
            last_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
            self.save(self.lightning_module.state_dict(), last_path)

        if self.global_rank == 0:
            # todo, pass complete checkpoint as state dictionary
            self.mp_queue.put(best_model_path)
            self.mp_queue.put(last_path)
            self.mp_queue.put(results)