diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml index 04d68491f..f5cffd248 100644 --- a/.github/workflows/test_pr_and_main.yml +++ b/.github/workflows/test_pr_and_main.yml @@ -174,6 +174,10 @@ jobs: run: | coverage run $COV_ARGS -m pytest mpisppy/tests/test_vss.py -v + - name: Test checkpoint/resume + run: | + coverage run $COV_ARGS -m pytest mpisppy/tests/test_checkpoint.py -v + - name: Test prox_approx end-to-end run: | coverage run $COV_ARGS -m pytest mpisppy/tests/test_prox_approx_e2e.py -v @@ -1234,6 +1238,7 @@ jobs: mpisppy/tests/test_feasible_xhat.py \ mpisppy/tests/test_proper_bundler.py \ mpisppy/tests/test_incumbent_writing.py \ + mpisppy/tests/test_checkpoint.py \ mpisppy/tests/test_iis_on_infeasible.py \ mpisppy/tests/test_grad_rho_bundles.py \ mpisppy/tests/test_sensi_rho_bundles.py \ diff --git a/doc/designs/checkpointing_design.md b/doc/designs/checkpointing_design.md new file mode 100644 index 000000000..9190cdab7 --- /dev/null +++ b/doc/designs/checkpointing_design.md @@ -0,0 +1,1038 @@ +# Checkpoint / Resume for mpi-sppy — Design + +Status: **draft** (framework and dill-reload backend both PoC-validated — the +latter on a serial MIP; multi-rank and cylinders pending). Scope: checkpoint a +running mpi-sppy job so it can be stopped and resumed later. Must work on multiple +MPI ranks and for cylinder (hub-and-spoke) runs. + +--- + +## 1. Goals and non-goals + +**Primary use case.** A long (multi-day) run that is **intentionally stopped and +resumed** on a schedule — e.g. a three-day study that ends each day and picks up +the next morning on the same cluster. Checkpoints are **infrequent** (a small +number over the whole run; roughly twice as many writes as resumes) and resumes +are planned, not crash-driven. The scenarios are **large MIPs**. + +That regime drives the design decisions below (dill the scenario models, restore +a MIP warm start), and it is worth stating up front because a *different* use +case — frequent checkpoints purely for hard-kill safety — would push toward a +lighter, leaf-data-only checkpoint — the leaf-rebuild backend, kept in this design +as a future option but **not currently planned** (§4, §11 Phase 6). + +**Goals** + +- Resume a Progressive Hedging (PH) run — serial, multi-rank, or full cylinders + (hub + spokes) — after a planned stop (and, as a bonus, after a crash). +- **Continue the optimization as if it had not stopped**, warm-started from the + last iterate, without losing the **best feasible solution found so far (the best + xhat), not just the best bound.** +- Survive a hard kill (`kill -9`, node failure, walltime) as a secondary benefit, + via atomic publication (§9). +- Add no measurable overhead when checkpointing is off, and — because checkpoints + are infrequent here — tolerate a heavier per-checkpoint cost in exchange for a + complete, warm-startable restore. + +**Non-goals (initially)** + +- Resuming across a *different* rank count or scenario-to-rank distribution + (cross-geometry remap). The first cut requires identical geometry and refuses a + mismatch with a clear error. +- **Bit-identical reproduction for MIPs.** Multi-threaded MIP solves are not + deterministic and admit multiple optima, so a resumed MIP run is *not* + bit-reproducible against a hypothetical uninterrupted run. The guarantee is + correct, warm-started continuation with the incumbent preserved (§7). (A + deterministic LP/QP solve *can* be bit-identical under the leaf-rebuild backend + — that is what the PoC showed — but it is not the target here.) +- Bit-identical reproduction of *bounds* (§7 — bounds are async and not + reproducible; carried forward as best-so-far). +- Robustness across a library/env upgrade between stop and resume. The use case + resumes in the **same environment** the next day, so a dill-based checkpoint + (welded to the current Pyomo/mpi-sppy/model code) is acceptable. Cross-version + resume is out of scope. +- APH. A C++ APH is expected to replace the Python `opt/aph.py`; PH-family only. +- **Catching external OS signals** (SIGTERM/SIGUSR1 from a scheduler, Ctrl-C) to + trigger a checkpoint. The terminal checkpoint (§8) fires on the run's *own* + termination — including hitting `--time-limit` — which covers the planned-stop + use case. A scheduler that hard-kills mid-solve is covered only insofar as the + previous checkpoint is preserved by atomic publication (§9), not by catching the + signal and checkpointing in response. Where a known walltime is the hazard, + `--checkpoint-before-seconds` (§8) is the answer instead of a signal handler: + it writes at the last iteration boundary preceding a user-supplied deadline. + +--- + +## 2. What to serialize, and what not to + +There are two very different "just dill it" ideas, and they get opposite answers. + +### 2.1 Do NOT dill the opt/hub object graph + +`dill.dump()` of the whole opt/hub object and reload is not viable — the core +objects are built around **live, non-serializable OS/MPI/solver handles**, not +data: + +- **MPI communicators** — `SPBase.mpicomm`, `SPBase.comms` (`spbase.py`), and the + `fullcomm`/`strata_comm`/`cylinder_comm` on the spcomm + (`cylinders/spcommunicator.py`). Handles into a running MPI runtime; meaningless + once the process exits. +- **MPI RMA windows and `MPI.Alloc_mem` buffers** — `SPWindow` + (`cylinders/spwindow.py`) and the `FieldArray` send/receive buffers. Kernel + shared-memory regions. +- **Persistent solver handles** — `s._solver_plugin` for gurobi/cplex/xpress + persistent interfaces (`spopt.py`). A C handle + license session. + +A resumed run launches a **fresh process** (new MPI job) anyway, so these must be +**reconstructed via normal startup** regardless of how anything else is restored. +That is not negotiable and not something a checkpoint can carry. + +### 2.2 DO dill the scenario models (the recommended backend here) + +The narrower idea — dill each **scenario Pyomo model** — is not only viable, it is +the right backend for this use case. The repo already dills *clean* scenario +models for the scenario-pickle path (§4); the only extension is to dill them +**mid-run**. The non-serializable attribute mpi-sppy itself puts on a scenario +model is `s._solver_plugin`, which we drop before writing and rebuild with +`set_instance` on resume (a dance we already do — the reconstruct step needs it +regardless); `_mpisppy_data` is a `pyo.Block` on the model with **no** +back-references to the comms or opt object. + +**But the model can also be made undillable by the user's own modeling code,** +and that is not something checkpointing can strip. The known case: a Pyomo rule +written as a nested function that closes over the `Config` object (`cfg`). The +closure drags the `Config` into the model's serialization graph, and Pyomo's +`ConfigDict` does not survive dill — although it *does* survive stdlib pickle. +`examples/stoch_distr/stoch_distr.py` has exactly this shape, so its scenario +models cannot be dilled at all, wrapper or no wrapper (issue #828). This is not +specific to checkpointing: the existing `--pickle-scenarios-dir` path would fail +the same way on such a model. Two consequences for this design: + +- The dill-reload backend cannot promise to checkpoint an *arbitrary* model; it + requires a dill-serializable one. The implementation therefore **probes one + local scenario at setup** and fails immediately with an actionable message + rather than discovering the problem at a terminal checkpoint hours in. +- The workaround lives in the model (hoist the value out of the closure before + defining the rule), so the user-facing docs must say so. + +Dilling the mid-run model captures, in one shot and mutually consistent, +everything that lives *on* the scenario model: + +- the dual weights `W`, `rho`, `xbars` (on `s._mpisppy_model`); +- nonant values **and fixedness** (from variable-fixing/forcing extensions); +- **second-stage (recourse) variable values → a MIP warm start** (§5.2); +- the proximal-approximation `xsqvar` and its accumulated `xsqvar_cuts`, plus the + `ProxApproxManager` bookkeeping on `_mpisppy_data` (§5.3) — the linearized prox + is likely in use for large MIPs (a MIQP prox is often intractable), and dill + brings the cuts back for free instead of replaying them; +- model-attached extension state, e.g. `fixer`'s per-variable `conv_iter_count` + on `s._mpisppy_data`. + +The costs that argued against this backend elsewhere — **version fragility** (dill +serializes by class/closure reference) and **per-checkpoint overhead** (full +models are large) — are both moot for this use case: the resume is same-environment +and the checkpoints are infrequent, so a heavy write paid a handful of times over +three days is negligible. And it *avoids re-running an expensive `scenario_creator`* +on every resume, which for large models is itself a real saving. (Exception: +the ADMM paths, where the wrapper re-runs the creator at startup regardless — +§8.2, item 2.) + +The alternative — rebuild each model via `scenario_creator` and overlay the state +as leaf data (arrays/name→value maps) — is the **low-cost backend**, designed here +but **not currently planned** (§11 Phase 6): +its checkpoints are tiny and fast (a handful of `O(first-stage)` arrays, no model +structure) and version-robust (plain numbers, not pickled classes). That makes it +the right choice for small scenarios, cheap creators, or *frequent* kill-safety +checkpoints where a per-write model dump would hurt — the mirror image of this +use case. It is also what the PoC validated (§6, §11 Phase 6). The two backends +share the same framework and manifest; only the scenario-model restore step +differs. + +--- + +## 3. Approach: reconstruct the scaffolding, restore the state + +A checkpoint is **not** a snapshot of the object graph. On resume: + +1. **Reconstruct the scaffolding** via the normal startup path — comms, RMA + windows, and persistent solvers. This is exactly what a fresh run already does. +2. **Restore the scenario-model state** via the chosen backend: + - **dill-reload (recommended here):** load each rank's dilled mid-run scenario + models and swap them in ahead of solver creation, so `_solver_plugin` is + attached (`set_instance`) to the reloaded model rather than to one about to + be discarded; mark `solution_available` so the first solve warm-starts + (§5.2). Because the reloaded model already carries the spliced W/prox + objective and the prox cuts, the deferred objective attach must be disarmed + so it does not run again downstream (§9, item 2). + - **leaf-rebuild (alternative):** rebuild each model via `scenario_creator`, + then overlay W / rho / nonant values+fixedness / prox `cut_values` from the + checkpoint. +3. **Restore the non-model state** — the pieces that do *not* live on a scenario + model and so are never captured by dilling models: the global iteration + counter, hub bounds/incumbent objective, the spoke incumbent (best xhat), the + internal state of extension *objects*, and cursor/RNG. These are always small + leaf data (§5.4–5.6). + +--- + +## 4. Building blocks already in the repo + +- **Scenario-model pickling (dill):** `utils/pickle_bundle.py` + (`dill_pickle`/`dill_unpickle`) and `generic/scenario_io.py` pickle each + scenario Pyomo model *alone*, driven by `--pickle-scenarios-dir` / + `--unpickle-scenarios-dir` (with `iter0_before_pickle` baking an iter-0 solve + into the pickle). The dill-reload backend (§2.2) is exactly this path + **generalized from iter-0 to iter-k** — pickle the model as it stands at the + checkpoint, not only after iter 0. +- **Warm-start plumbing:** `spopt.py` already supports warm-starting subproblem + solves — the `warmstart_subproblems` option plus `WarmstartStatus.PRIOR_SOLUTION` + use a warm start when `s._mpisppy_data.solution_available` is set + (`warmstart_subproblems` in `spopt.py`). Restoring the model's variable values and setting + `solution_available=True` feeds the restored MIP solution straight into this + path — no new solver code. +- **W / xbar persistence:** `utils/w_utils/wxbarwriter.py` (writes in + `post_everything`) / `wxbarreader.py` (reads in `pre_iter0`) round-trip `W` and + `xbar` as CSV. A precedent for the leaf-rebuild backend; unnecessary under + dill-reload (W/xbar ride in the model). +- **Incumbent-to-disk (spokes):** `cylinders/spoke.py` + `_maybe_write_incumbent_on_improvement` + (`--incumbent-on-improvement-filename-prefix`) already writes the first-stage + solution on each improvement — the reference for serializing the incumbent. + +--- + +## 5. State inventory + +For each piece: is it **reconstructed** (rebuilt by startup, no save), **carried +in the dilled model** (dill-reload backend), or **restored as non-model leaf +data** (always)? And if restored, is it **carried forward** as a valid +best-so-far value or (LP/QP only) potentially bit-reproducible? + +### 5.1 Hub PH primal state — *in the dilled model* + +The hub's primal trajectory is pure synchronous PH, independent of the spokes +(lagrangian only contributes an outer bound; xhatshuffle only an incumbent). +Per local scenario, all of the following live **on the scenario model** and are +therefore captured by dilling it: + +| State | Where it lives | Note | +|---|---|---| +| `W[ndn_i]` (accumulated duals) | `s._mpisppy_model.W` | `Update_W` accumulates — not recomputable; must be preserved | +| nonant values | nonant vardata `_value` | drive `Compute_Xbar`; `xbar` itself need not be separately saved | +| nonant **fixedness** (+ fixed value) | nonant vardata `.fixed` | `fixer`/`slammer` leave nonants fixed; must survive (§5.5) | +| `rho[ndn_i]` | `s._mpisppy_model.rho` | rho-updaters mutate it | +| `xbars[ndn_i]` | `s._mpisppy_model.xbars` | consensus target | +| smoothing `z/p/beta` | `s._mpisppy_model` | only if `--smoothing` | + +Under **leaf-rebuild**, this set is instead gathered/restored explicitly — helpers +already exist: `_populate_W_cache`/`W_from_flat_list` (`phbase.py`), +`_save_nonants`/`_restore_nonants` (`spopt.py`, which already captures fixedness in +`fixedness_cache`). + +**Restore point: a resume branch inside `Iter0`, in place of the iter-0 solve.** +`Iter0` (`phbase.py`) already contains the structural precedent: the +`iter0_from_pickle` option replaces the iter-0 `solve_loop` with +`_iter0_use_pickled_solution()`. Resume is a third branch in the same method, +*instead of* the solve loop, that loads the checkpoint and proceeds to +`iterk_loop`. This matters for the target use case: restoring in a *post*-iter0 +extension hook (what the PoC did, to avoid core changes) would first solve every +fresh model with `W = 0` and then throw those solutions away — one discarded +full MIP solve per scenario per resume, plausibly hours. The design therefore +specifies the core branch, not the hook (§9, item 2): the PoC's +`post_iter0_after_sync` restore was a zero-core-change validation crutch and is +**not** the design. + +**The model swap goes *before* `_create_solvers()`, not after.** +`_create_solvers` (`spopt.py`) walks `local_scenarios` attaching a +`_solver_plugin` and, for a persistent solver, calling `set_instance_retry` — +which builds the whole model inside the solver. If the reload happened after it, +every scenario would pay `set_instance` **twice** per resume: once on the fresh +model that is about to be discarded, once on the reloaded one. At this design's +target scale that is precisely the cost §2.2 sets out to avoid. Reloading first +lets the existing `_create_solvers` call attach the solver to the right model, +once, with no strip-and-rebuild dance. Nothing between the top of `Iter0` and +`_create_solvers` depends on the fresh models' identity. + +On the resume branch, the rest of `Iter0` adjusts as follows: the feasibility +check reads the restored per-scenario feasibility flags; `trivial_bound` / +`best_bound_obj_val` are restored from the checkpoint's leaf data rather than +recomputed via `Ebound`; the spoke sync still runs (publishing the *restored* +W/xbar/nonants to the spokes — they start from checkpointed state +immediately); the `rho_setter` is skipped (rho rides in the reloaded model); the +converger is constructed as usual and extension `restore_state` hooks (§9, item +3) fire before `iterk_loop`. + +**The deferred objective attach must be disarmed, not "skipped".** Iteration-0 +deferral (`_deferred_ph_attach`) splices the W/prox terms into the objective at +the *end* of `Iter0` — after the resume branch has already run — and re-runs +`set_instance` on persistent solvers. On a reloaded model that would duplicate +the prox components and double the W terms, so the resume branch clears the flag +rather than relying on control flow to miss it. Note that the *other* attach, +`attach_Ws_and_prox`, runs earlier still, in `PH_Prep`, and therefore cannot be +skipped from a branch inside `Iter0` at all; it is harmless (it decorates fresh +models that the swap discards), which is why disarming the deferred attach is +the requirement and the earlier call is not (§9, item 2). + +### 5.2 Recourse variable values — *warm start, in the dilled model* + +The second-stage (recourse) variable values are the bulk of a large scenario and +are **not** part of the algorithmic primal state (only the nonants and params +are). Their value is as a **MIP warm start**: restoring them and setting +`s._mpisppy_data.solution_available = True` makes the first resumed subproblem +solve start from the last iterate's solution via the existing +`warmstart_subproblems` path (§4). For large MIPs this can save substantial +branch-and-bound time on the first solve after each resume. + +Because they ride in the dilled model, they cost nothing extra here. Under the +leaf-rebuild backend they would be an **optional** all-var snapshot (off by +default): a per-checkpoint O(scenario) cost that only pays off for MIP/simplex +warm starts and is pure overhead for barrier solves — a bad trade when +checkpoints are frequent, which is why it is opt-in there. + +**Caveat (both backends):** an xhat/incumbent evaluation fixes the first stage and +re-solves, leaving recourse vars in the *eval* state; `_restore_nonants` restores +only nonants. So the model must be checkpointed (dilled) at a point where its +recourse values reflect the true last subproblem solve, not a mid-eval state — +i.e. snapshot before an eval corrupts them, or evaluate on a copy (§9, item 4). + +### 5.3 Proximal-approximation cuts (`--linearize-proximal-terms`) — *in the dilled model* + +When the linearized prox is on, `attach_PH_to_objective` builds, per scenario +(in `attach_PH_to_objective`, `phbase.py`): + +- `s._mpisppy_model.xsqvar` — the epigraph var for `x²`; +- `s._mpisppy_model.xsqvar_cuts` — a `Constraint` that accumulates one linear cut + per visited x-location (`add_cut` in `prox_approx.py`); +- `s._mpisppy_data.xsqvar_prox_approx[ndn_i]` — a `ProxApproxManager` whose + bookkeeping (`cut_index`, the sorted `cut_values` array; `ProxApproxManager` in `prox_approx.py`) + decides when a new cut is redundant. + +The cut *constraints* live on the model; the manager's bookkeeping lives on +`_mpisppy_data` (a Block on the model). **Dilling the model captures both, and +keeps them consistent** (the manager's references and the constraint set come back +in lockstep) — no replay, no re-binding. + +Under **leaf-rebuild**, neither survives (`attach_PH_to_objective` rebuilds +`xsqvar_cuts` empty). Each cut is fully determined by its x-location (continuous: +`xsqvar ≥ 2v·x − v²`; discrete: integer-keyed), so the checkpoint stores only the +per-nonant `cut_values` arrays and **replays** `add_cut` into the fresh model on +restore. Skipping this leaves resume correct (cuts regenerate lazily via +`check_tol_add_cut`) but coarser initially — fine for MIPs (not bit-reproducible +anyway), relevant only if an LP/QP run wants bit-identity. + +### 5.4 Hub bounds + incumbent, and the spoke incumbent — *non-model leaf data, carried forward* + +None of this lives on a hub scenario model, so it is restored as leaf data under +**both** backends: + +- `spcomm.BestInnerBound`, `spcomm.BestOuterBound`; `opt.best_bound_obj_val`, + `opt.best_solution_obj_val`. Products of **async** spoke interaction — their + timing is not reproducible, so they are carried forward as best-so-far. They + stay valid: a restored looser bound is improved again; a restored incumbent + objective is never regressed because `update_best_solution_if_improving` + (`spbase.py`) only accepts improvements. In cylinders the hub's + `best_solution_obj_val` is often `None` — the inner bound arrives as a scalar via + `receive_innerbounds` (`spcommunicator.py`) into `spcomm.BestInnerBound`. +- **The best xhat SOLUTION values live on the xhat spoke**, in + `spoke.opt.best_solution_cache` (a `ComponentMap` over all vars) + + `spoke.best_inner_bound`; `InnerBoundSpoke.finalize()` (`spoke.py`) loads + them back. So **"keep the best xhat" requires checkpointing the spoke + incumbent**, not just hub bounds. The spoke checkpoints its own cache **on its + own schedule** — on each improvement, reusing + `_maybe_write_incumbent_on_improvement`, independent of the hub checkpoint (§9, + item 6). Serialize the `ComponentMap` **by variable name** (`{var.name: value}`) + and rebuild by name lookup on the reconstructed model. +- **The initially-fixed-nonant baseline** (`opt._initial_fixed_varibles`), which + gates whether the outer bound may be updated at all. It lives on the opt object + and is keyed by variable *identity*, so it neither rides in the dill nor + survives the model swap; it is checkpointed **by variable name** and rebuilt + against the reloaded models. A plain PH hub happens to be insulated from + getting this wrong, but `Subgradient` and `FWPH` are not — see §9, item 11. + +### 5.5 Stateful extensions — *split: object state is leaf data, model state rides in the dill* + +Several extensions hold trajectory-driving state and **must** be restored or resume +diverges: + +- rho updaters (`mult_rho_updater`, `norm_rho_updater`, `grad_rho`), convergers — + multiplier / gradient / convergence history, kept on the **extension object**. +- variable-fixing/forcing extensions — `fixer.py` and `slammer.py` pin nonants and + then **skip what they already pinned**, so their tracking *is* the trajectory. + They span both storage locations: `slammer._slammed` is on the **extension + object**, while `fixer`'s per-variable `conv_iter_count` is on the **scenario + model** (`s._mpisppy_data`). + +Consequences: + +- **Model-attached tracker state (`fixer`) rides in the dilled model** for free — + consistent with the nonant fixedness it pairs with (§5.1). Under leaf-rebuild it + must be gathered explicitly. +- **Extension-object state is never on a model**, so it needs a serialization + contract regardless of backend. The `Extension` base has none today; add + `checkpoint_state()` / `restore_state()` (no-ops by default; implemented by rho + updaters, `fixer`, `slammer`, convergers), aggregated by the `Checkpointer` + (§9, item 3). The same contract serves hub and xhatter extensions + (`MultiExtension`). +- **The tracker and the actual variable state must agree.** Restoring "already + fixed X" without X's real `.fixed`/value (§5.1) makes the extension skip X while + the solver frees it — worse than no tracking. dill-reload gives this for free + (both come back together); leaf-rebuild must restore fixedness and the tracker as + one unit. + +### 5.6 RNG and spoke cursor — *non-model leaf data, partially restored* + +- xhatshuffle seeds its stream to a fixed `42` and samples **once** + (`main()` in `xhatshufflelooper_bounder.py`) — deterministic, no RNG state to save. +- The `ScenarioCycler` cursor and `xh_iter` are **local variables inside `main()`** + — unreachable. Exact spoke-cursor resume needs them hoisted onto `self` (Phase + 5). Without it the spoke restarts its cursor; this only changes *which* scenario + it tries next, not the preserved best (restored from §5.4). +- lagrangian / lagranger spokes use **no RNG**; their bound is deterministic given + the hub's `W`. State to carry: `_PHIter`, `trivial_bound`, last `bound`, received + `localWs`. + +### 5.7 Geometry / cfg fingerprint — *checkpoint metadata* + +Each per-rank file records `{n_proc, rank, local scenario list}` and a cfg hash. +Resume verifies the current layout matches and **refuses a mismatch with a clear +error** (validated — §6). + +--- + +## 6. PoC evidence (what is validated, and what is not) + +A throwaway PoC (serial + multi-rank + cylinders, farmer LP, gurobi_persistent) +validated the **framework and the leaf-rebuild backend**: + +- **Serial:** resume-from-iter-6 reproduced a full 12-iteration run with + `max|diff| = 0.000e+00` for W, nonants, rho (bit-identical — LP, deterministic + solver). Persistent solver survives the rebuild (Iter0 re-creates + + `set_instance`). +- **Multi-rank:** `-np 3` (1 scenario/rank) and uneven `-np 2` (2+1) resume + bit-identical on every rank; per-rank rank-tagged files, barrier + atomic + temp-then-rename write. Geometry mismatch fails with a clear error. +- **Cylinders (PH hub + lagrangian + xhatshuffle):** hub primal resumes + bit-identical inside `WheelSpinner`; the best xhat *solution* (on the spoke) is + preserved exactly; `BestInnerBound` carried exactly; `BestOuterBound` differed + run-to-run (async) but stayed valid. + +A second PoC then validated the **dill-reload backend on a MIP** (`sizes` SIZES3, +`gurobi_persistent`, single-thread `Threads=1`/`Seed=1`/`MIPGap=0` for a +deterministic solve — the §7 validation crutch): + +- **Mid-run model round-trip.** After a few PH iterations, a scenario model was + stripped of `_solver_plugin`, dilled, and reloaded **both in-process and in a + fresh process**; a new solver was attached with `set_instance` and the + subproblem re-solved. The reloaded model reproduced the original solve's + objective and **every decision variable exactly** — including the hardest case, + **linearized prox** (176 KB carrying **845 `xsqvar_cuts` + 65 + `ProxApproxManager`s** on `_mpisppy_data`), which came back structurally + identical and self-consistent. The only difference was the x² epigraph auxiliary + `xsqvar` wobbling ~1.5e-6 at solver feasibility tolerance (immaterial; MIQP was + exact). This is the load-bearing assumption — that a mid-run MIP model, cuts and + all, survives dill — and it **holds**. +- **Stop → reload → continue, bit-identical.** Stopping PH at iteration 3, dilling + the mid-run models, then rebuilding the scaffolding and continuing through the + reload branch reproduced an uninterrupted 6-iteration run with + `max|dW| = max|d nonant| = 0.0` — for **both** quadratic and linearized prox. + Under the deterministic single-thread solve this is exact bit-identity, the + strong "nothing was lost" check. + +Still to prove in later phases (this PoC was serial and focused on the model +round-trip + continuation): the dill-reload backend under **multi-rank** and +**cylinders**; carrying the **incumbent** across a dill-reload stop; a measured +warm-start speedup; the disk/time footprint at true model scale; and the +mid-run dill round-trip of a **stoch-ADMM wrapper-mutated model** (§8.2, item +4), which is structurally stranger than anything this PoC dilled. Note also +that both PoCs restored in the `post_iter0_after_sync` hook; the design +replaces that with the in-core resume branch (§5.1), which is itself unproven. + +--- + +## 7. Determinism contract (what resume guarantees) + +- **For the target MIP use case:** resume **continues the optimization correctly + and warm-started**, and **never loses or regresses the best xhat**. It is *not* + bit-reproducible — multi-threaded MIP solves are nondeterministic and admit + multiple optima, so the resumed iterates may differ from a hypothetical + uninterrupted run. That is expected, not a bug. +- **Bounds and incumbent:** valid and best-so-far, not bit-reproducible (async, + timing-dependent). Resume never reports a *worse* best-so-far than the + checkpoint. +- **Leaf-rebuild on a deterministic LP/QP solver:** the primal trajectory (W, + nonants, rho, xbar) *can* be bit-identical — this is what the PoC showed — but it + is a bonus, not the target guarantee. + +State this in user docs so a differing (but valid) trajectory or bound after +resuming a MIP is not mistaken for a bug. + +--- + +## 8. Configuration and semantics + +Checkpointing is **opt-in** and adds nothing when off. It is enabled by +`--checkpoint-dir`; with a directory set, the triggers below decide *when* a +checkpoint is written. They **compose** — whichever fires writes a checkpoint, all +sharing the same atomic publish (§9). With no `--checkpoint-dir` the `Checkpointer` +extension is not attached at all — zero overhead, no files. + +**Triggers** + +- **`--checkpoint-at-termination` (terminal checkpoint; default on).** Write one + complete, resumable checkpoint when the run terminates for *any* internal reason + — convergence, `--max-iterations`, cylinder convergence, or hitting + `--time-limit`. This is the primary trigger for the planned-stop use case: set + `--time-limit` to the daily budget and the run stops itself and checkpoints, + ready to resume the next morning. Nearly free — it fires in the hub's existing + `post_everything` hook (`phbase.py`), which runs once after the PH loop + regardless of *why* it exited, capturing the state of the last completed solve. + Turn it off (`--checkpoint-at-termination=False`) for a run that only wants + periodic insurance. It is *not* driven by external OS signals (a non-goal, §1); + the run's own termination — including `--time-limit` — is the trigger. +- **`--checkpoint-every-seconds S` (optional insurance).** Also checkpoint roughly + every `S` wall-clock seconds, for crash coverage during a long run. Checked at + each `enditer` as `allreduce_or(now − last_checkpoint ≥ S)` — the same + collective-decision pattern the existing `--time-limit` termination uses + (`phbase.py`), so every rank agrees and none writes at the barrier while others + sail past (a deadlock). Because it is tested only at iteration boundaries, the + guarantee is "the first boundary at least `S` seconds after the previous + checkpoint"; for large MIPs one solve can exceed `S`, which is expected — a + checkpoint cannot be taken mid-solve. Distinct from `--time-limit`, which *stops* + the run: this keeps it running and snapshots. +- **`--checkpoint-before-seconds S` (one-shot, anticipated).** Write **one** + checkpoint at the last iteration boundary that precedes `S` wall-clock seconds, + and keep running. Where `--checkpoint-every-seconds` is *reactive* — it fires at + the first boundary at or after its interval, so the write lands up to a full + iteration **late** — this trigger is *anticipatory*: at each `enditer` it asks + whether another boundary will arrive before `S`, estimating the next iteration + by the duration of the most recently completed one, and writes now if the answer + is no: + + ``` + allreduce_or(elapsed + last_iteration_seconds >= S) + ``` + + Same collective-decision pattern as the `time_limit` check in `phbase.py`, so + every rank decides together and none writes at the barrier while others sail + past. The motivating case is a scheduler walltime: the run will be hard-killed at + a known wall-clock time, and the user wants the checkpoint as late as possible + while still landing *before* the axe. This trigger **never terminates + anything**; the run continues and is killed, converges, or hits `--time-limit` + on its own. + + **One-shot.** Once the checkpoint is written the trigger is disarmed for the rest + of the run — it does not re-arm at `2S`. If the run outlives `S`, any later + checkpoints come from the other triggers, which are OR'd with this one as usual + (two triggers firing at the same `enditer` produce one write, §9 item 7). + + **No implicit safety margin.** `S` is used exactly as given: mpi-sppy applies no + fudge factor to the iteration estimate and does not attempt to predict how long + the checkpoint write itself will take. Both effects run the same direction and + matter to the user's choice of `S`: PH iteration times *sometimes grow* + (prox-approx cuts accumulate, `W` tightens, subproblem MIPs get harder), so the + last iteration can be an optimistic estimate of the next; and the write — + dilling large MIP models under the `dill-reload` backend — begins at the decision + point and is not free. The user supplies an `S` that already discounts for both. + To calibrate: run once + with checkpointing enabled and read the write duration off the bracketing `toc` + lines (§9 item 10), then set `S` to the walltime minus that duration minus the + user's own margin. **This must be prominent in the user docs** — an `S` chosen as + if it were the raw walltime is the one way this option quietly fails to do its + job. + + **What the clock measures.** `elapsed` is `time.perf_counter() − self.start_time`, + and `start_time` is stamped in `SPBase.__init__` (`spbase.py`) — *not* at job + submission and not at process launch. Everything before the SP object is + constructed is outside `S`: queue wait, interpreter startup, imports (over a + shared filesystem on a cluster, not always fast), MPI initialization at high rank + counts, and module/`cfg` setup. Scenario construction *is* inside it + (`_create_scenarios` is called later in the same `__init__`), as is the dill reload + on a resumed run (it happens in `Iter0`, §5.1). So `S` is measured on mpi-sppy's + clock, which starts *after* the scheduler's; a user aiming at a walltime must + subtract that startup gap as well — the same way they subtract the write cost. + The `toc` line mpi-sppy already emits at import (`global_toc("Initializing + mpi-sppy")`) versus the first checkpoint-related `toc` gives a usable read on how + large the gap is for a given cluster. Cylinders each construct their own + `SPBase`, so their `start_time`s differ slightly; this trigger is decided over + the hub's ranks alone, where the skew is a construction-time difference, not a + concern. + + **First arming.** At the `enditer` of iteration 1 the most recently *completed* + iteration is iteration 0, so its duration is the seed; no special case is needed + provided PH records iteration 0's duration alongside the later ones (§9 item 9). +- **`--checkpoint-every-iterations K` (optional insurance).** Also checkpoint every + `K` PH iterations. Checked at `enditer`. (This is the former + `--checkpoint-every k`, renamed for symmetry with + `--checkpoint-every-seconds`.) There is no special iteration-0 checkpoint: + resume is an in-core branch that restores any checkpointed iteration directly + (§5.1), so iteration 0 is not a privileged baseline. Given infrequent planned + stops, most runs leave the periodic triggers off and rely on the terminal + checkpoint alone. + +**Other options** + +- **`--checkpoint-dir `** — where per-rank files and the manifest are written + (§10); its presence is what enables checkpointing. +- **`--checkpoint-backend {dill-reload, leaf}`** — how scenario-model state is + restored (§2.2). `dill-reload` is the default and the only backend implemented in + the planned phases (captures the warm start + cuts, dodges an expensive + `scenario_creator` re-run). `leaf` — the **low-cost** option (tiny, fast, + version-robust checkpoints, for small/cheap-creator runs or frequent kill-safety + writes) — is designed but **not currently planned** (§11 Phase 6); until that + phase lands, `dill-reload` is the only valid value. +- **`--resume-from `** (or `--resume`, auto-selecting the latest *complete* + checkpoint from the manifest) — reconstruct the wheel and restore. Resume + requires identical geometry (§5.7); a mismatch is refused with a clear error. + +Each checkpoint is published atomically (§9, item 7; §10), so a kill *during* a +write leaves the previous complete checkpoint intact and referenced — never a +half-written one. + +### 8.1 Bundles + +mpi-sppy has **only proper bundles** now — loose bundling was removed in 2026 +(`spbase.py`; `doc/src/properbundles.rst`). A proper bundle is a **first-class +subproblem**: it appears in `local_scenarios` with its own `nonant_indices`, and +is itself a Pyomo model. So checkpointing **applies uniformly** — dilling +`local_scenarios` dills bundles exactly as it dills plain scenarios, and the +leaf-rebuild path iterates `nonant_indices` identically. Holds whether bundles are +in memory (`--scenarios-per-bundle`) or pickled (`--pickle-bundles-dir` / +`--unpickle-bundles-dir`). + +One cleanup: `_restore_nonants` still carries a 2019 comment that it "will not work +on bundles" (`spopt.py`). That predates proper bundles and refers to the removed +loose mechanism; re-verify and refresh it when bundle checkpointing is validated +(Phase 2). + +### 8.2 ADMM (deterministic and stochastic) + +`--admm` / `--stoch-admm` runs (`generic/admm.py`, `utils/admmWrapper.py`, +`utils/stoch_admmWrapper.py`) are plain PH hubs over *wrapped* scenarios, so +the checkpoint machinery applies in principle — but the wrapper path breaks +several assumptions made elsewhere in this design. Each of the following must +be honored or checkpointing will not work for ADMM: + +1. **Scenario naming and file discovery.** The existing pickle paths that §4 + builds on are *hard-refused* for ADMM (`_check_admm_compatibility`, + `generic/admm.py`) because `scenario_io.py` derives file names from + `module.scenario_names_creator` and `sputils.extract_num` — wrapped names + (`ADMM_STOCH__ADMM____ADMM__`) come from the wrapper, not the + module, and `extract_num` scrapes trailing digits, colliding across ADMM + subproblems that share a stochastic scenario. The checkpoint code must + therefore (a) enumerate `opt.local_scenarios.keys()` — never a module name + creator — for both write and restore, (b) never use `extract_num` in file + names (§10), and (c) not be swept into the ADMM incompatibility checks the + way the pickle flags were: checkpointing is *supposed* to work here, and a + test should pin that it does. +2. **The creator-cost saving does not apply, and a naive resume doubles model + memory.** §2.2 counts "avoids re-running an expensive `scenario_creator`" + as a dill-reload benefit. Not for ADMM: `Stoch_AdmmWrapper.__init__` runs + the user's `scenario_creator` for every local wrapped scenario (plus probe + scenarios) during normal startup — it needs the built models to assemble + consensus lists, `varprob_dict`, node names, and objective scaling — so an + ADMM resume pays the full creator cost regardless. Worse, after the reload + branch swaps the dilled models into `local_scenarios`, the fresh models + remain referenced by `wrapper.local_admm_stoch_subproblem_scenarios` and by + the `cfg._admm_variable_probability` bound method — a *persistent* 2× + per-rank model footprint for large MIPs. The reload branch must release or + replace the wrapper-held fresh models (§9, item 2). +3. **`variable_probability` is object-identity-keyed.** The wrapper's + `varprob_dict` maps scenario *object* → `(id(var), prob)` pairs + (`stoch_admmWrapper.py`; `AdmmBundler._bundle_varprob` likewise). This is + safe today only because `_use_variable_probability_setter` runs exactly + once, in `SPBase.__init__`, against the wrapper's own model objects, and + its results land on the model itself (`s._mpisppy_data.prob_coeff` / + `prob0_mask`) — which the dilled model carries back. The reload branch + depends on that invariant: **variable probabilities are consumed only at + construction; after the swap, the reloaded model's `_mpisppy_data` masks + and fixed-at-0 dummy vars are authoritative, and `var_prob_list` must never + be called with a reloaded model** (it would `KeyError` — or silently + mismatch if the dict were rebuilt with new ids). mpi-sppy masks `W` (not + prox) for zero-probability nonants and assumes each surrogate/dummy var is + fixed at 0; dill-reload preserves the mask and the fixedness together, and + the ADMM resume test must assert both survive. (A leaf-rebuild ADMM resume + would have to re-apply the mask and re-fix the dummies explicitly — one + more reason that backend is deferred, §11 Phase 6.) +4. **The dill round-trip is unvalidated for a wrapper-mutated model, and the + intended vehicle cannot currently test it.** `stoch_distr` scenario models + do not dill *at all* — not because of the wrapper, but because the model + defines a Pyomo rule closing over `cfg` (§2.2, issue #828). A bare + `scenario_creator` result fails identically, so the wrapper is exonerated + and simultaneously untested. Validating this item needs either a fix to + `stoch_distr` or a different stoch-ADMM model, and that is now a + prerequisite of Phase 2 rather than a step within it. The rest of this item + describes what still has to be proven once a vehicle exists. The MIP + PoC (§6) dilled a plain `sizes` model. A stoch-ADMM scenario is stranger: + inline dummy `pyo.Var()`s added post-construction with bracket-mangled + names, rewritten `ScenarioNode`s carrying *unattached* + `pyo.Expression(expr=0)` cost expressions and `surrogate_vardatas` sets of + vardata references, a rescaled objective, an appended ADMM stage + (a multistage tree even for a 2-stage-origin problem), and + probability-mask arrays on `_mpisppy_data`. dill should handle the cycles, + but this is a load-bearing assumption of the same kind §6 insisted on + PoC-ing — validate a stoch-ADMM mid-run round-trip early + (`mpisppy/tests/examples/stoch_distr` is the vehicle, + `test_stoch_admmWrapper.py` the harness; §11 Phases 2 and 4). +5. **Bundled stoch-ADMM.** `--stoch-admm --scenarios-per-bundle` + (`AdmmBundler`) creates bundles on the fly as EFs; they are first-class + subproblems and should dill like other proper bundles (§8.1). Its + `var_prob_list` has the same identity keying as item 3. +6. **The spoke set differs.** For stoch-ADMM cylinders: FWPH is refused, + `xhatshuffle` requires `--stage2-ef-solver-name`, and `xhatxbar` is the + variable-probability-native inner bounder. The Phase 4 test matrix must + include a stoch-ADMM configuration (§11). + +--- + +## 9. Core changes required + +Touch-points an implementation needs beyond the PoC's extension/subclass hacks: + +1. **Global iteration counter / resume offset.** `iterk_loop` hardcodes + `for _PHIter in range(1, max+1)` (in `iterk_loop`, `phbase.py`), so a resumed run renumbers + from 1 and its checkpoints collide with the pre-crash ones. Add a resume offset + so checkpoint numbering is the global iteration and termination honors the + original `max_iterations`. +2. **A reload-model resume branch, in `Iter0`, replacing the iter-0 solve.** + The branch lives where `iter0_from_pickle` already branches (§5.1), instead + of the iter-0 `solve_loop` — so a resume never pays a throwaway `W = 0` solve + of the fresh models — and the **model swap itself happens before + `_create_solvers()`**, so the existing solver creation attaches a + `_solver_plugin` (and calls `set_instance` once) to the reloaded model rather + than to a model about to be discarded (§5.1). Set `solution_available` for + the warm start (§5.2). + + Two attaches must be handled, and they are not symmetric. The deferred + objective attach (`_attach_PH_to_objective_after_iter0`, driven by + `_deferred_ph_attach`) runs at the **end** of `Iter0`, downstream of the + branch, and on a reloaded model would duplicate the prox components and + double the W terms — the branch **clears the flag** rather than relying on + control flow to miss it. `attach_Ws_and_prox`, by contrast, runs in + `PH_Prep`, *upstream* of `Iter0` entirely, so it cannot be skipped from this + branch; it is also harmless, since it decorates the fresh models that the + swap discards. + + Details the PoC and the ADMM analysis surfaced: **refresh + `saved_objectives[sname]`** for each reloaded model — `Eobjective` reads + those objective handles (populated by `_save_active_objectives` in + `SPOpt.__init__`) and they otherwise dangle to the discarded fresh model; + swap the reloaded model into `local_scenarios` (which `SPOpt.solve_loop` + iterates); where the `local_subproblems` alias exists, refresh it too — a + plain PH keeps no `local_subproblems`, but the **generic file-based path** + the dill-reload backend builds on (`scenario_io.py` sets + `sp.local_subproblems = sp.local_scenarios`) maintains it, and + `CGBase.solve_loop` iterates it; and on ADMM runs, **release or replace the + fresh models held by the wrapper** (`local_admm_stoch_subproblem_scenarios` + and the `cfg._admm_variable_probability` closure), or the run keeps two + copies of every local scenario alive for its whole life (§8.2, item 2). See + also item 11 for the one piece of *opt-object* state that is keyed by + variable identity and so cannot survive the swap untouched. This is a + distinct branch from the leaf-rebuild "build fresh, overlay values" path; + the `Checkpointer` picks the branch from `--checkpoint-backend`. +3. **Extension `checkpoint_state` / `restore_state` contract** on `Extension` + (no-ops by default; implemented by rho updaters, `fixer`, `slammer`, + convergers). Covers **extension-object** state under both backends; + model-attached state (`fixer`'s `conv_iter_count`) rides in the dill under + dill-reload but must be gathered explicitly under leaf-rebuild (§5.5). The + `Checkpointer` aggregates the dicts into the per-rank file. +4. **Clean-point model snapshot (xhat/incumbent eval).** Evaluating an xhat fixes + the first stage and re-solves, corrupting recourse vars (§5.2). The model must + be dilled (or its values gathered) when recourse values reflect the true last + solve — snapshot before an eval, or evaluate on a copy. +5. **Geometry / cfg fingerprint** (§5.7) with a clear refusal on mismatch. +6. **Async per-spoke incumbent checkpoints — no hub↔spoke coordination.** Each + spoke serializes its *own* best incumbent (the best xhat solution values, §5.4) + and bound whenever its incumbent improves — reusing + `_maybe_write_incumbent_on_improvement` — to its own rank-tagged file with the + same atomic write (item 7). Spokes are **not** synchronized to the hub's + checkpoint iteration: the determinism contract (§7) makes bounds/incumbent + best-so-far, not bit-reproducible, so a globally-consistent "snapshot at + iteration `k`" across cylinders is unnecessary. On resume the hub restores its + primal state while each spoke reloads its latest incumbent/bound, all accepted + only if improving (`update_best_solution_if_improving` in `spbase.py`). This + also avoids a hub-triggered snapshot barrier and its stall/deadlock risk. +7. **Atomic writes with a single published generation.** Each rank writes only its + local state (dilled models + leaf non-model data) to rank-tagged temp files and + renames them into place; the set of per-rank files is then published as one + checkpoint by atomically rewriting `manifest.json` (itself temp-then-rename) to + point at the new complete generation (§10). That flip is the single commit + point, so **one committed generation is enough**: a kill before it keeps the + previous checkpoint, a kill after it keeps the new one. The prior generation is + deleted once the manifest is in place. Retaining more than one checkpoint is + **not supported**: exactly one committed generation exists at any time (plus + the in-progress one transiently during a publish). +8. **A `Checkpointer` extension** that writes on its active triggers; restore + itself is the in-core resume branch (item 2), with extension + `restore_state` hooks (item 3) fired from it before `iterk_loop`: + - *periodic* (`--checkpoint-every-iterations` / `--checkpoint-every-seconds`) — + at `enditer`. The seconds trigger tests + `allreduce_or(now − last_checkpoint ≥ S)` so all ranks decide together + (mirroring the `time_limit` check in `phbase.py`), avoiding a rank-skew + deadlock at the write barrier. + - *anticipated one-shot* (`--checkpoint-before-seconds`) — also at `enditer`, + testing `allreduce_or(elapsed + last_iteration_seconds ≥ S)` with the same + collective pattern, then latching so it fires at most once (§8). It needs the + most-recent iteration duration (item 9); everything else it shares with the + periodic path. + - *terminal* (`--checkpoint-at-termination`, default on) — in `post_everything` + (`phbase.py`), which fires once after the PH loop however it exited + (convergence, `--max-iterations`, `--time-limit`), capturing the last + completed solve. Caveat: `post_everything` runs *after* `scenario_denouement`; + standard denouements only report, but one that re-solves or mutates a model + would be captured — call this out in user docs. + + For spokes, the xhatter `main()` loop calls no per-iteration extension hook — + add a single `self.opt.extobject.enditer()` (or a dedicated checkpoint hook) + inside it so **one `Checkpointer` serves hub and xhatter uniformly** (restore + already has a home: `pre_iter0`/`post_iter0` fire once in `xhat_prep` in + `xhatbase.py`). +9. **Most-recent iteration duration kept on `self`.** `iterk_loop` (`phbase.py`) + times each iteration into a *local* `iteration_start_time`, used only by the + `display_progress` print. `--checkpoint-before-seconds` needs that duration at + `enditer`, so record it on the object (e.g. `self._last_iteration_seconds`) as + each iteration completes — and record iteration 0's duration the same way, since + it is the seed the first time the trigger is tested (§8). Nothing else in PH + changes: no new hook, no change to the loop's control flow. +10. **`toc` on both ends of every checkpoint write.** The `Checkpointer` emits a + `global_toc` when a write begins and another when it completes — on every + trigger, hub and spokes alike, gated on `cylinder_rank == 0` so a multi-rank + cylinder prints one pair rather than one per rank. Because `tt_timer.toc` + stamps absolute elapsed time, the pair *is* the measured write duration, which + is what a user needs to choose `S` for `--checkpoint-before-seconds` — mpi-sppy + deliberately does not estimate that cost for them (§8). It also makes an + otherwise invisible multi-minute stall in a long run legible in the log. +11. **Restore the initially-fixed-nonant baseline, by name.** `SPOpt.__init__` + builds `_initial_fixed_varibles`, a `ComponentSet` of the nonant *vardata + objects* that were already fixed when the run started (`spopt.py`), and + `_can_update_best_bound` refuses to update `best_bound_obj_val` whenever a + nonant is fixed that is not in that set — because fixing a nonant mid-run + invalidates the outer bound. This is **opt-object state keyed by variable + identity**, so the model swap (item 2) breaks it in both directions and + neither is acceptable: + + - *Left alone*, the set holds vardata from the discarded fresh models. Every + reloaded nonant is a different object, so a fixed nonant reads as + unrecognized and the gate refuses to update the bound. + - *Naively rebuilt after the swap*, it absorbs whatever `fixer`/`slammer` + pinned before the stop, so those mid-run fixings look original and the + gate admits a bound the uninterrupted run would have refused. + + **How much this bites depends on the hub, and for a plain PH hub it is + currently masked.** `PHBase._can_update_best_bound` first returns `False` + whenever the proximal term is enabled, so PH consults the fixedness check + only with prox off — which happens at exactly one place, the iteration-0 + trivial bound, and that is the path the resume branch replaces anyway. The + hubs that consult the fixedness gate on their own terms are `Subgradient` + (which calls the `SPOpt` version directly, every iteration, bypassing the + prox short-circuit) and `FWPH` (its own override over the same set). Those + are where a stale baseline would actually change results. + + The checkpoint therefore records the originally-fixed nonants **by variable + name** (the same by-name discipline §5.4 uses for the incumbent cache), and + the resume branch rebuilds `_initial_fixed_varibles` from those names against + the reloaded models. This is the same identity-keying hazard §8.2 item 3 + catches for ADMM's `varprob_dict`. It belongs in the first phase not because + plain PH is currently broken by it, but because it is the correct restore of + opt-object state, it costs nothing, and it is load-bearing the moment resume + covers a hub that consults the gate per iteration — leaving a knowingly stale + ComponentSet behind for a later phase to trip over is the worse trade. + +--- + +## 10. File layout (proposed) + +``` +/ + manifest.json # cfg hash, n_proc, backend, cylinder map, latest complete hub generation + hub/ + gen_/ # NNNN = global PH iteration at the checkpoint + hub_rank_.pkl # non-model leaf state: iter counter, bounds, extension-object state + hub_rank__scen_.dill # dilled scenario model(s) for this rank (dill-reload backend) + spokes/ + spoke__rank_.pkl # each spoke's latest incumbent (best xhat, by name) + bound, + # overwritten asynchronously on improvement (§9, item 6) +``` + +The hub writes each checkpoint as an iteration-tagged generation under `hub/`, +deleting the prior one after the manifest flip (§9, item 7); each spoke keeps a +single latest-wins file under `spokes/` that it overwrites atomically on +improvement — the two are deliberately *not* aligned (§9, item 6). +`manifest.json` is the single commit point: it names the latest *complete* hub +generation and records the backend so resume loads the right way. Under the `leaf` +backend the `.dill` model files are replaced by numeric arrays inside the +`hub_rank_*.pkl`. Use plain `pickle` for the numeric/leaf state; `dill` for the +scenario models. + +`` is the scenario's full name sanitized for the filesystem (or its index in +the rank's local scenario list) — **never** `sputils.extract_num`, which is not +unique for ADMM wrapped names (§8.2, item 1). More generally, both write and +restore enumerate `opt.local_scenarios.keys()`, not a module name creator. + +**Disk footprint.** Dilled large MIP models × scenarios/rank can be large. The +single-generation policy (§9, item 7) keeps exactly one checkpoint live, but +the peak is **two generations transiently during a publish** (the new one is +fully written before the manifest flip deletes the old one) — state the peak in +user docs so disk quotas are sized for it. + +--- + +## 11. Phased rollout + +Each phase is a review-sized PR that is green on its own and adds user-visible +value. New tests are wired into `run_coverage.bash` **and** +`test_pr_and_main.yml` in the same commit. + +### 11.1 The A/B resume harness (every phase's acceptance test) + +The core CI test shape, reused by every phase, is an **A/B comparison**: + +- **Run A (reference):** an uninterrupted run of `N` iterations on a small + instance. +- **Run B (checkpointed):** the same instance stopped at iteration `k < N` + with a checkpoint written, then resumed in a **fresh process** and run to + `N`. +- **Compare A and B** under the §7 determinism contract: + - **Deterministic LP instances** (farmer; farmer + CVaR): `W`, nonants, + `rho`, `xbar` at each common iteration and the final objective must be + **bit-identical** (`max|diff| == 0.0`), and final bounds equal. + - **MIP instances** (`sizes`): with single-thread deterministic solver + settings (`Threads=1`, fixed seed, `MIPGap=0` — the §7 validation crutch) + the same bit-identity check applies; under default settings assert instead + that the run **continues** (global iteration numbering, no re-attach / + duplicate-component errors), the **incumbent never regresses**, bounds + stay valid, and the final objective agrees within a stated tolerance. +- Also assert the negatives: run B performs **no iter-0 subproblem solve** on + resume (§5.1), and a geometry/cfg mismatch is refused with a clear error + (§5.7). + +Instances — all small enough for the pip-installed, size-limited CPLEX/Xpress +CI solvers: + +- **farmer** — deterministic-LP baseline, serial and cylinders. +- **farmer + `--cvar`** (`utils/cvar.py`) — a mutate-after-creation transform: + the deactivated risk-neutral objective, the active `WITH_CVAR` objective, + and the eta var appended to the root nonants must all survive the dill + round-trip, and the resume branch's `saved_objectives` refresh (§9, item 2) + must resolve to `WITH_CVAR`, not the deactivated original. +- **stoch-distr (`--stoch-admm`)** — intended to exercise everything in §8.2: + wrapped names in file discovery, variable-probability masks and fixed-at-0 + dummy vars, the wrapper-mutated model dill round-trip, and release of the + wrapper-held fresh models. **Blocked**: `stoch_distr`'s models are not + dill-serializable (§2.2, issue #828), so this instance cannot be used until + the model is fixed or another stoch-ADMM model is chosen. +- **`sizes`** — the MIP target: warm start taken on resume, incumbent carried. + +The phase bullets below say where each instance enters (Phase 1a: serial +farmer; Phase 1b: farmer+CVaR and `sizes`; Phase 2: multi-rank, bundles, +stoch-ADMM; Phase 4: cylinders, including a stoch-ADMM configuration). + +Phase 1 is split into two review-sized PRs. Phase 1a is the whole serial +stop-and-resume story with the single trigger the primary use case actually +needs; Phase 1b adds the optional triggers and the harder test instances on top. +Each is green on its own, and 1a is independently useful — a run that stops at +`--time-limit` and resumes the next morning needs nothing from 1b. + +- **Phase 1a — Serial hub checkpoint/resume, terminal trigger only.** The + framework: `Checkpointer` extension; global iteration counter / resume offset + (§9 item 1); reload-model resume branch **in `Iter0`, replacing the iter-0 + solve**, with the swap ahead of `_create_solvers()`, the deferred objective + attach disarmed, `saved_objectives` refreshed, and the warm start set (§5.1, §9 + item 2); restore of the initially-fixed-nonant baseline by name (§9 item 11); + geometry+cfg fingerprint (§5.7); atomic per-rank writes + manifest publish (§9 + item 7); the terminal checkpoint (`--checkpoint-at-termination`, default on) + from `post_everything`; `toc` on both ends of every write (§9 item 10). CLI + flags `--checkpoint-dir`, `--checkpoint-at-termination`, + `--checkpoint-backend`, `--resume-from`/`--resume`, with a clear error when + `dill` is not installed (it is an optional `extras` dependency). Tests (the + §11.1 A/B harness, serial): **farmer** bit-identical A vs B; no iter-0 + subproblem solve occurs on resume; geometry/cfg mismatch refused. +- **Phase 1b — The optional triggers and the harder instances.** + `--checkpoint-every-iterations` / `--checkpoint-every-seconds` and the one-shot + anticipated `--checkpoint-before-seconds` (§8), all decided collectively at + `enditer`, with the most-recent iteration duration hoisted onto `self` (§9 item + 9). Tests: trigger semantics (fires once and latches; two triggers at one + `enditer` produce one write); **farmer + `--cvar`** bit-identical A vs B — the + mutate-after-creation case, where the resume branch's `saved_objectives` + refresh must resolve to `WITH_CVAR` and not the deactivated original; + **`sizes`** (MIP) — bit-identical under deterministic solver settings, and + under default settings run continues correctly, incumbent preserved, warm start + taken (mid-run model dill round-trip proven, §6). +- **Phase 2 — Multi-rank + bundles + stoch-ADMM.** Barriers, rank-tagged files, + single-generation atomic publish. Validate with **proper bundles** (§8.1) and + refresh the stale `_restore_nonants` comment. Validate **stoch-ADMM** (§8.2): + the wrapper-mutated model dill round-trip (item 4), file naming with wrapped + scenario names (item 1), wrapper-held fresh models released on resume (item + 2), and the probability mask + dummy-var fixedness surviving restore (item + 3) — `mpisppy/tests/examples/stoch_distr` is the vehicle. Tests (the §11.1 + A/B harness under `mpiexec`): MIP stop+resume compared on every rank, incl. + uneven distribution and `--scenarios-per-bundle`; a stoch-distr + (`--stoch-admm`) A/B stop+resume; mismatch refusal. +- **Phase 3 — Extension-object state contract.** `checkpoint_state`/`restore_state` + on `Extension`; implement for rho updaters, `fixer`, `slammer`, convergers. + (Model-attached `fixer` counter and nonant fixedness ride in the dill.) Test: PH + + norm-rho-updater, PH + `fixer`, PH + `slammer` each resume with state intact + and consistent with variable fixedness. +- **Phase 4 — Cylinders / spokes.** One-line xhatter write hook; unified + `Checkpointer` on spoke opts; each spoke checkpoints its own **best xhat** (by + name) asynchronously on improvement — no hub↔spoke coordination (§9, item 6). + Tests (the §11.1 A/B harness on cylinders): farmer/`sizes` + (hub+lagrangian+xhatshuffle) stop+resume — hub primal trajectory compared A + vs B (bit-identical for farmer, per the §6 PoC), best xhat preserved, bounds + valid best-so-far — **plus a stoch-ADMM cylinders configuration** (§8.2, + item 6: no FWPH; `xhatshuffle` with `--stage2-ef-solver-name`, or + `xhatxbar`). +- **Phase 5 — Exact spoke continuity (optional).** Hoist `ScenarioCycler`/`xh_iter` + onto `self`; checkpoint the cursor (+ RNG getstate if a stream becomes stateful). +- **Phase 6 — Leaf-rebuild backend + broader coverage (not currently planned).** + A possible future phase, deferred: the primary use case is fully served by the + dill-reload backend (Phases 1–4), so this is recorded for when a lighter, + version-robust checkpoint is actually needed rather than scheduled now. It would + add the `--checkpoint-backend leaf` path (rebuild via `scenario_creator`, overlay + W/rho/nonants/fixedness, replay prox `cut_values`, optional all-var warm start — + what the PoC prototyped) plus lagranger, FWPH, and subgradient spoke coverage. + The design deliberately keeps this backend's hooks and the shared + framework/manifest (§2.2) so it can be added later without disturbing the shipped + dill-reload path. + +--- + +## 12. Design decisions (resolved) and deferrals + +Resolved (given the §1 use case): + +- **Backend choice.** dill the scenario models (§2.2): overhead is negligible at a + few checkpoints, version robustness is unneeded (same-environment resume next + day), and it captures the warm start + prox cuts + model-attached state for free + while avoiding an expensive `scenario_creator` re-run (except on ADMM paths — + §8.2, item 2). +- **Warm start.** Worthwhile for MIPs (branch-and-bound benefits), free via the + dilled model, fed through the existing `warmstart_subproblems` / + `solution_available` path. +- **Restore point.** An in-core resume branch in `Iter0` replacing the iter-0 + solve (§5.1; §9, item 2) — no throwaway `W = 0` solve on resume, and + consequently no special iteration-0 checkpoint (§8). The PoCs' extension-hook + restore was a validation crutch, not the design. +- **Checkpoint retention** (§9, item 7): exactly one manifest-published + generation is kept; retaining older generations is not supported. The disk + peak is two generations transiently during a publish (§10) — a documented + cost, not an open question. +- **Spoke snapshot coordination** (§9, item 6): resolved by *not* coordinating. +- **variable_probability / surrogate vars (incl. ADMM).** Resolved by the §8.2 + contract: probabilities are consumed only at `SPBase` construction; after the + reload swap, the reloaded model's `_mpisppy_data` masks and fixed-at-0 + dummy/surrogate vars are authoritative, and `var_prob_list` is never called + with a reloaded model. Validation is scheduled (Phase 2). +- **Mid-run MIP model dill round-trip** — was the load-bearing unvalidated + assumption; **validated by the MIP dill-reload PoC** (§6), including the + linearized-prox cuts, in-process and cross-process, with serial stop→reload→ + continue bit-identical under a deterministic solver. The stoch-ADMM + wrapper-mutated variant of the same assumption is not yet validated — that is + a scheduled validation item (§6; §8.2, item 4; §11 Phase 2), not an open + design question. + +Deferred: + +- **Cross-geometry resume** (different rank count or scenario-to-rank + distribution) — a §1 non-goal; revisit if HPC users need to resume on a + different node count. +- **Leaf-rebuild backend** — designed (§2.2) but not scheduled (§11, Phase 6). diff --git a/doc/src/checkpointing.rst b/doc/src/checkpointing.rst new file mode 100644 index 000000000..3dc570318 --- /dev/null +++ b/doc/src/checkpointing.rst @@ -0,0 +1,127 @@ +.. _checkpointing: + +Checkpointing and Resuming a Run +================================ + +A long Progressive Hedging run can be stopped and picked up later. The intended +use is a planned stop: a multi-day study that ends each day and resumes the next +morning on the same cluster, without losing the work done so far. + +Checkpointing is entirely opt-in. With no ``--checkpoint-dir`` the machinery is +not attached at all, and a run that does not ask for it pays nothing. + +.. note:: + The current implementation covers a **serial PH hub**. Multi-rank runs, + bundles, and cylinder (hub-and-spoke) runs are planned but not yet + supported. See ``doc/designs/checkpointing_design.md`` for the full design + and the phased rollout. + +Writing a checkpoint +-------------------- + +Give a directory and the run writes one checkpoint when it terminates:: + + python -m mpisppy.generic_cylinders --module-name farmer --num-scens 3 \ + --solver-name cplex --max-iterations 100 --default-rho 1.0 \ + --time-limit 28800 \ + --checkpoint-dir ./ckpt + +``--checkpoint-at-termination`` is on by default, so the checkpoint is written +when the run ends for *any* internal reason: convergence, ``--max-iterations``, +or hitting ``--time-limit``. Pairing it with ``--time-limit`` is the planned-stop +recipe above -- set the day's budget and the run stops itself and checkpoints. +Turn it off with ``--disable-checkpoint-at-termination``. + +Each write is bracketed by a pair of timestamped ``toc`` lines, so the log shows +how long it took:: + + [ 1234.56] Writing checkpoint (termination) at iteration 42 to ./ckpt + [ 1261.03] Checkpoint written (termination) at iteration 42 + +Resuming +-------- + +Point a new run at the directory:: + + python -m mpisppy.generic_cylinders --module-name farmer --num-scens 3 \ + --solver-name cplex --max-iterations 100 --default-rho 1.0 \ + --resume-from ./ckpt + +The resumed run continues from the checkpointed iterate rather than starting +over. It does **not** re-solve the subproblems at iteration 0 -- for large MIPs +that solve is often the most expensive in the run, and its answer would be +thrown away. Iteration numbering continues where it left off, so +``--max-iterations`` bounds the run as a whole rather than each leg of it. + +What must match, and what may change +------------------------------------ + +A checkpoint records the layout it was written with and refuses to load into a +run that does not match, rather than producing a subtly wrong answer. Resuming +requires: + +* the same number of MPI ranks, and the same scenarios on each rank; +* the same **structural** options -- the ones that change the shape of the + scenario models or the meaning of the state stored in them: + ``--default-rho``, ``--linearize-proximal-terms``, + ``--linearize-binary-proximal-terms``, + ``--proximal-linearization-tolerance``, the ``--smoothing`` settings, the + ``--cvar`` settings, ``--module-name``, ``--num-scens``, + ``--branching-factors``, and ``--scenarios-per-bundle``. + +Everything else is free to change. In particular **the iteration limit, the +time limit, and the display/verbosity options may all differ** on a resume -- +picking a run back up with a different budget is the whole point, so those are +deliberately outside the check. + +A mismatch is reported with an explicit message naming what differs. + +What resume guarantees +---------------------- + +For the target case -- large MIP subproblems -- a resumed run **continues +correctly and warm-started**, and never loses or regresses the best solution +found so far. It is **not** bit-for-bit reproducible against a hypothetical +uninterrupted run: multi-threaded MIP solves are not deterministic and admit +multiple optima, so the resumed iterates may differ. That is expected, not a +bug. + +Bounds and the incumbent are carried forward as valid best-so-far values. A +resumed run never reports a worse best-so-far than its checkpoint. + +On a deterministic LP or QP solve the primal trajectory can come back +bit-identical, but that is a bonus rather than the guarantee. + +Disk usage +---------- + +Exactly one checkpoint is kept. Retaining older generations is not supported. +The new checkpoint is written in full before the old one is deleted, so the +**peak** on disk is two generations -- size disk quotas for that, not for one. +Under the default ``dill-reload`` backend a checkpoint holds a serialized copy +of every local scenario model, which for large MIPs is not small. + +Publication is atomic: the files are written, then a manifest is rewritten to +point at them. A run killed during a write leaves the previous complete +checkpoint intact and referenced, never a half-written one. + +Requirements and limitations +---------------------------- + +**dill is required.** ``--checkpoint-backend dill-reload`` is the default and +currently the only implemented backend, and it needs the optional ``dill`` +package:: + + pip install mpi-sppy[extras] + +**Your scenario models must be serializable.** A model can be made +unserializable by what its ``scenario_creator`` closes over -- most commonly a +Pyomo rule written as a nested function that reads ``cfg`` directly, which pulls +the whole configuration object into the model. See :ref:`scenario_creator` for +the pattern and the fix. Checkpointing checks one scenario at setup rather than +discovering the problem hours later at the terminal checkpoint, and the error +names the offending rule. + +**The terminal checkpoint runs after ``scenario_denouement``.** Standard +denouement functions only report, but one that re-solves or mutates a model +would have those changes captured in the checkpoint. diff --git a/doc/src/index.rst b/doc/src/index.rst index e49f7ff78..1a12a6b4f 100644 --- a/doc/src/index.rst +++ b/doc/src/index.rst @@ -58,6 +58,7 @@ MPI is used. properbundles.rst pickling.rst + checkpointing.rst jensens.rst vss.rst feasible_xhat.rst diff --git a/mpisppy/extensions/checkpointer.py b/mpisppy/extensions/checkpointer.py new file mode 100644 index 000000000..d0e76bbaa --- /dev/null +++ b/mpisppy/extensions/checkpointer.py @@ -0,0 +1,83 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""Write checkpoints so a run can be stopped and resumed later. + +Attached only when ``--checkpoint-dir`` is given, so a run that does not ask +for checkpointing pays nothing. This extension decides *when* to write; +``mpisppy/utils/checkpointing.py`` owns the on-disk format, and the resume +branch lives in ``PHBase.Iter0`` (restoring has to happen mid-startup, before +solvers are created). + +The trigger implemented here is the terminal checkpoint +(``--checkpoint-at-termination``, on by default): one complete, resumable +checkpoint written when the run ends for any internal reason -- convergence, +the iteration limit, or hitting ``--time-limit``. That is the planned-stop path +for a multi-day study: set ``--time-limit`` to the day's budget and the run +stops itself and checkpoints, ready to resume the next morning. + +See ``doc/designs/checkpointing_design.md``. +""" + +import os + +from mpisppy import global_toc +from mpisppy.extensions.extension import Extension +import mpisppy.utils.checkpointing as ckpt + + +class Checkpointer(Extension): + """Write a resumable checkpoint on the run's own termination.""" + + def __init__(self, opt): + super().__init__(opt) + options = opt.options + self.ckpt_dir = options.get("checkpoint_dir", None) + self.backend = options.get("checkpoint_backend", + ckpt.DILL_RELOAD_BACKEND) + self.at_termination = options.get("checkpoint_at_termination", True) + + if self.ckpt_dir is None: + raise RuntimeError( + "Checkpointer was attached without a checkpoint directory. " + "It should only be attached when --checkpoint-dir is set." + ) + # Fail at setup rather than after a multi-hour run reaches its first + # write and discovers the backend is unusable. + ckpt.require_dill(self.backend) + + def _write(self, why): + """Write one generation, bracketed by toc so the cost is legible. + + The pair of timestamps *is* the measured write duration, which is what + a user needs in order to choose a deadline for the anticipated + trigger -- mpi-sppy deliberately does not estimate that cost for them. + """ + rank0 = self.opt.cylinder_rank == 0 + generation = int(getattr(self.opt, "_PHIter", 0)) + os.makedirs(self.ckpt_dir, exist_ok=True) + global_toc(f"Writing checkpoint ({why}) at iteration {generation} " + f"to {self.ckpt_dir}", rank0) + ckpt.write_checkpoint(self.opt, self.ckpt_dir, generation, + backend=self.backend) + global_toc(f"Checkpoint written ({why}) at iteration {generation}", + rank0) + + def pre_iter0(self): + # Prove now that this run's models can actually be checkpointed. A run + # that only finds out at its terminal checkpoint would lose exactly the + # state it was trying to preserve. + ckpt.probe_model_is_dillable(self.opt) + + def post_everything(self): + if not self.at_termination: + return + # post_everything runs after scenario_denouement. Standard denouements + # only report, but one that re-solves or mutates a model would be + # captured in this checkpoint. + self._write("termination") diff --git a/mpisppy/generic/parsing.py b/mpisppy/generic/parsing.py index 1e9523066..487078334 100644 --- a/mpisppy/generic/parsing.py +++ b/mpisppy/generic/parsing.py @@ -117,6 +117,7 @@ def add_decomp_args(cfg): cfg.relaxed_ph_fixer_args() cfg.integer_relax_then_enforce_args() cfg.slamming_args() + cfg.checkpoint_args() cfg.w_oscillation_args() cfg.gapper_args() cfg.gapper_args(name="lagrangian") diff --git a/mpisppy/phbase.py b/mpisppy/phbase.py index 59a942120..d78148891 100644 --- a/mpisppy/phbase.py +++ b/mpisppy/phbase.py @@ -18,6 +18,7 @@ import mpisppy.utils.sputils as sputils import mpisppy.spopt +import mpisppy.utils.checkpointing as checkpointing from mpisppy.utils.prox_approx import ProxApproxManager from mpisppy.utils.rho_utils import check_rhos_positive @@ -251,6 +252,14 @@ class PHBase(mpisppy.spopt.SPOpt): Function to set variable specific probabilities. """ + + # Resume state, set by _restore_from_checkpoint_if_resuming when + # --resume-from is given. Class-level defaults so every other code path + # -- including runs with no checkpointing at all -- reads sane values. + _resumed_from_checkpoint = False + _resume_iteration = 0 + _checkpoint_leaf_state = None + def __init__( self, options, @@ -1151,6 +1160,54 @@ def _iter0_use_pickled_solution(self): s._mpisppy_data.outer_bound = md["iter0_outer_bound"] s._mpisppy_data.inner_bound = md["iter0_inner_bound"] + def _restore_from_checkpoint_if_resuming(self): + """Splice a checkpoint's scenario models into this run, if resuming. + + Called from ``Iter0`` before ``_create_solvers``. A no-op unless + ``--resume-from`` was given. See + ``doc/designs/checkpointing_design.md`` sections 5.1 and 9. + """ + ckpt_dir = self.options.get("resume_from", None) + if not ckpt_dir: + return + + leaf, models = checkpointing.load_checkpoint(self, ckpt_dir) + + for sname, model in models.items(): + self.local_scenarios[sname] = model + # Eobjective/Ebound read these objective handles; left alone they + # would dangle to the fresh model we just replaced. + self.saved_objectives[sname] = sputils.find_active_objective(model) + # The recourse values that came back with the model are a warm + # start for the first resumed solve. + model._mpisppy_data.solution_available = True + # The generic file-based path keeps this alias and CGBase.solve_loop + # iterates it; a plain PH has no such attribute. + if getattr(self, "local_subproblems", None) is not None: + self.local_subproblems = self.local_scenarios + + # _initial_fixed_varibles is a ComponentSet of vardata belonging to the + # models we just discarded, so every reloaded nonant would look + # unrecognized and _can_update_best_bound would refuse to update the + # bound for the rest of the run. Rebuilding it from the checkpointed + # *names* also keeps nonants that a fixing extension pinned mid-run + # from passing as originally fixed, which would let through a bound the + # uninterrupted run would have refused. + self._restore_fixed_nonant_baseline(leaf["initially_fixed_nonants"]) + + # The deferred attach runs at the end of Iter0, downstream of here. The + # reloaded model already carries the spliced objective and the prox + # cuts, so letting it run would double the W terms and duplicate the + # prox components. + self._deferred_ph_attach = False + + self._checkpoint_leaf_state = leaf + self._resumed_from_checkpoint = True + self._resume_iteration = int(leaf["generation"]) + global_toc(f"Resuming from checkpoint in {ckpt_dir} " + f"(iteration {self._resume_iteration})", + self.cylinder_rank == 0) + def Iter0(self): """ Create solvers and perform the initial PH solve (with no dual weights or prox terms). @@ -1185,6 +1242,14 @@ def _vb(msg): self._PHIter = 0 self._save_original_nonants() + # Resume, if asked: swap the checkpointed models in *before* solvers + # are created, so _create_solvers attaches a solver (and, for a + # persistent solver, calls set_instance) to the reloaded model rather + # than to one we are about to discard. Doing it the other way round + # makes every scenario pay set_instance twice per resume, which for + # large MIPs is exactly the cost checkpointing is meant to avoid. + self._restore_from_checkpoint_if_resuming() + global_toc("Creating solvers") self._create_solvers() @@ -1196,7 +1261,13 @@ def _vb(msg): and self.cylinder_rank == 0 ) - if self.options.get("iter0_from_pickle", False): + if self._resumed_from_checkpoint: + # The reloaded models already carry a solved iterate; re-solving + # them here with W = 0 would discard it. Bookkeeping that + # solve_loop would have set was restored with the models. + global_toc("Skipping PHBase.Iter0 solve loop (--resume-from); " + "continuing from the checkpointed iterate") + elif self.options.get("iter0_from_pickle", False): self._iter0_use_pickled_solution() else: if self.options["verbose"]: @@ -1233,9 +1304,18 @@ def _vb(msg): if have_extensions: self.extobject.post_iter0() - self.trivial_bound = self.Ebound(verbose) - if self.trivial_bound is not None and self._can_update_best_bound(): - self.best_bound_obj_val = self.trivial_bound + if self._resumed_from_checkpoint: + # The trivial bound belongs to iteration 0 of the original run; + # recomputing it here would use the checkpointed (W-laden) iterate + # and produce something that is not the trivial bound at all. + self.trivial_bound = self._checkpoint_leaf_state["trivial_bound"] + restored_bound = self._checkpoint_leaf_state["best_bound_obj_val"] + if restored_bound is not None: + self.best_bound_obj_val = restored_bound + else: + self.trivial_bound = self.Ebound(verbose) + if self.trivial_bound is not None and self._can_update_best_bound(): + self.best_bound_obj_val = self.trivial_bound if hasattr(self.spcomm, "sync_nonants"): self.spcomm.sync_nonants() @@ -1246,7 +1326,10 @@ def _vb(msg): if have_extensions: self.extobject.post_iter0_after_sync() - if self.rho_setter is not None: + # On a resume, rho rides in the reloaded model -- including whatever a + # rho updater had done to it by the time of the checkpoint -- so + # re-running the setter would silently undo that. + if self.rho_setter is not None and not self._resumed_from_checkpoint: if self.cylinder_rank == 0: self._use_rho_setter(verbose) else: @@ -1330,7 +1413,11 @@ def iterk_loop(self): global_toc("Cylinder convergence", self.cylinder_rank == 0) return - for self._PHIter in range(1, max_iterations+1): + # _PHIter is the *global* iteration number: on a resume it picks up + # where the checkpoint left off, so checkpoint generations do not + # collide with the pre-stop ones and max_iterations stays a limit on + # the run as a whole rather than on this leg of it. + for self._PHIter in range(self._resume_iteration + 1, max_iterations+1): iteration_start_time = time.time() if dprogress: diff --git a/mpisppy/spopt.py b/mpisppy/spopt.py index f81eb475a..531479583 100644 --- a/mpisppy/spopt.py +++ b/mpisppy/spopt.py @@ -1262,6 +1262,25 @@ def _create_fixed_nonant_cache(self): if v.fixed: self._initial_fixed_varibles.add(v) + def _restore_fixed_nonant_baseline(self, names): + """Rebuild `_initial_fixed_varibles` from checkpointed variable names. + + `_create_fixed_nonant_cache` records vardata *objects*, so the cache + does not survive a resume that replaces the scenario models: every + reloaded nonant is a different object, and `_can_update_best_bound` + would treat each one as newly fixed and refuse to update the bound for + the rest of the run. Rebuilding by name restores the original meaning + -- which nonants were already fixed when the run first started -- so a + nonant that a fixing extension pinned mid-run is still correctly + recognized as *not* part of the baseline. + """ + wanted = set(names) + self._initial_fixed_varibles = ComponentSet() + for s in self.local_scenarios.values(): + for v in s._mpisppy_data.nonant_indices.values(): + if v.name in wanted: + self._initial_fixed_varibles.add(v) + def _can_update_best_bound(self): for s in self.local_scenarios.values(): for v in s._mpisppy_data.nonant_indices.values(): diff --git a/mpisppy/tests/test_checkpoint.py b/mpisppy/tests/test_checkpoint.py new file mode 100644 index 000000000..4da80d00e --- /dev/null +++ b/mpisppy/tests/test_checkpoint.py @@ -0,0 +1,372 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""Tests for checkpoint/resume (doc/designs/checkpointing_design.md), phase 1a. + +The load-bearing test is the A/B harness: run A is an uninterrupted run of N +iterations; run B stops at k < N with a checkpoint, then resumes and continues +to N. On farmer -- a deterministic LP -- the two must agree bit-for-bit, which +is the strong "nothing was lost" check. + +The rest pin the things that would make a resume quietly wrong rather than +loudly broken: that resume does not solve the fresh models at iteration 0 +(which would throw away the checkpointed iterate and, for a large MIP, cost +hours), that a geometry or structural-option mismatch is refused instead of +producing nonsense, and that the initially-fixed-nonant baseline survives the +model swap -- without it a resumed run silently stops updating its best bound. +""" + +import os +import tempfile +import unittest + +import mpisppy.utils.checkpointing as checkpointing +import mpisppy.tests.examples.farmer as farmer +from mpisppy.extensions.checkpointer import Checkpointer +from mpisppy.opt.ph import PH +from mpisppy.tests.utils import get_solver +from mpisppy.utils.config import Config + +solver_available, solver_name, persistent_available, persistent_solver_name = \ + get_solver() + +SCENARIO_NAMES = ["scen0", "scen1", "scen2"] +CREATOR_KWARGS = {"use_integer": False, "crops_multiplier": 1} + + +def _options(max_iters, ckpt_dir=None, resume_from=None, **overrides): + options = { + "solver_name": solver_name, + "PHIterLimit": max_iters, + "defaultPHrho": 1.0, + # Never converge early: the A/B comparison needs a fixed iteration + # count on both sides. + "convthresh": -1.0, + "verbose": False, + "display_progress": False, + "display_timing": False, + "display_convergence_detail": False, + "iter0_solver_options": None, + "iterk_solver_options": None, + "tee-rank0-solves": False, + "smoothed": 0, + "time_limit": None, + } + if ckpt_dir is not None: + options["checkpoint_dir"] = ckpt_dir + options["checkpoint_at_termination"] = True + options["checkpoint_backend"] = checkpointing.DILL_RELOAD_BACKEND + if resume_from is not None: + options["resume_from"] = resume_from + options.update(overrides) + return options + + +def _make_ph(options, scenario_names=None): + extensions = Checkpointer if "checkpoint_dir" in options else None + return PH( + options, + scenario_names if scenario_names is not None else SCENARIO_NAMES, + farmer.scenario_creator, + farmer.scenario_denouement, + scenario_creator_kwargs=CREATOR_KWARGS, + extensions=extensions, + ) + + +def _primal_snapshot(ph): + """W, rho and nonant values for every local scenario, keyed by name.""" + snap = {} + for sname, s in ph.local_scenarios.items(): + for ndn_i, v in s._mpisppy_data.nonant_indices.items(): + snap[(sname, "x", v.name)] = v._value + snap[(sname, "W", str(ndn_i))] = \ + float(s._mpisppy_model.W[ndn_i]._value) + snap[(sname, "rho", str(ndn_i))] = \ + float(s._mpisppy_model.rho[ndn_i]._value) + return snap + + +@unittest.skipIf(not solver_available, + "no solver is available for the A/B resume harness") +class TestResumeABFarmer(unittest.TestCase): + """Uninterrupted vs stop-and-resume on a deterministic LP.""" + + N = 6 + STOP = 3 + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.ckpt_dir = os.path.join(self._tmp.name, "ckpt") + + def tearDown(self): + self._tmp.cleanup() + + def _run_b(self): + """Stop at STOP with a checkpoint, then resume and finish.""" + stopped = _make_ph(_options(self.STOP, ckpt_dir=self.ckpt_dir)) + stopped.ph_main() + resumed = _make_ph(_options(self.N, resume_from=self.ckpt_dir)) + resumed.ph_main() + return stopped, resumed + + def test_resume_is_bit_identical(self): + reference = _make_ph(_options(self.N)) + reference.ph_main() + _, resumed = self._run_b() + + want = _primal_snapshot(reference) + got = _primal_snapshot(resumed) + self.assertEqual(set(want), set(got)) + for key in want: + self.assertEqual( + want[key], got[key], + msg=f"{key} differs after resume: {want[key]} vs {got[key]}") + + def test_iteration_numbering_is_global(self): + """A resumed run continues the count instead of restarting at 1.""" + stopped, resumed = self._run_b() + self.assertEqual(stopped._PHIter, self.STOP) + self.assertTrue(resumed._resumed_from_checkpoint) + self.assertEqual(resumed._resume_iteration, self.STOP) + self.assertEqual(resumed._PHIter, self.N) + + def test_resume_performs_no_iter0_solve(self): + """The whole point of the in-core branch: no throwaway W = 0 solve. + + For a large MIP that solve is the most expensive in the run -- cold, + unregularized, no warm start -- and its answer is discarded. + """ + _make_ph(_options(self.STOP, ckpt_dir=self.ckpt_dir)).ph_main() + + resumed = _make_ph(_options(self.N, resume_from=self.ckpt_dir)) + calls = [] + original = resumed.solve_loop + + def counting_solve_loop(*args, **kwargs): + calls.append(resumed._PHIter) + return original(*args, **kwargs) + + resumed.solve_loop = counting_solve_loop + resumed.ph_main() + + self.assertNotIn( + 0, calls, + msg="resume solved the fresh models at iteration 0; the " + "checkpointed iterate would have been discarded") + self.assertEqual(calls, list(range(self.STOP + 1, self.N + 1))) + + def test_trivial_bound_is_restored_not_recomputed(self): + """The trivial bound belongs to iteration 0 of the original run.""" + reference = _make_ph(_options(self.N)) + reference.ph_main() + _, resumed = self._run_b() + self.assertEqual(reference.trivial_bound, resumed.trivial_bound) + + def test_writes_one_generation_and_a_manifest(self): + """Retention is exactly one published generation.""" + _make_ph(_options(self.STOP, ckpt_dir=self.ckpt_dir)).ph_main() + self.assertTrue( + os.path.exists(os.path.join(self.ckpt_dir, "manifest.json"))) + generations = os.listdir(os.path.join(self.ckpt_dir, "hub")) + self.assertEqual(generations, [f"gen_{self.STOP:04d}"]) + + +@unittest.skipIf(not solver_available, + "no solver is available to write a checkpoint to refuse") +class TestResumeRefusesMismatch(unittest.TestCase): + """A checkpoint that does not fit the current run must be refused.""" + + STOP = 2 + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.ckpt_dir = os.path.join(self._tmp.name, "ckpt") + _make_ph(_options(self.STOP, ckpt_dir=self.ckpt_dir)).ph_main() + + def tearDown(self): + self._tmp.cleanup() + + def test_structural_option_mismatch_is_refused(self): + """Changing rho changes the meaning of the state in the checkpoint.""" + options = _options(4, resume_from=self.ckpt_dir, defaultPHrho=2.0) + with self.assertRaises(checkpointing.CheckpointMismatch) as ctx: + _make_ph(options).ph_main() + self.assertIn("structural options", str(ctx.exception)) + + def test_scenario_distribution_mismatch_is_refused(self): + """Resuming with a different scenario set is refused, not guessed at.""" + options = _options(4, resume_from=self.ckpt_dir) + with self.assertRaises(checkpointing.CheckpointMismatch) as ctx: + _make_ph(options, scenario_names=["scen0", "scen1"]).ph_main() + self.assertIn("scenario", str(ctx.exception).lower()) + + def test_missing_manifest_is_refused_clearly(self): + options = _options(4, resume_from=os.path.join(self._tmp.name, "nope")) + with self.assertRaises(checkpointing.CheckpointMismatch) as ctx: + _make_ph(options).ph_main() + self.assertIn("manifest", str(ctx.exception)) + + def test_iteration_limit_may_change_on_resume(self): + """The limit and the clock are deliberately outside the fingerprint. + + Picking a run back up the next morning with a different budget is the + primary use case, so these must not be treated as a mismatch. + """ + options = _options(4, resume_from=self.ckpt_dir, time_limit=3600) + resumed = _make_ph(options) + resumed.ph_main() + self.assertEqual(resumed._PHIter, 4) + + +class TestStructuralFingerprint(unittest.TestCase): + """Which option changes block a resume, and which do not.""" + + def _fingerprint(self, **options): + base = {"defaultPHrho": 1.0, "linearize_proximal_terms": False} + base.update(options) + return checkpointing.structural_fingerprint(base) + + def test_identical_options_match(self): + self.assertEqual(self._fingerprint(), self._fingerprint()) + + def test_structural_change_is_detected(self): + self.assertNotEqual(self._fingerprint(), + self._fingerprint(defaultPHrho=2.0)) + self.assertNotEqual(self._fingerprint(), + self._fingerprint(linearize_proximal_terms=True)) + + def test_non_structural_change_is_ignored(self): + """A named subset, so harmless flags do not block a resume.""" + self.assertEqual( + self._fingerprint(), + self._fingerprint(PHIterLimit=999, time_limit=60, + display_progress=True, verbose=True, + solver_name="some_other_solver")) + + def test_structural_cfg_extras_are_covered(self): + """Settings PH never reads, but which reshape the model, still count.""" + with_cvar = self._fingerprint( + checkpoint_structural_cfg={"cvar": True, "cvar_alpha": 0.95}) + without = self._fingerprint( + checkpoint_structural_cfg={"cvar": False, "cvar_alpha": 0.95}) + self.assertNotEqual(with_cvar, without) + + +class TestConfigRegistration(unittest.TestCase): + """Config.checkpoint_args registers the phase-1a flags with sane defaults.""" + + def setUp(self): + self.cfg = Config() + self.cfg.checkpoint_args() + + def test_flags_are_registered(self): + for name in ("checkpoint_dir", "checkpoint_at_termination", + "checkpoint_backend", "resume_from"): + self.assertIn(name, self.cfg) + + def test_checkpointing_is_off_by_default(self): + self.assertIsNone(self.cfg.checkpoint_dir) + self.assertIsNone(self.cfg.resume_from) + + def test_terminal_checkpoint_defaults_on(self): + self.assertTrue(self.cfg.checkpoint_at_termination) + + def test_backend_defaults_to_dill_reload(self): + self.assertEqual(self.cfg.checkpoint_backend, + checkpointing.DILL_RELOAD_BACKEND) + + +class TestFilenameSanitizing(unittest.TestCase): + """File names must never go through extract_num (not unique for ADMM).""" + + def test_wrapped_admm_names_stay_distinct(self): + first = checkpointing.sanitize_for_filename( + "ADMM_STOCH__ADMM__region1__ADMM__scen3") + second = checkpointing.sanitize_for_filename( + "ADMM_STOCH__ADMM__region2__ADMM__scen3") + self.assertNotEqual(first, second) + + def test_path_separators_are_removed(self): + self.assertNotIn("/", checkpointing.sanitize_for_filename("a/b c")) + + +class TestFixedNonantBaseline(unittest.TestCase): + """The initially-fixed baseline must survive the model swap, by name. + + `_initial_fixed_varibles` is a ComponentSet of vardata, so a resume that + replaces the scenario models invalidates it by identity. Both failure + directions are pinned here: lose the baseline and the gate refuses to + update the bound; rebuild it from the *current* fixedness and a nonant that + a fixing extension pinned mid-run passes as original, admitting a bound the + uninterrupted run would have refused. + + These call the gate directly. A plain PH hub is insulated in practice -- + `PHBase._can_update_best_bound` short-circuits whenever prox is enabled, and + the one consultation with prox off is the iteration-0 trivial bound, which + the resume branch replaces -- but `Subgradient` and `FWPH` consult the same + baseline per iteration, so restoring it correctly is what keeps this from + becoming a bug the moment resume covers them. See design section 9, item 11. + """ + + def setUp(self): + # No solve, so no solver is needed -- this is pure bookkeeping. + # PH_Prep attaches the W/prox parameters that the PHBase override of + # _can_update_best_bound inspects before delegating to the fixedness + # check; with the attach deferred, prox is off, which is the state the + # gate is actually consulted in. + self.ph = _make_ph(_options(1)) + self.ph.PH_Prep() + scenario = next(iter(self.ph.local_scenarios.values())) + self.nonant = next(iter(scenario._mpisppy_data.nonant_indices.values())) + self.nonant.fix(self.nonant._value if self.nonant._value else 0.0) + + def test_baseline_by_name_allows_bound_updates(self): + self.ph._restore_fixed_nonant_baseline([self.nonant.name]) + self.assertTrue( + self.ph._can_update_best_bound(), + msg="a nonant fixed before the run started must stay part of the " + "baseline, or the resumed run stops updating its bound") + + def test_lost_baseline_would_block_bound_updates(self): + """What an identity-keyed cache degrades to after a swap.""" + self.ph._restore_fixed_nonant_baseline([]) + self.assertFalse(self.ph._can_update_best_bound()) + + def test_midrun_fixings_are_not_absorbed_into_the_baseline(self): + """A nonant pinned after the start must not pass as original.""" + others = [v for s in self.ph.local_scenarios.values() + for v in s._mpisppy_data.nonant_indices.values() + if v is not self.nonant] + midrun = others[0] + midrun.fix(midrun._value if midrun._value else 0.0) + + # Only the original is in the checkpointed baseline. + self.ph._restore_fixed_nonant_baseline([self.nonant.name]) + self.assertFalse( + self.ph._can_update_best_bound(), + msg="a mid-run fixing was treated as originally fixed, which " + "would admit a bound the uninterrupted run would refuse") + + def test_rebuilt_baseline_holds_current_model_objects(self): + """Rebuilt by name means the objects belong to the live models.""" + self.ph._restore_fixed_nonant_baseline([self.nonant.name]) + live = {id(v) for s in self.ph.local_scenarios.values() + for v in s._mpisppy_data.nonant_indices.values()} + for v in self.ph._initial_fixed_varibles: + self.assertIn(id(v), live) + + +class TestUnknownBackend(unittest.TestCase): + def test_require_dill_ignores_other_backends(self): + # Only the dill-reload backend needs dill; nothing should raise here. + checkpointing.require_dill(checkpointing.LEAF_BACKEND) + + +if __name__ == "__main__": + unittest.main() diff --git a/mpisppy/utils/cfg_vanilla.py b/mpisppy/utils/cfg_vanilla.py index d5a858899..3dec04330 100644 --- a/mpisppy/utils/cfg_vanilla.py +++ b/mpisppy/utils/cfg_vanilla.py @@ -386,6 +386,7 @@ def ph_hub( add_wxbar_read_write(hub_dict, cfg) add_ph_tracking(hub_dict, cfg) add_timed_mipgap(hub_dict, cfg) + add_checkpointing(hub_dict, cfg) return hub_dict def cg_hub( @@ -933,6 +934,42 @@ def add_wxbar_read_write(hub_dict, cfg): }) return hub_dict +def add_checkpointing(hub_dict, cfg): + """Attach the Checkpointer extension and forward its options. + + Both --checkpoint-dir (write) and --resume-from (read) are handled here. + Resuming does not need the extension -- the resume branch lives in + PHBase.Iter0 -- but it does need resume_from in the options dict, so a + resume-only run (stop today, resume tomorrow with a fresh checkpoint dir) + still works. + + See doc/designs/checkpointing_design.md. + """ + from mpisppy.utils.checkpointing import STRUCTURAL_CFG_KEYS + + if _hasit(cfg, 'checkpoint_dir'): + from mpisppy.extensions.checkpointer import Checkpointer + hub_dict = extension_adder(hub_dict, Checkpointer) + hub_dict["opt_kwargs"]["options"].update( + {"checkpoint_dir": cfg.checkpoint_dir, + "checkpoint_at_termination": cfg.checkpoint_at_termination, + "checkpoint_backend": cfg.checkpoint_backend, + }) + + if _hasit(cfg, 'resume_from'): + hub_dict["opt_kwargs"]["options"]["resume_from"] = cfg.resume_from + + if _hasit(cfg, 'checkpoint_dir') or _hasit(cfg, 'resume_from'): + # Structural settings PH itself never reads -- they act on the model + # before PH sees it, or describe the scenario tree -- but which must + # match for a checkpoint to be resumable. Collected here so the + # fingerprint can cover them without the extension needing the cfg. + hub_dict["opt_kwargs"]["options"]["checkpoint_structural_cfg"] = { + k: cfg.get(k, None) for k in STRUCTURAL_CFG_KEYS if k in cfg + } + + return hub_dict + def add_ph_tracking(cylinder_dict, cfg, spoke=False): """ Manage the phtracker extension and bridge gap between config and ph options dict Args: diff --git a/mpisppy/utils/checkpointing.py b/mpisppy/utils/checkpointing.py new file mode 100644 index 000000000..3bde267f2 --- /dev/null +++ b/mpisppy/utils/checkpointing.py @@ -0,0 +1,415 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""Read/write primitives for checkpointing a run so it can be resumed later. + +See ``doc/designs/checkpointing_design.md``. The division of labor is: + +- This module owns the *file format*: what a checkpoint generation looks like + on disk, how it is published atomically, and the fingerprints that decide + whether a checkpoint may be resumed into the current run. +- ``mpisppy/extensions/checkpointer.py`` owns *when* a checkpoint is written. +- ``PHBase.Iter0`` owns the resume branch itself, because restoring has to + happen in the middle of startup (the reloaded models must be in place before + solvers are created). + +The ``dill-reload`` backend dills each mid-run scenario model, which brings back +the dual weights, rho, nonant values and fixedness, the recourse values that +serve as a MIP warm start, and the proximal-approximation cuts, all mutually +consistent. Everything that does *not* live on a scenario model -- the global +iteration counter, bounds, and the initially-fixed-nonant baseline -- is written +alongside as a small pickle of plain data. +""" + +import json +import os +import pickle +import re +import shutil +import hashlib + +from pyomo.common.dependencies import attempt_import + +import mpisppy.utils.pickle_bundle as pickle_bundle + +dill, dill_available = attempt_import("dill") + +# Bump when the on-disk layout changes in a way older readers cannot handle. +FORMAT_VERSION = 1 + +DILL_RELOAD_BACKEND = "dill-reload" +LEAF_BACKEND = "leaf" + +MANIFEST_NAME = "manifest.json" +HUB_SUBDIR = "hub" + +# Option keys that must match for a checkpoint to be resumable. Deliberately a +# named subset rather than the whole configuration: these are the entries that +# change the *structure* of the scenario models or the meaning of the state +# riding in them, so a mismatch means the checkpoint cannot be restored into +# this run. Everything else is free to change between a stop and a resume -- +# notably the iteration limit and the time limit, which a user legitimately +# adjusts when picking a run back up the next morning, and the display/verbosity +# options, which have no bearing on the state at all. +STRUCTURAL_OPTION_KEYS = ( + "defaultPHrho", + "linearize_proximal_terms", + "linearize_binary_proximal_terms", + "proximal_linearization_tolerance", + "smoothed", + "defaultPHp", + "defaultPHbeta", +) + +# Structural values that do not appear in ``opt.options`` because PH itself +# never reads them -- they act on the model before PH sees it, or describe the +# scenario tree. ``cfg_vanilla.add_checkpointing`` collects these into +# ``options["checkpoint_structural_cfg"]`` so the fingerprint can cover them. +STRUCTURAL_CFG_KEYS = ( + "module_name", + "num_scens", + "branching_factors", + "scenarios_per_bundle", + "cvar", + "cvar_weight", + "cvar_alpha", + "cvar_mean_weight", +) + + +class CheckpointMismatch(RuntimeError): + """A checkpoint exists but cannot be resumed into the current run.""" + + +def _canonical(value): + """Render an option value as something JSON can hash reproducibly.""" + if isinstance(value, (list, tuple)): + return [_canonical(v) for v in value] + if isinstance(value, dict): + return {str(k): _canonical(v) for k, v in sorted(value.items())} + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + +def structural_fingerprint(options): + """Hash the structural subset of ``options`` (see STRUCTURAL_OPTION_KEYS).""" + payload = {k: _canonical(options.get(k)) for k in STRUCTURAL_OPTION_KEYS} + extras = options.get("checkpoint_structural_cfg") or {} + for k in STRUCTURAL_CFG_KEYS: + if k in extras: + payload[f"cfg:{k}"] = _canonical(extras[k]) + blob = json.dumps(payload, sort_keys=True).encode("utf-8") + return hashlib.sha256(blob).hexdigest() + + +def sanitize_for_filename(name): + """Make a scenario name safe to embed in a file name. + + Never use ``sputils.extract_num`` here: it scrapes trailing digits, which + are not unique for ADMM's wrapped scenario names. + """ + return re.sub(r"[^A-Za-z0-9_.-]", "_", str(name)) + + +def _generation_dirname(generation): + return f"gen_{generation:04d}" + + +def _leaf_filename(rank): + return f"hub_rank_{rank:04d}.pkl" + + +def _model_filename(rank, sname): + return f"hub_rank_{rank:04d}_scen_{sanitize_for_filename(sname)}.dill" + + +def _atomic_write_bytes(path, write_callback): + """Write via a temp file in the same directory, then rename into place.""" + tmp = f"{path}.tmp" + with open(tmp, "wb") as f: + write_callback(f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def require_dill(backend): + if backend == DILL_RELOAD_BACKEND and not dill_available: + raise RuntimeError( + "The '{}' checkpoint backend requires dill, which is not " + "installed. Install the optional dependencies with " + "'pip install mpi-sppy[extras]' (or 'pip install dill'), or " + "choose a different --checkpoint-backend.".format(backend) + ) + + + + +def probe_model_is_dillable(opt): + """Serialize one local scenario to memory to prove checkpointing will work. + + Called once at setup. A run that only discovers at its terminal checkpoint + -- possibly many hours in -- that its models cannot be dilled would lose + exactly the state checkpointing exists to preserve, so this trades one + model serialization up front for a failure that arrives immediately and + says what to do about it. The probe runs at iteration 0, when the model is + at its smallest (no accumulated prox-approximation cuts). + """ + if not opt.local_scenarios: + return + sname, s = next(iter(opt.local_scenarios.items())) + solver_plugin = getattr(s, "_solver_plugin", None) + if solver_plugin is not None: + del s._solver_plugin + try: + dill.dumps(s) + except Exception as exc: + raise RuntimeError( + "Checkpointing is enabled, but no checkpoint could ever be " + "written.\n\n" + + pickle_bundle.describe_dill_failure( + s, exc, what=f"scenario '{sname}'") + ) from exc + finally: + if solver_plugin is not None: + s._solver_plugin = solver_plugin + + +def geometry(opt): + """The rank layout a resume must reproduce (see design section 5.7).""" + return { + "n_proc": int(opt.n_proc), + "rank": int(opt.cylinder_rank), + "scenario_names": sorted(opt.local_scenarios.keys()), + } + + +def initially_fixed_nonant_names(opt): + """Names of the nonants that were already fixed when the run first started. + + This is the baseline ``_can_update_best_bound`` compares against, and it is + the one piece of opt-object state that is keyed by variable *identity* -- + a ``ComponentSet`` of vardata belonging to models that a resume replaces. + Recording it by name is what lets the resume rebuild it correctly; see + design section 9, item 11, for what goes wrong otherwise. + """ + baseline = getattr(opt, "_initial_fixed_varibles", None) + if baseline is None: + return [] + return sorted(v.name for v in baseline) + + +def write_checkpoint(opt, ckpt_dir, generation, backend=DILL_RELOAD_BACKEND): + """Write and atomically publish one checkpoint generation. + + The rank writes its own files into a temporary generation directory, which + is renamed into place; the manifest is then rewritten (itself + temp-then-rename) to point at the new generation. That manifest flip is the + single commit point, so a kill before it leaves the previous checkpoint + intact and a kill after it leaves the new one. The prior generation is + deleted once the manifest names its replacement. + """ + require_dill(backend) + + rank = int(opt.cylinder_rank) + hub_dir = os.path.join(ckpt_dir, HUB_SUBDIR) + final_dir = os.path.join(hub_dir, _generation_dirname(generation)) + staging_dir = f"{final_dir}.tmp" + + if os.path.isdir(staging_dir): + shutil.rmtree(staging_dir) + os.makedirs(staging_dir, exist_ok=True) + + try: + model_files = _write_models(opt, staging_dir, rank, backend) + except Exception as exc: + # Leave no half-written generation behind; the previous checkpoint (if + # any) stays published, since the manifest was never touched. + shutil.rmtree(staging_dir, ignore_errors=True) + if isinstance(exc, ValueError): + raise + first = next(iter(opt.local_scenarios.values()), None) + detail = (pickle_bundle.describe_dill_failure(first, exc, + what="scenario model") + if first is not None + else f"{type(exc).__name__}: {exc}") + raise RuntimeError( + f"Failed to write the checkpoint to '{ckpt_dir}'. Any previously " + f"published checkpoint is untouched.\n\n" + detail + ) from exc + + leaf = { + "format_version": FORMAT_VERSION, + "backend": backend, + "generation": int(generation), + "geometry": geometry(opt), + "structural_fingerprint": structural_fingerprint(opt.options), + "model_files": model_files, + "initially_fixed_nonants": initially_fixed_nonant_names(opt), + "trivial_bound": _as_float_or_none(getattr(opt, "trivial_bound", None)), + "best_bound_obj_val": _as_float_or_none( + getattr(opt, "best_bound_obj_val", None)), + "best_solution_obj_val": _as_float_or_none( + getattr(opt, "best_solution_obj_val", None)), + } + _atomic_write_bytes( + os.path.join(staging_dir, _leaf_filename(rank)), + lambda f: pickle.dump(leaf, f), + ) + + if os.path.isdir(final_dir): + shutil.rmtree(final_dir) + os.replace(staging_dir, final_dir) + + previous = _read_manifest(ckpt_dir, missing_ok=True) + _publish_manifest(ckpt_dir, { + "format_version": FORMAT_VERSION, + "backend": backend, + "generation": int(generation), + "n_proc": int(opt.n_proc), + "structural_fingerprint": structural_fingerprint(opt.options), + }) + + # Exactly one committed generation is kept; retaining more is not + # supported. Two exist only transiently, between the generation rename + # above and this delete. + if previous is not None and previous.get("generation") != int(generation): + stale = os.path.join(hub_dir, _generation_dirname(previous["generation"])) + if os.path.isdir(stale): + shutil.rmtree(stale, ignore_errors=True) + + return final_dir + + +def _write_models(opt, staging_dir, rank, backend): + """Dill each local scenario model into the staging directory.""" + if backend != DILL_RELOAD_BACKEND: + raise ValueError( + f"Unknown checkpoint backend '{backend}'. The only implemented " + f"backend is '{DILL_RELOAD_BACKEND}'." + ) + model_files = {} + for sname, s in opt.local_scenarios.items(): + fname = _model_filename(rank, sname) + # The solver plugin is a live C handle plus a license session; it + # cannot be serialized and is rebuilt by _create_solvers on resume. + solver_plugin = getattr(s, "_solver_plugin", None) + if solver_plugin is not None: + del s._solver_plugin + try: + _atomic_write_bytes( + os.path.join(staging_dir, fname), + lambda f, model=s: dill.dump(model, f), + ) + finally: + if solver_plugin is not None: + s._solver_plugin = solver_plugin + model_files[sname] = fname + return model_files + + +def _as_float_or_none(value): + return None if value is None else float(value) + + +def _publish_manifest(ckpt_dir, manifest): + path = os.path.join(ckpt_dir, MANIFEST_NAME) + tmp = f"{path}.tmp" + with open(tmp, "w") as f: + json.dump(manifest, f, indent=2, sort_keys=True) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def _read_manifest(ckpt_dir, missing_ok=False): + path = os.path.join(ckpt_dir, MANIFEST_NAME) + if not os.path.exists(path): + if missing_ok: + return None + raise CheckpointMismatch( + f"No checkpoint manifest at '{path}'. --resume-from expects a " + f"directory that a previous run wrote with --checkpoint-dir." + ) + with open(path) as f: + return json.load(f) + + +def load_checkpoint(opt, ckpt_dir): + """Load this rank's checkpoint, refusing a mismatch with a clear error. + + Returns ``(leaf_state, {scenario_name: reloaded_model})``. The caller is + responsible for splicing the models into the run (see + ``PHBase._resume_from_checkpoint``). + """ + manifest = _read_manifest(ckpt_dir) + + if manifest.get("format_version") != FORMAT_VERSION: + raise CheckpointMismatch( + f"Checkpoint in '{ckpt_dir}' has format version " + f"{manifest.get('format_version')}, but this mpi-sppy writes and " + f"reads version {FORMAT_VERSION}. Checkpoints are not portable " + f"across format versions." + ) + + backend = manifest.get("backend") + require_dill(backend) + + expected_fp = structural_fingerprint(opt.options) + if manifest.get("structural_fingerprint") != expected_fp: + raise CheckpointMismatch( + f"The checkpoint in '{ckpt_dir}' was written by a run whose " + f"structural options differ from this one. The options that must " + f"match are {', '.join(STRUCTURAL_OPTION_KEYS)} plus " + f"{', '.join(STRUCTURAL_CFG_KEYS)}; the iteration limit, time " + f"limit, and display options may be changed freely on a resume." + ) + + if int(manifest.get("n_proc", -1)) != int(opt.n_proc): + raise CheckpointMismatch( + f"The checkpoint in '{ckpt_dir}' was written on " + f"{manifest.get('n_proc')} rank(s) but this run has " + f"{opt.n_proc}. Resuming across a different rank count is not " + f"supported; rerun with the original rank count." + ) + + generation = manifest["generation"] + gen_dir = os.path.join(ckpt_dir, HUB_SUBDIR, _generation_dirname(generation)) + rank = int(opt.cylinder_rank) + + leaf_path = os.path.join(gen_dir, _leaf_filename(rank)) + if not os.path.exists(leaf_path): + raise CheckpointMismatch( + f"The checkpoint in '{ckpt_dir}' has no state for rank {rank} " + f"(expected '{leaf_path}')." + ) + with open(leaf_path, "rb") as f: + leaf = pickle.load(f) + + have = sorted(opt.local_scenarios.keys()) + want = leaf["geometry"]["scenario_names"] + if have != want: + raise CheckpointMismatch( + f"Rank {rank} now owns scenarios {have}, but the checkpoint in " + f"'{ckpt_dir}' was written with {want} on that rank. Resuming " + f"requires an identical scenario-to-rank distribution." + ) + + models = {} + for sname, fname in leaf["model_files"].items(): + path = os.path.join(gen_dir, fname) + if not os.path.exists(path): + raise CheckpointMismatch( + f"The checkpoint in '{ckpt_dir}' is missing the model file " + f"'{fname}' for scenario '{sname}'." + ) + with open(path, "rb") as f: + models[sname] = dill.load(f) + + return leaf, models diff --git a/mpisppy/utils/config.py b/mpisppy/utils/config.py index 555693e61..659984678 100644 --- a/mpisppy/utils/config.py +++ b/mpisppy/utils/config.py @@ -771,6 +771,39 @@ def integer_relax_then_enforce_args(self): domain=float, default=0.5) + def checkpoint_args(self): + # Checkpoint/resume (see doc/designs/checkpointing_design.md). The + # Checkpointer extension is attached iff checkpoint_dir is set, so a + # run that does not ask for checkpointing pays nothing. + self.add_to_config("checkpoint_dir", + description="directory for checkpoint files; its " + "presence enables checkpointing (default None)", + domain=str, + default=None) + + self.add_to_config("checkpoint_at_termination", + description="write a checkpoint when the run ends " + "for any internal reason, including reaching " + "--time-limit (default True); requires " + "--checkpoint-dir", + domain=bool, + default=True) + + self.add_to_config("checkpoint_backend", + description="how scenario-model state is saved and " + "restored; 'dill-reload' is the only implemented " + "backend (default dill-reload)", + domain=str, + default="dill-reload") + + self.add_to_config("resume_from", + description="resume from the checkpoint in this " + "directory; requires the same rank count and " + "scenario-to-rank distribution as the run that " + "wrote it (default None)", + domain=str, + default=None) + def slamming_args(self): # Phase-1 preference-driven slamming (see doc/designs/slamming_design.md). # The Slammer extension is activated iff slamming_directives_file is set; diff --git a/run_coverage.bash b/run_coverage.bash index cbac43991..84b4d52c4 100755 --- a/run_coverage.bash +++ b/run_coverage.bash @@ -191,6 +191,9 @@ run_phase "test_xhat_feasibility_cuts (serial)" \ run_phase "test_incumbent_writing (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_incumbent_writing.py -v +run_phase "test_checkpoint (serial)" \ + coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_checkpoint.py -v + run_phase "test_iis_on_infeasible (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_iis_on_infeasible.py -v