Skip to content

Add configurable checkpointing to the larch2 optimization loop - #40

Open
matsen wants to merge 3 commits into
mainfrom
39-configurable-checkpointing
Open

Add configurable checkpointing to the larch2 optimization loop#40
matsen wants to merge 3 commits into
mainfrom
39-configurable-checkpointing

Conversation

@matsen

@matsen matsen commented Apr 27, 2026

Copy link
Copy Markdown
Contributor

Closes #39

Summary

  • New CLI flags --checkpoint-after <N>, --checkpoint-prefix <P>, and --resume <path> add configurable checkpoint/resume to the larch2 optimization loop.
  • Phased per the issue: Phase 1 writes DAG-only checkpoints (--resume is functionally equivalent to --dag-pb); Phase 2 writes a full-state larch_checkpoint (DAG + RNG + patience + results + git SHA + args fingerprint) so a resumed run is canonical-equivalent to an uninterrupted run with the same seeds.
  • Args fingerprint is SHA-256 over a JSON payload of the trajectory-affecting flags. The payload itself is also stored alongside the digest so a mismatch on resume names the conflicting field. Resume rejects every input flag (--dag-pb, --tree-pb, --fasta, --newick, --refseq, --vcf) since the checkpoint is self-describing.

Implementation notes

  • proto/checkpoint.proto + include/larch/checkpoint.hpp define the larch_checkpoint message family. C++ proto-mirror structs use the existing reflection-based encoder, plus new std::optional<int64_t> overloads so optional uint64 fields distinguish 0 from "unset".
  • mt19937 (de)serialization uses std::ostringstream / operator<< (the C++ standard guarantees byte-exact round-trip).
  • SHA-256 is a self-contained FIPS 180-4 implementation in src/checkpoint.cpp (no new dep).
  • Atomic write: <path>.tmp then rename.
  • Format detection: a larch_checkpoint always starts with the byte 0x08 (tag for schema_version, varint), distinguishing it from raw dag_data files (0x0a or other length-delimited fields).
  • Determinism fix to build_phylo_dag_from_proto: walk proto IDs in sorted order and append edges in edge_id order so a save+load round trip preserves the in-memory phylo_dag's node/edge layout. The previous unordered_map iteration and parent-sort silently broke determinism for resume because sampling traverses the DAG by node-index and children-list order.

Tests added

  • test/checkpoint_test.cpp — mt19937 round-trip, larch_checkpoint proto round-trip (incl. nested dag_data, results vector, patience state, RNG state), SHA-256 known-answer tests, optional<int64_t> zero-vs-absent presence semantics.
  • test/checkpoint_integration.sh — end-to-end CLI: Phase 2 deterministic resume for K∈{1,2,3} with canonical-equal final output; multi-resume chain (catches RNG-state non-idempotency); Phase 1 score-not-worse-than-baseline; args fingerprint mismatch refused with the conflicting field named; --resume + --vcf refused; pre-drift checkpoint emission with --patience 2 --drift 1; future schema_version rejected with a clear error.
  • test/dag_canonical_equal.cpp — content-keyed canonical comparison helper used by the integration test (uses sample_id for leaves and the compact-genome string for inner nodes; edge tuples as a multiset). Stronger than the issue's suggested proto-id sort, which would fail when merge state differs across runs and reorders proto-id assignments even though the DAGs are logically equivalent.

Test plan

  • ctest --test-dir build — all 21 tests pass (19 existing + 2 new).
  • Manual verification: 6-iter run with --checkpoint-after 1 writes out.ckpt-{1..6}.pb.gz and the final output canonical-equals a baseline uninterrupted run for every checkpoint K.
  • Pre-drift checkpoint manually verified with --patience 2 --drift 1 on data/testcase/full_dag.pb.gz.

🤖 Generated with Claude Code

matsen and others added 3 commits April 27, 2026 07:03
New CLI flags --checkpoint-after, --checkpoint-prefix, and --resume
let a long-running optimization be interrupted and resumed without
losing progress. The implementation ships in two phases as the issue
requires: Phase 1 writes DAG-only checkpoints (resume == --dag-pb),
Phase 2 writes a full-state larch_checkpoint message (RNG, patience,
results, args fingerprint) that resumes byte-equivalent to an
uninterrupted run with the same seeds.

Determinism fix in build_phylo_dag_from_proto: walk proto IDs in
sorted order and append edges in edge_id order so a save+load round
trip preserves the in-memory phylo_dag's node/edge layout. Sampling
traverses by node-index and children-list order, so the previous
unordered_map iteration broke deterministic resume.

Includes:
- proto/checkpoint.proto and the C++ proto-mirror in
  include/larch/checkpoint.hpp.
- Self-contained SHA-256 (FIPS 180-4) for the args fingerprint.
- mt19937 round-trip via std::ostringstream + operator<<.
- Atomic checkpoint write via "<path>.tmp" + rename.
- Unit tests (RNG, proto round-trip, SHA-256 KATs, optional<int64>
  presence) and an integration test that asserts canonical-equal
  resume across K∈{1,2,3}, multi-resume chain, fingerprint mismatch,
  forbidden-input rejection, pre-drift checkpoint emission, and
  schema_version-too-new rejection.

Closes #39

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Extract encode_dag_to_proto() helper in save_proto_dag.hpp so the
  larch_checkpoint writer no longer duplicates dag-encoding logic.
- Reorder main() resume path to call resume_state_from_msg(ckpt) before
  consuming ckpt.dag, so no field is read after a partial move.
- Rename build_checkpoint_msg's completed_iter parameter to
  next_iteration to match the proto field semantics.
- static_assert that checkpoint_schema_version > 0, guarding the
  looks_like_checkpoint magic-byte test against future field reorders.
- Extend integration test K coverage to {1..5} (every-iteration
  interrupt), test --dag-pb as a forbidden flag with --resume, and
  guard extract_scores against an empty result.
- Strengthen checkpoint_test round-trip to assert every result/radius
  field, and add a non-empty drift_rng round-trip case.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Rename radius_result_msg_pb / optimize_result_msg_pb to plain _msg
  to match the existing patience_state_msg / rng_state_msg /
  larch_checkpoint_msg naming (the _pb suffix wrongly suggested
  these were codegen-produced).
- Collapse the three opt_u / opt_z / opt_i lambdas in
  args_fingerprint_payload into a single auto-deduced opt() lambda.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@matsen

matsen commented Apr 27, 2026

Copy link
Copy Markdown
Contributor Author

Issue Lead (final review)

Category: completed
Scope check: on-track. Every CLI flag, proto field, integration scenario, and success criterion in the issue body has a corresponding implementation and test. No scope drift, no premature deferral.

Assessment:

  • Phase 1 and Phase 2 both ship as specified. The issue's Phase 2 gate ("two runs with identical --seed, one uninterrupted and one interrupted+resumed, produce final outputs that match under canonical comparison and identical per-iteration parsimony scores") is verified by checkpoint_integration_test for K in {1..5} on a 6-iteration run, which is the full N-1 every-iteration interrupt property test from the issue's test plan Fix sampled trees dropping leaves from inconsistent DAGs #9 (stronger than the {3,7,13} sample originally suggested).
  • All 9 items in the issue's test plan are covered: RNG round-trip, checkpoint round-trip, Phase 1 score-not-worse, Phase 2 determinism, pre-drift checkpoint, args fingerprint mismatch (names the conflicting field), multi-resume chain, forbidden input flag, every-iteration interrupt.
  • The determinism fix to build_phylo_dag_from_proto (sorted proto-id walk, edges in edge_id order) is the right kind of root-cause fix — it identifies and resolves the unordered_map iteration that would have silently broken Phase 2 byte-equivalence even with otherwise-correct save/load.
  • Args fingerprint stores both the SHA-256 digest and the underlying JSON payload so a mismatch can name the conflicting field. Schema version is guarded by a static_assert against the magic-byte detection. Atomic write via .tmp + rename. mt19937 via operator<<. All match the issue's design choices.
  • 21/21 tests pass (full ctest run, integration test 38.6s).

Follow-ups filed: none. The PR body has no DEFERRED section and the changed files have no TODO/FIXME/"defer"/"future work" markers. The reviewer-flagged A4/N3 nits in load_proto_dag.hpp and the drift_escape switch are pre-existing code untouched by this PR — fixing them here would be scope drift.

Action: Phase set to completed. PR #40 is ready to merge.

@matsen
matsen requested a review from ognian- April 27, 2026 23:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add configurable checkpointing to the larch2 optimization loop

1 participant