feat: add persistent state methods to Model #522
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Add
write_state()andread_state()methods toart.Modelthat persist arbitrary JSON state to the model's output directory (state.json).Features
StateTypetype parameter with defaultdict[str, Any]for backward compatibilitywrite_state(state: StateType)- persists state as JSON to{output_dir}/state.jsonread_state() -> StateType | None- reads state, returnsNoneif not foundModel[MyConfig]syntax still worksModel[MyConfig, MyState]enforces state type with type checkersUse Case
This enables filesystem-based state tracking for training resumption, dataset position, and other metadata. For the Tinker RL training migration, this replaces wandb-based state tracking with filesystem-based persistence for better compatibility.
Usage Examples
File Storage
State is stored at:
.art/{project}/models/{model_name}/state.json