Merge pull request #499 from The-Strategy-Unit/multiprocessing_spawn

switches to using spawn multiprocessing

Author: tomjemmett

src/nhp/model/run.py

@@ -1,9 +1,9 @@
 """Run the model."""
 
 import logging
+import multiprocessing
 import os
 import time
-from multiprocessing import Pool
 from typing import Any, Callable, Tuple, Type
 
 from tqdm.auto import tqdm as base_tqdm
@@ -84,7 +84,8 @@ def _run_model(
     cpus = os.cpu_count()
     batch_size = int(os.getenv("BATCH_SIZE", "1"))
 
-    with Pool(cpus) as pool:
+    ctx = multiprocessing.get_context("spawn")
+    with ctx.Pool(cpus) as pool:
         baseline = model.go(0)  # baseline
         model_results: list[ModelRunResult] = list(
             tqdm(

tests/unit/nhp/model/test_run.py

@@ -51,13 +51,16 @@ def test_run_model(mocker):
     params = {"start_year": 2020, "end_year": 2022, "model_runs": 2}
     mocker.patch("os.cpu_count", return_value=2)
 
-    pool_mock = mocker.patch("nhp.model.run.Pool")
+    pool_ctx_mock = mocker.patch("multiprocessing.get_context")
+    pool_mock = pool_ctx_mock().Pool
     pool_ctm = pool_mock.return_value.__enter__.return_value
     pool_ctm.name = "pool"
     pool_ctm.imap = Mock(wraps=lambda f, i, **kwargs: map(f, i))
 
     pc_m = Mock()
 
+    pool_ctx_mock.reset_mock()
+
     # act
     actual = _run_model(model_m, params, "data", "hsa", "run_params", pc_m, False)  # type: ignore
 
@@ -66,6 +69,9 @@ def test_run_model(mocker):
     assert actual == [model_m().go()] * 3
     pc_m.assert_called_once_with(2)
 
+    pool_ctx_mock.assert_called_once_with("spawn")
+    pool_mock.assert_called_once_with(2)
+
 
 def test_noop_progress_callback():
     # arrange, act & assert