stoch-admm phase B: consensus_vars ergonomics + first-stage hook extensions

[DLWoodruff]

Summary

Phase B of the stoch-admm user-API automation plan (doc/designs/admm_user_api_automation_design.md), three closely related items that build on Phase A's first_stage_cost / first_stage_varlist module hooks:

• B.1 — consensus_vars_creator accepts Pyomo Var / VarData in addition to strings; the wrapper normalizes via .name at construction time. Mixed lists OK for migration. Eliminates the \"flow[(DC1, DC2)]\" vs \"flow['DC1', 'DC2']\" formatting foot-gun.
• B.2 — The wrapper auto-merges each ADMM subproblem's first-stage Vars into its consensus_vars entry, so consensus_vars_creator can return only the admm-consensus Vars. Eliminates the throwaway-model trick stoch_distr.py's consensus_vars_creator used to walk first_stage_varlist from a freshly built before-wrap scenario.
• B.3 — Optional first_stage_surrogate_nonant_list / first_stage_nonant_ef_suppl_list module hooks that forward to sputils.attach_root_node's advanced parameters. Users of surrogates or EF-supplemental nonants are no longer forced back to the legacy manual-attach path.

ADMM vocabulary glossary

Writing B.2's docstrings forced us to settle on consistent terms across the ADMM family. The full glossary now lives in the module docstring at the top of mpisppy/utils/admmWrapper.py and the three admm test files have header comments referencing it:

• before-wrap scenario / wrapped scenario — the Pyomo model the user's scenario_creator returns vs. the same model after the wrapper has applied the wrap operation; what mpi-sppy iterates over.
• wrap (verb, narrow scope) — append the ADMM consensus stage (Stoch_AdmmWrapper) or flatten consensus Vars into ROOT (AdmmBundler), plus surrogate / probability bookkeeping. Does not cover sputils.attach_root_node (that step is setup, not wrap).
• the wrapper — any of AdmmWrapper, Stoch_AdmmWrapper, AdmmBundler.
• stochastic scenario ($\xi \in \Xi$) — one realization of the random data. "scenario" alone is now ambiguous in prose; always qualify.
• ADMM subproblem ($a \in \mathcal{A}$) — one partition / region. "ADMM" capitalized in prose; code identifier is admm_subproblem.
• first-stage Vars (user-API / algorithm-level) vs. root-node Vars (Vars attached to _mpisppy_node_list[0], the implementation-level term).

All B.1 / B.2 / B.3 docstrings and comments — plus examples/stoch_distr/stoch_distr.py and the three admm test files — use this vocabulary consistently.

Notable findings during implementation (recorded in the design doc)

• B.2: first-stage Vars are per-ADMM-subproblem, not global (significant divergence from the original draft). Different ADMM subproblems may carry different first-stage Vars (stoch_distr's per-region factories). The helper takes a per-subproblem dict; the wrapper gathers names from already-built local before-wrap scenarios where possible and probes for ADMM subproblems the rank does not host locally.
• B.1: Pyomo VarData weakref gotcha — .name only resolves to a real name if the source before-wrap scenario is alive at wrapper-construction time. Documented in the helper docstring; callers building throwaway scenarios in helpers must keep them alive across the wrapper call.
• B.3: bundle ROOT does not propagate constituent surrogates — the advanced lists live on the per-constituent before-wrap scenario roots; the EF reads from there when assembling. Bundler test verifies forwarding via a spy on sputils.attach_root_node.

Files

Source

• mpisppy/utils/admmWrapper.py — module docstring with vocabulary glossary; new helpers _admm_normalize_consensus_vars(consensus_vars, *, tuple_form) (B.1) and _merge_first_stage_into_consensus_vars(consensus_vars, first_stage_var_names_per_sub, root_stage=1) (B.2); AdmmWrapper.__init__ calls the normalizer.
• mpisppy/utils/stoch_admmWrapper.py — B.1 normalizer + B.2 per-subproblem first-stage merge + B.3 advanced hooks (first_stage_surrogate_nonant_list, first_stage_nonant_ef_suppl_list) forwarded to attach_root_node.
• mpisppy/utils/admm_bundler.py — same B.1 / B.2 / B.3 plumbing. Builds one probe before-wrap scenario per ADMM subproblem at __init__ since it does not pre-create scenarios there.
• mpisppy/generic/admm.py — extracted both-or-neither / advanced-needs-core validation into a shared _discover_first_stage_hooks(module) used by both setup_stoch_admm and setup_stoch_admm_with_bundles.

Example

• examples/stoch_distr/stoch_distr.py — consensus_vars_creator no longer instantiates a throwaway scenario; just returns the admm-consensus Vars. The wrapper merges first-stage Vars at construction time.

Tests

• mpisppy/tests/test_admmWrapper.py — TestAdmmConsensusVarsNormalize (6 unit tests of the B.1 normalizer), TestAdmmWrapperVarConsensusInputs (end-to-end B.1 on AdmmWrapper), TestMergeFirstStageIntoConsensusVars (4 unit tests of the B.2 merger).
• mpisppy/tests/test_stoch_admmWrapper.py — test_consensus_vars_accepts_var_objects (B.1), test_b2_auto_merge_first_stage_into_consensus_vars and test_b2_per_subproblem_first_stage_vars (B.2), test_advanced_hooks_forwarded / test_advanced_hook_alone_ok / test_advanced_hook_without_core_hooks_errors / test_setup_stoch_admm_advanced_without_core_errors (B.3). test_hooks_and_legacy_paths_produce_same_node_list updated to pre-merge first-stage Vars into the legacy consensus_vars (apples-to-apples post-B.2).
• mpisppy/tests/test_admm_bundler.py — TestAdmmBundlerB2AutoMerge and TestAdmmBundlerB3AdvancedHooks.

Docs

• doc/designs/admm_user_api_automation_design.md — "Findings during implementation" subsections recording divergences from the original Phase B draft.
• doc/src/generic_admm.rst — new "Advanced first-stage hooks (optional)" subsection covering B.3, with :ref: cross-references to the existing "EF Supplement List" / "Surrogate Nonant List" sections in scenario_creator.rst (which gained Sphinx labels in this PR). The legacy-path subsection now explicitly mentions passing surrogate_nonant_list / nonant_ef_suppl_list directly to attach_root_node.
• doc/src/scenario_creator.rst — added Sphinx labels _ef_supplement_list and _surrogate_nonant_list so the cross-references resolve.

Test plan

• `python -m pytest mpisppy/tests/test_admmWrapper.py mpisppy/tests/test_admm_bundler.py mpisppy/tests/test_stoch_admmWrapper.py --deselect mpisppy/tests/test_stoch_admmWrapper.py::TestStochAdmmWrapper::test_values` — 71/71 pass (the deselected `test_values` is the pre-existing flaky test from PR stoch-admm phase A: first_stage_cost / first_stage_varlist module hooks #686, unrelated).
• `ruff check` on the changed Python files — clean.
• `cd doc && make html` — clean (the four pre-existing warnings are unrelated).
• `SOLVERNAME=xpress bash examples/stoch_distr/stoch_admm_stage2ef.bash` — converges (rel gap 0.884%).

Follow-up

Phases C, D, E in the design doc remain and can branch off main after this PR merges. None of them depend on Phase B's specific shape — they target combining_names / split_admm_stoch_subproblem_scenario_name defaults (C), admm_stoch_subproblem_scenario_names_creator defaults (D), and inparser_adder boilerplate (E).

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 98.96907% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 72.71%. Comparing base (f083396) to head (b1cdb7b).

Files with missing lines	Patch %	Lines
mpisppy/generic/admm.py	93.75%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #722      +/-   ##
==========================================
+ Coverage   72.60%   72.71%   +0.11%     
==========================================
  Files         162      162              
  Lines       20643    20720      +77     
==========================================
+ Hits        14987    15066      +79     
+ Misses       5656     5654       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

This PR implements Phase B of the stochastic-ADMM user-API automation plan by improving consensus_vars ergonomics, removing the need for throwaway scenarios to discover first-stage variables, and extending the first-stage hook mechanism to support surrogate and EF-supplemental nonants.

Changes:

• Accept Pyomo Var / VarData objects (in addition to strings) in consensus_vars, normalizing to .name at wrapper construction time.
• Auto-merge each ADMM subproblem’s first-stage variable names into its consensus_vars entry when the Phase A first_stage_* hooks are used.
• Add optional advanced first-stage hooks to forward surrogate_nonant_list / nonant_ef_suppl_list into sputils.attach_root_node, with shared validation in setup_stoch_admm*.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
mpisppy/utils/stoch_admmWrapper.py	Normalizes consensus var identifiers, auto-merges first-stage vars per subproblem, and forwards advanced attach-root kwargs via new optional hooks.
mpisppy/utils/admmWrapper.py	Adds shared helpers for consensus-var normalization and first-stage auto-merge; applies normalization in AdmmWrapper.
mpisppy/utils/admm_bundler.py	Mirrors normalization/auto-merge behavior and forwards advanced attach-root kwargs for bundled stoch-ADMM.
mpisppy/generic/admm.py	Extracts shared hook discovery/validation into _discover_first_stage_hooks and forwards results into wrapper/bundler setup.
examples/stoch_distr/stoch_distr.py	Removes throwaway-scenario first-stage discovery in consensus_vars_creator, relying on wrapper auto-merge.
mpisppy/tests/test_stoch_admmWrapper.py	Adds/updates tests covering B.1 normalization, B.2 auto-merge (including per-subproblem + probe fallback), and B.3 advanced hook forwarding/validation.
mpisppy/tests/test_admmWrapper.py	Adds unit tests for normalization and first-stage merge helpers plus an end-to-end var-object consensus input test for AdmmWrapper.
mpisppy/tests/test_admm_bundler.py	Adds tests for bundler B.2 auto-merge semantics and B.3 advanced hook forwarding/validation.
doc/src/scenario_creator.rst	Adds Sphinx labels to enable cross-references for EF supplement / surrogate nonant documentation.
doc/src/generic_admm.rst	Documents the two new advanced first-stage hooks and clarifies legacy/manual attach behavior.
doc/designs/admm_user_api_automation_design.md	Records implementation findings and clarifies Phase B semantics (per-subproblem first-stage vars, probe behavior, weakref caveat).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.