parallel snstop (#417)

* parallel snstop

* formatting fixes

* added comments

* added test - does it make sense?

* fixed MPI check

* actually fixing MPI check

* iSort fix

* maybe this time the test will be skipped?

* maybe like this?

* what about this

* cleanup

* add timeout option to testflo

* rerun tests

* updated test with send/receive

---------

Co-authored-by: Marco Mangano <[email protected]>

.github/test_real.sh

@@ -11,4 +11,4 @@ cd tests
 # we have to copy over the coveragerc file to make sure it's in the
 # same directory where codecov is run
 cp ../.coveragerc .
-testflo --pre_announce --disallow_deprecations -v --coverage --coverpkg pyoptsparse $EXTRA_FLAGS
+testflo --pre_announce --disallow_deprecations -v --coverage --coverpkg pyoptsparse $EXTRA_FLAGS --timeout 60

.github/windows.yaml

@@ -37,5 +37,5 @@ jobs:
 
     - script: |
         cd tests
-        testflo -n 1 .
+        testflo -n 1 --timeout 60 .
       displayName: Run tests

.github/workflows/windows-build.yml

@@ -42,4 +42,4 @@ jobs:
         run: |
           conda activate pyos-build
           cd tests
-          testflo --pre_announce -v -n 1 .
+          testflo --pre_announce -v -n 1 --timeout 60 .

pyoptsparse/pyOpt_optimizer.py

@@ -704,15 +704,19 @@ def _waitLoop(self):
 
             # Receive mode and quit if mode is -1:
             mode = self.optProb.comm.bcast(mode, root=0)
-            if mode == -1:
+            # mode = 0 call masterfunc2 as broadcast by root in masterfunc
+            if mode == 0:
+                # Receive info from shell function
+                info = self.optProb.comm.bcast(info, root=0)
+
+                # Call the generic internal function. We don't care
+                # about return values on these procs
+                self._masterFunc2(*info)
+            # mode = -1 exit wait loop
+            elif mode == -1:
                 break
-
-            # Otherwise receive info from shell function
-            info = self.optProb.comm.bcast(info, root=0)
-
-            # Call the generic internal function. We don't care
-            # about return values on these procs
-            self._masterFunc2(*info)
+            else:
+                raise Error("Wait loop recieved code %d must be -1 or 0" % mode)
 
     def _setInitialCacheValues(self):
         """

pyoptsparse/pySNOPT/pySNOPT.py

@@ -536,6 +536,46 @@ def __call__(
         else:
             return commSol
 
+    def _waitLoop(self):
+        """Non-root processors go into this waiting loop while the
+        root proc does all the work in the optimization algorithm
+
+        This function overwrites the namesake in the Optimizer class to add a new mode enabling parallel snstop function
+        """
+
+        mode = None
+        info = None
+        while True:
+            # * Note*: No checks for MPI here since this code is
+            # * only run in parallel, which assumes mpi4py is working
+
+            # Receive mode and quit if mode is -1:
+            mode = self.optProb.comm.bcast(mode, root=0)
+
+            # mode = 0 call masterfunc2 as broadcast by root in masterfunc
+            if mode == 0:
+                # Receive info from shell function
+                info = self.optProb.comm.bcast(info, root=0)
+
+                # Call the generic internal function. We don't care
+                # about return values on these procs
+                self._masterFunc2(*info)
+
+            # mode = -1 exit wait loop
+            elif mode == -1:
+                break
+
+            # mode = 1 call user snSTOP function
+            elif mode == 1:
+                # Receive function arguments from root
+                info = self.optProb.comm.bcast(info, root=0)
+                # Get function handle and make call
+                snstop_handle = self.getOption("snSTOP function handle")
+                if snstop_handle is not None:
+                    snstop_handle(*info)
+            else:
+                raise Error("Wait loop recieved code %d must be -1, 0, or 1 " % mode)
+
     def _userfg_wrap(self, mode, nnJac, x, fobj, gobj, fcon, gcon, nState, cu, iu, ru):
         """
         The snopt user function. This is what is actually called from snopt.
@@ -703,6 +743,10 @@ def _snstop(self, ktcond, mjrprtlvl, minimize, n, nncon, nnobj, ns, itn, nmajor,
 
             if not self.storeHistory:
                 raise Error("snSTOP function handle must be used with storeHistory=True")
+
+            # Broadcasting flag to call user snstop function
+            self.optProb.comm.bcast(1, root=0)
+            self.optProb.comm.bcast(snstopArgs, root=0)
             iabort = snstop_handle(*snstopArgs)
             # write iterDict again if anything was inserted
             if self.storeHistory and callCounter is not None:

tests/test_hs015_parallel.py

@@ -0,0 +1,139 @@
+"""Test solution of problem HS15 from the Hock & Schittkowski collection"""
+
+# Standard Python modules
+import unittest
+
+# External modules
+import numpy as np
+
+try:
+    HAS_MPI = True
+    # External modules
+    from mpi4py import MPI
+
+    # Setting up MPI communicators
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+except ImportError:
+    HAS_MPI = False
+
+# First party modules
+from pyoptsparse import Optimization
+
+# Local modules
+from testing_utils import OptTest
+
+
+@unittest.skipIf(not HAS_MPI, "MPI not available")
+class TestHS15(OptTest):
+    ## Solve test problem HS15 from the Hock & Schittkowski collection.
+    #
+    #  min   100 (x2 - x1^2)^2 + (1 - x1)^2
+    #  s.t.  x1 x2 >= 1
+    #        x1 + x2^2 >= 0
+    #        x1 <= 0.5
+    #
+    #  The standard start point (-2, 1) usually converges to the standard
+    #  minimum at (0.5, 2.0), with final objective = 306.5.
+    #  Sometimes the solver converges to another local minimum
+    #  at (-0.79212, -1.26243), with final objective = 360.4.
+    ##
+
+    N_PROCS = 2  # Run case on two procs
+
+    name = "HS015"
+    DVs = {"xvars"}
+    cons = {"con"}
+    objs = {"obj"}
+    extras = {"extra1", "extra2"}
+    fStar = [
+        306.5,
+        360.379767,
+    ]
+    xStar = [
+        {"xvars": (0.5, 2.0)},
+        {"xvars": (-0.79212322, -1.26242985)},
+    ]
+    optOptions = {}
+
+    def objfunc(self, xdict):
+        self.nf += 1
+        x = xdict["xvars"]
+        funcs = {}
+        funcs["obj"] = [100 * (x[1] - x[0] ** 2) ** 2 + (1 - x[0]) ** 2]
+        conval = np.zeros(2, "D")
+        conval[0] = x[0] * x[1]
+        conval[1] = x[0] + x[1] ** 2
+        funcs["con"] = conval
+        # extra keys
+        funcs["extra1"] = 0.0
+        funcs["extra2"] = 1.0
+        fail = False
+        return funcs, fail
+
+    def sens(self, xdict, funcs):
+        self.ng += 1
+        x = xdict["xvars"]
+        funcsSens = {}
+        funcsSens["obj"] = {
+            "xvars": [2 * 100 * (x[1] - x[0] ** 2) * (-2 * x[0]) - 2 * (1 - x[0]), 2 * 100 * (x[1] - x[0] ** 2)]
+        }
+        funcsSens["con"] = {"xvars": [[x[1], x[0]], [1, 2 * x[1]]]}
+        fail = False
+        return funcsSens, fail
+
+    def setup_optProb(self):
+        # Optimization Object
+        self.optProb = Optimization("HS15 Constraint Problem", self.objfunc)
+
+        # Design Variables
+        lower = [-5.0, -5.0]
+        upper = [0.5, 5.0]
+        value = [-2, 1.0]
+        self.optProb.addVarGroup("xvars", 2, lower=lower, upper=upper, value=value)
+
+        # Constraints
+        lower = [1.0, 0.0]
+        upper = [None, None]
+        self.optProb.addConGroup("con", 2, lower=lower, upper=upper)
+
+        # Objective
+        self.optProb.addObj("obj")
+
+    @staticmethod
+    def my_snstop(iterDict):
+        """manually terminate SNOPT after 1 major iteration if"""
+
+        return_idx = 0
+        if iterDict["nMajor"] == 1:
+            if comm.rank == 1:
+                comm.send(1, dest=0, tag=comm.rank)
+            elif comm.rank == 0:
+                return_idx = comm.recv(source=1)
+        return return_idx
+
+    def test_optimization(self):
+        self.optName = "SNOPT"
+        self.setup_optProb()
+        sol = self.optimize()
+        # Check Solution
+        self.assert_solution_allclose(sol, 1e-12)
+        # Check informs
+        self.assert_inform_equal(sol)
+
+    def test_snopt_snstop(self):
+        self.optName = "SNOPT"
+        self.setup_optProb()
+        optOptions = {
+            "snSTOP function handle": self.my_snstop,
+        }
+        sol = self.optimize(optOptions=optOptions, storeHistory=True)
+        # Check informs
+        # we should get 70/74
+        self.assert_inform_equal(sol, optInform=74)
+
+
+if __name__ == "__main__":
+    unittest.main()