Fix for multiprocessing on Linux (#927)

blue-marble · Jun 10, 2022 · 57a43d9 · 57a43d9
1 parent 1d0adb6
commit 57a43d9
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 33 deletions.
diff --git a/gridpath/get_scenario_inputs.py b/gridpath/get_scenario_inputs.py
@@ -23,7 +23,7 @@
 
 from argparse import ArgumentParser
 import csv
-from multiprocessing import Pool
+from multiprocessing import get_context
 import os.path
 import pandas as pd
 import sys
@@ -132,7 +132,8 @@ def write_model_inputs(
             ]
         )
 
-        pool = Pool(n_parallel_subproblems)
+        # Pool must use spawn to work properly on Linux
+        pool = get_context("spawn").Pool(n_parallel_subproblems)
         pool.map(get_inputs_for_subproblem_pool, pool_data)
         pool.close()
 

diff --git a/gridpath/run_scenario.py b/gridpath/run_scenario.py
@@ -23,7 +23,7 @@
 import argparse
 from csv import reader, writer
 import datetime
-from multiprocessing import Pool, Manager
+from multiprocessing import get_context, Manager
 import os.path
 from pyomo.environ import (
     AbstractModel,
@@ -406,7 +406,8 @@ def run_scenario(scenario_directory, subproblem_structure, parsed_arguments):
             for subproblem in subproblem_structure.SUBPROBLEM_STAGES.keys():
                 objective_values[subproblem] = manager.dict()
 
-            pool = Pool(n_parallel_subproblems)
+            # Pool must use spawn to work properly on Linux
+            pool = get_context("spawn").Pool(n_parallel_subproblems)
             pool_data = tuple(
                 [
                     [

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -530,38 +530,30 @@ def test_example_single_stage_prod_cost_cycle_select(self):
     def test_example_multi_stage_prod_cost_parallel(self):
         """
         Check validation and objective function values of
-        "multi_stage_prod_cost" example
+        "multi_stage_prod_cost" example running subproblems in parallel
         :return:
         """
-        # TODO: figure why run_e2e processed gets terminated on linux when
-        #  using parallel processing; skip test on linux for the time being
-        if platform.system() == "Linux":
-            print(
-                "Skipping test_example_multi_stage_prod_cost_parallel on ",
-                platform.system(),
-            )
-        else:
-            self.run_and_check_objective(
-                "multi_stage_prod_cost",
-                {
-                    1: {
-                        1: -1265436373826.0408,
-                        2: -1265436373826.0408,
-                        3: -1265436373826.099,
-                    },
-                    2: {
-                        1: -1265436373826.0408,
-                        2: -1265436373826.0408,
-                        3: -1265436373826.099,
-                    },
-                    3: {
-                        1: -1265436373826.0408,
-                        2: -1265436373826.0408,
-                        3: -1265436373826.099,
-                    },
+        self.run_and_check_objective(
+            "multi_stage_prod_cost",
+            {
+                1: {
+                    1: -1265436373826.0408,
+                    2: -1265436373826.0408,
+                    3: -1265436373826.099,
                 },
-                parallel=3,
-            )
+                2: {
+                    1: -1265436373826.0408,
+                    2: -1265436373826.0408,
+                    3: -1265436373826.099,
+                },
+                3: {
+                    1: -1265436373826.0408,
+                    2: -1265436373826.0408,
+                    3: -1265436373826.099,
+                },
+            },
+            parallel=3,
+        )
 
     def test_example_multi_stage_prod_cost_w_hydro(self):
         """