stoch-admm phase A: first_stage_cost / first_stage_varlist module hooks#686
Open
DLWoodruff wants to merge 7 commits into
Open
stoch-admm phase A: first_stage_cost / first_stage_varlist module hooks#686DLWoodruff wants to merge 7 commits into
DLWoodruff wants to merge 7 commits into
Conversation
Eliminate the asymmetry where --admm users must NOT call
sputils.attach_root_node in scenario_creator but --stoch-admm users
MUST. When the model module defines first_stage_cost(scenario) and
first_stage_varlist(scenario) callables, Stoch_AdmmWrapper (and
AdmmBundler, for the bundled path) call attach_root_node themselves;
scenario_creator just returns the constructed model. The legacy
path (where scenario_creator calls attach_root_node itself) still
works for back-compat.
Strict both-or-neither contract. setup_stoch_admm (and
setup_stoch_admm_with_bundles) inspect the module for both hooks;
mixing produces a clear error naming the missing function.
Per-scenario error matrix in the wrappers, checked right after
scenario_creator returns:
| both hooks | neither hook
--------------------------------------+-----------------
attach_root_node called | RuntimeError| OK (legacy)
not called | OK (hooks) | RuntimeError
The "hooks + manual call -> error" choice is deliberately strict
(rather than silently overwriting the user's call): users learn the
contract instead of half-migrating undetected. The "neither hook +
no call" error replaces a cryptic deep assertion in
assign_variable_probs with a message pointing at both remediation
options.
Migrate examples/stoch_distr/stoch_distr.py to use the hooks:
scenario_creator stashes the first-stage varlist on the model and
the hook reads it back. consensus_vars_creator is refactored to
read first-stage vars via first_stage_varlist() instead of walking
model._mpisppy_node_list (which no longer exists at that point,
because the wrapper hasn't attached the root yet).
Update the custom driver stoch_distr_admm_cylinders.py to forward
the hooks to Stoch_AdmmWrapper so the stoch_distr_ef.py
extensive-form path still works.
Tests: 6 new unit tests in test_stoch_admmWrapper.py covering both
paths, all four error-matrix cells, and the setup_stoch_admm-level
half-hooks error. The legacy-path tests
(test_admm_bundler.py + the pre-existing test_stoch_admmWrapper
suite + test_admm_bundler tests via the tests/examples copy of
stoch_distr.py which keeps the manual attach_root_node) all
continue to pass — 42/42 (the pre-existing test_values flakiness
under subprocess capture is unrelated and skipped).
Doc: generic_admm.rst "Extending to Stochastic ADMM" rewritten
around the new hook-based API as the recommended path, with the
legacy manual-attach pattern preserved as a fallback section.
This commit supersedes the doc clarification from commit 24bf580
("clarify attach_root_node contract differs for --admm vs
--stoch-admm").
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #686 +/- ##
==========================================
+ Coverage 70.80% 70.89% +0.08%
==========================================
Files 154 154
Lines 19252 19290 +38
==========================================
+ Hits 13631 13675 +44
+ Misses 5621 5615 -6 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- "the function must split it" — name scenario_creator explicitly (the prior pronoun read ambiguously next to split_admm_stoch_subproblem_scenario_name). - "Under the hood, Stoch_AdmmWrapper reads…" — flag the paragraph as implementation detail rather than user-facing API. - "When both hooks (first_stage_cost and first_stage_varlist) are defined" — name the hooks inline. - Drop the parenthetical rationale on the both-or-neither error (the rule is what matters; the design doc explains the why). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
"Phase A" is a development-cycle label that won't make sense to future readers. Strip it from inline comments across the wrappers, generic_cylinders setup, the migrated example, and the test class name (TestStochAdmmWrapperPhaseAHooks -> TestStochAdmmWrapperFirstStageHooks). The substance of each comment is preserved; only the temporal label goes away. The design doc and commit messages still reference "Phase A" appropriately (they're about the phased plan). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This was referenced May 13, 2026
Adds the design doc that motivates this PR's phase A implementation.
The doc covers all five phases (A-E) but only phase A is implemented
here; B-E will follow as separate PRs.
A. Auto-build the scenario tree for --stoch-admm via new
first_stage_cost / first_stage_varlist module hooks (this PR).
B. Accept Pyomo Var/VarData in consensus_vars_creator.
C. Default combining_names / split_admm_stoch_subproblem_scenario_name
pair in mpisppy.utils.stoch_admmWrapper, with __ADMM__ as the
delimiter.
D. Default admm_stoch_subproblem_scenario_names_creator with the
correct nesting (outer = stoch, inner = admm).
E. Drop redundant ADMM-arg registrations from example
inparser_adder functions; admm_args is canonical.
Item F (merging AdmmWrapper and Stoch_AdmmWrapper) and
multistage-origin stoch-admm are explicitly out of scope.
Each phase is independently shippable, preserves the old API as a
fallback, and migrates the canonical example. Phase A's per-scenario
error matrix (hooks-x-attached) and its both-or-neither contract are
spelled out in section 2 of the doc, and the implementation in this
PR matches.
The design doc was originally drafted on a separate branch and went
through review (closed PR Pyomo#684) before phase A was implemented; the
"files" list in section 2 was updated after implementation to
reflect findings (AdmmBundler needed the same plumbing,
consensus_vars_creator required a refactor away from
_mpisppy_node_list, the custom driver had to forward hooks too).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Symmetric with the existing test_setup_stoch_admm_half_hooks_errors and test_wrapper_half_hooks_errors: bumps patch coverage on the new both-or-neither check in mpisppy/utils/admm_bundler.py (AdmmBundler __init__) and mpisppy/generic/admm.py (setup_stoch_admm_with_bundles). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase A of the ADMM user-API automation plan. The phased plan A–E itself is folded into this PR as
doc/designs/admm_user_api_automation_design.md(see thedesign:commit at the end of the branch).Eliminate the asymmetry where
--admmusers must NOT callsputils.attach_root_nodeinscenario_creatorbut--stoch-admmusers MUST. When the model module definesfirst_stage_cost(scenario)andfirst_stage_varlist(scenario)callables,Stoch_AdmmWrapper(andAdmmBundler, for the bundled path) callattach_root_nodethemselves;scenario_creatorjust returns the constructed model. The legacy path is preserved for backward compatibility.Strict both-or-neither contract. Per-scenario error matrix, checked right after
scenario_creatorreturns:attach_root_nodecalledRuntimeErrorRuntimeErrorThe "hooks + manual call → error" choice is deliberately strict (rather than silently overwriting the user's call): users learn the contract instead of half-migrating undetected. The "neither hook + no call" error replaces a cryptic deep assertion in
assign_variable_probswith a message pointing at both remediation options.Files
doc/designs/admm_user_api_automation_design.md: phased plan A–E with rationale, per-phase scope, back-compat strategy, and resolved open questions.mpisppy/utils/stoch_admmWrapper.py: hook parameters + per-scenario error matrix.mpisppy/utils/admm_bundler.py: same plumbing for the bundled path (constructs scenarios in its ownscenario_creatormethod; runs the matrix per scenario inside the bundle).mpisppy/generic/admm.py:setup_stoch_admmandsetup_stoch_admm_with_bundlesdiscover hooks on the module and enforce both-or-neither.examples/stoch_distr/stoch_distr.py: migrate to the new API.scenario_creatorstashes the first-stage varlist on the model; new hook functions read it back.consensus_vars_creatoris refactored to read first-stage vars viafirst_stage_varlist()instead of walkingmodel._mpisppy_node_list(which no longer exists at that point — the wrapper hasn't attached the root yet).examples/stoch_distr/stoch_distr_admm_cylinders.py: the custom driver forwards the module's hooks sostoch_distr_ef.py(which uses this driver) continues to work after the migration.doc/src/generic_admm.rst: rewrites "Extending to Stochastic ADMM" around the hook-based API as the recommended path; preserves the legacy manual-attach pattern as a fallback subsection.mpisppy/tests/test_stoch_admmWrapper.py: 6 new tests covering both paths, all four error-matrix cells, and thesetup_stoch_admm-level half-hooks error.Test plan
Follow-up
🤖 Generated with Claude Code