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] `pyarrow.fs` persistence (11/n): Support pausing trials (and certain schedulers) #38355

Merged

krfricke merged 41 commits into ray-project:master from justinvyu:air/persistence/pause_trials

Aug 15, 2023

Contributor

justinvyu commented Aug 11, 2023 •

edited

Loading

Why are these changes needed?

This PR gets pausing trials to work without in-memory checkpoints. This is needed for pausing scheduler support (PBT, BOHB).

Instead of placing a save_to_object future in the trial data and immediately stopping the trial, this PR now waits to handle the save future before stopping the trial. Otherwise, the save future gets immediately cleared when schedule_trial_stop is called, causing no checkpoint to be saved --> possibly starting training from scratch.

Main part to review

This is the workaround I have for pausing trials w/o memory checkpoints:

TODO

This is stacked on [air] pyarrow.fs persistence (10/n): Unify Tune and Train sessions to support new persistence path in FunctionTrainable #38284
May need to pre-emptively set trial status = paused. Otherwise will new scheduler decisions come in for the saving/pausing trial?

Related issue number

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 added 30 commits

August 8, 2023 11:51


          Move _TrainingResult to session.py

1d8f49b

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Prototype unified session

c88df86

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

7bca696

…persistence/unify_sessions

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix incorrect merge

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Implement reset with unified session (for actor reuse)

d87100d

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Eager mode in session (difference in tune/train behavior)

5455c2c

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Working for train again

e1bd135

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Remove unused ckpt index code

60ea1ee

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix lint

3d23988

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

4d209e7

…persistence/unify_sessions


          Add dict checkpoint utils for tests

67a0721

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Add env var as a constant

2a35b3e

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Remove prints

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix tune.run sync config = None issue

74a2f19

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Pause trials w/o memory checkpoints

5fa0bfd

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Update some test_actor_reuse tests to verify new codepath

01b76a7

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

dfff147

…persistence/unify_sessions


          Rename eager_mode -> synchronous_result_reporting

35aaf2c

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Improve some comments + some cleanups

d08e6e4

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Some more cleanups (remove unused code)

c2e33cc

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Remove reference to global session

9b475ec

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Remove shared storage context (store in global session instead)

bf32992

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Update tuner e2e test to test restoration / checkpointing

b59e0b6

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix lint

86b5520

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          More cleanups + docstrings

6c3b4e5

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          synch result reporting logic is flipped...

ce9920b

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Handle trainable outputs correctly

a97025f

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Propagate storage on trial.reset (for restarting upon restore)

b14bce5

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix lint

81bdc2e

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Merge branch 'air/persistence/unify_sessions' into air/persistence/pa…

26697b2

…use_trials

ericl added the @author-action-required label

Contributor

ericl commented Aug 11, 2023

(can do final review once dependency is merged)

justinvyu added 5 commits

August 12, 2023 10:51


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

ead9e81

…persistence/pause_trials

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Only skip checkpointing if decision=PAUSE for new codepath

e606033

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Update logdir -> local_path

656294d

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Update test_actor_reuse

c8f0ae6

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          Fix lint

34e0c8e

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

justinvyu force-pushed the air/persistence/pause_trials branch from 2e1b691 to 34e0c8e Compare

August 12, 2023 19:29

justinvyu marked this pull request as ready for review

August 12, 2023 19:29


          incorrect logic for skipping save

1e8cb72

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

ericl approved these changes

View reviewed changes

krfricke requested changes

View reviewed changes

Contributor

krfricke left a comment

Generally lgtm but requesting minimal changes for the tune controller changes

python/ray/tune/execution/tune_controller.py Outdated Show resolved Hide resolved

justinvyu added 3 commits

August 14, 2023 10:19


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

20f5155

…persistence/pause_trials


          Increase size of test actor reuse

edd261a

Signed-off-by: Justin Yu <justinvyu@anyscale.com>


          address comments

5698f34

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

justinvyu commented

View reviewed changes

python/ray/tune/execution/tune_controller.py

+                              def stop_after_save_result(*args, **kwargs):
+                                  self._on_saving_result(*args, **kwargs)
+                                  self._schedule_trial_stop(trial)
+                                  self._set_trial_status(trial, Trial.PAUSED)

Contributor Author

justinvyu Aug 14, 2023

@krfricke Is it ok that there's some time between the schedule save and the save result coming in, where the trial is still RUNNING?

Previously, the memory checkpoint path would just return a future and immediately set to PAUSED.

Contributor

krfricke Aug 15, 2023

Ah good point.

I think it is a problem. If a pause instruction doesn't immediately pause the trial, we may send multiple instructions.

How about we set the status to PAUSED immediately, and again after the saving result resolved?

Contributor Author

justinvyu Aug 15, 2023

Ok, done

justinvyu commented

View reviewed changes

python/ray/tune/tests/test_actor_reuse.py

Comment on lines +154 to +160

+                  if trainable_type == "function":
+                      monkeypatch.setenv(RAY_AIR_NEW_PERSISTENCE_MODE, "1")
+                      yield train_fn
+                      monkeypatch.setenv(RAY_AIR_NEW_PERSISTENCE_MODE, "0")
+                  elif trainable_type == "class":
+                      yield MyResettableClass
+                  else:

Contributor Author

justinvyu Aug 14, 2023

This is just for now, before the other class trainable PR is in.

justinvyu requested a review from krfricke

August 14, 2023 17:43

justinvyu removed the @author-action-required label

justinvyu added 2 commits

August 15, 2023 08:57


          Merge branch 'master' of https://github.com/ray-project/ray into air/…

cce3395

…persistence/pause_trials


          set trial to paused in the interim + fix the should_checkpoint=False …

f47a653

…case

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

justinvyu added the tests-ok label

krfricke approved these changes

View reviewed changes

Contributor

krfricke left a comment

LGTM

krfricke merged commit 921ec47 into ray-project:master

58 of 62 checks passed

justinvyu deleted the air/persistence/pause_trials branch

August 15, 2023 19:09

harborn pushed a commit to harborn/ray that referenced this pull request


          [air] pyarrow.fs persistence (11/n): Support pausing trials (and ce…

f571de8

…rtain schedulers) (ray-project#38355)

This PR gets pausing trials to work without in-memory checkpoints. This is needed for pausing scheduler support (PBT, BOHB).

Instead of placing a `save_to_object` future in the trial data and immediately stopping the trial, this PR now waits to handle the `save` future before stopping the trial. Otherwise, the save future gets immediately cleared when `schedule_trial_stop` is called, causing no checkpoint to be saved --> possibly starting training from scratch.

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


          [air] pyarrow.fs persistence (11/n): Support pausing trials (and ce…

…rtain schedulers) (ray-project#38355)

This PR gets pausing trials to work without in-memory checkpoints. This is needed for pausing scheduler support (PBT, BOHB).

Instead of placing a `save_to_object` future in the trial data and immediately stopping the trial, this PR now waits to handle the `save` future before stopping the trial. Otherwise, the save future gets immediately cleared when `schedule_trial_stop` is called, causing no checkpoint to be saved --> possibly starting training from scratch.

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

matthewdeng mentioned this pull request

[tune/train] Implement new persistence strategy and roll out as default option #38294

Closed

arvind-chandra pushed a commit to lmco/ray that referenced this pull request


          [air] pyarrow.fs persistence (11/n): Support pausing trials (and ce…

…rtain schedulers) (ray-project#38355)

This PR gets pausing trials to work without in-memory checkpoints. This is needed for pausing scheduler support (PBT, BOHB).

Instead of placing a `save_to_object` future in the trial data and immediately stopping the trial, this PR now waits to handle the `save` future before stopping the trial. Otherwise, the save future gets immediately cleared when `schedule_trial_stop` is called, causing no checkpoint to be saved --> possibly starting training from scratch.

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


          [air] pyarrow.fs persistence (11/n): Support pausing trials (and ce…

49ba4a5

…rtain schedulers) (ray-project#38355)

This PR gets pausing trials to work without in-memory checkpoints. This is needed for pausing scheduler support (PBT, BOHB).

Instead of placing a `save_to_object` future in the trial data and immediately stopping the trial, this PR now waits to handle the `save` future before stopping the trial. Otherwise, the save future gets immediately cleared when `schedule_trial_stop` is called, causing no checkpoint to be saved --> possibly starting training from scratch.

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