Merge pull request #5325 from hakonhagland/python_mpi

Allow Python bindings to control MPI setup in the OPM::Main constructor

opm/simulators/flow/Main.cpp

@@ -49,7 +49,9 @@ Main::Main(int argc, char** argv, bool ownMPI)
     }
 }
 
-Main::Main(const std::string& filename)
+Main::Main(const std::string& filename, bool mpi_init, bool mpi_finalize)
+    : mpi_init_{mpi_init}
+    , mpi_finalize_{mpi_finalize}
 {
     setArgvArgc_(filename);
     initMPI();
@@ -58,10 +60,14 @@ Main::Main(const std::string& filename)
 Main::Main(const std::string& filename,
            std::shared_ptr<EclipseState> eclipseState,
            std::shared_ptr<Schedule> schedule,
-           std::shared_ptr<SummaryConfig> summaryConfig)
+           std::shared_ptr<SummaryConfig> summaryConfig,
+           bool mpi_init,
+           bool mpi_finalize)
     : eclipseState_{std::move(eclipseState)}
     , schedule_{std::move(schedule)}
     , summaryConfig_{std::move(summaryConfig)}
+    , mpi_init_{mpi_init}
+    , mpi_finalize_{mpi_finalize}
 {
     setArgvArgc_(filename);
     initMPI();
@@ -107,7 +113,7 @@ Main::~Main()
 #endif // HAVE_DAMARIS
 
 #if HAVE_MPI && !HAVE_DUNE_FEM
-    if (ownMPI_) {
+    if (ownMPI_ && this->mpi_finalize_) {
         MPI_Finalize();
     }
 #endif
@@ -132,9 +138,15 @@ void Main::setArgvArgc_(const std::string& filename)
 void Main::initMPI()
 {
 #if HAVE_DUNE_FEM
-    Dune::Fem::MPIManager::initialize(argc_, argv_);
+    // The instance() method already checks if MPI has been initialized so we may
+    // not need to check mpi_init_ here.
+    if (this->mpi_init_) {
+        Dune::MPIHelper::instance(argc_, argv_);
+    }
 #elif HAVE_MPI
-    MPI_Init(&argc_, &argv_);
+    if (this->mpi_init_) {
+        MPI_Init(&argc_, &argv_);
+    }
 #endif
     FlowGenericVanguard::setCommunication(std::make_unique<Parallel::Communication>());
 

opm/simulators/flow/Main.hpp

@@ -130,14 +130,16 @@ public:
     Main(int argc, char** argv, bool ownMPI = true);
 
     // This constructor can be called from Python
-    Main(const std::string& filename);
+    Main(const std::string& filename, bool mpi_init = true, bool mpi_finalize = true);
 
     // This constructor can be called from Python when Python has
     // already parsed a deck
     Main(const std::string& filename,
          std::shared_ptr<EclipseState> eclipseState,
          std::shared_ptr<Schedule> schedule,
-         std::shared_ptr<SummaryConfig> summaryConfig);
+         std::shared_ptr<SummaryConfig> summaryConfig,
+         bool mpi_init = true,
+         bool mpi_finalize = true);
 
     ~Main();
 
@@ -748,6 +750,8 @@ private:
     std::shared_ptr<EclipseState> eclipseState_{};
     std::shared_ptr<Schedule> schedule_{};
     std::shared_ptr<SummaryConfig> summaryConfig_{};
+    bool mpi_init_{true}; //!< True if MPI_Init should be called
+    bool mpi_finalize_{true}; //!< True if MPI_Finalize should be called
 
     // To demonstrate run with non_world_comm
     bool test_split_comm_ = false;

opm/simulators/flow/python/PyBlackOilSimulator.hpp

@@ -62,6 +62,7 @@ public:
         const std::string &idx_name,
         py::array_t<double,
         py::array::c_style | py::array::forcecast> array);
+    void setupMpi(bool init_mpi, bool finalize_mpi);
     int step();
     int stepCleanup();
     int stepInit();
@@ -74,6 +75,8 @@ private:
     const std::string deck_filename_;
     bool has_run_init_ = false;
     bool has_run_cleanup_ = false;
+    bool mpi_init_ = true;
+    bool mpi_finalize_ = true;
     //bool debug_ = false;
     // This *must* be declared before other pointers
     // to simulator objects. This in order to deinitialize

python/docstrings.json

@@ -27,9 +27,9 @@
         "signature": "get_dt() -> float",
         "doc": "Get the timestep size of the last completed step.\n\n:return: Timestep size in days.\n:type return: float"
     },
-    "getFluidStateVariables": {
-        "signature": "get_fluid_state_variables(name: str) -> NDArray[float]",
-        "doc": "Retrieve a fluid state variable for the simulation grid.\n\n:para name: The name of the variable. Valid names are 'pw' (pressure water), 'pg' (pressure gas), 'po' (pressure oil), 'rho_w' (density water), 'rho_g' (density gas), 'rho_o' (density oil)'Rs' (soultion gas-oil ratio), 'Rv' (volatile gas-oil ratio), 'Sw' (water saturation), 'Sg' (gas saturation), 'So' (oil saturation), and 'T' (temperature).\n:type name: str\n\n:return: An array of fluid state variables.\n:type return: NDArray[float]"
+    "getFluidStateVariable": {
+        "signature": "get_fluid_state_variable(name: str) -> NDArray[float]",
+        "doc": "Retrieve a fluid state variable for the simulation grid.\n\n:param name: The name of the variable. Valid names are 'pw' (pressure water), 'pg' (pressure gas), 'po' (pressure oil), 'rho_w' (density water), 'rho_g' (density gas), 'rho_o' (density oil)'Rs' (soultion gas-oil ratio), 'Rv' (volatile gas-oil ratio), 'Sw' (water saturation), 'Sg' (gas saturation), 'So' (oil saturation), and 'T' (temperature).\n:type name: str\n\n:return: An array of fluid state variables.\n:type return: NDArray[float]"
     },
     "getPorosity": {
         "signature": "get_porosity() -> numpy.ndarray",
@@ -37,15 +37,15 @@
     },
     "getPrimaryVarMeaning": {
         "signature": "get_primary_var_meaning(variable: str) -> NDArray[int]",
-        "doc": "Retrieve the primary variable meaning of the simulation grid.\n\n:para variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: An array of primary variable meanings. See ``get_primary_variable_meaning_map()`` for more information.\n:type return: NDArray[int]"
+        "doc": "Retrieve the primary variable meaning of the simulation grid.\n\n:param variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: An array of primary variable meanings. See ``get_primary_variable_meaning_map()`` for more information.\n:type return: NDArray[int]"
     },
     "getPrimaryVarMeaningMap": {
         "signature": "get_primary_var_meaning_map(variable: str) -> dict[str, int]",
-        "doc": "Retrieve the primary variable meaning map for each primary variable.\n\n:para variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: A dictionary of primary variable meanings. The keys are the primary variable meanings and the values are the corresponding integer codes. The integer codes are used to represent the primary variable meanings in the simulation grid. For variable name 'pressure', the valid keys are: 'Po', 'Pg', and 'Pw', for variable name 'water', the valid keys are: 'Sw', 'Rvw', 'Rsw', and 'Disabled', for variable name 'gas', the valid keys are: 'Sg', 'Rs', 'Rv', and 'Disabled', for variable name 'brine', the valid keys are: 'Cs', 'Sp', and 'Disabled'.\n:type return: dict[str, int]"
+        "doc": "Retrieve the primary variable meaning map for each primary variable.\n\n:param variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: A dictionary of primary variable meanings. The keys are the primary variable meanings and the values are the corresponding integer codes. The integer codes are used to represent the primary variable meanings in the simulation grid. For variable name 'pressure', the valid keys are: 'Po', 'Pg', and 'Pw', for variable name 'water', the valid keys are: 'Sw', 'Rvw', 'Rsw', and 'Disabled', for variable name 'gas', the valid keys are: 'Sg', 'Rs', 'Rv', and 'Disabled', for variable name 'brine', the valid keys are: 'Cs', 'Sp', and 'Disabled'.\n:type return: dict[str, int]"
     },
     "getPrimaryVariable": {
         "signature": "get_primary_variable(variable: str) -> NDArray[float]",
-        "doc": "Retrieve the primary variable's values for the simulation grid.\n\n:para variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: An array of primary variable values. See ``get_primary_variable_meaning()`` for more information.\n:type return: NDArray[float]"
+        "doc": "Retrieve the primary variable's values for the simulation grid.\n\n:param variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n\n:return: An array of primary variable values. See ``get_primary_variable_meaning()`` for more information.\n:type return: NDArray[float]"
     },
     "run": {
         "signature": "run() -> int",
@@ -57,7 +57,11 @@
     },
     "setPrimaryVariable": {
         "signature": "set_primary_variable(variable: str, value: NDArray[float]) -> None",
-        "doc": "Set the primary variable's values for the simulation grid.\n\n:para variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n:para value: An array of primary variable values to be set. See ``get_primary_variable()`` for more information.\n:type value: NDArray[float]"
+        "doc": "Set the primary variable's values for the simulation grid.\n\n:param variable: The name of the variable. Valid names are 'pressure', 'water', 'gas', and 'brine'.\n:type variable: str\n:param value: An array of primary variable values to be set. See ``get_primary_variable()`` for more information.\n:type value: NDArray[float]"
+    },
+    "setupMpi": {
+        "signature": "mpi_init(init: bool, finalize: bool) -> None",
+        "doc": "Setup MPI for parallel simulation. This method should be called before any other method.\n:param init: Whether to call ``MPI_Init()`` or not.\n:param finalize:Whether to call ``MPI_Finalize()```when the simulator object goes out of scope.\n\n:return: None"
     },
     "step": {
         "signature": "step() -> int",

python/simulators/CMakeLists.txt

@@ -60,7 +60,7 @@ if(OPM_ENABLE_PYTHON_TESTS)
   #   splitting the python tests into multiple add_test() tests instead
   #   of having a single "python -m unittest" test call that will run all
   #   the tests in the "test" sub directory.
-  foreach(case_name IN ITEMS basic fluidstate_variables primary_variables schedule throw)
+  foreach(case_name IN ITEMS basic fluidstate_variables mpi primary_variables schedule throw)
     add_test(NAME python_${case_name}
         WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/python
         COMMAND ${CMAKE_COMMAND}

python/simulators/PyBlackOilSimulator.cpp

@@ -171,6 +171,15 @@ setPrimaryVariable(
     getFluidState().setPrimaryVariable(variable, data, size_);
 }
 
+void PyBlackOilSimulator::setupMpi(bool mpi_init, bool mpi_finalize)
+{
+    if (this->has_run_init_) {
+        throw std::logic_error("mpi_init() called after step_init()");
+    }
+    this->mpi_init_ = mpi_init;
+    this->mpi_finalize_ = mpi_finalize;
+}
+
 int PyBlackOilSimulator::step()
 {
     if (!this->has_run_init_) {
@@ -211,11 +220,17 @@ int PyBlackOilSimulator::stepInit()
             this->deck_->getDataFile(),
             this->eclipse_state_,
             this->schedule_,
-            this->summary_config_
+            this->summary_config_,
+            this->mpi_init_,
+            this->mpi_finalize_
         );
     }
     else {
-        this->main_ = std::make_unique<Opm::Main>( this->deck_filename_ );
+        this->main_ = std::make_unique<Opm::Main>(
+            this->deck_filename_,
+            this->mpi_init_,
+            this->mpi_finalize_
+        );
     }
     int exit_code = EXIT_SUCCESS;
     this->flow_main_ = this->main_->initFlowBlackoil(exit_code);
@@ -293,18 +308,19 @@ void export_PyBlackOilSimulator(py::module& m)
         .def("get_cell_volumes", &PyBlackOilSimulator::getCellVolumes, getCellVolumes_docstring)
         .def("get_dt", &PyBlackOilSimulator::getDT, getDT_docstring)
         .def("get_fluidstate_variable", &PyBlackOilSimulator::getFluidStateVariable,
-            py::return_value_policy::copy, py::arg("name"))
+            py::return_value_policy::copy, getFluidStateVariable_docstring, py::arg("name"))
         .def("get_porosity", &PyBlackOilSimulator::getPorosity, getPorosity_docstring)
         .def("get_primary_variable_meaning", &PyBlackOilSimulator::getPrimaryVarMeaning,
-            py::return_value_policy::copy, py::arg("variable"))
+            py::return_value_policy::copy, getPrimaryVarMeaning_docstring, py::arg("variable"))
         .def("get_primary_variable_meaning_map", &PyBlackOilSimulator::getPrimaryVarMeaningMap,
-            py::return_value_policy::copy, py::arg("variable"))
+            py::return_value_policy::copy, getPrimaryVarMeaningMap_docstring, py::arg("variable"))
         .def("get_primary_variable", &PyBlackOilSimulator::getPrimaryVariable,
-            py::return_value_policy::copy, py::arg("variable"))
+            py::return_value_policy::copy, getPrimaryVariable_docstring, py::arg("variable"))
         .def("run", &PyBlackOilSimulator::run, run_docstring)
         .def("set_porosity", &PyBlackOilSimulator::setPorosity, setPorosity_docstring, py::arg("array"))
         .def("set_primary_variable", &PyBlackOilSimulator::setPrimaryVariable,
-            py::arg("variable"), py::arg("value"))
+            py::arg("variable"), setPrimaryVariable_docstring, py::arg("value"))
+        .def("setup_mpi", &PyBlackOilSimulator::setupMpi, setupMpi_docstring, py::arg("init"), py::arg("finalize"))
         .def("step", &PyBlackOilSimulator::step, step_docstring)
         .def("step_cleanup", &PyBlackOilSimulator::stepCleanup, stepCleanup_docstring)
         .def("step_init", &PyBlackOilSimulator::stepInit, stepInit_docstring);

python/test/test_basic.py

@@ -71,3 +71,4 @@ class TestBasic(unittest.TestCase):
             poro2 = sim.get_porosity()
             self.assertAlmostEqual(poro2[0], 0.285, places=7, msg='value of porosity 2')
 
+

python/test/test_mpi.py

@@ -0,0 +1,34 @@
+import os
+import unittest
+from pathlib import Path
+from opm.simulators import BlackOilSimulator
+from .pytest_common import pushd
+
+class TestBasic(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # NOTE: See comment in test_basic.py for the reason why we are
+        #   only using a single test_all() function instead of splitting
+        #   it up in multiple test functions
+        test_dir = Path(os.path.dirname(__file__))
+        cls.data_dir = test_dir.parent.joinpath("test_data/SPE1CASE1a")
+
+    # IMPORTANT: Tests are run alphabetically, so we need to make sure that
+    #  the the first test calls MPI_Init(), therefore the name of the tests
+    #  have a numeric label like "01" in test_01_mpi_init to ensure that they
+    #  are run in a given order.
+    def test_01_mpi_init(self):
+        with pushd(self.data_dir):
+            sim = BlackOilSimulator("SPE1CASE1.DATA")
+            sim.setup_mpi(init=True, finalize=False)
+            sim.step_init()  # This will create the OPM::Main() object which will call MPI_Init()
+            assert True
+
+    def test_02_mpi_no_init(self):
+        with pushd(self.data_dir):
+            sim = BlackOilSimulator("SPE1CASE1.DATA")
+            sim.setup_mpi(init=False, finalize=True)
+            sim.step_init()  # This will create the OPM::Main() object which will not call MPI_Init()
+            # That this test runs shows that the simulator does not call
+            # MPI_Init() a second time
+            assert True

python/test/test_throw.py

@@ -20,3 +20,4 @@ class TestBasic(unittest.TestCase):
             #   has not been initialized
             with self.assertRaises(RuntimeError):
                 sim.get_dt()
+