checkpoint_manager#

Checkpoint manager for activation hook scoring with periodic saves and resume support.

Classes

ScoringCheckpointManager

Manages checkpointing for activation hook scoring with periodic saves.
class ScoringCheckpointManager#

Bases: object

Manages checkpointing for activation hook scoring with periodic saves.

__init__(checkpoint_dir, activation_hooks=None, checkpoint_interval=100)#

Initialize checkpoint manager.

Parameters:

  • checkpoint_dir (str) – Directory to save checkpoints

  • activation_hooks – Dictionary of activation hooks to manage

  • checkpoint_interval (int) – Save checkpoint every N batches

finalize()#

Mark scoring as completed.

load_checkpoint()#

Load existing checkpoint if available, including hook states.

Returns:

Dict with checkpoint info or None if no checkpoint exists

Return type:

dict[str, Any] | None

load_hook_states(activation_hooks)#

Load hook states from checkpoint files.

Parameters:

activation_hooks – Hook objects to load states into

Returns:

True if hook states were successfully loaded, False otherwise

Return type:

bool

save_checkpoint()#

Save current checkpoint to disk (progress info only). Hook states are saved separately in update_progress.

should_skip_batch(batch_idx)#

Check if we should skip this batch (already processed in previous run).

Parameters:

batch_idx (int)

Return type:

bool

update_progress(batch_idx, total_batches)#

Update progress and potentially save checkpoint.

Parameters:

  • batch_idx (int) – Current batch index

  • total_batches (int) – Total number of batches