data_juicer.utils.ckpt_utils module

class data_juicer.utils.ckpt_utils.CheckpointManager(ckpt_dir, original_process_list, num_proc=1)

Bases: object

This class is used to save the latest version of dataset to checkpoint directory or load it from checkpoint directory, a bit like cache management Rerun the same config will reload the checkpoint and skip ops before it.

If any args of operator in process list is changed, all ops will be rerun from the beginning.

__init__(ckpt_dir, original_process_list, num_proc=1)

Initialization method.

Parameters:

• ckpt_dir – path to save and load checkpoint

• original_process_list – process list in config

• num_proc – number of process workers when saving dataset

get_left_process_list()

Get left process list of ops for processing dataset, when checkpoint is available, remove some ops from process list, otherwise keep it unchanged.

Returns:

process list of left ops

check_ckpt()

Check if checkpoint is available.

Returns:

True when checkpoint is available, else False

record(op_cfg: dict)

Save op name and args to op record, which is used to compare with the process list from config to decide if a checkpoint is available.

check_ops_to_skip()

Check which ops need to be skipped in the process list.

If op record list from checkpoint are the same as the prefix part of process list, then skip these ops and start processing from the checkpoint. Otherwise, process the original dataset from scratch.

Returns:

whether to skip some ops or not

save_ckpt(ds)

Save dataset to checkpoint directory and dump processed ops list.

Parameters:

ds – input dataset to save

load_ckpt()

Load dataset from a checkpoint file.

Returns:

a dataset stored in checkpoint file.