ai2_kit.core.checkpoint module

class ai2_kit.core.checkpoint.CheckpointService

Bases: object

apply_checkpoint(key_fn: str | Callable[[FnInfo], str], disable=False)

set_checkpoint_dir(checkpoint_dir: str)

class ai2_kit.core.checkpoint.FnInfo(fn_name, args, kwargs, call_site)

Bases: NamedTuple

args: tuple

Alias for field number 1

call_site: str

Alias for field number 3

fn_name: str

Alias for field number 0

kwargs: dict

Alias for field number 2

ai2_kit.core.checkpoint.apply_checkpoint(key_fn: str | Callable[[FnInfo], str], disable=False)

apply checkpoint for function.

Note: This checkpoint implementation doesn’t support multiprocess. To support multiple process we need to have a dedicated background process to read/write checkpoint, which will require message queue (e.g. nanomsg or nng) to implement it.

Example:

>>> set_checkpoint_file('/tmp/test.ckpt')
>>> task_fn = lambda a, b: a + b
>>> checkpoint('task_1+2')(task_fn)(1, 2)

ai2_kit.core.checkpoint.set_checkpoint_dir(checkpoint_dir: str)