diff --git a/tests/utils/test_checkpoint.py b/tests/utils/test_checkpoint.py index 7bd074cbac..edb90518f4 100644 --- a/tests/utils/test_checkpoint.py +++ b/tests/utils/test_checkpoint.py @@ -258,6 +258,21 @@ def test_from_str(self) -> None: metric_data=MetricData("eval_loss", 6.486097566010406e18), ), ), + # phase-naive checkpoints under a dirpath that contains a phase + # marker are still phase-naive + ( + "foo/eval_step_sweeps/epoch_2_step_1_acc=0.98", + CheckpointPath( + "foo/eval_step_sweeps", + epoch=2, + step={Phase.NONE: 1}, + metric_data=MetricData("acc", 0.98), + ), + ), + ( + "foo/train_step_dir/epoch_2_step_1", + CheckpointPath("foo/train_step_dir", epoch=2, step=1), + ), ] for path, expected_ckpt in valid_paths: parsed_ckpt = CheckpointPath.from_str(path) diff --git a/torchtnt/utils/checkpoint.py b/torchtnt/utils/checkpoint.py index c297957b51..f853b48144 100644 --- a/torchtnt/utils/checkpoint.py +++ b/torchtnt/utils/checkpoint.py @@ -160,9 +160,13 @@ def _populate_from_str(self, checkpoint_path: str) -> None: Raises: ValueError: If the path is malformed (either non-parsable, or contains wrong data types) """ + # Only the checkpoint name (the final path component) decides this. If we + # looked at the whole path, a dirpath that happens to contain + # "train_step"/"eval_step"/"predict_step" would flip a phase-naive + # checkpoint onto the phase-aware regex. + ckpt_name = checkpoint_path.rstrip("/").rsplit("/", 1)[-1] is_phase_aware = any( - phase in checkpoint_path - for phase in ["train_step", "eval_step", "predict_step"] + phase in ckpt_name for phase in ["train_step", "eval_step", "predict_step"] ) regex = self.PHASE_AWARE_REGEX if is_phase_aware else self.PHASE_NAIVE_REGEX path_match = regex.match(checkpoint_path)