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

doc/designs/admm_user_api_automation_design.md

@@ -213,6 +213,24 @@ manual name formatting", not "skip the find_component step".
   `consensus_vars_creator` produces the same `varprob_dict` as the
   string-based form.
 
+#### Findings during implementation
+
+- **Pyomo VarData weakref**: a `VarData` holds its parent block via
+  a weakref, so the wrapper's `.name` lookup only resolves to a real
+  name if the before-wrap scenario the Var was taken from is still
+  alive at wrapper-construction time.  Callers that build a
+  throwaway before-wrap scenario inside a helper must keep it alive
+  across the wrapper call, or snapshot the names eagerly before
+  letting it go out of scope.  Documented in the helper docstring;
+  test_admmWrapper.py's integration test has to hold a
+  `live_scenarios` list across the wrapper call.
+- **Normalization location**: rather than pre-normalizing inside
+  `assign_variable_probs` (as originally sketched), implementation
+  normalizes once in `__init__` so `_consensus_vars_number_creator`
+  and the B.2 / B.3 plumbing all see canonical string identifiers.
+  Helper: `_admm_normalize_consensus_vars(consensus_vars, *,
+  tuple_form)` in `mpisppy/utils/admmWrapper.py`.
+
 ### B.2 Auto-merge first-stage Vars into `consensus_vars`
 
 #### Today's contract
@@ -260,6 +278,35 @@ already-merged list *and* has the hooks defined, to avoid a
 double-merge — easiest is to de-duplicate the merged list by Var
 identity.
 
+#### Findings during implementation
+
+- **First-stage Vars are per-ADMM-subproblem, not global**
+  (significant divergence from the original "append to every
+  subproblem entry" wording above).  In stoch_distr each region
+  owns its own factory production decisions, so `first_stage_varlist`
+  returns *different* Var names for different ADMM subproblems
+  even though they share Var-name conventions.  The implemented
+  helper takes a per-subproblem dict, not a flat list:
+  `_merge_first_stage_into_consensus_vars(consensus_vars,
+  first_stage_var_names_per_sub, root_stage=1)`.  De-dup is
+  string-name based (no Var-identity comparison needed once the
+  consensus_vars dict is normalized).
+- **Per-rank gathering**: every rank must end up with the same
+  `consensus_vars` so `_consensus_vars_number_creator` and
+  `assign_variable_probs` stay consistent.  `Stoch_AdmmWrapper`
+  pulls first-stage Var names from already-built local before-wrap
+  scenarios where possible, then builds one *probe* before-wrap
+  scenario per ADMM subproblem the local rank does not host.
+  `AdmmBundler` always builds one probe per ADMM subproblem in
+  `__init__` because it pre-builds no before-wrap scenarios there.
+- **Vocabulary glossary established**: the prose for B.2's
+  docstrings forced us to settle on consistent terms across the
+  ADMM family — before-wrap scenario, wrapped scenario, wrap (verb,
+  narrow scope), ADMM subproblem, first-stage Vars vs. root-node
+  Vars.  The full glossary now lives in the module docstring at the
+  top of `mpisppy/utils/admmWrapper.py` and is referenced by header
+  comments at the top of the three admm test files.
+
 ### B.3 Optional surrogate / EF-supplemental hooks
 
 #### Today's contract
@@ -306,6 +353,28 @@ message).
 - Test: a synthetic scenario with a surrogate Var; verify the
   attached node carries `surrogate_nonant_list` correctly.
 
+#### Findings during implementation
+
+- **Bundle ROOT does not propagate constituent surrogates.**
+  `AdmmBundler` builds the bundle's own ROOT by flattening every
+  consensus Var into a fresh `attach_root_node(bundle, 0,
+  nonantlist)` call with no advanced kwargs.  The surrogate /
+  EF-supplemental Vars instead live on the *per-constituent*
+  before-wrap scenario root nodes — the EF construction reads them
+  from there when assembling the bundle.  The B.3 plumbing
+  therefore forwards the advanced lists only to the constituent
+  attach_root_node calls inside `_process_scenario`, not to the
+  bundle-level attach.  This is correct (EF reads from
+  constituents) but worth flagging because a reader might expect
+  the bundle ROOT to carry the surrogate annotation directly.
+  Bundler test verifies forwarding via a spy on
+  `sputils.attach_root_node`.
+- **Shared discovery helper**: discovery and validation moved into
+  `_discover_first_stage_hooks(module)` in `mpisppy/generic/admm.py`
+  so `setup_stoch_admm` and `setup_stoch_admm_with_bundles` no
+  longer duplicate the both-or-neither / advanced-needs-core
+  checks.
+
 ### Open questions
 
 None significant.

doc/src/generic_admm.rst

@@ -426,13 +426,49 @@ inside ``scenario_creator`` so the hook can find it.
    ``RuntimeError`` at ``setup_stoch_admm`` time.  Mixing the hooks
    with a manual ``attach_root_node`` call also raises.
 
+Advanced first-stage hooks (optional)
+"""""""""""""""""""""""""""""""""""""""
+
+``sputils.attach_root_node`` accepts two further optional parameters,
+``surrogate_nonant_list`` and ``nonant_ef_suppl_list`` (see
+:ref:`surrogate_nonant_list` and :ref:`ef_supplement_list` for what
+each does), for problems that need to mark some first-stage Vars as
+surrogates (EF skips their nonant equality) or as EF-supplemental
+nonants (extra Vars carried through the EF construction).  If your
+problem needs either, define the corresponding optional module-level
+hook:
+
+.. code-block:: python
+
+   def first_stage_surrogate_nonant_list(scenario):
+       """Optional. Forwarded to attach_root_node's surrogate_nonant_list."""
+       return scenario._surrogate_nonants   # stashed in scenario_creator
+
+   def first_stage_nonant_ef_suppl_list(scenario):
+       """Optional. Forwarded to attach_root_node's nonant_ef_suppl_list."""
+       return scenario._ef_suppl_nonants
+
+Each advanced hook is independent of the other — defining either one
+alone is fine — but both depend on the two core hooks
+(``first_stage_cost`` and ``first_stage_varlist``) also being defined,
+because there is nothing for the wrapper to attach the advanced lists
+onto otherwise.  Defining an advanced hook without the core hooks
+raises ``RuntimeError`` at ``setup_stoch_admm`` time.
+
+On the legacy path (no core hooks), pass ``surrogate_nonant_list`` and
+``nonant_ef_suppl_list`` directly to your own ``sputils.attach_root_node``
+call inside ``scenario_creator``; the wrapper inherits whatever you
+attached.
+
 First-stage attachment via manual ``attach_root_node`` (legacy)
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 
 If you omit both hooks, ``scenario_creator`` must itself call
 ``sputils.attach_root_node`` with the original problem's first-stage
-cost and varlist.  Skipping the call (when no hooks are defined)
-raises ``RuntimeError`` with a message pointing at both options.
+cost and varlist (and ``surrogate_nonant_list`` /
+``nonant_ef_suppl_list`` if you need them).  Skipping the call (when
+no hooks are defined) raises ``RuntimeError`` with a message pointing
+at both options.
 
 This path is preserved for backward compatibility with model modules
 written before the hooks existed (and for direct uses of

doc/src/scenario_creator.rst

@@ -63,6 +63,8 @@ equally likely, you can avoid a warning by assigning the string
 "uniform" to the ``_mpisppy_probability`` attribute of the scenario
 model.
 
+.. _ef_supplement_list:
+
 EF Supplement List
 ------------------
 
@@ -72,6 +74,8 @@ constraints when an EF is formed, either to solve the EF or when bundles are
 formed. For some problems, with the appropriate solver, adding redundant nonanticipativity constraints
 for auxiliary variables to the bundle/EF will result in a (much) smaller pre-solved model.
 
+.. _surrogate_nonant_list:
+
 Surrogate Nonant List
 ---------------------
 

examples/stoch_distr/stoch_distr.py

@@ -268,19 +268,9 @@ def consensus_vars_creator(admm_subproblem_names, stoch_scenario_name, inter_reg
         if admm_subproblem_name not in consensus_vars:
             print(f"WARNING: {admm_subproblem_name} has no consensus_vars")
             consensus_vars[admm_subproblem_name] = list()
-
-    # Now add the original problem's first-stage variables (factory
-    # production decisions) to each subproblem's consensus list.
-    # Read them directly from the first_stage_varlist hook —
-    # scenario_creator no longer attaches the root node, so
-    # model._mpisppy_node_list does not exist at this point.
-    # The original first stage is stage 1 in the augmented tree.
-    for admm_subproblem_name in admm_subproblem_names:
-        admm_stoch_subproblem_scenario_name = combining_names(admm_subproblem_name,stoch_scenario_name)
-        model = scenario_creator(admm_stoch_subproblem_scenario_name, inter_region_dict=inter_region_dict, cfg=cfg, data_params=data_params, all_nodes_dict=all_nodes_dict)
-        for var in first_stage_varlist(model):
-            if (var.name, 1) not in consensus_vars[admm_subproblem_name]:
-                consensus_vars[admm_subproblem_name].append((var.name, 1))
+    # Each ADMM subproblem's first-stage Vars are added to its
+    # consensus list by the wrapper (Stoch_AdmmWrapper / AdmmBundler)
+    # at construction time, driven by the first_stage_varlist hook.
     return consensus_vars
 
 

mpisppy/generic/admm.py

@@ -146,6 +146,64 @@ def setup_admm(module, cfg, n_cylinders):
             all_scenario_names, None)
 
 
+def _discover_first_stage_hooks(module):
+    """Discover the four optional first-stage hooks on a stoch-admm
+    model module and validate the both-or-neither contract.
+
+    The two core hooks (first_stage_cost, first_stage_varlist) must
+    be defined together or both omitted; mixing them is an error.
+
+    The two advanced hooks (first_stage_surrogate_nonant_list,
+    first_stage_nonant_ef_suppl_list) are each independent of the
+    other; defining either alone is fine, but only when the two
+    core hooks are also defined -- they forward to
+    sputils.attach_root_node's optional surrogate_nonant_list /
+    nonant_ef_suppl_list parameters and have nothing to attach to
+    otherwise.
+
+    Returns:
+        dict: keyword arguments suitable for forwarding to
+        Stoch_AdmmWrapper / AdmmBundler.  Hooks that the module did
+        not define are passed as None.
+    """
+    first_stage_cost = getattr(module, "first_stage_cost", None)
+    first_stage_varlist = getattr(module, "first_stage_varlist", None)
+    if (first_stage_cost is None) != (first_stage_varlist is None):
+        present = "first_stage_cost" if first_stage_cost is not None else "first_stage_varlist"
+        missing = "first_stage_varlist" if first_stage_cost is not None else "first_stage_cost"
+        raise RuntimeError(
+            f"Module {module.__name__!r} defines {present} but not "
+            f"{missing}.  These hooks must be defined together "
+            f"(or both omitted).  See doc/src/generic_admm.rst."
+        )
+
+    first_stage_surrogate_nonant_list = getattr(
+        module, "first_stage_surrogate_nonant_list", None)
+    first_stage_nonant_ef_suppl_list = getattr(
+        module, "first_stage_nonant_ef_suppl_list", None)
+    advanced = {
+        "first_stage_surrogate_nonant_list": first_stage_surrogate_nonant_list,
+        "first_stage_nonant_ef_suppl_list": first_stage_nonant_ef_suppl_list,
+    }
+    present_advanced = [n for n, h in advanced.items() if h is not None]
+    if present_advanced and first_stage_cost is None:
+        raise RuntimeError(
+            f"Module {module.__name__!r} defines the advanced hook(s) "
+            f"{present_advanced} but not first_stage_cost / "
+            f"first_stage_varlist.  Advanced hooks forward to "
+            f"sputils.attach_root_node's optional parameters and only "
+            f"make sense when the core hooks are also defined.  See "
+            f"doc/src/generic_admm.rst."
+        )
+
+    return {
+        "first_stage_cost": first_stage_cost,
+        "first_stage_varlist": first_stage_varlist,
+        "first_stage_surrogate_nonant_list": first_stage_surrogate_nonant_list,
+        "first_stage_nonant_ef_suppl_list": first_stage_nonant_ef_suppl_list,
+    }
+
+
 def setup_stoch_admm(module, cfg, n_cylinders):
     """Create Stoch_AdmmWrapper for stochastic ADMM.
 
@@ -165,19 +223,7 @@ def setup_stoch_admm(module, cfg, n_cylinders):
     consensus_vars = module.consensus_vars_creator(
         admm_subproblem_names, stoch_scenario_name, **scenario_creator_kwargs)
 
-    # Discover optional first-stage hooks on the module.  Both must be
-    # defined together or both omitted; mixing produces a clear error
-    # here (rather than half-migrating silently).
-    first_stage_cost = getattr(module, "first_stage_cost", None)
-    first_stage_varlist = getattr(module, "first_stage_varlist", None)
-    if (first_stage_cost is None) != (first_stage_varlist is None):
-        present = "first_stage_cost" if first_stage_cost is not None else "first_stage_varlist"
-        missing = "first_stage_varlist" if first_stage_cost is not None else "first_stage_cost"
-        raise RuntimeError(
-            f"Module {module.__name__!r} defines {present} but not "
-            f"{missing}.  These hooks must be defined together "
-            f"(or both omitted).  See doc/src/generic_admm.rst."
-        )
+    first_stage_hooks = _discover_first_stage_hooks(module)
 
     admm = Stoch_AdmmWrapper(
         options={},
@@ -191,8 +237,7 @@ def setup_stoch_admm(module, cfg, n_cylinders):
         mpicomm=MPI.COMM_WORLD,
         scenario_creator_kwargs=scenario_creator_kwargs,
         BFs=cfg.get("branching_factors"),
-        first_stage_cost=first_stage_cost,
-        first_stage_varlist=first_stage_varlist,
+        **first_stage_hooks,
     )
 
     # Store on cfg as plain attributes (Pyomo Config can't handle these types)
@@ -232,18 +277,7 @@ def setup_stoch_admm_with_bundles(module, cfg, n_cylinders):
     consensus_vars = module.consensus_vars_creator(
         admm_subproblem_names, stoch_scenario_name, **scenario_creator_kwargs)
 
-    # Discover optional first-stage hooks on the module.  Same
-    # both-or-neither contract as the non-bundled path.
-    first_stage_cost = getattr(module, "first_stage_cost", None)
-    first_stage_varlist = getattr(module, "first_stage_varlist", None)
-    if (first_stage_cost is None) != (first_stage_varlist is None):
-        present = "first_stage_cost" if first_stage_cost is not None else "first_stage_varlist"
-        missing = "first_stage_varlist" if first_stage_cost is not None else "first_stage_cost"
-        raise RuntimeError(
-            f"Module {module.__name__!r} defines {present} but not "
-            f"{missing}.  These hooks must be defined together "
-            f"(or both omitted).  See doc/src/generic_admm.rst."
-        )
+    first_stage_hooks = _discover_first_stage_hooks(module)
 
     bundler = AdmmBundler(
         module=module,
@@ -254,8 +288,7 @@ def setup_stoch_admm_with_bundles(module, cfg, n_cylinders):
         combining_fn=module.combining_names,
         split_fn=module.split_admm_stoch_subproblem_scenario_name,
         scenario_creator_kwargs=scenario_creator_kwargs,
-        first_stage_cost=first_stage_cost,
-        first_stage_varlist=first_stage_varlist,
+        **first_stage_hooks,
     )
     bundle_names = bundler.bundle_names_creator()