hyperstate.StateManager

class hyperstate.StateManager(config_cls: Type[C], state_cls: Type[S], initial_state: Callable[[C, Dict[str, Any]], S], init_path: str | Path | None = None, checkpoint_dir: str | Path | None = None, overrides: List[str] | None = None, ignore_extra_fields: bool = False, checkpoint_key: str = 'step')

The StateManager manages the full lifecycle of loading configs, creating the initial state, applying hyperparameter schedules, persisting checkpoints, and restoring state from checkpoints.

__init__(config_cls: Type[C], state_cls: Type[S], initial_state: Callable[[C, Dict[str, Any]], S], init_path: str | Path | None = None, checkpoint_dir: str | Path | None = None, overrides: List[str] | None = None, ignore_extra_fields: bool = False, checkpoint_key: str = 'step') None

Instantiate a StateManager.

Parameters:
  • config_cls – The type of the config object.

  • state_cls – The type of the state object.

  • initial_state – A function that takes the config object and returns the initial state.

  • init_path – Path to a config file or checkpoint.

  • checkpoint_dir – Directory to store checkpoints. If the directory contains a valid checkpoint, the latest checkpoint will be loaded and init_path will be ignored.

  • overrides – A list of overrides to apply to the config. (Example: [“optimizer.lr=0.1”])

  • ignore_extra_fields – If True, ignore extra fields in the config file.

  • checkpoint_key – The name of a field in the state object that will have a unique value for each checkpoint.

Methods

__init__(config_cls, state_cls, initial_state)

Instantiate a StateManager.

checkpoint(target_dir)

Persist the current state to a checkpoint.

config_dict()

Returns the config as a nested dictionary.

set_deserialize_ctx(key, value)

Add a value to the deserialization context.

step()

Updates all hyperparameter schedules and persists the current state to the checkpoint directory.

Attributes

checkpoint_dir

config

state