Skip to content

Conversation

@corbt
Copy link
Contributor

@corbt corbt commented Jan 20, 2026

Summary

Add write_state() and read_state() methods to art.Model that persist arbitrary JSON state to the model's output directory (state.json).

Features

  • StateType type parameter with default dict[str, Any] for backward compatibility
  • write_state(state: StateType) - persists state as JSON to {output_dir}/state.json
  • read_state() -> StateType | None - reads state, returns None if not found
  • Full backward compatibility: existing Model[MyConfig] syntax still works
  • Optional type safety: Model[MyConfig, MyState] enforces state type with type checkers

Use 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

# Without type parameter (backward compatible)
model = art.TrainableModel(name="my-model", project="my-project", base_model="...")
model.write_state({"step": 5, "dataset_offset": 100})
state = model.read_state()  # Returns dict[str, Any] | None

# With type parameter for type safety
class TrainingState(TypedDict):
    step: int
    dataset_offset: int
    policy_version: int

model = art.TrainableModel[TinkerConfig, TrainingState](...)
model.write_state({"step": 5, "dataset_offset": 100, "policy_version": 3})
state = model.read_state()  # Returns TrainingState | None

File Storage

State is stored at: .art/{project}/models/{model_name}/state.json

Cursor Bot added 2 commits January 20, 2026 01:53
Add write_state() and read_state() methods to art.Model that persist
arbitrary JSON state to the model's output directory (state.json).

Features:
- StateType type parameter with default dict[str, Any] for type safety
- write_state(state: StateType) - persists state as JSON
- read_state() -> StateType | None - reads state, returns None if not found
- Full backward compatibility: existing Model[MyConfig] syntax still works
- Optional type safety: Model[MyConfig, MyState] enforces state type

This enables filesystem-based state tracking for training resumption,
dataset position, and other metadata that was previously tracked via wandb.
- Add pyright: ignore[reportInconsistentOverload] to __new__ overloads
- Update Backend method signatures to use AnyModel/AnyTrainableModel type aliases
- Add type: ignore comment for return value in Model.__new__
@corbt corbt requested a review from bradhilton January 20, 2026 02:12
Copy link
Collaborator

@bradhilton bradhilton left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@corbt corbt merged commit 1775016 into main Jan 20, 2026
2 checks passed
corbt added a commit that referenced this pull request Jan 20, 2026
* feat: add persistent state methods to Model

Add write_state() and read_state() methods to art.Model that persist
arbitrary JSON state to the model's output directory (state.json).

Features:
- StateType type parameter with default dict[str, Any] for type safety
- write_state(state: StateType) - persists state as JSON
- read_state() -> StateType | None - reads state, returns None if not found
- Full backward compatibility: existing Model[MyConfig] syntax still works
- Optional type safety: Model[MyConfig, MyState] enforces state type

This enables filesystem-based state tracking for training resumption,
dataset position, and other metadata that was previously tracked via wandb.

* fix: resolve pyright type checking errors

- Add pyright: ignore[reportInconsistentOverload] to __new__ overloads
- Update Backend method signatures to use AnyModel/AnyTrainableModel type aliases
- Add type: ignore comment for return value in Model.__new__

---------

Co-authored-by: Cursor Bot <bot@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants