wehub-resource-sync
240 lines
8.0 KiB
Python
240 lines
8.0 KiB
Python
import random
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from ray.train import Checkpoint, CheckpointConfig
|
|
from ray.train._internal.checkpoint_manager import _CheckpointManager, _TrainingResult
|
|
from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE
|
|
|
|
|
|
@pytest.fixture
|
|
def checkpoint_paths(tmp_path):
|
|
checkpoint_paths = []
|
|
for i in range(10):
|
|
checkpoint_path = tmp_path / f"ckpt_{i}"
|
|
checkpoint_path.mkdir()
|
|
(checkpoint_path / "dummy.txt").write_text(f"{i}")
|
|
checkpoint_paths.append(checkpoint_path)
|
|
|
|
yield [str(path) for path in checkpoint_paths]
|
|
|
|
|
|
def test_unlimited_checkpoints(checkpoint_paths: List[str]):
|
|
manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=None))
|
|
|
|
for i in range(10):
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
|
|
metrics={"iter": i},
|
|
)
|
|
)
|
|
|
|
assert len(manager.best_checkpoint_results) == 10
|
|
|
|
|
|
def test_limited_checkpoints(checkpoint_paths: List[str]):
|
|
manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=2))
|
|
|
|
for i in range(10):
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
|
|
metrics={"iter": i},
|
|
)
|
|
)
|
|
|
|
assert len(manager.best_checkpoint_results) == 2
|
|
|
|
# Keep the latest checkpoints if no metric is given.
|
|
assert {
|
|
tracked_checkpoint.metrics["iter"]
|
|
for tracked_checkpoint in manager.best_checkpoint_results
|
|
} == {8, 9}
|
|
|
|
# The first 8 checkpoints should be deleted.
|
|
for i in range(8):
|
|
assert not Path(checkpoint_paths[i]).exists()
|
|
|
|
assert Path(checkpoint_paths[8]).exists()
|
|
assert Path(checkpoint_paths[9]).exists()
|
|
|
|
|
|
@pytest.mark.parametrize("order", ["min", "max"])
|
|
def test_keep_checkpoints_by_score(order, checkpoint_paths):
|
|
num_to_keep = 2
|
|
score_attribute = "score"
|
|
|
|
manager = _CheckpointManager(
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=num_to_keep,
|
|
checkpoint_score_attribute=score_attribute,
|
|
checkpoint_score_order=order,
|
|
)
|
|
)
|
|
|
|
scores = []
|
|
for i in range(10):
|
|
score = random.random()
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
|
|
metrics={"iter": i, score_attribute: score},
|
|
)
|
|
)
|
|
scores.append(score)
|
|
|
|
sorted_scores = sorted(scores, reverse=order == "max")
|
|
assert set(sorted_scores[:num_to_keep]) == {
|
|
tracked_checkpoint.metrics[score_attribute]
|
|
for tracked_checkpoint in manager.best_checkpoint_results
|
|
}
|
|
|
|
# Make sure the bottom checkpoints are deleted.
|
|
best_checkpoint_iters = {
|
|
tracked_checkpoint.metrics["iter"]
|
|
for tracked_checkpoint in manager.best_checkpoint_results
|
|
}
|
|
for i, checkpoint_path in enumerate(checkpoint_paths):
|
|
if i in best_checkpoint_iters or i == 9:
|
|
# The checkpoint should only exist if it's one of the top K or the latest.
|
|
assert Path(checkpoint_path).exists()
|
|
else:
|
|
assert not Path(checkpoint_path).exists()
|
|
|
|
|
|
def test_keep_latest_checkpoint(checkpoint_paths):
|
|
manager = _CheckpointManager(
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=2,
|
|
checkpoint_score_attribute="score",
|
|
checkpoint_score_order="max",
|
|
)
|
|
)
|
|
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
|
|
metrics={"score": 3.0},
|
|
)
|
|
)
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
|
|
metrics={"score": 2.0},
|
|
)
|
|
)
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
|
|
metrics={"score": 1.0},
|
|
)
|
|
)
|
|
|
|
assert len(manager.best_checkpoint_results) == 2
|
|
|
|
# The latest checkpoint with the lowest score should not be deleted yet.
|
|
assert manager.latest_checkpoint_result.metrics["score"] == 1.0
|
|
|
|
# The latest checkpoint with the lowest score should not be deleted yet.
|
|
assert Path(checkpoint_paths[2]).exists()
|
|
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[3]),
|
|
metrics={"score": 0.0},
|
|
)
|
|
)
|
|
# A newer checkpoint came in. Even though the new one has a lower score, there are
|
|
# already num_to_keep better checkpoints, so the previous one should be deleted.
|
|
assert not Path(checkpoint_paths[2]).exists()
|
|
|
|
# Quick sanity check to make sure that the new checkpoint is kept.
|
|
assert manager.latest_checkpoint_result.metrics["score"] == 0.0
|
|
assert Path(checkpoint_paths[3]).exists()
|
|
|
|
# The original 2 checkpoints should still exist
|
|
assert Path(checkpoint_paths[0]).exists()
|
|
assert Path(checkpoint_paths[1]).exists()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"metrics",
|
|
[
|
|
{"nested": {"sub": {"attr": 5}}},
|
|
{"nested": {"sub/attr": 5}},
|
|
{"nested/sub": {"attr": 5}},
|
|
{"nested/sub/attr": 5},
|
|
],
|
|
)
|
|
def test_nested_get_checkpoint_score(metrics):
|
|
manager = _CheckpointManager(
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=2,
|
|
checkpoint_score_attribute="nested/sub/attr",
|
|
checkpoint_score_order="max",
|
|
)
|
|
)
|
|
|
|
tracked_checkpoint = _TrainingResult(checkpoint=None, metrics=metrics)
|
|
assert manager._get_checkpoint_score(tracked_checkpoint) == (True, 5.0)
|
|
|
|
|
|
@pytest.mark.parametrize("has_score_attr", [True, False])
|
|
def test_only_store_score_attr(has_score_attr, checkpoint_paths, monkeypatch):
|
|
monkeypatch.setenv(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, "1")
|
|
|
|
# Set up CheckpointManager.
|
|
if has_score_attr:
|
|
checkpoint_config = CheckpointConfig(
|
|
num_to_keep=None,
|
|
checkpoint_score_attribute="score",
|
|
checkpoint_score_order="max",
|
|
)
|
|
else:
|
|
checkpoint_config = CheckpointConfig(num_to_keep=None)
|
|
manager = _CheckpointManager(checkpoint_config=checkpoint_config)
|
|
|
|
# Ensure we insert TrainingResults with score in the right order.
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
|
|
metrics={"score": 3.0},
|
|
)
|
|
)
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
|
|
metrics={"score": 1.0, "another_unsaved_metric": 6.0},
|
|
)
|
|
)
|
|
manager.register_checkpoint(
|
|
_TrainingResult(
|
|
checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
|
|
metrics={"another_unsaved_metric": 1.0},
|
|
)
|
|
)
|
|
assert len(manager.best_checkpoint_results) == 3
|
|
if has_score_attr:
|
|
assert manager.best_checkpoint_results[0].metrics == {"score": 1.0}
|
|
assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[1]
|
|
assert manager.best_checkpoint_results[1].metrics == {"score": 3.0}
|
|
assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[0]
|
|
assert manager.best_checkpoint_results[2].metrics == {}
|
|
assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
|
|
else:
|
|
assert manager.best_checkpoint_results[0].metrics == {}
|
|
assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[0]
|
|
assert manager.best_checkpoint_results[1].metrics == {}
|
|
assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[1]
|
|
assert manager.best_checkpoint_results[2].metrics == {}
|
|
assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|