Merge branch 'main' into sync_locs

bknueven · web-flow · commit 79af1ec4c971 · 2025-02-28T12:09:14.000-07:00
diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml
@@ -407,7 +407,7 @@ jobs:
         run: |
           cd mpisppy/tests
           python test_gradient_rho.py
-          python test_w_writer.py
+          python test_xbar_w_reader_writer.py
 
   test-headers:
     name: header test
diff --git a/examples/farmer/farmer_cylinders.py b/examples/farmer/farmer_cylinders.py
@@ -112,6 +112,8 @@ def main():
                                   ph_converger=ph_converger,
                                   rho_setter = rho_setter)
 
+    hub_dict['opt_kwargs']['options']['cfg'] = cfg                    
+
     if cfg.primal_dual_converger:
         hub_dict['opt_kwargs']['options']\
             ['primal_dual_converger_options'] = {
diff --git a/examples/generic_cylinders.bash b/examples/generic_cylinders.bash
@@ -4,6 +4,12 @@
 SOLVER="cplex"
 SPB=1
 
+echo "^^^ hub only with w-writer (smoke) ^^^"
+python -m mpi4py ../mpisppy/generic_cylinders.py --module-name farmer/farmer --num-scens 3 --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --W-writer --W-fname w_values.csv 
+
+echo "^^^ hub only with w-reader (smoke) ^^^"
+python -m mpi4py ../mpisppy/generic_cylinders.py --module-name farmer/farmer --num-scens 3 --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --W-reader --init-W-fname w_values.csv
+
 echo "^^^ Multi-stage AirCond ^^^"
 mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name mpisppy.tests.examples.aircond --branching-factors "3 3 3" --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --lagrangian --xhatxbar --rel-gap 0.01 --solution-base-name aircond_nonants
 # --xhatshuffle --stag2EFsolvern
diff --git a/mpisppy/generic_cylinders.py b/mpisppy/generic_cylinders.py
@@ -14,22 +14,31 @@
 import shutil
 import numpy as np
 import pyomo.environ as pyo
+import pyomo.common.config as pyofig
+
 from mpisppy.spin_the_wheel import WheelSpinner
+
 import mpisppy.utils.cfg_vanilla as vanilla
 import mpisppy.utils.config as config
 import mpisppy.utils.sputils as sputils
+
 from mpisppy.convergers.norm_rho_converger import NormRhoConverger
 from mpisppy.convergers.primal_dual_converger import PrimalDualConverger
+
 from mpisppy.extensions.extension import MultiExtension
 from mpisppy.extensions.fixer import Fixer
 from mpisppy.extensions.mipgapper import Gapper
 from mpisppy.extensions.gradient_extension import Gradient_extension
 from mpisppy.extensions.scenario_lpfiles import Scenario_lpfiles
+
+from mpisppy.utils.wxbarwriter import WXBarWriter
+from mpisppy.utils.wxbarreader import WXBarReader
+
 import mpisppy.utils.solver_spec as solver_spec
+
 from mpisppy import global_toc
 from mpisppy import MPI
 
-
 def _parse_args(m):
     # m is the model file module
     cfg = config.Config()
@@ -82,6 +91,18 @@ def _parse_args(m):
     cfg.coeff_rho_args()
     cfg.sensi_rho_args()
     cfg.reduced_costs_rho_args()
+
+    cfg.add_to_config("user_defined_extensions",
+                      description="Space-delimited module names for user extensions",
+                      domain=pyofig.ListOf(str),
+                      default=None)
+    # TBD - think about adding directory for json options files
+
+    cfg.add_to_config("hub_and_spoke_dict_callback",
+                      description="[FOR EXPERTS ONLY] Module that contains the function hub_and_spoke_dict_callback that will be passed the hubdict and list of spokedicts prior to spin-the-wheel (last chance for intervention)",
+                      domain=str,
+                      default=None)    
+    
     cfg.parse_command_line(f"mpi-sppy for {cfg.module_name}")
     
     cfg.checker()  # looks for inconsistencies 
@@ -163,7 +184,11 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
                                   rho_setter = rho_setter,
                                   all_nodenames = all_nodenames,
                                   )
-    
+
+    # the intent of the following is to transition to strictly
+    # cfg-based option passing, as opposed to dictionary-based processing.
+    hub_dict['opt_kwargs']['options']['cfg'] = cfg                
+        
     # Extend and/or correct the vanilla dictionary
     ext_classes = list()
     # TBD: add cross_scenario_cuts, which also needs a cylinder
@@ -198,6 +223,24 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
     if cfg.scenario_lpfiles:
         ext_classes.append(Scenario_lpfiles)
 
+    if cfg.W_and_xbar_reader:
+        ext_classes.append(WXBarReader)
+
+    if cfg.W_and_xbar_writer:
+        ext_classes.append(WXBarWriter)
+
+    if cfg.user_defined_extensions is not None:
+        for ext_name in cfg.user_defined_extensions:
+            module = sputils.module_name_to_module(ext_name)
+            vanilla.extension_adder(module)
+            # grab JSON for this module's option dictionary
+            json_filename = ext_name+".json"
+            if os.path.exists(json_filename):
+                ext_options= json.load(json_filename)
+                hub_dict['opt_kwargs']['options'][ext_name] = ext_options
+            else:
+                raise RuntimeError(f"JSON options file {json_filename} for user defined extension not found")
+
     if cfg.sep_rho:
         vanilla.add_sep_rho(hub_dict, cfg)
 
@@ -322,7 +365,11 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
         list_of_spoke_dict.append(xhatxbar_spoke)
     if cfg.reduced_costs:
         list_of_spoke_dict.append(reduced_costs_spoke)
-        
+
+    # if the user dares, let them mess with the hubdict prior to solve
+    if cfg.hub_and_spoke_dict_callback is not None:
+        module = sputils.module_name_to_module(cfg.hub_and_spoke_dict_callback)
+        module.hub_and_spoke_dict_callback(hub_dict, list_of_spoke_dict)
 
     wheel = WheelSpinner(hub_dict, list_of_spoke_dict)
     wheel.spin()
diff --git a/mpisppy/phbase.py b/mpisppy/phbase.py
@@ -239,8 +239,6 @@ class PHBase(mpisppy.spopt.SPOpt):
                 Function to set rho values throughout the PH algorithm.
             variable_probability (callable, optional):
                 Function to set variable specific probabilities.
-            cfg (config object, optional?)  controls (mainly from user)
-                (Maybe this should move up to spbase)
 
     """
     def __init__(
diff --git a/mpisppy/tests/test_gradient_rho.py b/mpisppy/tests/test_gradient_rho.py
@@ -57,6 +57,7 @@ def _create_ph_farmer(self):
         scenario_creator_kwargs = farmer.kw_creator(self.cfg)
         beans = (self.cfg, scenario_creator, scenario_denouement, all_scenario_names)
         hub_dict = vanilla.ph_hub(*beans, scenario_creator_kwargs=scenario_creator_kwargs)
+        hub_dict['opt_kwargs']['options']['cfg'] = self.cfg                            
         list_of_spoke_dict = list()
         wheel = WheelSpinner(hub_dict, list_of_spoke_dict)
         wheel.spin()
diff --git a/mpisppy/tests/test_xbar_w_reader_writer.py b/mpisppy/tests/test_xbar_w_reader_writer.py
@@ -22,16 +22,17 @@
 import mpisppy.tests.examples.farmer as farmer
 from mpisppy.spin_the_wheel import WheelSpinner
 from mpisppy.tests.utils import get_solver
-from mpisppy.utils.wxbarwriter import WXBarWriter
-from mpisppy.utils.wxbarreader import WXBarReader
+import mpisppy.utils.wxbarreader as wxbarreader
+import mpisppy.utils.wxbarwriter as wxbarwriter
 
-
-__version__ = 0.1
+__version__ = 0.2
 
 solver_available,solver_name, persistent_available, persistent_solver_name= get_solver()
 
 def _create_cfg():
     cfg = config.Config()
+    wxbarreader.add_options_to_config(cfg)        
+    wxbarwriter.add_options_to_config(cfg)
     cfg.add_branching_factors()
     cfg.num_scens_required()
     cfg.popular_args()
@@ -43,7 +44,7 @@ def _create_cfg():
 
 #*****************************************************************************
 
-class Test_w_writer_farmer(unittest.TestCase):
+class Test_xbar_w_reader_writer_farmer(unittest.TestCase):
     """ Test the gradient code using farmer."""
 
     def _create_ph_farmer(self, ph_extensions=None, max_iter=100):
@@ -59,14 +60,15 @@ def _create_ph_farmer(self, ph_extensions=None, max_iter=100):
         self.cfg.max_iterations = max_iter
         beans = (self.cfg, scenario_creator, scenario_denouement, all_scenario_names)
         hub_dict = vanilla.ph_hub(*beans, scenario_creator_kwargs=scenario_creator_kwargs, ph_extensions=ph_extensions)
-        if ph_extensions==WXBarWriter: #tbd
-            hub_dict['opt_kwargs']['options']["W_and_xbar_writer"] =  {"Wcsvdir": "Wdir"}
-            hub_dict['opt_kwargs']['options']['W_fname'] = self.temp_w_file_name
-            hub_dict['opt_kwargs']['options']['Xbar_fname'] = self.temp_xbar_file_name
-        if ph_extensions==WXBarReader:
-            hub_dict['opt_kwargs']['options']["W_and_xbar_reader"] =  {"Wcsvdir": "Wdir"}
-            hub_dict['opt_kwargs']['options']['init_W_fname'] = self.w_file_name
-            hub_dict['opt_kwargs']['options']['init_Xbar_fname'] = self.xbar_file_name
+        hub_dict['opt_kwargs']['options']['cfg'] = self.cfg
+        if ph_extensions==wxbarwriter.WXBarWriter:
+            self.cfg.W_and_xbar_writer = True
+            self.cfg.W_fname = self.temp_w_file_name
+            self.cfg.Xbar_fname = self.temp_xbar_file_name
+        if ph_extensions==wxbarreader.WXBarReader:
+            self.cfg.W_and_xbar_reader = True
+            self.cfg.init_W_fname = self.w_file_name
+            self.cfg.init_Xbar_fname = self.xbar_file_name
         list_of_spoke_dict = list()
         wheel = WheelSpinner(hub_dict, list_of_spoke_dict)
         wheel.spin()
@@ -79,7 +81,7 @@ def setUp(self):
         self.ph_object = None
     
     def test_wwriter(self):
-        self.ph_object = self._create_ph_farmer(ph_extensions=WXBarWriter, max_iter=5)
+        self.ph_object = self._create_ph_farmer(ph_extensions=wxbarwriter.WXBarWriter, max_iter=5)
         with open(self.temp_w_file_name, 'r') as f:
             read = csv.reader(f)
             rows = list(read)
@@ -88,7 +90,7 @@ def test_wwriter(self):
         os.remove(self.temp_w_file_name)
 
     def test_xbarwriter(self):
-        self.ph_object = self._create_ph_farmer(ph_extensions=WXBarWriter, max_iter=5)
+        self.ph_object = self._create_ph_farmer(ph_extensions=wxbarwriter.WXBarWriter, max_iter=5)
         with open(self.temp_xbar_file_name, 'r') as f:
             read = csv.reader(f)
             rows = list(read)
@@ -97,15 +99,15 @@ def test_xbarwriter(self):
         os.remove(self.temp_xbar_file_name)
 
     def test_wreader(self):
-        self.ph_object = self._create_ph_farmer(ph_extensions=WXBarReader, max_iter=1)
+        self.ph_object = self._create_ph_farmer(ph_extensions=wxbarreader.WXBarReader, max_iter=1)
         for sname, scenario in self.ph_object.local_scenarios.items():
             if sname == 'scen0':
                 self.assertAlmostEqual(scenario._mpisppy_model.W[("ROOT", 1)]._value, 70.84705093609978)
             if sname == 'scen1':
                 self.assertAlmostEqual(scenario._mpisppy_model.W[("ROOT", 0)]._value, -41.104251445950844)
 
     def test_xbarreader(self):
-        self.ph_object = self._create_ph_farmer(ph_extensions=WXBarReader, max_iter=1)
+        self.ph_object = self._create_ph_farmer(ph_extensions=wxbarreader.WXBarReader, max_iter=1)
         for sname, scenario in self.ph_object.local_scenarios.items():
             if sname == 'scen0':
                 self.assertAlmostEqual(scenario._mpisppy_model.xbars[("ROOT", 1)]._value, 274.2239371483933)
diff --git a/mpisppy/utils/config.py b/mpisppy/utils/config.py
@@ -964,32 +964,12 @@ def tracking_args(self):
                             domain=int,
                             default=0)
 
-
     def wxbar_read_write_args(self):
-        self.add_to_config("init_W_fname",
-                                description="Path of initial W file (default None)",
-                                domain=str,
-                                default=None)
-        self.add_to_config("init_Xbar_fname",
-                                description="Path of initial Xbar file (default None)",
-                                domain=str,
-                                default=None)
-        self.add_to_config("init_separate_W_files",
-                                description="If True, W is read from separate files (default False)",
-                                domain=bool,
-                                default=False)
-        self.add_to_config("W_fname",
-                                description="Path of final W file (default None)",
-                                domain=str,
-                                default=None)
-        self.add_to_config("Xbar_fname",
-                                description="Path of final Xbar file (default None)",
-                                domain=str,
-                                default=None)
-        self.add_to_config("separate_W_files",
-                                description="If True, writes W to separate files (default False)",
-                                domain=bool,
-                                default=False)
+        import mpisppy.utils.wxbarreader as wxbarreader
+        wxbarreader.add_options_to_config(self)
+
+        import mpisppy.utils.wxbarwriter as wxbarwriter
+        wxbarwriter.add_options_to_config(self)        
 
     def proper_bundle_config(self):
         self.add_to_config('pickle_bundles_dir',
diff --git a/mpisppy/utils/gradient.py b/mpisppy/utils/gradient.py
@@ -15,7 +15,6 @@
 import csv
 import copy
 
-from mpisppy.utils import config
 import mpisppy.utils.cfg_vanilla as vanilla
 from mpisppy.utils.wxbarwriter import WXBarWriter
 from mpisppy.spin_the_wheel import WheelSpinner
@@ -190,29 +189,6 @@ def write_grad_rho(self):
 ###################################################################################
 
 
-def _parser_setup():
-    """ Set up config object and return it, but don't parse 
-
-    Returns:
-       cfg (Config): config object
-
-    Notes:
-       parsers for the non-model-specific arguments; but the model_module_name will be pulled off first
-
-    """
-
-    cfg = config.Config()
-    cfg.add_branching_factors()
-    cfg.num_scens_required()
-    cfg.popular_args()
-    cfg.two_sided_args()
-    cfg.ph_args()
-    
-    cfg.gradient_args()
-
-    return cfg
-
-
 def grad_cost_and_rho(mname, original_cfg):
     """ Creates a ph object from cfg and using the module 'mname' functions. Then computes the corresponding grad cost and rho.
 
@@ -244,10 +220,12 @@ def grad_cost_and_rho(mname, original_cfg):
     if hasattr(model_module, '_variable_probability'):
         variable_probability = model_module._variable_probability
     beans = (cfg, scenario_creator, scenario_denouement, all_scenario_names)
+    # Note: the callers need to set w_writer cfg options
     hub_dict = vanilla.ph_hub(*beans,
                               scenario_creator_kwargs=scenario_creator_kwargs,
                               ph_extensions=WXBarWriter,
                               variable_probability=variable_probability)
+    hub_dict['opt_kwargs']['options']['cfg'] = cfg                
     list_of_spoke_dict = list()
     wheel = WheelSpinner(hub_dict, list_of_spoke_dict)
     wheel.spin() #TODO: steal only what's needed in  WheelSpinner
diff --git a/mpisppy/utils/wxbarreader.py b/mpisppy/utils/wxbarreader.py
@@ -39,18 +39,42 @@
 n_proc = MPI.COMM_WORLD.Get_size()
 rank = MPI.COMM_WORLD.Get_rank()
 
+def add_options_to_config(cfg):
+
+    cfg.add_to_config("W_and_xbar_reader",
+                      description="Enables the w and xbar reader (default False)",
+                      domain=bool,
+                      default=False)
+    cfg.add_to_config("init_W_fname",
+                      description="Path of initial W file (default None)",
+                      domain=str,
+                      default=None)
+    cfg.add_to_config("init_Xbar_fname",
+                      description="Path of initial Xbar file (default None)",
+                      domain=str,
+                      default=None)
+    cfg.add_to_config("init_separate_W_files",
+                      description="If True, W is read from separate files (default False)",
+                      domain=bool,
+                      default=False)
+
 class WXBarReader(mpisppy.extensions.extension.Extension):
     """ Extension class for reading W values
     """
     def __init__(self, ph):
 
+        assert 'cfg' in ph.options
+        self.cfg = ph.options['cfg']
+        if self.cfg.get("W_and_xbar_reader") is None or not self.cfg.W_and_xbar_reader:
+            self.not_active = True
+            return  # nothing to do here
+        else:
+            self.not_active = False
+
         ''' Do a bunch of checking if files exist '''
-        w_fname, x_fname, sep_files = None, None, False
-        if ('init_separate_W_files' in ph.options):
-            sep_files = ph.options['init_separate_W_files']
+        w_fname, x_fname, sep_files = self.cfg.init_W_fname, self.cfg.init_Xbar_fname, self.cfg.init_separate_W_files
 
-        if ('init_W_fname' in ph.options):
-            w_fname = ph.options['init_W_fname']
+        if w_fname is not None:
             if (not os.path.exists(w_fname)):
                 if (rank == 0):
                     if (sep_files):
@@ -59,8 +83,7 @@ def __init__(self, ph):
                         print('Cannot find file', w_fname)
                 quit()
 
-        if ('init_Xbar_fname' in ph.options):
-            x_fname = ph.options['init_Xbar_fname']
+        if x_fname is not None:
             if (not os.path.exists(x_fname)):
                 if (rank == 0):
                     print('Cannot find file', x_fname)
@@ -85,6 +108,8 @@ def post_iter0(self):
 
     def miditer(self):
         ''' Called before the solveloop is called '''
+        if self.not_active:
+            return  # nothing to do.
         if self.PHB._PHIter == 1:
             if self.w_fname:
                 mpisppy.utils.wxbarutils.set_W_from_file(
diff --git a/mpisppy/utils/wxbarwriter.py b/mpisppy/utils/wxbarwriter.py
