diff --git a/diffct/_solver_state.py b/diffct/_solver_state.py index d9cede4..6d4ec8c 100644 --- a/diffct/_solver_state.py +++ b/diffct/_solver_state.py @@ -288,6 +288,11 @@ def save_checkpoint(operator, directory, states, metadata, iteration, history, m step = Path(directory).resolve() / f"step-{iteration:06d}" step.mkdir(parents=True, exist_ok=True) rank = step / f"rank-{getattr(operator, 'rank', 0):05d}" + if rank.exists(): + # Replacing a published step in place could destroy it on a + # failed write, and ranks cannot roll back together. + raise ValueError(f"checkpoint {rank} already exists; use a new checkpoint_directory " + "or pass resume= to continue from it") rank.mkdir() for name, source in states.items(): source = TensorStore(source.detach()) if isinstance(source, torch.Tensor) else source diff --git a/tests/test_solver_checkpoint_edges.py b/tests/test_solver_checkpoint_edges.py index 5c691e5..cc1dd6a 100644 --- a/tests/test_solver_checkpoint_edges.py +++ b/tests/test_solver_checkpoint_edges.py @@ -285,3 +285,80 @@ def test_four_rank_atomic_metadata_failure_rejects_collectively_before_more_oper assert len(record["calls"]) == record["calls_at_publication"], "rank executed another operator after publication failure" assert record["complete_state_observations"] >= 2, "missing all-rank write/flush evidence before step2 publication" assert record["elapsed"] < 15, "publication failure waited for the Gloo timeout" + + +@pytest.fixture +def checkpoint_problem(): + from diffct.geometry import circular_trajectory_2d_parallel + + trajectory = circular_trajectory_2d_parallel(8, device="cpu", dtype=torch.float32) + projector = diffct.Projector(trajectory, volume_shape=(8, 8), detector_shape=(12,), + beam="parallel", backend="torch") + truth = torch.arange(64, dtype=torch.float32).reshape(8, 8) / 64 + return projector, projector(truth) + + +@pytest.mark.parametrize("solver", [diffct.cgls, diffct.sirt], ids=["cgls", "sirt"]) +@pytest.mark.parametrize("resume_step", [None, 1], ids=["fresh", "resumed"]) +def test_existing_checkpoint_rejects_next_step_and_preserves_resumable_state( + solver, resume_step, checkpoint_problem, tmp_path, +): + projector, sinogram = checkpoint_problem + directory = tmp_path / "checkpoint" + options = dict(rtol=0., atol=0.) + expected = solver(projector, sinogram, 4, **options) + partial = solver(projector, sinogram, 2, checkpoint_directory=directory, **options) + assert partial.iterations == 2 + before = {path.relative_to(directory): path.read_bytes() + for path in directory.rglob("*") if path.is_file()} + assert before + resume = None if resume_step is None else _checkpoint(directory, resume_step) + next_step = 1 if resume_step is None else resume_step + 1 + existing = _checkpoint(directory, next_step) + with pytest.raises(ValueError) as caught: + solver(projector, sinogram, 4, checkpoint_directory=directory, resume=resume, **options) + message = str(caught.value) + assert str(existing) in message + assert "new checkpoint_directory" in message + assert "resume=" in message + after = {path.relative_to(directory): path.read_bytes() + for path in directory.rglob("*") if path.is_file()} + assert after == before + resumed = solver(projector, sinogram, 4, resume=existing, **options) + assert resumed.iterations == expected.iterations == 4 + torch.testing.assert_close(resumed.volume, expected.volume) + + +@pytest.mark.parametrize("solver", [diffct.cgls, diffct.sirt], ids=["cgls", "sirt"]) +def test_empty_checkpoint_directory_publishes_each_step(solver, checkpoint_problem, tmp_path): + projector, sinogram = checkpoint_problem + options = dict(rtol=0., atol=0.) + expected = solver(projector, sinogram, 3, **options) + result = solver(projector, sinogram, 3, checkpoint_directory=tmp_path, **options) + assert result.iterations == expected.iterations == 3 + assert sorted(path.name for path in tmp_path.iterdir()) == [ + "step-000001", "step-000002", "step-000003", + ] + torch.testing.assert_close(result.volume, expected.volume) + for step in range(1, 4): + resumed = solver(projector, sinogram, 3, resume=_checkpoint(tmp_path, step), **options) + assert resumed.iterations == 3 + torch.testing.assert_close(resumed.volume, expected.volume) + + +@pytest.mark.parametrize("solver", [diffct.cgls, diffct.sirt], ids=["cgls", "sirt"]) +def test_resume_publishes_remaining_steps_in_new_empty_directory(solver, checkpoint_problem, tmp_path): + projector, sinogram = checkpoint_problem + original, destination = tmp_path / "original", tmp_path / "new" + destination.mkdir() + options = dict(rtol=0., atol=0.) + expected = solver(projector, sinogram, 4, **options) + partial = solver(projector, sinogram, 2, checkpoint_directory=original, **options) + assert partial.iterations == 2 + resumed = solver(projector, sinogram, 4, resume=_checkpoint(original, 2), + checkpoint_directory=destination, **options) + assert resumed.iterations == expected.iterations == 4 + assert sorted(path.name for path in destination.iterdir()) == ["step-000003", "step-000004"] + torch.testing.assert_close(resumed.volume, expected.volume) + reloaded = solver(projector, sinogram, 4, resume=_checkpoint(destination, 4), **options) + torch.testing.assert_close(reloaded.volume, expected.volume)