Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/plugins/training_type/tpu_spawn.py (Python)

Code before:
        if self.local_rank == 0:
            time.sleep(2)

        # ensure that spawned processes go through teardown before joining
        trainer._call_teardown_hook()

    def model_to_device(self) -> None:
        self.model = self.wrapped_model.to(self.root_device)

    def barrier(self, name: Optional[str] = None) -> None:
        if self.is_distributed:
            rendezvous(name)

    def __transfer_distrib_spawn_state_on_fit_end(self, trainer: "pl.Trainer", results: Any) -> None:
        rank_zero_info("Transferring results back to main process ...")
        checkpoint_callback = trainer.checkpoint_callback
        best_model_path = checkpoint_callback.best_model_path if checkpoint_callback else None

        # requires to compute the state_dict on all processes in case Metrics are present
        state_dict = self.lightning_module.state_dict()

        # save the last weights
        last_path = None
        if trainer.state.fn == TrainerFn.FITTING 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.checkpoint_io.save_checkpoint(state_dict, last_path)

        if self.local_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)
            self.lightning_module.add_to_queue(self.mp_queue)  # adds the `callback_metrics` to the queue

    def broadcast(self, obj: object, src: int = 0) -> object:
        if not self.is_distributed:
            return obj
        buffer = io.BytesIO()
        torch.save(obj, buffer)
        data = bytearray(buffer.getbuffer())
        data_tensor = torch.tensor(data, device=self.root_device, dtype=torch.float)
        data = xm.all_gather(data_tensor)
        buffer = io.BytesIO(data.cpu().byte().numpy())
        obj = torch.load(buffer)
        return obj

    def reduce_boolean_decision(self, decision: bool) -> bool:
        decision = torch.tensor(int(decision), device=self.lightning_module.device)
        decision = self.reduce(decision, reduce_op="sum")
        decision = bool(decision == self.world_size)
        return decision

Review comment (suggestion):
```suggestion
            tmp_ckpt_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
            self.checkpoint_io.save_checkpoint(state_dict, tmp_ckpt_path)
```

Code after:
        if self.local_rank == 0:
            time.sleep(2)

        # ensure that spawned processes go through teardown before joining
        trainer._call_teardown_hook()

    def model_to_device(self) -> None:
        self.model = self.wrapped_model.to(self.root_device)

    def barrier(self, name: Optional[str] = None) -> None:
        if self.is_distributed:
            rendezvous(name)

    def __transfer_distrib_spawn_state_on_fit_end(self, trainer: "pl.Trainer", results: Any) -> None:
        rank_zero_warn("cleaning up tpu spawn environment...")
        checkpoint_callback = trainer.checkpoint_callback
        best_model_path = checkpoint_callback.best_model_path if checkpoint_callback else None

        # requires to compute the state_dict on all processes in case Metrics are present
        state_dict = self.lightning_module.state_dict()

        # save the last weights
        last_path = None
        if trainer.state.fn == TrainerFn.FITTING 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.checkpoint_io.save_checkpoint(state_dict, last_path)

        # We use `local_rank` here as separate filesystems are used for each VM for TPU Pod Training
        if self.local_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)
            # adds the `callback_metrics` to the queue
            # TODO: Remove the if in v1.7
            if is_overridden("add_to_queue", self.lightning_module):
                self.lightning_module.add_to_queue(self.mp_queue)
            else:
                self.add_to_queue(trainer, self.mp_queue)

    def broadcast(self, obj: object, src: int = 0) -> object:
        if not self.is_distributed:
            return obj
        buffer = io.BytesIO()
        torch.save(obj, buffer)
        data = bytearray(buffer.getbuffer())
        data_tensor = torch.tensor(data, device=self.root_device, dtype=torch.float)
        data = xm.all_gather(data_tensor)
        buffer = io.BytesIO(data.cpu().byte().numpy())
        obj = torch.load(buffer)
        return obj