Checkpoint recovery refactoring #439

melanibe · 2021-04-19T15:01:46Z

Changes the way we handle checkpoints for run recoveries, for the following use cases:

If AML fails for unkown reason and we want to (manually) restart the run: provide --run_recovery_id=old_run_id tag -> this should download all checkpoints available in the old_run checkpoint folder and copy them in the new run's checkpoint folder.
If an AML job gets preempted, the job will automatically restart with all file system in place as interrupted -> we can find the last recovery checkpoint in the local checkpoint folder and provide it as recovery checkpoint to PL. This will avoid restarting training from scratch in case of pre-emption in the future.
For some use cases, we may want to keep more than one recovery checkpoint -> allow the user to specify the number of recovery checkpoint to keep.

This PR also gets rid of deprecated start_epoch config argument as PL will automatically restore epoch flag, optimizer, lr scheduler state from the loaded recovery checkpoint itself

# Conflicts: # InnerEye/ML/deep_learning_config.py # InnerEye/ML/model_training.py # InnerEye/ML/utils/checkpoint_handling.py # InnerEye/ML/utils/run_recovery.py # Tests/ML/models/test_scalar_model.py # Tests/ML/utils/test_checkpoint_handling.py

…ning into melanibe/checkpoint-saving � Conflicts: � CHANGELOG.md

melanibe · 2021-04-20T11:04:50Z

Closes #311

InnerEye/ML/deep_learning_config.py

InnerEye/ML/common.py

InnerEye/ML/deep_learning_config.py

InnerEye/ML/model_training.py

InnerEye/ML/utils/checkpoint_handling.py

InnerEye/ML/utils/run_recovery.py

Tests/ML/models/test_scalar_model.py

Tests/ML/test_model_training.py

melanibe added 17 commits April 19, 2021 16:12

Add auto-restart

2e52d58

Change handling of checkpoints and clean-up

be04f44

Save last k recovery checkpoints

c939279

Log epoch for keeping last ckpt

5ae4f28

Keeping k last checkpoints

8bebc7c

Add possibility to recover from particular checkpoint

2c8afdb

Update tests

bfc3975

Check k recovery

096fdcd

Re-add skipif

020a8c1

Correct pick up of recovery runs and add test

aea1b0b

Correct pick up of recovery runs and add test

9ce3e54

Remove all start epochs

0657b61

Remove all start epochs

0c22b53

Spimplify run recovery logic

422a859

Fix it

9dcfef3

Merge conflicts import errors

fcba209

melanibe changed the title ~~Melanibe/checkpoint saving~~ Checkpoint recovery refractoring Apr 19, 2021

melanibe added 12 commits April 19, 2021 18:26

Fix it

b50dcd8

Fix tests in test_scalar_model.py

7f3f44a

Fix tests in test_model_util.py

4702280

Fix tests in test_scalar_model.py

6ff6492

Fix tests in test_model_training.py

06852fc

Avoid forcing the user to log epoch

d98d41b

Fix test_get_checkpoints

3aebe38

Fix test_checkpoint_handling.py

0693fce

Fix callback

0bb57cc

Update CHANGELOG.md

23740a4

Self PR review comments

310a408

Fix more tests

ad488ad

melanibe added 3 commits April 20, 2021 10:53

Merge branch 'main' of https://github.com/microsoft/InnerEye-DeepLear…

c877bc9

…ning into melanibe/checkpoint-saving � Conflicts: � CHANGELOG.md

Fix argument in test

b9ba095

Mypy

5ec80e2

melanibe marked this pull request as ready for review April 20, 2021 10:11

melanibe added 2 commits April 20, 2021 12:12

Update InnerEye-DeepLearning.iml

c915317

Update InnerEye-DeepLearning.iml

e8b6e41

melanibe requested review from ant0nsc and Shruthi42 and removed request for ant0nsc April 20, 2021 10:16

Fix mypy errors

9c02781

ant0nsc suggested changes Apr 20, 2021

View reviewed changes

ant0nsc changed the title ~~Checkpoint recovery refractoring~~ Checkpoint recovery refactoring Apr 20, 2021

melanibe added 4 commits April 20, 2021 20:23

Address PR comment

f56155f

Typo

7302282

mypy fix

68abb81

just style

96e9c83

melanibe requested review from ant0nsc and javier-alvarez April 21, 2021 08:37

ant0nsc approved these changes Apr 21, 2021

View reviewed changes

melanibe enabled auto-merge (squash) April 21, 2021 09:44

ozan-oktay approved these changes Apr 21, 2021

View reviewed changes

melanibe merged commit adffa95 into main Apr 21, 2021

melanibe deleted the melanibe/checkpoint-saving branch April 21, 2021 14:40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Checkpoint recovery refactoring #439

Checkpoint recovery refactoring #439

melanibe commented Apr 19, 2021 •

edited

Loading

melanibe commented Apr 20, 2021

Checkpoint recovery refactoring #439

Checkpoint recovery refactoring #439

Conversation

melanibe commented Apr 19, 2021 • edited Loading

melanibe commented Apr 20, 2021

melanibe commented Apr 19, 2021 •

edited

Loading