wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-14 12:40:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[tune] Prevent MEMORY checkpoints from breaking trial FT (#6691)

Browse Source 
* Prevent MEMORY checkpoints from breaking FT

* Add save/pause/resume/restore test

* change checkpoint return value based on status

* Fix test_checkpoint_manager_tests.

* Fix test + checkpoint manager bug

* lint

* Add docstring

* Add docstring to checkpoint_manager constructor

* Change variable name for clarity

* Revert on_checkpoint docstring wording

* Break after success

* nit: more informative warning

* Quarantine test
This commit is contained in:
Ujval Misra
2020-01-22 23:17:09 -08:00
committed by Richard Liaw
parent 0834bda8c1
commit 1558307ac4
10 changed files with 122 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/tune/checkpoint_manager.py
+24 -12
View File
@@ -45,6 +45,9 @@ class CheckpointManager:
def __init__(self, keep_checkpoints_num, checkpoint_score_attr, delete_fn):
"""Initializes a new CheckpointManager.
`newest_persistent_checkpoint` and `newest_memory_checkpoint` are
initialized to Checkpoint objects with values of None.
Args:
keep_checkpoints_num (int): Keep at least this many checkpoints.
checkpoint_score_attr (str): Attribute to use to determine which
@@ -60,28 +63,38 @@ class CheckpointManager:
self._checkpoint_score_attr = checkpoint_score_attr[4:]
else:
self._checkpoint_score_attr = checkpoint_score_attr
self.delete = delete_fn
self.newest_checkpoint = Checkpoint(Checkpoint.MEMORY, None)
self.newest_persistent_checkpoint = Checkpoint(Checkpoint.PERSISTENT,
None)
self.newest_memory_checkpoint = Checkpoint(Checkpoint.MEMORY, None)
self._best_checkpoints = []
self._membership = set()
def on_checkpoint(self, checkpoint):
"""Starts tracking checkpoint metadata on checkpoint.
Sets newest checkpoint. Deletes previous checkpoint as long as it isn't
one of the best ones. Also deletes the worst checkpoint if at capacity.
Sets the newest checkpoint. For PERSISTENT checkpoints: Deletes
previous checkpoint as long as it isn't one of the best ones. Also
deletes the worst checkpoint if at capacity.
Args:
checkpoint (Checkpoint): Trial state checkpoint.
"""
old_checkpoint = self.newest_checkpoint
self.newest_checkpoint = checkpoint
if checkpoint.storage == Checkpoint.MEMORY:
self.newest_memory_checkpoint = checkpoint
return
old_checkpoint = self.newest_persistent_checkpoint
self.newest_persistent_checkpoint = checkpoint
# Remove the old checkpoint if it isn't one of the best ones.
if old_checkpoint.value and old_checkpoint not in self._membership:
self.delete(old_checkpoint)
try:
queue_item = QueueItem(self._priority(checkpoint), checkpoint)
except KeyError:
if old_checkpoint not in self._membership:
self.delete(old_checkpoint)
logger.error("Result dict has no key: {}. "
"checkpoint_score_attr must be set to a key in the "
"result dict.".format(self._checkpoint_score_attr))
@@ -95,11 +108,10 @@ class CheckpointManager:
self._membership.add(checkpoint)
if worst in self._membership:
self._membership.remove(worst)
self.delete(worst)
# Remove the old checkpoint if it isn't one of the best ones.
if old_checkpoint.value and old_checkpoint not in self._membership:
self.delete(old_checkpoint)
# Don't delete the newest checkpoint. It will be deleted on the
# next on_checkpoint() call since it isn't in self._membership.
if worst != checkpoint:
self.delete(worst)
def best_checkpoints(self):
"""Returns best checkpoints, sorted by score."""