L-shaped: fix subproblem indexing when root_scenarios is used

.github/workflows/test_pr_and_main.yml

@@ -1227,6 +1227,7 @@ jobs:
             mpisppy/tests/test_flexible_rank_cli.py \
             mpisppy/tests/test_flex_xhat_assembly.py \
             mpisppy/tests/test_ciutils.py \
+            mpisppy/tests/test_lshaped_cuts.py \
             mpisppy/tests/test_prox_approx.py \
             mpisppy/tests/test_sep_rho.py \
             mpisppy/tests/test_reduced_costs_fixer.py \

mpisppy/tests/straight_tests.py

@@ -183,6 +183,28 @@ def _doone(cmdstr: str) -> bool:
 _doone(cmdstr)
 
 
+#####################################################
+# L-shaped with a multi-rank hub (regression test for issue #551).
+# farmer_lshapedhub.py builds the middle scenario into the root problem
+# (root_scenarios), so with one scenario per rank this exercises the
+# subproblem index bookkeeping in mpisppy.utils.lshaped_cuts.set_ls.
+lshaped_path = os.path.abspath(
+    os.path.join(_tests_dir, "..", "..", "examples", "farmer", "farmer_lshapedhub.py")
+)
+
+cmdstr = (
+    f"mpiexec -np 3 {pyexe} {python_args} -m mpi4py {shlex.quote(lshaped_path)} "
+    f"--num-scens 3 "
+    f"--bundles-per-rank=0 "
+    f"--max-iterations=50 "
+    f"--solver-name={shlex.quote(solver_name)} "
+    f"--max-solver-threads 1 "
+    f"--rel-gap=0.0"
+)
+
+_doone(cmdstr)
+
+
 #######################################################
 if badguys:
     print("\nstraight_tests.py failed commands:")

mpisppy/tests/test_lshaped_cuts.py

@@ -0,0 +1,115 @@
+###############################################################################
+# 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.
+###############################################################################
+# Serial unit tests for the subproblem index bookkeeping in
+# LShapedCutGeneratorData.set_ls (regression tests for issue #551: scenarios
+# built into the root problem have no eta variable and no Benders subproblem,
+# so the global subproblem indexing must skip them).
+
+import unittest
+
+import pyomo.environ as pyo
+import pyomo.contrib.benders.benders_cuts as bc
+
+from mpisppy.utils.lshaped_cuts import LShapedCutGenerator
+
+
+class _FakeLShaped:
+    """Just enough of LShapedMethod for LShapedCutGeneratorData.set_ls:
+    one MPI rank's view of the scenario assignment and the root model.
+    """
+    def __init__(self, all_scenario_names, local_scenario_names,
+                 root_scenarios=None):
+        self.all_scenario_names = all_scenario_names
+        self.local_scenario_names = local_scenario_names
+        self.has_root_scens = root_scenarios is not None
+        self.root_scenarios = root_scenarios
+        self.root = pyo.ConcreteModel()
+        eta_names = [s for s in all_scenario_names
+                     if root_scenarios is None or s not in root_scenarios]
+        self.root.eta = pyo.Var(eta_names)
+
+
+def _make_bender(ls):
+    m = pyo.ConcreteModel()
+    m.bender = LShapedCutGenerator()
+    m.bender.set_input(root_vars=[], tol=1e-8)
+    m.bender.set_ls(ls)
+    return m.bender
+
+
+@unittest.skipUnless(bc.mpi4py_available and bc.numpy_available,
+                     "LShapedCutGenerator requires mpi4py and numpy")
+class TestSetLSIndexing(unittest.TestCase):
+
+    _all_names = ["scen0", "scen1", "scen2"]
+
+    def test_no_root_scenarios_single_rank(self):
+        ls = _FakeLShaped(self._all_names, list(self._all_names))
+        bender = _make_bender(ls)
+        self.assertEqual(bender.global_num_subproblems(), 3)
+        self.assertEqual(bender._subproblem_ndx_map, {0: 0, 1: 1, 2: 2})
+        self.assertEqual(bender.all_root_etas, list(ls.root.eta.values()))
+
+    def test_no_root_scenarios_rank_view(self):
+        # the middle rank of a three-rank hub
+        ls = _FakeLShaped(self._all_names, ["scen1"])
+        bender = _make_bender(ls)
+        self.assertEqual(bender.global_num_subproblems(), 3)
+        self.assertEqual(bender._subproblem_ndx_map, {0: 1})
+
+    def test_no_has_root_scens_attribute(self):
+        # cross_scen_spoke.py calls set_ls with an opt object that has
+        # no has_root_scens attribute; all scenarios get subproblems
+        ls = _FakeLShaped(self._all_names, list(self._all_names))
+        del ls.has_root_scens
+        bender = _make_bender(ls)
+        self.assertEqual(bender.global_num_subproblems(), 3)
+        self.assertEqual(bender._subproblem_ndx_map, {0: 0, 1: 1, 2: 2})
+
+    def test_root_scenario_global_count(self):
+        ls = _FakeLShaped(self._all_names, list(self._all_names),
+                          root_scenarios=["scen1"])
+        bender = _make_bender(ls)
+        self.assertEqual(bender.global_num_subproblems(), 2)
+        self.assertEqual(len(bender.all_root_etas), 2)
+        # the root scenario contributes no subproblem, and the scenario
+        # after it maps to the second (not third) global slot
+        self.assertEqual(bender._subproblem_ndx_map, {0: 0, 1: 1})
+
+    def test_root_scenario_rank_views(self):
+        # issue #551: three-rank hub, one scenario per rank, the middle
+        # scenario built into the root problem; the rank owning scen2
+        # must map its subproblem inside range(global_num_subproblems())
+        rank_locals = [["scen0"], ["scen1"], ["scen2"]]
+        expected_maps = [{0: 0}, {}, {0: 1}]
+        seen_global_ndxs = []
+        for local_names, expected in zip(rank_locals, expected_maps):
+            ls = _FakeLShaped(self._all_names, local_names,
+                              root_scenarios=["scen1"])
+            bender = _make_bender(ls)
+            self.assertEqual(bender.global_num_subproblems(), 2)
+            self.assertEqual(bender._subproblem_ndx_map, expected)
+            for global_ndx in bender._subproblem_ndx_map.values():
+                self.assertLess(global_ndx, len(bender.all_root_etas))
+            seen_global_ndxs.extend(bender._subproblem_ndx_map.values())
+        # across the ranks, every global slot is covered exactly once
+        self.assertEqual(sorted(seen_global_ndxs), [0, 1])
+
+    def test_root_scenario_shared_rank_view(self):
+        # two-rank hub: the rank owning both scen0 and the root scenario
+        # has exactly one subproblem
+        ls = _FakeLShaped(self._all_names, ["scen0", "scen1"],
+                          root_scenarios=["scen1"])
+        bender = _make_bender(ls)
+        self.assertEqual(bender.global_num_subproblems(), 2)
+        self.assertEqual(bender._subproblem_ndx_map, {0: 0})
+
+
+if __name__ == "__main__":
+    unittest.main()

mpisppy/utils/lshaped_cuts.py

@@ -42,10 +42,22 @@ def __init__(self, component):
 
     def set_ls(self, ls):
         self.ls = ls
-        self.global_subproblem_count = len(self.ls.all_scenario_names)
-        self._subproblem_ndx_map = dict.fromkeys(range(len(self.ls.local_scenario_names)))
-        for s in self._subproblem_ndx_map.keys():
-            self._subproblem_ndx_map[s] = self.ls.all_scenario_names.index(self.ls.local_scenario_names[s])
+        # scenarios built into the root problem (root_scenarios) have no eta
+        # variable and no Benders subproblem, so all of the global indexing
+        # here is over the non-root scenarios only
+        if getattr(self.ls, "has_root_scens", False):
+            root_scenarios = self.ls.root_scenarios
+        else:
+            root_scenarios = ()
+        sub_scenario_names = [s for s in self.ls.all_scenario_names
+                              if s not in root_scenarios]
+        self.global_subproblem_count = len(sub_scenario_names)
+        local_sub_scenario_names = [s for s in self.ls.local_scenario_names
+                                    if s not in root_scenarios]
+        self._subproblem_ndx_map = {
+            ndx: sub_scenario_names.index(name)
+            for ndx, name in enumerate(local_sub_scenario_names)
+        }
         # print(self._subproblem_ndx_map)
         self.all_root_etas = list(self.ls.root.eta.values())
 

run_coverage.bash

@@ -203,6 +203,7 @@ run_phase "serial unit tests (serial)" \
         mpisppy/tests/test_buffer_inspect.py \
         mpisppy/tests/test_comm_lor_check.py \
         mpisppy/tests/test_ciutils.py \
+        mpisppy/tests/test_lshaped_cuts.py \
         mpisppy/tests/test_prox_approx.py \
         mpisppy/tests/test_sep_rho.py \
         mpisppy/tests/test_reduced_costs_fixer.py \