From 4e0c0a9c30aae0c295e93b4258e2709488d684c9 Mon Sep 17 00:00:00 2001 From: cslht11 Date: Sun, 1 Nov 2026 05:22:27 -0800 Subject: [PATCH 1/3] fix: let a reused checkpoint directory replace earlier steps Saving an iteration whose files already exist crashed with FileExistsError from the exclusive rank-file creation. Re-saving now replaces only this rank's previous files, so ranks writing the same step concurrently never touch each other; a regression test pins the directory contents and the deterministic results. --- diffct/_solver_state.py | 9 +++++++-- tests/test_solver_checkpoint_edges.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/diffct/_solver_state.py b/diffct/_solver_state.py index d9cede4..edca7b3 100644 --- a/diffct/_solver_state.py +++ b/diffct/_solver_state.py @@ -288,10 +288,15 @@ 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}" - rank.mkdir() + rank.mkdir(exist_ok=True) for name, source in states.items(): source = TensorStore(source.detach()) if isinstance(source, torch.Tensor) else source - target = NpyStore.create(rank / f"{name}.npy", source.shape, dtype=source.dtype) + path = rank / f"{name}.npy" + if path.exists(): + # Re-saving replaces only this rank's previous file; other + # ranks may be writing the same step concurrently. + path.unlink() + target = NpyStore.create(path, source.shape, dtype=source.dtype) copy_blocks(source, target, elements) target.flush() manifest[name] = dict(shape=list(source.shape), dtype=str(source.dtype), fingerprint=fingerprint(target, elements)) diff --git a/tests/test_solver_checkpoint_edges.py b/tests/test_solver_checkpoint_edges.py index 5c691e5..4e5a352 100644 --- a/tests/test_solver_checkpoint_edges.py +++ b/tests/test_solver_checkpoint_edges.py @@ -285,3 +285,17 @@ 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" + + +def test_reusing_a_checkpoint_directory_replaces_steps(tmp_path): + from diffct import Projector, cgls + from diffct.geometry import circular_trajectory_2d_parallel + + trajectory = circular_trajectory_2d_parallel(8, device="cpu", dtype=torch.float32) + projector = Projector(trajectory, volume_shape=(16, 16), detector_shape=(24,), + beam="parallel", backend="torch") + sinogram = projector(torch.rand(16, 16)) + first = cgls(projector, sinogram, iterations=2, checkpoint_directory=str(tmp_path)) + second = cgls(projector, sinogram, iterations=2, checkpoint_directory=str(tmp_path)) + assert [p.name for p in sorted(tmp_path.iterdir())] == ["step-000001", "step-000002"] + torch.testing.assert_close(first.volume.float(), second.volume.float()) From 1211e28023ab620a6369cf234ad2f5c7b2e2a693 Mon Sep 17 00:00:00 2001 From: Yipeng Sun Date: Mon, 2 Nov 2026 18:29:10 +0100 Subject: [PATCH 2/3] fix: reject an existing checkpoint step with a clear error instead of replacing it --- diffct/_solver_state.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/diffct/_solver_state.py b/diffct/_solver_state.py index edca7b3..6d4ec8c 100644 --- a/diffct/_solver_state.py +++ b/diffct/_solver_state.py @@ -288,15 +288,15 @@ 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}" - rank.mkdir(exist_ok=True) + 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 - path = rank / f"{name}.npy" - if path.exists(): - # Re-saving replaces only this rank's previous file; other - # ranks may be writing the same step concurrently. - path.unlink() - target = NpyStore.create(path, source.shape, dtype=source.dtype) + target = NpyStore.create(rank / f"{name}.npy", source.shape, dtype=source.dtype) copy_blocks(source, target, elements) target.flush() manifest[name] = dict(shape=list(source.shape), dtype=str(source.dtype), fingerprint=fingerprint(target, elements)) From 4548f0026b91ff13dd0a49922993e05a8417a17e Mon Sep 17 00:00:00 2001 From: Yipeng Sun Date: Mon, 2 Nov 2026 18:31:06 +0100 Subject: [PATCH 3/3] Test checkpoint collision rejection and solver resume preservation --- tests/test_solver_checkpoint_edges.py | 81 ++++++++++++++++++++++++--- 1 file changed, 72 insertions(+), 9 deletions(-) diff --git a/tests/test_solver_checkpoint_edges.py b/tests/test_solver_checkpoint_edges.py index 4e5a352..cc1dd6a 100644 --- a/tests/test_solver_checkpoint_edges.py +++ b/tests/test_solver_checkpoint_edges.py @@ -287,15 +287,78 @@ def test_four_rank_atomic_metadata_failure_rejects_collectively_before_more_oper assert record["elapsed"] < 15, "publication failure waited for the Gloo timeout" -def test_reusing_a_checkpoint_directory_replaces_steps(tmp_path): - from diffct import Projector, cgls +@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 = Projector(trajectory, volume_shape=(16, 16), detector_shape=(24,), - beam="parallel", backend="torch") - sinogram = projector(torch.rand(16, 16)) - first = cgls(projector, sinogram, iterations=2, checkpoint_directory=str(tmp_path)) - second = cgls(projector, sinogram, iterations=2, checkpoint_directory=str(tmp_path)) - assert [p.name for p in sorted(tmp_path.iterdir())] == ["step-000001", "step-000002"] - torch.testing.assert_close(first.volume.float(), second.volume.float()) + 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)