diff --git a/flytekit/core/checkpointer.py b/flytekit/core/checkpointer.py index bd17b7748b..bed776ed4c 100644 --- a/flytekit/core/checkpointer.py +++ b/flytekit/core/checkpointer.py @@ -146,7 +146,9 @@ def read(self) -> typing.Optional[bytes]: if p is None: return None files = list(p.iterdir()) - if len(files) == 0 or len(files) > 1: + if len(files) == 0: + return None + if len(files) > 1: raise ValueError(f"Expected exactly one checkpoint - found {len(files)}") f = files[0] return f.read_bytes() diff --git a/tests/flytekit/unit/core/test_checkpoint.py b/tests/flytekit/unit/core/test_checkpoint.py index 0199737335..f1dbbbd5ef 100644 --- a/tests/flytekit/unit/core/test_checkpoint.py +++ b/tests/flytekit/unit/core/test_checkpoint.py @@ -83,6 +83,16 @@ def test_sync_checkpoint_restore_default_path(tmpdir): assert cp.restore() == cp._prev_download_path +def test_sync_checkpoint_read_empty_dir(tmpdir): + td_path = Path(tmpdir) + dest = td_path.joinpath("dest") + dest.mkdir() + src = td_path.joinpath("src") + src.mkdir() + cp = SyncCheckpoint(checkpoint_dest=str(dest), checkpoint_src=str(src)) + assert cp.read() is None + + def test_sync_checkpoint_read_multiple_files(tmpdir): """ Read can only work with one file.