def save_checkpoint()

in jcm/checkpoints.py [0:0]

• 18 lines of code
• 4 McCabe index (conditional complexity)

def save_checkpoint(ckpt_dir, target, step, prefix="checkpoint_", keep=1):
    """Save a checkpoint of the model.

    Attempts to be pre-emption safe by writing to temporary before
    a final rename and cleanup of past files.

    Args:
      ckpt_dir: str: path to store checkpoint files in.
      target: serializable flax object, usually a flax optimizer.
      step: int or float: training step number or other metric number.
      prefix: str: checkpoint file name prefix.
      keep: number of past checkpoint files to keep.

    Returns:
      Filename of saved checkpoint.
    """
    # Write temporary checkpoint file.
    logging.info("Saving checkpoint at step: %s", step)
    ckpt_tmp_path = _checkpoint_path(ckpt_dir, "tmp", prefix)
    ckpt_path = _checkpoint_path(ckpt_dir, step, prefix)
    blobfile.makedirs(os.path.dirname(ckpt_path))
    with blobfile.BlobFile(ckpt_tmp_path, "wb") as fp:
        fp.write(serialization.to_bytes(target))

    # Rename once serialization and writing finished.
    blobfile.copy(ckpt_tmp_path, ckpt_path, overwrite=True)
    blobfile.remove(ckpt_tmp_path)
    logging.info("Saved checkpoint at %s", ckpt_path)

    # Remove old checkpoint files.
    base_path = os.path.join(ckpt_dir, f"{prefix}")
    checkpoint_files = natural_sort(blobfile.glob(base_path + "*"))
    if len(checkpoint_files) > keep:
        old_ckpts = checkpoint_files[:-keep]
        for path in old_ckpts:
            logging.info("Removing checkpoint at %s", path)
            blobfile.remove(path)

    return ckpt_path