Merge pull request #451 from The-Strategy-Unit/improve_exception_handling_in_containers

improve exception handling in containers

Author: tomjemmett

src/nhp/docker/__main__.py

@@ -10,10 +10,9 @@
 from nhp.model.run import run_all
 
 
-# %%
 def _exit_container():
     logging.error("\nTimed out, killing container")
-    os._exit(1)
+    os._exit(2)
 
 
 def parse_args():
@@ -38,7 +37,6 @@ def parse_args():
     return parser.parse_args()
 
 
-# %%
 def main():
     """The main method."""
     args = parse_args()
@@ -71,18 +69,25 @@ def main():
     logging.info("complete")
 
 
-# %%
 def init():
     """Method for calling main."""
     if __name__ == "__main__":
-        # start a timer to kill the container if we reach a timeout
-        t = threading.Timer(config.CONTAINER_TIMEOUT_SECONDS, _exit_container)
-        t.start()
-        # run the model
-        main()
-        # cancel the timer
-        t.cancel()
+        exc = None
+        try:
+            # start a timer to kill the container if we reach a timeout
+            t = threading.Timer(config.CONTAINER_TIMEOUT_SECONDS, _exit_container)
+            t.start()
+            # run the model
+            main()
+        except Exception as e:
+            logging.error("An error occurred: %s", str(e))
+            exc = e
+        finally:
+            # cancel the timer
+            t.cancel()
+            # if there was an exception, raise it
+            if exc is not None:
+                raise exc
 
 
-# %%
 init()

src/nhp/model/run.py

@@ -6,7 +6,6 @@
 from multiprocessing import Pool
 from typing import Any, Callable, Tuple, Type
 
-import pandas as pd
 from tqdm.auto import tqdm as base_tqdm
 
 from nhp.model.aae import AaEModel
@@ -80,13 +79,14 @@ def _run_model(
 
     # model run 0 is the baseline
     # model run 1:n are the monte carlo sims
-    model_runs = list(range(params["model_runs"] + 1))
+    model_runs = [i + 1 for i in range(params["model_runs"])]
 
     cpus = os.cpu_count()
     batch_size = int(os.getenv("BATCH_SIZE", "1"))
 
     with Pool(cpus) as pool:
-        results: list[ModelRunResult] = list(
+        baseline = model.go(0)  # baseline
+        model_results: list[ModelRunResult] = list(
             tqdm(
                 pool.imap(
                     model.go,
@@ -101,7 +101,7 @@ def _run_model(
     # ensure that the callback reports all model runs are complete
     progress_callback(params["model_runs"])
 
-    return results
+    return [baseline, *model_results]
 
 
 def noop_progress_callback(_: Any) -> Callable[[Any], None]:

tests/unit/nhp/docker/test___main__.py

@@ -15,7 +15,7 @@ def test_exit_container(mocker):
 
     r._exit_container()
 
-    m.assert_called_once_with(1)
+    m.assert_called_once_with(2)
 
 
 @pytest.mark.parametrize(
@@ -139,12 +139,12 @@ def test_init_timeout_call_exit(mocker):
     import nhp.docker.__main__ as r
 
     main_mock = mocker.patch("nhp.docker.__main__.main")
-    exit_mock = mocker.patch("nhp.docker.__main__._exit_container")
+    exit_container_mock = mocker.patch("nhp.docker.__main__._exit_container")
     main_mock.side_effect = lambda: time.sleep(0.2)
     with patch.object(r, "__name__", "__main__"):
         r.init()
 
-    exit_mock.assert_called_once()
+    exit_container_mock.assert_called_once()
 
 
 def test_init_timeout_dont_call_exit(mocker):
@@ -153,9 +153,25 @@ def test_init_timeout_dont_call_exit(mocker):
     import nhp.docker.__main__ as r
 
     main_mock = mocker.patch("nhp.docker.__main__.main")
-    exit_mock = mocker.patch("nhp.docker.__main__._exit_container")
+    exit_container_mock = mocker.patch("nhp.docker.__main__._exit_container")
     main_mock.side_effect = lambda: time.sleep(0.02)
     with patch.object(r, "__name__", "__main__"):
         r.init()
 
-    exit_mock.assert_not_called()
+    exit_container_mock.assert_not_called()
+
+
+def test_init_catches_exception(mocker):
+    # arrange
+    mocker.patch("nhp.docker.__main__.main", side_effect=Exception("Test error"))
+    import nhp.docker.__main__ as r
+
+    m = mocker.patch("logging.error")
+
+    # act
+    with patch.object(r, "__name__", "__main__"):
+        with pytest.raises(Exception, match="Test error"):
+            r.init()
+
+    # assert
+    m.assert_called_once_with("An error occurred: %s", "Test error")

tests/unit/nhp/model/test_run.py

@@ -62,7 +62,7 @@ def test_run_model(mocker):
     actual = _run_model(model_m, params, "data", "hsa", "run_params", pc_m, False)  # type: ignore
 
     # assert
-    pool_ctm.imap.assert_called_once_with(model_m().go, [0, 1, 2], chunksize=1)
+    pool_ctm.imap.assert_called_once_with(model_m().go, [1, 2], chunksize=1)
     assert actual == [model_m().go()] * 3
     pc_m.assert_called_once_with(2)