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

Code before:
            self.ckpt_process = None
            self.local_rank = None
            print(
                f'[{ENGINE_NAME}]: No checkpoint process self.global_rank={self.global_rank}  self.dp_writer_config={self.dp_writer_config}'
            )
        else:
            self.save_event = mp.Event()
            self.save_queue = mp.SimpleQueue()
            engine_args = (config_params, dp_writer_config, self.save_event, self.save_queue, self.optimize_dp_state)
            self.ckpt_process = mp.Process(target=init_decoupled_checkpoint, args=engine_args)
            self.ckpt_process.start()
            self.local_rank = dp_writer_config.local_rank
            print(
                f'[{ENGINE_NAME}]: Create checkpoint process self.global_rank={self.global_rank}  self.ckpt_process.pid={self.ckpt_process.pid} self.dp_writer_config={self.dp_writer_config}'
            )

    def __del__(self):
        try:
            self.cleanup()
        except Exception:
            # Suppress exceptions in destructor to avoid crashes during shutdown
            pass

    def _check_process_alive(self):
        """Check if the checkpoint process is still alive."""
        if self.ckpt_process is None:
            return False
        return self.ckpt_process.is_alive()

    def _wait_for_event_with_timeout(self, timeout_seconds=DEFAULT_CHECKPOINT_TIMEOUT_SECONDS):
        """Wait for save_event with timeout and process health checks.

        Returns True if event was set, raises RuntimeError if process died or timeout occurred.
        """
        elapsed = 0
        while elapsed < timeout_seconds:
            if self.save_event.wait(timeout=PROCESS_HEALTH_CHECK_INTERVAL_SECONDS):
                return True
            elapsed += PROCESS_HEALTH_CHECK_INTERVAL_SECONDS

            # Check if process is still alive
            if not self._check_process_alive():
                raise RuntimeError(f"[{ENGINE_NAME}] Checkpoint process died unexpectedly. "
                                   f"Check logs for OOM or other errors in the checkpoint subprocess.")

        raise RuntimeError(f"[{ENGINE_NAME}] Checkpoint commit timed out after {timeout_seconds} seconds. "
                           f"Process alive: {self._check_process_alive()}")

    def create(self, info: CheckpointCommitInfo):
        self.commit_info = info
        if self.checkpoint_size.gb_size() is None:

Review comment (bug):
Actually, `self.ckpt_process == None` should not be a failing condition. The reason is that not all engines are [expected](https://github.com/deepspeedai/DeepSpeed/blob/6d9c3dc0b8383832d181cd70784f4d7f22348d0d/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py#L80-L83) to have `self.ckpt_process`.  For more details see Figure 6 of the [paper](https://arxiv.org/pdf/2406.13768).

This is why in many code paths `self.ckpt_process is None` is used to guard stateful changes. So, `_check_process_alive()` can be reduced to `self.ckpt_process.is_alive()`. 

Hope that makes sense. 

Code after:
        self.checkpoint_size = CheckpointSize()
        self.global_rank = dist.get_rank()
        self.optimize_dp_state = optimize_dp_state
        self._cleanup_called = False
        if dp_writer_config is None:
            self.save_event = None
            self.save_queue = None
            self.ckpt_process = None
            self.local_rank = None
            print(
                f'[{ENGINE_NAME}]: No checkpoint process self.global_rank={self.global_rank}  self.dp_writer_config={self.dp_writer_config}'
            )
        else:
            self.save_event = mp.Event()
            self.save_queue = mp.SimpleQueue()
            engine_args = (config_params, dp_writer_config, self.save_event, self.save_queue, self.optimize_dp_state)
            self.ckpt_process = mp.Process(target=init_decoupled_checkpoint, args=engine_args)
            self.ckpt_process.start()
            self.local_rank = dp_writer_config.local_rank
            print(
                f'[{ENGINE_NAME}]: Create checkpoint process self.global_rank={self.global_rank}  self.ckpt_process.pid={self.ckpt_process.pid} self.dp_writer_config={self.dp_writer_config}'
            )

    def __del__(self):
        try:
            self.cleanup()
        except Exception:
            # Suppress exceptions in destructor to avoid crashes during shutdown
            pass

    def _check_process_alive(self):
        """Check if the checkpoint process is still alive.

        Note: Only call this when self.ckpt_process is not None.
        Some ranks don't have a checkpoint process by design (see Figure 6 in paper).
        """
        return self.ckpt_process.is_alive()

    def _wait_for_event_with_timeout(self, timeout_seconds=DEFAULT_CHECKPOINT_TIMEOUT_SECONDS):
        """Wait for save_event with timeout and process health checks.

        Returns True if event was set, raises RuntimeError if process died or timeout occurred.
        """
        elapsed = 0
        while elapsed < timeout_seconds:
            if self.save_event.wait(timeout=PROCESS_HEALTH_CHECK_INTERVAL_SECONDS):
                return True
            elapsed += PROCESS_HEALTH_CHECK_INTERVAL_SECONDS

            # Check if process is still alive
            if not self._check_process_alive():