feat(training): add max_time support to Trainer.fit - #10
Merged
Conversation
- Implement StageTimer to handle time-based termination of training. - Add max_time parameter to Fit stage and integrate the timer as a listener. - Expose max_time parameter in Trainer to allow passing it down to Fit. - Update __init__.py to export StageTimer. - Add unit tests for StageTimer and verify integration with Trainer.fit. - Add reax_config fixture to test/conftest.py for standardizing jax config in tests.
Member
Author
|
@mattia-spider , could you please review this? Are you happy with it? |
mattia-spider
approved these changes
Jul 6, 2026
mattia-spider
left a comment
mattia-spider
left a comment
Collaborator
There was a problem hiding this comment.
Looks perfect to me. I did some tests locally and seems to work smoothly.
Member
Author
|
Thanks! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
This PR introduces time-based termination for training in reax. Users can now specify a max_time argument when calling trainer.fit(...) to limit training duration, allowing for graceful interruption mid-epoch.
This includes:
Changes: