Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[air] New train.Checkpoint API: Add new framework-specific checkpoint implementations #38452

Merged
merged 37 commits into from
Aug 16, 2023
Merged

[air] New train.Checkpoint API: Add new framework-specific checkpoint implementations #38452

merged 37 commits into from
Aug 16, 2023

Conversation

justinvyu
Copy link
Contributor

@justinvyu justinvyu commented Aug 15, 2023 •
edited

Why are these changes needed?

This PR:

  • Introduces a FrameworkCheckpoint base class that implements preprocessor accessors for easy migration from the legacy framework checkpoints.
  • Introduces new implementations of framework specific checkpoints based on FrameworkCheckpoint / the new train.Checkpoint.
  • Renames all existing usage (except for unit tests for the new checkpoints) to Legacy*Checkpoint.

For reviewers

See framework_checkpoint.py and all new framework specific ckpt classes.

For now, I'm renaming all existing usage of checkpoints to Legacy*Checkpoint, except for the framework-specific ckpt unit tests. Next step is to transition all of that usage to Checkpoint and then delete the legacy classes.

Related issue number

#38292

Checks

  • I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
    • I've added any new APIs to the API Reference. For example, if I added a
      method in Tune, I've added it in doc/source/tune/api/ under the
      corresponding .rst file.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

@justinvyu
Add legacy framework ckpt with preprocessor accessors
77e5895 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
New torch checkpoint + rename old one to legacy
9837a03 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
Rename existing torch ckpt usage -> legacy
8d6f5da 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
switch back for torch ckpt unit test
0350587 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
Change Checkpoint.from_directory to classmethod
166fba3 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
Add docstrings to legacy ckpt
00f406e 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[torch] fix to use new preprocessor methods
554302c 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[lightning] new lightning ckpt
85f9d28 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[torch] fix lint
e925f33 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[lightning] rename existing usgae -> legacy
82a0eb8 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[lightning] fix test
cec7536 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>

[lightning] fix lint

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[transformers] add new transformers ckpt + update existing usage to l…
e7e00c6 
…egacy

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[xgboost] new xgb ckpt
7d380ec 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[xgboost] rename existing usage to legacy
4e8977a 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[xgboost] add back xgboost ckpt unit test
ee35e9a 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[xgboost] fix unit test
48267fe 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[lightgbm] new checkpoint
e904034 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[lightgbm] update existing to legacy
cfbf1a8 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[sklearn] new checkpoint
46bd6fb 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[sklearn] update legacy usage
67fd1d2 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[sklearn] fix lint
9acbd0f 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[tf] new checkpoint
d8c64ff 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[tf] existing usage to legacy
13eede7 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
Merge branch 'master' of https://github.com/ray-project/ray into air/…
7f96879 
…new_framework_ckpts

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[xgboost] fix batch_inference.rst
121ccad 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[transformers] allow new ckpt to be imported from ray.train.hf
7095c37 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
[torch] Return torch ckpt
a99f5f3 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
Fix return type cross reference for doc CI
56cd593 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
ericl
ericl approved these changes Aug 15, 2023
View reviewed changes
Copy link
Contributor

@ericl ericl left a comment

Choose a reason for hiding this comment

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

LGTM. I didn't review the new checkpoint code, as I assume we'll be testing/validating those in the next PR that flips those to the used impl.

python/ray/train/_internal/framework_checkpoint.py Outdated Show resolved Hide resolved
@ericl ericl added the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Aug 15, 2023
@justinvyu
Merge branch 'master' of https://github.com/ray-project/ray into air/…
ce0cc17 
…new_framework_ckpts
@justinvyu
Small fixes
43bee86 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@anyscalesam anyscalesam mentioned this pull request Aug 15, 2023
Closed
4 tasks
@justinvyu
Fix torch ckpt doctest
862049a 
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu justinvyu added tests-ok The tagger certifies test failures are unrelated and assumes personal liability. and removed @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. labels Aug 16, 2023
@justinvyu justinvyu mentioned this pull request Aug 16, 2023
Merged
8 tasks
@matthewdeng matthewdeng merged commit 2fbd5ff into ray-project:master Aug 16, 2023
94 of 98 checks passed
@justinvyu justinvyu deleted the air/new_framework_ckpts branch August 16, 2023 05:51
harborn pushed a commit to harborn/ray that referenced this pull request Aug 17, 2023
@justinvyu @harborn
[air] New train.Checkpoint API: Add new framework-specific checkpoi…
725367f 
…nt implementations (ray-project#38452)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: harborn <gangsheng.wu@intel.com>
harborn pushed a commit to harborn/ray that referenced this pull request Aug 17, 2023
@justinvyu @harborn
[air] New train.Checkpoint API: Add new framework-specific checkpoi…
01f16de 
…nt implementations (ray-project#38452)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
arvind-chandra pushed a commit to lmco/ray that referenced this pull request Aug 31, 2023
@justinvyu @arvind-chandra
[air] New train.Checkpoint API: Add new framework-specific checkpoi…
9568fad 
…nt implementations (ray-project#38452)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: e428265 <arvind.chandramouli@lmco.com>
vymao pushed a commit to vymao/ray that referenced this pull request Oct 11, 2023
@justinvyu
[air] New train.Checkpoint API: Add new framework-specific checkpoi…
709afe2 
…nt implementations (ray-project#38452)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Victor <vctr.y.m@example.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ericl ericl ericl approved these changes

@matthewdeng matthewdeng matthewdeng approved these changes

@richardliaw richardliaw Awaiting requested review from richardliaw

@gjoliver gjoliver Awaiting requested review from gjoliver

@krfricke krfricke Awaiting requested review from krfricke

@xwjiang2010 xwjiang2010 Awaiting requested review from xwjiang2010

@amogkam amogkam Awaiting requested review from amogkam amogkam is a code owner

@Yard1 Yard1 Awaiting requested review from Yard1

@maxpumperla maxpumperla Awaiting requested review from maxpumperla

@scv119 scv119 Awaiting requested review from scv119 scv119 is a code owner

@c21 c21 Awaiting requested review from c21 c21 is a code owner

@scottjlee scottjlee Awaiting requested review from scottjlee scottjlee is a code owner

@bveeramani bveeramani Awaiting requested review from bveeramani bveeramani is a code owner

@raulchen raulchen Awaiting requested review from raulchen raulchen is a code owner

Assignees

@ericl ericl

@krfricke krfricke

Labels
tests-ok The tagger certifies test failures are unrelated and assumes personal liability.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
4 participants
@justinvyu @ericl @matthewdeng @krfricke