Merge pull request #442 from fast-aircraft-design/mpi-fix

Fix for crash in multiprocessing environment

.github/workflows/tests.yml

@@ -20,6 +20,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
+      - uses: mpi4py/setup-mpi@v1
 
       - name: Install poetry
         run: pipx install poetry
@@ -33,7 +34,7 @@ jobs:
       - name: Activate environment and install dependencies
         run: |
           poetry env use ${{ matrix.python-version }}
-          poetry install
+          poetry install -E mpi4py
 
       - name: Check with Black
         run: |

pyproject.toml

@@ -66,6 +66,10 @@ stdatm = "0.*"
 Deprecated = "^1.2.13"
 click = "^8.0.3"
 importlib-metadata = { version = "^4.2", python = "<3.10" }
+mpi4py = {version = "^3", optional = true}
+
+[tool.poetry.extras]
+mpi4py = ["mpi4py"]
 
 [tool.poetry.dev-dependencies]
 fast-oad-cs25 = "0.*"

src/fastoad/openmdao/_utils.py

@@ -2,7 +2,7 @@
 Utility functions for OpenMDAO classes/instances
 """
 #  This file is part of FAST-OAD : A framework for rapid Overall Aircraft Design
-#  Copyright (C) 2021 ONERA & ISAE-SUPAERO
+#  Copyright (C) 2022 ONERA & ISAE-SUPAERO
 #  FAST is free software: you can redistribute it and/or modify
 #  it under the terms of the GNU General Public License as published by
 #  the Free Software Foundation, either version 3 of the License, or
@@ -14,11 +14,45 @@
 #  You should have received a copy of the GNU General Public License
 #  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+from contextlib import contextmanager
+from copy import deepcopy
 from typing import List, Tuple
-from deprecated import deprecated
 
 import numpy as np
 import openmdao.api as om
+from deprecated import deprecated
+from openmdao.utils.mpi import FakeComm
+
+
+@contextmanager
+def problem_without_mpi(problem: om.Problem) -> om.Problem:
+    """
+    Context manager that delivers a copy of the given OpenMDAO problem.
+
+    A deepcopy operation may crash if problem.comm is not pickle-able, like a
+    mpi4py.MPI.Intracomm object.
+
+    This context manager temporarily sets a FakeComm object as problem.comm and
+    does the copy.
+
+    It ensures the original problem gets back its original communicator after
+    the `with` block is ended.
+
+    :param problem: any openMDAO problem
+    :return: A copy of the given problem with a FakeComm object as problem.comm
+    """
+    # An actual MPI communicator will make the deepcopy crash if an MPI
+    # library is installed.
+
+    actual_comm = problem.comm
+    problem.comm = FakeComm()
+
+    try:
+        problem_copy = deepcopy(problem)
+        problem_copy.comm = problem.comm
+        yield problem_copy
+    finally:
+        problem.comm = actual_comm
 
 
 @deprecated(

src/fastoad/openmdao/problem.py

@@ -26,6 +26,7 @@
 from fastoad.module_management.service_registry import RegisterSubmodel
 from fastoad.openmdao.validity_checker import ValidityDomainChecker
 from fastoad.openmdao.variables import VariableList
+from ._utils import problem_without_mpi
 from .exceptions import FASTOpenMDAONanInInputFile
 
 # Name of IVC that will contain input values
@@ -73,20 +74,21 @@ def setup(self, *args, **kwargs):
         """
         Set up the problem before run.
         """
-        problem_copy = deepcopy(self)
-        try:
-            super(FASTOADProblem, problem_copy).setup(*args, **kwargs)
-        except RuntimeError:
-            vars_metadata = self._get_undetermined_dynamic_vars_metadata(problem_copy)
-            if vars_metadata:
-                # If vars_metadata is empty, it means the RuntimeError was not because
-                # of dynamic shapes, and the incoming self.setup() will raise it.
-                ivc = om.IndepVarComp()
-                for name, meta in vars_metadata.items():
-                    # We use a (2,)-shaped array as value here. This way, it will be easier to
-                    # identify dynamic-shaped data in an input file generated from current problem.
-                    ivc.add_output(name, [np.nan, np.nan], units=meta["units"])
-                self.model.add_subsystem(SHAPER_SYSTEM_NAME, ivc, promotes=["*"])
+        with problem_without_mpi(self) as problem_copy:
+            try:
+                super(FASTOADProblem, problem_copy).setup(*args, **kwargs)
+            except RuntimeError:
+                vars_metadata = self._get_undetermined_dynamic_vars_metadata(problem_copy)
+                if vars_metadata:
+                    # If vars_metadata is empty, it means the RuntimeError was not because
+                    # of dynamic shapes, and the incoming self.setup() will raise it.
+                    ivc = om.IndepVarComp()
+                    for name, meta in vars_metadata.items():
+                        # We use a (2,)-shaped array as value here. This way, it will be easier
+                        # to identify dynamic-shaped data in an input file generated from current
+                        # problem.
+                        ivc.add_output(name, [np.nan, np.nan], units=meta["units"])
+                    self.model.add_subsystem(SHAPER_SYSTEM_NAME, ivc, promotes=["*"])
 
         super().setup(*args, **kwargs)
 
@@ -170,14 +172,15 @@ def _read_inputs_without_setup_done(self):
         """
         input_variables, unused_variables = self._get_problem_inputs()
         self.additional_variables = unused_variables
-        tmp_prob = deepcopy(self)
-        tmp_prob.setup()
-        # At this point, there may be non-fed dynamically shaped inputs, so the setup may
-        # create the "shaper" IVC, but we ignore it because we need to redefine these variables
-        # in input file.
-        ivc_vars = tmp_prob.model.get_io_metadata(
-            "output", tags="indep_var", excludes=f"{SHAPER_SYSTEM_NAME}.*"
-        )
+        with problem_without_mpi(self):
+            tmp_prob = deepcopy(self)
+            tmp_prob.setup()
+            # At this point, there may be non-fed dynamically shaped inputs, so the setup may
+            # create the "shaper" IVC, but we ignore it because we need to redefine these variables
+            # in input file.
+            ivc_vars = tmp_prob.model.get_io_metadata(
+                "output", tags="indep_var", excludes=f"{SHAPER_SYSTEM_NAME}.*"
+            )
         for meta in ivc_vars.values():
             try:
                 del input_variables[meta["prom_name"]]
@@ -187,15 +190,16 @@ def _read_inputs_without_setup_done(self):
             self._insert_input_ivc(input_variables.to_ivc())
 
     def _insert_input_ivc(self, ivc: om.IndepVarComp, subsystem_name=INPUT_SYSTEM_NAME):
-        tmp_prob = deepcopy(self)
-        tmp_prob.setup()
-
-        # We get order from copied problem, but we have to ignore the "shaper" and the auto IVCs.
-        previous_order = [
-            system.name
-            for system in tmp_prob.model.system_iter(recurse=False)
-            if system.name != "_auto_ivc" and system.name != SHAPER_SYSTEM_NAME
-        ]
+        with problem_without_mpi(self) as tmp_prob:
+            tmp_prob.setup()
+
+            # We get order from copied problem, but we have to ignore the "shaper"
+            # and the auto IVCs.
+            previous_order = [
+                system.name
+                for system in tmp_prob.model.system_iter(recurse=False)
+                if system.name != "_auto_ivc" and system.name != SHAPER_SYSTEM_NAME
+            ]
 
         self.model.add_subsystem(subsystem_name, ivc, promotes=["*"])
         self.model.set_order([subsystem_name] + previous_order)

src/fastoad/openmdao/variables/variable_list.py

@@ -24,7 +24,7 @@
 from deprecated import deprecated
 from openmdao.core.constants import _SetupStatus
 
-from fastoad.openmdao._utils import get_unconnected_input_names
+from fastoad.openmdao._utils import get_unconnected_input_names, problem_without_mpi
 from .variable import METADATA_TO_IGNORE, Variable
 
 
@@ -279,8 +279,9 @@ def from_problem(
         """
 
         if not problem._metadata or problem._metadata["setup_status"] < _SetupStatus.POST_SETUP:
-            problem = deepcopy(problem)
-            problem.setup()
+            with problem_without_mpi(problem) as problem_copy:
+                problem_copy.setup()
+                problem = problem_copy
 
         # Get inputs and outputs
         metadata_keys = (