Repository: microsoft/DeepSpeed
File: deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py (Python)

Code before:
        if not self._check_process_alive():
            return

        save_info = (state_dict, path)
        self.save_queue.put((save_info, DecoupledEvent.SAVE_EVENT))

    def commit(self, info: CheckpointCommitInfo):
        # Use proper validation instead of assert (assert is disabled with python -O)
        if info != self.commit_info:
            raise ValueError(f"[{ENGINE_NAME}] Checkpoint commit info mismatch: "
                             f"expected {self.commit_info}, got {info}")

        if self._check_process_alive():
            self.save_queue.put((None, DecoupledEvent.COMMIT_EVENT))
            # Wait with timeout and health checks instead of blocking forever
            self._wait_for_event_with_timeout()
            self.save_event.clear()

        self.commit_info = None

        if self.checkpoint_size.gb_size() is None:
            dist.barrier()
            post_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank)
            self.checkpoint_size.set_post_size(post_size)

        if self.global_rank == 0 and self.checkpoint_size.gb_size() is not None:
            print(
                f'{self.name} self.global_rank={self.global_rank} created checkpoint of {round(self.checkpoint_size.gb_size(), 2)} GB'
            )

        return True

    def get_commit_info(self):
        # print(f'getting commit info {self.commit_info=}')
        return self.commit_info

    def is_decoupled(self):
        return True

    def cleanup(self):
        # Prevent multiple cleanup calls (especially from __del__)
        if self._cleanup_called:
            return
        self._cleanup_called = True

        try:
            if self.get_commit_info() is not None:
                self.commit(self.commit_info)
        except Exception as e:
            logger.warning(f"[{ENGINE_NAME}] Error during commit in cleanup: {e}")


Review comment (suggestion):
I think it is better to add an assert that `self.checkpoint_size.gb_size() is not None` after line 180. 

Code after:
    def save(self, state_dict, path: str):
        if self.ckpt_process is None:
            return

        # Check process health before attempting to save
        if not self._check_process_alive():
            return

        save_info = (state_dict, path)
        self.save_queue.put((save_info, DecoupledEvent.SAVE_EVENT))

    def commit(self, info: CheckpointCommitInfo):
        # Use proper validation instead of assert (assert is disabled with python -O)
        if info != self.commit_info:
            raise ValueError(f"[{ENGINE_NAME}] Checkpoint commit info mismatch: "
                             f"expected {self.commit_info}, got {info}")

        if self.ckpt_process is not None:
            # Check process health before waiting
            if not self._check_process_alive():
                raise RuntimeError(f"[{ENGINE_NAME}] Cannot commit checkpoint: checkpoint process is not running.")

            self.save_queue.put((None, DecoupledEvent.COMMIT_EVENT))
            # Wait with timeout and health checks instead of blocking forever
            self._wait_for_event_with_timeout()
            self.save_event.clear()

        self.commit_info = None

        if self.checkpoint_size.gb_size() is None:
            dist.barrier()
            post_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank)
            self.checkpoint_size.set_post_size(post_size)

        assert self.checkpoint_size.gb_size() is not None, "Checkpoint size should be set after commit"

        if self.global_rank == 0:
            print(
                f'{self.name} self.global_rank={self.global_rank} created checkpoint of {round(self.checkpoint_size.gb_size(), 2)} GB'
            )

        return True

    def get_commit_info(self):
        # print(f'getting commit info {self.commit_info=}')
        return self.commit_info

    def is_decoupled(self):
        return True

    def cleanup(self):