Separate reconstructed and in-place PETSc checkpoint reloads - #674
Merged
lmoresi merged 2 commits intoSep 4, 2026
Merged
Conversation
Restore reconstructed checkpoint meshes through DMPlex section/local-vector metadata and the topology migration SF, because raw global-vector ordering is not stable across mesh reconstruction. Preserve an explicit same-layout vector path for in-place disk-snapshot restoration and reject snapshot rank-count changes. Add serial format coverage and a four-rank regression that verifies both in-place snapshot restore and fresh reconstructed-mesh reload exactly reproduce coordinate-defined fields.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
There are a couple of concrete robustness issues in the new MPI error-path handling/metadata validation that can lead to hangs or misleading behavior.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR separates two PETSc checkpoint restore modes for MeshVariable.read_checkpoint(): a safe default that remaps onto a reconstructed mesh via DMPlex topology migration, and an explicit same_layout=True mode for exact in-place reload onto the original live mesh layout. This aligns disk snapshots with the in-place case (and enforces same MPI rank count), while keeping production “reconstructed restart” behavior safe by default.
Changes:
- Add
same_layoutoption toMeshVariable.read_checkpoint()and write an explicit owned global-vector payload under/uw_checkpointfor in-place restores. - Make disk snapshot restore use
same_layout=Trueand reject a changed MPI rank count; also invalidate the mesh-wide packed auxiliary vector after reload. - Extend snapshot tests (serial + MPI ptest) to validate both in-place and reconstructed-mesh reload correctness and to assert the new on-disk HDF5 structure.
File summaries
| File | Description |
|---|---|
tests/test_0010_snapshot_disk_format.py |
Extends snapshot disk-format and restore behavior tests, including auxiliary-vector invalidation and rank-count rejection. |
tests/parallel/ptest_0010_snapshot_disk.py |
Strengthens MPI snapshot ptest to validate analytic DOF correctness and reconstructed-mesh reload correctness. |
src/underworld3/discretisation/discretisation_mesh.py |
Writes the new /uw_checkpoint global-vector payload for explicit same-layout reloads. |
src/underworld3/discretisation/discretisation_mesh_variables.py |
Implements same_layout branching in read_checkpoint() and invalidates mesh._stale_lvec after reload. |
src/underworld3/checkpoint/disk_snapshot.py |
Enforces same MPI rank count for exact snapshot reload and routes snapshot restore through same_layout=True. |
Review details
- Files reviewed: 5/5 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1475
to
+1494
| if same_layout: | ||
| import h5py | ||
|
|
||
| if uw.mpi.rank == 0: | ||
| with h5py.File(filename, "r") as checkpoint_h5: | ||
| has_direct_vector = ( | ||
| "uw_checkpoint" in checkpoint_h5 | ||
| and data_name in checkpoint_h5["uw_checkpoint"] | ||
| ) | ||
| else: | ||
| has_direct_vector = None | ||
| has_direct_vector = uw.mpi.comm.bcast( | ||
| has_direct_vector, | ||
| root=0, | ||
| ) | ||
| if not has_direct_vector: | ||
| raise RuntimeError( | ||
| f"{filename} has no in-place checkpoint vector for " | ||
| f"{data_name!r}. Reload it with same_layout=False." | ||
| ) |
Comment on lines
487
to
+489
| md = read_snapshot_metadata(path) | ||
| write_size = int(md.get("mpi_ranks_at_write", 1)) | ||
| if write_size != int(uw.mpi.size): |
5 tasks
gthyagi
added a commit
to gthyagi/underworld3
that referenced
this pull request
Sep 5, 2026
Detect the same-layout checkpoint API added by upstream underworldcode#674 before asserting MPI disk replay. Older transport-branch checkouts continue to test serial disk restore and MPI in-memory restore; rebased development checkouts automatically exercise the full distributed disk path. This keeps the SUPG PR focused while making its dependency on the already-merged checkpoint fix explicit. Underworld development team with AI support from Claude Code.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
PETSc checkpoint fields have two distinct restore contexts:
checkpoint.mesh.00000.h5. Its global DOF ordering may differ from the writer even with the same MPI rank count.model.load_state()restores into the exact live mesh object that wrote the disk snapshot.A raw global vector is exact for case 2 but is unsafe for case 1. In the 48-rank Zhong A1 restart, applying the raw vector to a reconstructed mesh kept the scalar range at
[0,1]while attaching values to the wrong coordinates. The first Stokes solve then increasedvelocity_maxfrom about3.5to75.8.Change
MeshVariable.read_checkpoint()on the DMPlex section/local-vector path by default. PETSc's topology migration SF maps the saved point data onto a reconstructed mesh.same_layout=Truefor the exact in-place restore case only./uw_checkpointfor that explicit path.same_layout=Trueinternally and reject a changed MPI rank count.Benchmark scripts do not select a reload mode. A normal reconstructed production restart continues to call
read_checkpoint()with its default behavior.48-rank Gadi validation
All rows use the same
cellsize=1/32, quadrature-degree-4 Zhong A1 checkpoint.1.00.00.953050.0velocity_max75.783.49[-0.0418, 1.0288][0,1]11.010.7910The corrected step-2 job also retained the expected early transient diagnostics: surface/CMB Nusselt numbers
0.99/1.01, mean temperature0.31202, and51.49 GiBPBS peak memory.Tests
27 passed:tests/test_0010_snapshot_disk_format.pytests/parallel/ptest_0010_snapshot_disk.py: in-place field max error0.0; reconstructed-mesh field max error0.00.00, finite U/P/T/Tdot checkpoint written