text: import os
  import os.path as osp
  import re
  from typing import Tuple
  
  import torch
  
  from colossalai.context import Config
  from colossalai.context.parallel_mode import ParallelMode
  from colossalai.core import global_context as gpc
  
  __all__ = [
      'get_checkpoint_path',
      'get_latest_checkpoint_path',
      'get_latest_checkpoint_pattern',
      'save_checkpoint',
      'load_checkpoint'
  ]
  
  
  def unwrap_config(config: Config):
      '''
      unwrap Config objects to normal dicts
      '''
      config_dict = dict()
      for k, v in config.items():
          if isinstance(v, dict):
              config_dict[k] = unwrap_config(v)
          else:
              config_dict[k] = v
  
      return config_dict
  
  
  def _get_ranks_name():
      # tensor parallel
      tp_local_rank = 0
      if gpc.is_initialized(ParallelMode.TENSOR):
          tp_local_rank = gpc.get_local_rank(ParallelMode.TENSOR)
  
      # pipeline parallel
      pp_local_rank = 0
      if gpc.is_initialized(ParallelMode.PIPELINE):
          pp_local_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
  
      ranks_name = f'tp{tp_local_rank}-pp{pp_local_rank}'
      return ranks_name
  
  
  def _get_standard_checkpoint_filename(epoch: int, suffix: str = ''):
      ranks_name = _get_ranks_name()
      return f'epoch{epoch}-{ranks_name}{suffix}.pt'
  
  
  def get_checkpoint_path(checkpoint_dir: str, epoch: int, suffix: str = ''):
      '''This is a function to generate the checkpoint path from the (checkpoint_dir, epoch, suffix, gpu_parallel_rank) tuple.
      This is useful during generation and recuperation of the checkpoint.
  
      :param checkpoint_dir: set up a directory for saving checkpoints
      :type checkpoint_dir: str
      :param epoch: epoch number (indicate how many epochs have you trained this model)
      :type epoch: int
      :param suffix: additional notation to specify the model or checkpoint, defaults to ''
      :type suffix: str, optional
      :return: checkpoint path to be generated 
      :rtype: path
      '''
      ckpt_filename = _get_standard_checkpoint_filename(epoch, suffix)
      return os.path.join(checkpoint_dir, ckpt_filename)
  
  
  def _ensure_directory_exists(filename: str):
      # ensure the directory exists
      dir = os.path.dirname(filename)
      if not os.path.exists(dir):
          os.makedirs(dir)
  
  
  def get_latest_checkpoint_pattern(suffix: str = ''):
      '''Generate Regular expression of latest checkpoint's pattern
  
      :param suffix: additional notation to specify the model or checkpoint, defaults to ''
      :type suffix: str, optional
      :return: checkpoint pattern
      :rtype: regular expression
      '''
      ranks_name = _get_ranks_name()
      ckpt_pattern = re.compile(f'epoch(\d+)-{ranks_name}{suffix}\.pt')
      return ckpt_pattern
  
  
  def get_latest_checkpoint_path(checkpoint_dir: str, suffix: str = ''):
      '''This is a function to retrieve the latest checkpoint path from the (checkpoint_dir, suffix, gpu_parallel_rank) tuple.
      This is useful during recuperation of the checkpoint, especially when you do not know the epoch number.
  
      :param checkpoint_dir: directory for saving checkpoints
      :type checkpoint_dir: str
      :param suffix: additional notation to specify the model or checkpoint, defaults to ''
      :type suffix: str, optional
      :raises FileNotFoundError: raise error when we cannot find the latest checkpoint file with inputs given
      :return: the latest checkpoint path to be retrieved 
      :rtype: path
      '''
      CKPT_NAME_PAT = get_latest_checkpoint_pattern(suffix=suffix)
  
      last_epoch = -1
      assert osp.isdir(checkpoint_dir), f'{checkpoint_dir} is not a directory'
  
      for filename in os.listdir(checkpoint_dir):
          ret = CKPT_NAME_PAT.match(filename)
          if ret:
              epoch = int(ret[0].split('-')[0].lstrip('epoch'))
              if epoch > last_epoch:
                  last_epoch = epoch
  
      if last_epoch == -1:
          ranks_name = _get_ranks_name()
          raise FileNotFoundError(f"Cannot find the latest checkpoint file for {ranks_name} in {checkpoint_dir}")
      else:
          target_file = _get_standard_checkpoint_filename(last_epoch, suffix=suffix)
          path = osp.join(checkpoint_dir, target_file)
          return path
  
  
  def save_checkpoint(checkpoint_path: str,
                      epoch: int,
                      model: torch.nn.Module,
                      optimizer: torch.optim.Optimizer,
                      lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
                      **kwargs):
      '''Given a directory to store the checkpoints, saves all the training components' parameters or buffers, such as model, optimizer, lr_scheduler and etc. into a checkpoint dictionary. 
      
      This method can be used for both colosalai nn.BaseModel and normal pytorch nn.Module.
  
  
      :param checkpoint_path: set up a directory for saving checkpoints
      :type checkpoint_path: str
      :param epoch: epoch number (indicate how many epochs have you trained this model)
      :type epoch: int
      :param model: model to be registered
      :type model: torch.nn.Module
      :param optimizer: optimizer to be registered
      :type optimizer: torch.optim.Optimizer
      :param lr_scheduler: lr_scheduler to be registered, defaults to None
      :type lr_scheduler: torch.optim.lr_scheduler._LRScheduler, optional
      '''
      # for compatibility with normal pytorch nn.Module
      if hasattr(model, 'state_dict_for_save_checkpoint'):
          model_sd = model.state_dict_for_save_checkpoint()
      else:
          model_sd = model.state_dict()
  
      # ckpt container
      checkpoint = {
          'epoch': epoch,
          'model': model_sd,
          'optimizer': optimizer.state_dict(),
          **kwargs
      }
      if lr_scheduler is not None:
          checkpoint['lr_scheduler'] = lr_scheduler.state_dict()
  
      _ensure_directory_exists(checkpoint_path)
      torch.save(checkpoint, checkpoint_path)
  
  
  def load_checkpoint(checkpoint_path: str,
                      model: torch.nn.Module,
                      optimizer: torch.optim.Optimizer,
                      lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
                      finetune: bool = False,
                      strict: bool = True) -> Tuple:
      '''Loads the checkpoint file. 
      If finetune is False, then we intend to continue/resume the training process from the checkpoint given. 
      So we copy parameters and buffers from state_dict into these modules(model, optimizer,lr_scheduler) and its descendants. 
      If finetune is True, then only the weights and buffers of model should be reload.
      If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
      
      :param checkpoint_path: the exact and matched checkpoint_path directory to retrieve appropriate state_dict
      :type checkpoint_path: str
      :param model: model to reload parameters and buffers
      :type model: torch.nn.Module
      :param optimizer: optimizer to recuperate
      :type optimizer: torch.optim.Optimizer 
      :param lr_scheduler: lr_scheduler to recuperate, defaults to None
      :type lr_scheduler: torch.optim.lr_scheduler._LRScheduler, optional
      :param finetune: whether to finetune the model with new dataset or continue the pre-training, defaults to False
      :type finetune: bool, optional
      :param strict: whether to strictly enforce that the keys in
          :attr:`state_dict` of the checkpoint match the names of
          parameters and buffers in model., defaults to True
      :type strict: bool, optional
      :raises ValueError: raise error if the model/optimizer cannot successfully be recuperated
      :return: (the epoch number of the checkpoint retrieved, the checkpoint retrieved)
      :rtype: Tuple
  
      '''
      # Load the checkpoint.
      checkpoint = torch.load(checkpoint_path, map_location='cpu')
      try:
          last_epoch = checkpoint.pop('epoch') if not finetune else 0
          model.load_state_dict(checkpoint.pop('model'), strict=strict)
      except KeyError:
          raise ValueError('Checkpoint is corrupted')
  
      if not finetune:
          try:
              optimizer.load_state_dict(checkpoint.pop('optimizer'))
          except KeyError:
              raise ValueError('Checkpoint is corrupted')
  
          if lr_scheduler is not None and 'lr_scheduler' in checkpoint:
              lr_scheduler.load_state_dict(checkpoint.pop('lr_scheduler'))
  
      return last_epoch, checkpoint
  