diff --git a/doc/designs/admm_user_api_automation_design.md b/doc/designs/admm_user_api_automation_design.md new file mode 100644 index 000000000..2f528a543 --- /dev/null +++ b/doc/designs/admm_user_api_automation_design.md @@ -0,0 +1,359 @@ +# ADMM user-API automation — design + +Status: draft, awaiting review. + +Related: ongoing ADMM doc/usability work on the `admm-doc-crossref` +branch (cross-references to `generic_cylinders`, BFs flow-through, +xhatshuffle stage2ef hard error, `attach_root_node` contract +clarification). + +--- + +## 0. Goals + +Reduce the boilerplate and foot-guns a user must navigate when writing +an ADMM model module for `generic_cylinders --admm` or `--stoch-admm`. +Today the user has to: + +1. Call `sputils.attach_root_node` in their `scenario_creator` — but + only for `--stoch-admm`; for `--admm` the wrapper overwrites the + call. Asymmetric and easy to get wrong (#4 in our recent chat). +2. Stringly-type variable names in `consensus_vars_creator` so they + exactly match `var.name`; typos / indexed-var formatting differences + fail with cryptic errors. +3. Hand-write a `combining_names` / `split_admm_stoch_subproblem_scenario_name` + inverse pair, with a documented underscore-collision warning. +4. Write `admm_stoch_subproblem_scenario_names_creator` and remember to + nest the loops correctly (outer = stoch, inner = admm) so MPI rank + assignment works. +5. Register `num_admm_subproblems` / `num_stoch_scens` in `inparser_adder` + (already partially automated, but examples still register them). + +Non-goals: + +- Merging `AdmmWrapper` and `Stoch_AdmmWrapper` (item F; explicitly + out of scope). +- Performance changes. +- Multistage-origin stoch-admm. Phase A targets 2-stage origin only. + Multistage origin is not currently exercised by any example or test + (every caller passes `BFs=None`); extending the automation there + belongs in a follow-up if/when a real multistage-origin model + appears. + +--- + +## 1. Phased plan + +Five phases (A–E), each independently shippable, each its own +review-sized PR. Each PR must: + +- Keep the existing string-typed / hand-written API working (the old + path is the back-compat fallback). +- Migrate the canonical example (`examples/distr/distr.py` for `--admm`, + `examples/stoch_distr/stoch_distr.py` for `--stoch-admm`) to the + new API so the doc has something to point at. +- Add a unit test exercising both the old and the new path. +- Update `doc/src/generic_admm.rst`. + +--- + +## 2. Phase A — auto-build the scenario tree for `--stoch-admm` + +### Today's contract + +- `--admm`: `AdmmWrapper.assign_variable_probs` calls + `sputils.attach_root_node(s, objfunc, varlist)` at + `mpisppy/utils/admmWrapper.py:148`, overwriting any prior call. User + must not call it; if they do, it's silently a no-op. +- `--stoch-admm`: `Stoch_AdmmWrapper` reads + `s._mpisppy_node_list[-1]` at `mpisppy/utils/stoch_admmWrapper.py:236` + and appends an ADMM-consensus node. User MUST call `attach_root_node` + with the original problem's first-stage cost + varlist or the + wrapper assertion-fails. + +### Proposed + +Add two optional module-level functions for `--stoch-admm`: + +```python +def first_stage_cost(scenario): + """Returns the Pyomo Expression for the original problem's + first-stage cost (e.g., scenario.FirstStageCost).""" + +def first_stage_varlist(scenario): + """Returns the list of first-stage Pyomo Vars / VarDatas (e.g., + [scenario.y[n] for n in factory_nodes]).""" +``` + +**Strict both-or-neither contract.** `setup_stoch_admm` inspects the +module for both attributes. If exactly one is present, raise +`RuntimeError` naming the missing function. This catches users who +half-migrate. + +When both are defined, `Stoch_AdmmWrapper` calls +`sputils.attach_root_node` itself for each local scenario *before* +`assign_variable_probs` runs (so the existing append-an-ADMM-stage +logic still works on top of the wrapper-attached node). + +**Per-scenario error matrix** (checked after `scenario_creator` +returns, before the wrapper appends the ADMM stage): + +| | Both hooks defined | Neither hook defined | +|---|---|---| +| `s._mpisppy_node_list` is set (user called `attach_root_node`) | **`RuntimeError`** — "hooks make `attach_root_node` redundant; remove the call from `scenario_creator`" | OK — legacy path; wrapper appends ADMM stage on top of the user's node list | +| `s._mpisppy_node_list` is **not** set | OK — wrapper attaches root via the hooks | **`RuntimeError`** — "scenario is missing `_mpisppy_node_list`. Either call `sputils.attach_root_node` in `scenario_creator`, or define `first_stage_cost` and `first_stage_varlist` module functions (recommended)" | + +The check fires per-scenario, not just once, so inconsistent branching +in `scenario_creator` (e.g., calling `attach_root_node` only on some +scenario names) is also caught. Both errors fire at +scenario-construction time, before `assign_variable_probs`, so users +get a clear message instead of a cryptic deep assertion. + +The "both hooks + user call → error" choice is deliberately strict: +the wrapper's call would silently overwrite the user's, but erroring +instead teaches the user that the hooks obviate manual attachment. + +### Files + +- `mpisppy/utils/stoch_admmWrapper.py`: in `__init__`, after + `scenario_creator` returns, apply the error matrix and (if hooks + supplied) call `attach_root_node`. Add `first_stage_cost` and + `first_stage_varlist` parameters with both-or-neither validation. +- `mpisppy/utils/admm_bundler.py`: same hook plumbing for the + bundled path (`AdmmBundler` constructs scenarios in its own + `scenario_creator` method and runs the same error matrix + per-scenario). This was missed in the original design and only + surfaced when the bundled-admm tests started failing during + implementation. +- `mpisppy/generic/admm.py:setup_stoch_admm`: pass module hooks + through. Both-or-neither check at module level (so the error + message can name the user's module). +- `mpisppy/generic/admm.py:setup_stoch_admm_with_bundles`: same + hook-discovery and forwarding as the non-bundled path. +- `examples/stoch_distr/stoch_distr.py`: migrate. Stash the + first-stage varlist on the model in `scenario_creator` (so the + hook can find it); remove the `attach_root_node` call; add the + two new hook functions. **Also refactor + `consensus_vars_creator`**: it used to walk + `model._mpisppy_node_list` to derive first-stage nonants, which + now doesn't exist at that call site (the wrapper hasn't attached + the root yet). Read first-stage vars via `first_stage_varlist()` + directly instead. +- `examples/stoch_distr/stoch_distr_admm_cylinders.py`: the custom + driver constructs `Stoch_AdmmWrapper` directly; update it to + forward the module's hooks so `stoch_distr_ef.py` (which goes + through this driver) continues to work after the example + migration. +- `doc/src/generic_admm.rst`: rewrite the "Extending to Stochastic + ADMM" note around the new API; preserve the old `attach_root_node` + call as a "legacy path" subsection. +- `mpisppy/tests/test_stoch_admmWrapper.py`: a new test class + exercising both paths (hooks + legacy) and all four error-matrix + cells, plus the `setup_stoch_admm`-level half-hooks error. + +### Resolved + +- **Hook signature: scenario form** (`first_stage_cost(scenario)` / + `first_stage_varlist(scenario)`). The function gets the constructed + scenario and reads attributes directly. Phase B will fix the + parallel string-vs-Var issue in `consensus_vars_creator`. +- Naming: `first_stage_cost` / `first_stage_varlist`, matching + `attach_root_node`'s parameter names (`firstobj`, `varlist`). + +--- + +## 3. Phase B — accept Pyomo Var/VarData in `consensus_vars_creator` + +### Today's contract + +- `--admm`: `consensus_vars_creator` returns + `{subproblem_name: [varname_str, ...]}`. Wrapper calls + `s.find_component(varname_str)`. +- `--stoch-admm`: returns `{subproblem_name: [(varname_str, stage), ...]}`. + +Mismatches (e.g., `"flow[(DC1, DC2)]"` vs `"flow['DC1', 'DC2']"`) fail +with a list-of-pairs RuntimeError deep in `assign_variable_probs`. + +### Proposed + +Accept Pyomo Var / VarData objects *or* strings in the lists. At wrapper +construction, normalize each item to a string via `obj.name` before the +existing `find_component` path. Mixed lists OK (e.g., for migration). + +Caveat: `consensus_vars_creator` for `--stoch-admm` is called *once* +with one stochastic scenario. Vars passed in are tied to that +scenario instance; the wrapper still has to look them up by name in +the other scenarios via `find_component`. The user benefit is "no +manual name formatting", not "skip the find_component step". + +### Files + +- `mpisppy/utils/admmWrapper.py:assign_variable_probs`: pre-normalize + the lists. +- `mpisppy/utils/stoch_admmWrapper.py:assign_variable_probs`: same. +- Examples: migrate `consensus_vars_creator` in `distr.py` and + `stoch_distr.py` to return Vars. +- Doc. +- Test: a unit test in `test_admm_bundler.py` or a new file that + constructs a small model and verifies a Var-based + `consensus_vars_creator` produces the same `varprob_dict` as the + string-based form. + +### Open questions + +None significant. + +--- + +## 4. Phase C — default `combining_names` / `split_admm_stoch_subproblem_scenario_name` + +### Today's contract + +User-defined inverse pair, e.g. from `stoch_distr.py`: + +```python +ADMM_STOCH_PREFIX = "ADMM_STOCH_" + +def combining_names(admm_sub, stoch_scen): + return f"{ADMM_STOCH_PREFIX}{admm_sub}_{stoch_scen}" + +def split_admm_stoch_subproblem_scenario_name(name): + parts = name.split("_") + return parts[2], parts[3] +``` + +`split_...` is fragile if any subproblem or stoch-scenario name +contains an underscore (today's doc has a Warning about this). + +### Proposed + +Provide a default pair in `mpisppy.utils.stoch_admmWrapper`: + +```python +_DEFAULT_DELIM = "__ADMM__" # verbose but unambiguous; safe for shells, + # filenames, and CSV writers + +def default_combining_names(admm_sub, stoch_scen): + return f"ADMM_STOCH{_DEFAULT_DELIM}{admm_sub}{_DEFAULT_DELIM}{stoch_scen}" + +def default_split_admm_stoch_subproblem_scenario_name(name): + parts = name.split(_DEFAULT_DELIM) + return parts[1], parts[2] +``` + +In `setup_stoch_admm`, if the module does not define `combining_names` +*and* `split_admm_stoch_subproblem_scenario_name`, use the defaults. +If the module defines one, it must define the other (raise a clear +error otherwise — we cannot pair a default split with a custom combine +or vice versa). + +### Files + +- `mpisppy/utils/stoch_admmWrapper.py`: add defaults; export them. +- `mpisppy/generic/admm.py:setup_stoch_admm`: fallback logic. +- Examples: optionally drop the user-defined pair to demonstrate the + default; leave at least one example using a custom pair to show how. +- Doc: move the underscore-collision warning into a "Customizing + the naming convention" subsection; default-path users no longer + hit it. +- Test: unit test that with no module-defined pair, the wrapper builds + the same scenarios via the default pair as via the explicit pair. + +### Resolved + +- Default delimiter: `__ADMM__` (verbose but unambiguous; safe for + shells, filenames, and CSV writers). + +--- + +## 5. Phase D — default `admm_stoch_subproblem_scenario_names_creator` + +### Today's contract + +User function with required nesting order: + +```python +def admm_stoch_subproblem_scenario_names_creator(admm_subproblem_names, + stoch_scenario_names): + return [combining_names(sub, stoch) + for stoch in stoch_scenario_names # outer + for sub in admm_subproblem_names] # inner +``` + +If the user nests them backwards, MPI rank assignment goes wrong. + +### Proposed + +If the module does not define this function, use a default in +`mpisppy.utils.stoch_admmWrapper` that uses the module's +`combining_names` (or the phase-C default). + +### Files + +- `mpisppy/utils/stoch_admmWrapper.py`: add default. +- `mpisppy/generic/admm.py:setup_stoch_admm`: fallback logic. +- Examples: drop the user-defined version where appropriate. +- Doc. +- Test. + +### Open questions + +None. + +--- + +## 6. Phase E — `inparser_adder` boilerplate + +### Today's state + +`mpisppy/generic/admm.py:admm_args` already registers `admm`, +`stoch_admm`, `num_admm_subproblems`, `num_stoch_scens`, +`branching_factors`, `stage2_ef_solver_name`, with `if not already +present` guards. Example modules' `inparser_adder` typically also +register `num_admm_subproblems` / `num_stoch_scens`; the guards prevent +double-registration. + +### Proposed + +- Document `admm_args` as the canonical source for ADMM-specific args. +- Remove the redundant registrations from example modules. +- Optional: emit a one-time `DeprecationWarning` when a model module + redundantly registers an ADMM arg. (Maybe not — silent skip is + fine.) + +### Files + +- Examples: drop registrations. +- Doc: update "Standard functions" / "Additional functions" lists. + +### Open questions + +Whether the DeprecationWarning is worth the noise. Default: skip it. + +--- + +## 7. Roll-out + +After this design is approved: + +1. Phase A as its own branch + PR off main. +2. After A merges, phase B branches off main (or off A if D depends on + it — no expected dependency). +3. ...etc. + +Each phase's PR title prefixes `stoch-admm:` (matching recent commits) +or `admm:`. Each ships with green tests including the new automation +tests. + +--- + +## 8. Risks and rollback + +- Each phase preserves the old API as a fallback path; users with + existing model modules are unaffected. +- If a phase reveals a deeper issue (e.g., phase A interacting with + bundling), that phase can be reverted on its own without affecting + earlier phases. +- The doc note added in commit `24bf5802` ("clarify attach_root_node + contract differs for --admm vs --stoch-admm") will become obsolete + when phase A lands; phase A's PR explicitly supersedes it. diff --git a/doc/src/generic_admm.rst b/doc/src/generic_admm.rst index e85059a6c..b17784bde 100644 --- a/doc/src/generic_admm.rst +++ b/doc/src/generic_admm.rst @@ -379,27 +379,70 @@ To support ``--stoch-admm``, additionally implement: These functions define the naming convention for composite scenarios. See ``examples/stoch_distr/stoch_distr.py`` for a complete working example. -Your ``consensus_vars_creator`` returns ``(variable_name, stage)`` tuples -instead of plain strings. +.. Note:: + + ``scenario_creator`` for ``--stoch-admm`` differs from the + deterministic case in one way: + + **It receives a composite name.** The argument is e.g. + ``"ADMM_STOCH_Region1_StochasticScenario3"``; ``scenario_creator`` + must call ``split_admm_stoch_subproblem_scenario_name`` on it to + recover the ADMM subproblem name and the stochastic scenario name, + then build the corresponding model. + +First-stage attachment via module hooks (recommended) +"""""""""""""""""""""""""""""""""""""""""""""""""""""" + +Under the hood, ``Stoch_AdmmWrapper`` reads the user-supplied +``_mpisppy_node_list`` and *appends* an ADMM-consensus stage to it +(whereas ``AdmmWrapper`` overwrites the node list). The wrapper can +attach the root node for you if you provide two module-level hook +functions: + +.. code-block:: python + + def first_stage_cost(scenario): + """Original problem's first-stage cost expression.""" + return scenario.FirstStageCost + + def first_stage_varlist(scenario): + """Original problem's first-stage variables (NOT ADMM consensus vars).""" + return scenario._first_stage_vars # stashed in scenario_creator + +When both hooks (``first_stage_cost`` and ``first_stage_varlist``) are +defined on the module, the wrapper calls +``sputils.attach_root_node(scenario, first_stage_cost(scenario), +first_stage_varlist(scenario))`` itself for each scenario before +running its consensus-stage logic. ``scenario_creator`` no longer +needs to call ``attach_root_node`` (and must not — see error matrix +below). + +See ``examples/stoch_distr/stoch_distr.py`` for the canonical +pattern, including how to stash the varlist on the scenario from +inside ``scenario_creator`` so the hook can find it. .. Note:: + The hooks are **both-or-neither**: defining only one raises + ``RuntimeError`` at ``setup_stoch_admm`` time. Mixing the hooks + with a manual ``attach_root_node`` call also raises. + +First-stage attachment via manual ``attach_root_node`` (legacy) +""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""" - ``scenario_creator`` differs from the deterministic case in two ways: +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. - 1. **Receives a composite name.** The argument is e.g. - ``"ADMM_STOCH_Region1_StochasticScenario3"``; the function must split - it (using ``split_admm_stoch_subproblem_scenario_name``) to determine - both which ADMM subproblem and which stochastic scenario to build. +This path is preserved for backward compatibility with model modules +written before the hooks existed (and for direct uses of +``Stoch_AdmmWrapper`` that bypass ``setup_stoch_admm``). - 2. **Must call** ``sputils.attach_root_node`` with the *original* - problem's first-stage variables (i.e., not the ADMM consensus vars). - Unlike ``AdmmWrapper`` (which overwrites the node list), - ``Stoch_AdmmWrapper`` *reads* the user-supplied node list and - *appends* an ADMM-consensus stage to it. Skipping the call will - assertion-fail inside the wrapper. +Consensus vars +"""""""""""""" - See ``stoch_distr.scenario_creator`` for the canonical pattern: - ``sputils.attach_root_node(model, model.FirstStageCost, [...factory vars...])``. +Your ``consensus_vars_creator`` returns ``(variable_name, stage)`` +tuples instead of plain strings. .. _admm_bundling: diff --git a/examples/stoch_distr/stoch_distr.py b/examples/stoch_distr/stoch_distr.py index df2644f83..011aee494 100644 --- a/examples/stoch_distr/stoch_distr.py +++ b/examples/stoch_distr/stoch_distr.py @@ -175,11 +175,31 @@ def scenario_creator(admm_stoch_subproblem_scenario_name, inter_region_dict=None # Generating the model model = min_cost_distr_problem(local_dict, cfg, stoch_scenario_name, max_revenue=data_params["max revenue"]) - sputils.attach_root_node(model, model.FirstStageCost, [model.y[n] for n in local_dict["factory nodes"]]) - + # Stash the original first-stage variable list so the module-level + # first_stage_varlist hook can retrieve it. Stoch_AdmmWrapper will + # call sputils.attach_root_node on our behalf using the hooks. + model._first_stage_vars = [model.y[n] for n in local_dict["factory nodes"]] + return model +def first_stage_cost(scenario): + """Return the original problem's first-stage cost expression. + + Used by Stoch_AdmmWrapper to attach the root node automatically; + the user no longer needs to call sputils.attach_root_node in + scenario_creator. See doc/src/generic_admm.rst. + """ + return scenario.FirstStageCost + + +def first_stage_varlist(scenario): + """Return the original problem's first-stage variables (factory + production decisions for stoch_distr). Companion to first_stage_cost. + """ + return scenario._first_stage_vars + + def scenario_denouement(rank, admm_stoch_subproblem_scenario_name, scenario, eps=10**(-6)): """For each admm stochastic scenario subproblem prints its name and the final variable values @@ -249,16 +269,18 @@ def consensus_vars_creator(admm_subproblem_names, stoch_scenario_name, inter_reg print(f"WARNING: {admm_subproblem_name} has no consensus_vars") consensus_vars[admm_subproblem_name] = list() - # now add the parents. It doesn't depend on the stochastic scenario so we chose one and - # then we go through the models (created by scenario creator) for all the admm_stoch_subproblem_scenario - # which have this scenario as an ancestor (parent) in the tree + # 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 node in model._mpisppy_node_list: - for var in node.nonant_list: - if var.name not in consensus_vars[admm_subproblem_name]: - consensus_vars[admm_subproblem_name].append((var.name, node.stage)) + 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)) return consensus_vars diff --git a/examples/stoch_distr/stoch_distr_admm_cylinders.py b/examples/stoch_distr/stoch_distr_admm_cylinders.py index 952a00dd9..733cc1c2b 100644 --- a/examples/stoch_distr/stoch_distr_admm_cylinders.py +++ b/examples/stoch_distr/stoch_distr_admm_cylinders.py @@ -72,6 +72,9 @@ def _make_admm(cfg, n_cylinders, verbose=None): scenario_creator_kwargs = stoch_distr.kw_creator(cfg) stoch_scenario_name = stoch_scenario_names[0] # choice of any scenario consensus_vars = stoch_distr.consensus_vars_creator(admm_subproblem_names, stoch_scenario_name, **scenario_creator_kwargs) + # stoch_distr exposes first_stage_cost / first_stage_varlist + # module-level hooks so Stoch_AdmmWrapper attaches the root node + # itself; pass them through here too. admm = stoch_admmWrapper.Stoch_AdmmWrapper(options, all_admm_stoch_subproblem_scenario_names, split_admm_stoch_subproblem_scenario_name, @@ -84,6 +87,8 @@ def _make_admm(cfg, n_cylinders, verbose=None): scenario_creator_kwargs=scenario_creator_kwargs, verbose=verbose, BFs=None, + first_stage_cost=stoch_distr.first_stage_cost, + first_stage_varlist=stoch_distr.first_stage_varlist, ) return admm, all_admm_stoch_subproblem_scenario_names diff --git a/mpisppy/generic/admm.py b/mpisppy/generic/admm.py index 9c875f782..271abcaa3 100644 --- a/mpisppy/generic/admm.py +++ b/mpisppy/generic/admm.py @@ -165,6 +165,20 @@ 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." + ) + admm = Stoch_AdmmWrapper( options={}, all_admm_stoch_subproblem_scenario_names=all_names, @@ -177,6 +191,8 @@ 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, ) # Store on cfg as plain attributes (Pyomo Config can't handle these types) @@ -216,6 +232,19 @@ 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." + ) + bundler = AdmmBundler( module=module, scenarios_per_bundle=cfg.scenarios_per_bundle, @@ -225,6 +254,8 @@ 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, ) bundle_names = bundler.bundle_names_creator() diff --git a/mpisppy/tests/test_admm_bundler.py b/mpisppy/tests/test_admm_bundler.py index c0c91bfca..dc9986610 100644 --- a/mpisppy/tests/test_admm_bundler.py +++ b/mpisppy/tests/test_admm_bundler.py @@ -68,6 +68,38 @@ def test_bundle_names_three_subproblems(self): # 3 subproblems × 1 bundle each = 3 bundles self.assertEqual(len(names), 3) + def test_half_hooks_errors(self): + """Exactly one of first_stage_cost / first_stage_varlist passed + to AdmmBundler — both-or-neither contract violated.""" + cfg = config.Config() + stoch_distr.inparser_adder(cfg) + cfg.num_stoch_scens = 4 + cfg.num_admm_subproblems = 2 + admm_subproblem_names = stoch_distr.admm_subproblem_names_creator(cfg) + stoch_scenario_names = stoch_distr.stoch_scenario_names_creator(cfg) + scenario_creator_kwargs = stoch_distr.kw_creator(cfg) + consensus_vars = stoch_distr.consensus_vars_creator( + admm_subproblem_names, stoch_scenario_names[0], + **scenario_creator_kwargs) + + common = dict( + module=stoch_distr, + scenarios_per_bundle=4, + admm_subproblem_names=admm_subproblem_names, + stoch_scenario_names=stoch_scenario_names, + consensus_vars=consensus_vars, + combining_fn=stoch_distr.combining_names, + split_fn=stoch_distr.split_admm_stoch_subproblem_scenario_name, + scenario_creator_kwargs=scenario_creator_kwargs, + ) + for half in ( + {"first_stage_cost": lambda s: 0}, + {"first_stage_varlist": lambda s: []}, + ): + with self.assertRaises(RuntimeError) as cm: + AdmmBundler(**common, **half) + self.assertIn("must be defined together", str(cm.exception)) + @unittest.skipUnless(solver_available, "no solver available") def test_bundle_creation(self): """Test that a bundle can be created and has correct structure.""" diff --git a/mpisppy/tests/test_stoch_admmWrapper.py b/mpisppy/tests/test_stoch_admmWrapper.py index 9e4954c3c..0d0f6d84a 100644 --- a/mpisppy/tests/test_stoch_admmWrapper.py +++ b/mpisppy/tests/test_stoch_admmWrapper.py @@ -547,5 +547,275 @@ def test_ef_partial_consensus(self): os.chdir(original_dir) +class TestStochAdmmWrapperFirstStageHooks(unittest.TestCase): + """first_stage_cost / first_stage_varlist hooks on the + Stoch_AdmmWrapper API. + + When the hooks are supplied, the wrapper calls + sputils.attach_root_node itself for each scenario, and + scenario_creator must NOT call it. When the hooks are absent + (legacy path), scenario_creator must call attach_root_node + itself. Mixing the two (hooks + manual call, or no hooks + no + call) raises with a clear message at scenario-construction time. + """ + + @staticmethod + def _minimal_scenario_creator(call_attach): + """Return a scenario_creator that either does or does not call + sputils.attach_root_node, otherwise identical.""" + import pyomo.environ as pyo + import mpisppy.utils.sputils as sputils + + def scenario_creator(sname, **kwargs): + parts = sname.split("_") + admm_part = parts[2] + m = pyo.ConcreteModel() + # consensus var owned per admm subproblem + if admm_part == "A": + m.x = pyo.Var(bounds=(0, 1)) + own = m.x + else: + m.y = pyo.Var(bounds=(0, 1)) + own = m.y + m.fs = pyo.Var(bounds=(0, 1)) # original first-stage var + m.FirstStageCost = pyo.Expression(expr=m.fs) + m.obj = pyo.Objective(expr=own + m.fs, sense=pyo.minimize) + # Stash for the hook to find (mirrors examples/stoch_distr + # pattern): + m._first_stage_vars = [m.fs] + if call_attach: + sputils.attach_root_node(m, m.FirstStageCost, [m.fs]) + m._mpisppy_probability = "uniform" + return m + + return scenario_creator + + @staticmethod + def _hooks(): + def first_stage_cost(s): + return s.FirstStageCost + + def first_stage_varlist(s): + return s._first_stage_vars + + return first_stage_cost, first_stage_varlist + + @staticmethod + def _common_kwargs(): + admm_names = ["A", "B"] + stoch_names = ["S1", "S2"] + all_names = [f"ADMM_STOCH_{a}_{s}" + for s in stoch_names for a in admm_names] + + def split(sname): + parts = sname.split("_") + return parts[2], "_".join(parts[3:]) + + consensus_vars = {"A": [("x", 1)], "B": [("y", 1)]} + return { + "all_admm_stoch_subproblem_scenario_names": all_names, + "split_admm_stoch_subproblem_scenario_name": split, + "admm_subproblem_names": admm_names, + "stoch_scenario_names": stoch_names, + "consensus_vars": consensus_vars, + "n_cylinders": 1, + "scenario_creator_kwargs": {}, + } + + def test_hooks_path_attaches_root_node(self): + """With hooks defined and scenario_creator NOT calling + attach_root_node, the wrapper must attach the root itself and + produce a valid _mpisppy_node_list with the consensus stage + appended.""" + from mpisppy.utils.stoch_admmWrapper import Stoch_AdmmWrapper + from mpisppy import MPI + + fs_cost, fs_varlist = self._hooks() + admm = Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=False), + mpicomm=MPI.COMM_WORLD, + first_stage_cost=fs_cost, + first_stage_varlist=fs_varlist, + **self._common_kwargs(), + ) + for sname, s in admm.local_admm_stoch_subproblem_scenarios.items(): + self.assertTrue(hasattr(s, "_mpisppy_node_list"), + f"{sname}: node list missing") + # The wrapper appends an ADMM stage; expect 2 nodes for the + # 2-stage-origin case (root + admm-consensus). + self.assertEqual( + len(s._mpisppy_node_list), 2, + f"{sname}: expected 2-node list (root + admm), got " + f"{[n.name for n in s._mpisppy_node_list]}") + self.assertEqual(s._mpisppy_node_list[0].name, "ROOT") + + def test_hooks_and_legacy_paths_produce_same_node_list(self): + """The new hooks path must produce the same final + _mpisppy_node_list as the legacy path (where scenario_creator + called attach_root_node itself).""" + from mpisppy.utils.stoch_admmWrapper import Stoch_AdmmWrapper + from mpisppy import MPI + + fs_cost, fs_varlist = self._hooks() + admm_hooks = Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=False), + mpicomm=MPI.COMM_WORLD, + first_stage_cost=fs_cost, + first_stage_varlist=fs_varlist, + **self._common_kwargs(), + ) + admm_legacy = Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=True), + mpicomm=MPI.COMM_WORLD, + **self._common_kwargs(), + ) + for sname in admm_hooks.local_admm_stoch_subproblem_scenarios: + s_h = admm_hooks.local_admm_stoch_subproblem_scenarios[sname] + s_l = admm_legacy.local_admm_stoch_subproblem_scenarios[sname] + self.assertEqual( + [n.name for n in s_h._mpisppy_node_list], + [n.name for n in s_l._mpisppy_node_list], + f"{sname}: node-name list differs between hooks and legacy") + self.assertEqual( + len(s_h._mpisppy_node_list[0].nonant_vardata_list), + len(s_l._mpisppy_node_list[0].nonant_vardata_list), + f"{sname}: root nonant count differs") + + def test_hooks_plus_manual_attach_errors(self): + """Hooks defined AND scenario_creator also calls + attach_root_node — error so the user knows the hooks make the + manual call redundant (rather than silently overwriting it).""" + from mpisppy.utils.stoch_admmWrapper import Stoch_AdmmWrapper + from mpisppy import MPI + + fs_cost, fs_varlist = self._hooks() + with self.assertRaises(RuntimeError) as cm: + Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=True), + mpicomm=MPI.COMM_WORLD, + first_stage_cost=fs_cost, + first_stage_varlist=fs_varlist, + **self._common_kwargs(), + ) + self.assertIn("must NOT call", str(cm.exception)) + + def test_no_hooks_no_attach_errors(self): + """No hooks AND scenario_creator skips attach_root_node — error + with a message pointing at both remediation options.""" + from mpisppy.utils.stoch_admmWrapper import Stoch_AdmmWrapper + from mpisppy import MPI + + with self.assertRaises(RuntimeError) as cm: + Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=False), + mpicomm=MPI.COMM_WORLD, + **self._common_kwargs(), + ) + msg = str(cm.exception) + self.assertIn("first_stage_cost", msg) + self.assertIn("attach_root_node", msg) + + def test_wrapper_half_hooks_errors(self): + """Exactly one hook passed to the wrapper — both-or-neither + contract violated.""" + from mpisppy.utils.stoch_admmWrapper import Stoch_AdmmWrapper + from mpisppy import MPI + + fs_cost, fs_varlist = self._hooks() + for kw in ( + {"first_stage_cost": fs_cost}, + {"first_stage_varlist": fs_varlist}, + ): + with self.assertRaises(RuntimeError) as cm: + Stoch_AdmmWrapper( + options={}, + scenario_creator=self._minimal_scenario_creator(call_attach=False), + mpicomm=MPI.COMM_WORLD, + **kw, + **self._common_kwargs(), + ) + self.assertIn("must be defined together", str(cm.exception)) + + def test_setup_stoch_admm_half_hooks_errors(self): + """The setup_stoch_admm-level check fires before the wrapper-level + check, naming the user's module in the error message.""" + import types + from mpisppy.generic.admm import setup_stoch_admm + + fs_cost, _ = self._hooks() + + def admm_subproblem_names_creator(cfg): + return ["A", "B"] + + def stoch_scenario_names_creator(cfg): + return ["S1", "S2"] + + def admm_stoch_subproblem_scenario_names_creator(a_names, s_names): + return [f"ADMM_STOCH_{a}_{s}" for s in s_names for a in a_names] + + def split_admm_stoch_subproblem_scenario_name(name): + parts = name.split("_") + return parts[2], "_".join(parts[3:]) + + module = types.SimpleNamespace( + __name__="fake_module", + admm_subproblem_names_creator=admm_subproblem_names_creator, + stoch_scenario_names_creator=stoch_scenario_names_creator, + admm_stoch_subproblem_scenario_names_creator=admm_stoch_subproblem_scenario_names_creator, + split_admm_stoch_subproblem_scenario_name=split_admm_stoch_subproblem_scenario_name, + kw_creator=lambda cfg: {}, + consensus_vars_creator=lambda an, sn, **kw: {"A": [("x", 1)], "B": [("y", 1)]}, + scenario_creator=self._minimal_scenario_creator(call_attach=False), + first_stage_cost=fs_cost, + # first_stage_varlist intentionally missing + ) + cfg = config.Config() + cfg.add_to_config("branching_factors", description="", + domain=list, default=None) + with self.assertRaises(RuntimeError) as cm: + setup_stoch_admm(module, cfg, n_cylinders=1) + msg = str(cm.exception) + self.assertIn("fake_module", msg) + self.assertIn("first_stage_cost", msg) + self.assertIn("first_stage_varlist", msg) + + def test_setup_stoch_admm_with_bundles_half_hooks_errors(self): + """The bundled path enforces the same both-or-neither contract.""" + import types + from mpisppy.generic.admm import setup_stoch_admm_with_bundles + + fs_cost, _ = self._hooks() + + module = types.SimpleNamespace( + __name__="fake_module", + admm_subproblem_names_creator=lambda cfg: ["A", "B"], + stoch_scenario_names_creator=lambda cfg: ["S1", "S2"], + split_admm_stoch_subproblem_scenario_name=( + lambda name: (name.split("_")[2], + "_".join(name.split("_")[3:]))), + kw_creator=lambda cfg: {}, + consensus_vars_creator=( + lambda an, sn, **kw: {"A": [("x", 1)], "B": [("y", 1)]}), + combining_names=lambda a, s: f"ADMM_STOCH_{a}_{s}", + scenario_creator=self._minimal_scenario_creator(call_attach=False), + first_stage_cost=fs_cost, + # first_stage_varlist intentionally missing + ) + cfg = config.Config() + cfg.add_to_config("scenarios_per_bundle", description="", + domain=int, default=2) + with self.assertRaises(RuntimeError) as cm: + setup_stoch_admm_with_bundles(module, cfg, n_cylinders=1) + msg = str(cm.exception) + self.assertIn("fake_module", msg) + self.assertIn("first_stage_cost", msg) + self.assertIn("first_stage_varlist", msg) + + if __name__ == '__main__': unittest.main() diff --git a/mpisppy/utils/admm_bundler.py b/mpisppy/utils/admm_bundler.py index f33d3445e..e9b583cdd 100644 --- a/mpisppy/utils/admm_bundler.py +++ b/mpisppy/utils/admm_bundler.py @@ -54,7 +54,16 @@ class AdmmBundler: def __init__(self, module, scenarios_per_bundle, admm_subproblem_names, stoch_scenario_names, consensus_vars, combining_fn, split_fn, - scenario_creator_kwargs=None): + scenario_creator_kwargs=None, + first_stage_cost=None, first_stage_varlist=None): + # Same both-or-neither contract as Stoch_AdmmWrapper. + 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"AdmmBundler was given {present} but not {missing}. " + f"These hooks must be defined together (or both omitted)." + ) self.module = module self.scenarios_per_bundle = scenarios_per_bundle self.admm_subproblem_names = admm_subproblem_names @@ -63,6 +72,8 @@ def __init__(self, module, scenarios_per_bundle, self.combining_fn = combining_fn self.split_fn = split_fn self.scenario_creator_kwargs = scenario_creator_kwargs or {} + self.first_stage_cost = first_stage_cost + self.first_stage_varlist = first_stage_varlist self.number_admm_subproblems = len(admm_subproblem_names) self.consensus_vars_number = _consensus_vars_number_creator(consensus_vars) @@ -209,9 +220,34 @@ def scenario_creator(self, bundle_name, **kwargs): scen_dict = {} varprob_by_scenario = {} + has_first_stage_hooks = self.first_stage_cost is not None for sub_name, stoch_name in constituents: vsname = self.combining_fn(sub_name, stoch_name) s = self.module.scenario_creator(vsname, **self.scenario_creator_kwargs) + # Error matrix, same as Stoch_AdmmWrapper. See + # doc/src/generic_admm.rst. + already_attached = hasattr(s, "_mpisppy_node_list") + if has_first_stage_hooks: + if already_attached: + raise RuntimeError( + f"scenario {vsname!r}: first_stage_cost and " + f"first_stage_varlist hooks are defined, so " + f"scenario_creator must NOT call " + f"sputils.attach_root_node. Remove the " + f"attach_root_node call from scenario_creator." + ) + sputils.attach_root_node( + s, self.first_stage_cost(s), self.first_stage_varlist(s)) + else: + if not already_attached: + raise RuntimeError( + f"scenario {vsname!r}: no _mpisppy_node_list is " + f"set. Either call sputils.attach_root_node in " + f"scenario_creator (legacy), or define " + f"first_stage_cost(scenario) and " + f"first_stage_varlist(scenario) module functions " + f"(recommended). See doc/src/generic_admm.rst." + ) varprob = self._process_scenario(vsname, s, sub_name) scen_dict[vsname] = s varprob_by_scenario[vsname] = varprob diff --git a/mpisppy/utils/stoch_admmWrapper.py b/mpisppy/utils/stoch_admmWrapper.py index a3eeadb8a..201a5a767 100644 --- a/mpisppy/utils/stoch_admmWrapper.py +++ b/mpisppy/utils/stoch_admmWrapper.py @@ -92,20 +92,33 @@ def __init__(self, scenario_creator_kwargs=None, verbose=None, BFs=None, + first_stage_cost=None, + first_stage_varlist=None, ): assert len(options) == 0, "no options supported by stoch_admmWrapper" + # first_stage_cost / first_stage_varlist must be defined together + # or omitted together. Defensive check; setup_stoch_admm also + # enforces this, but the wrapper is callable directly. + 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"Stoch_AdmmWrapper was given {present} but not {missing}. " + f"These hooks must be defined together (or both omitted)." + ) + has_first_stage_hooks = first_stage_cost is not None # We need local_scenarios self.local_admm_stoch_subproblem_scenarios = {} scen_tree = sputils._ScenTree(["ROOT"], all_admm_stoch_subproblem_scenario_names) assert mpicomm.Get_size() % n_cylinders == 0, \ f"{mpicomm.Get_size()=} and {n_cylinders=}, but {mpicomm.Get_size() % n_cylinders=} should be 0" ranks_per_cylinder = mpicomm.Get_size() // n_cylinders - + scenario_names_to_rank, _rank_slices, _scenario_slices =\ scen_tree.scen_names_to_ranks(ranks_per_cylinder) cylinder_rank = mpicomm.Get_rank() // n_cylinders - + # taken from spbase self.local_admm_stoch_subproblem_scenarios_names = [ all_admm_stoch_subproblem_scenario_names[i] for i in _rank_slices[cylinder_rank] @@ -113,6 +126,31 @@ def __init__(self, for sname in self.local_admm_stoch_subproblem_scenarios_names: s = scenario_creator(sname, **scenario_creator_kwargs) self.local_admm_stoch_subproblem_scenarios[sname] = s + # Error matrix (hooks defined? x _mpisppy_node_list set?). + # Catch half-migrations at scenario-construction time before + # the deep assertion in assign_variable_probs. + already_attached = hasattr(s, "_mpisppy_node_list") + if has_first_stage_hooks: + if already_attached: + raise RuntimeError( + f"scenario {sname!r}: first_stage_cost and " + f"first_stage_varlist hooks are defined, so " + f"scenario_creator must NOT call " + f"sputils.attach_root_node. Remove the " + f"attach_root_node call from scenario_creator." + ) + sputils.attach_root_node( + s, first_stage_cost(s), first_stage_varlist(s)) + else: + if not already_attached: + raise RuntimeError( + f"scenario {sname!r}: no _mpisppy_node_list is " + f"set. Either call sputils.attach_root_node in " + f"scenario_creator (legacy), or define " + f"first_stage_cost(scenario) and " + f"first_stage_varlist(scenario) module functions " + f"(recommended). See doc/src/generic_admm.rst." + ) # we are not collecting instantiation time self.split_admm_stoch_subproblem_scenario_name = split_admm_stoch_subproblem_scenario_name