Harden nnUNetV2Runner against code execution

monai/apps/nnunet/nnunetv2_runner.py

@@ -17,6 +17,8 @@
 import re
 import shlex
 import subprocess
+import warnings
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 import monai
@@ -710,7 +712,11 @@ def train_parallel(
         **kwargs: Any,
     ) -> None:
         """
-        Create the line command for subprocess call for parallel training.
+        Launch subprocesses for parallel training.
+
+        The commands for each GPU run sequentially on that device, while different devices run in
+        parallel. Each stage waits for all of its devices to finish before the next stage starts.
+
         Note: to set the number of GPUs to use, use ``gpu_id_for_all`` instead of the `CUDA_VISIBLE_DEVICES`
         environment variable.
 
@@ -741,17 +747,19 @@ def train_parallel(
                     f"log '.txt' inside '{os.path.join(self.nnunet_results, self.dataset_name)}'"
                 )
         for stage in all_cmds:
-            processes = []
-            for device_id in stage:
-                if not stage[device_id]:
-                    continue
-                cmd_str = "; ".join(shlex.join(cmd) for cmd, _ in stage[device_id])
-                env = stage[device_id][0][1]
-                logger.info(f"Current running command on GPU device {device_id}:\n{cmd_str}\n")
-                processes.append(subprocess.Popen(cmd_str, shell=True, env=env, stdout=subprocess.DEVNULL))
-            # finish this stage first
-            for p in processes:
-                p.wait()
+            device_cmds = [(device_id, gpu_cmds) for device_id, gpu_cmds in stage.items() if gpu_cmds]
+            if not device_cmds:
+                continue
+
+            def _run_device_commands(item):
+                device_id, gpu_cmds = item
+                for cmd, env in gpu_cmds:
+                    cmd_str = shlex.join(cmd)
+                    logger.info(f"Current running command on GPU device {device_id}:\n{cmd_str}\n")
+                    subprocess.Popen(cmd, shell=False, env=env, stdout=subprocess.DEVNULL).wait()
+
+            with ThreadPoolExecutor(max_workers=len(device_cmds)) as executor:
+                list(executor.map(_run_device_commands, device_cmds))
 
     def validate_single_model(self, config: str, fold: int, **kwargs: Any) -> None:
         """
@@ -996,7 +1004,16 @@ def predict_ensemble_postprocessing(
 
         # apply postprocessing
         if run_postprocessing:
-            pp_fns, pp_fn_kwargs = load_pickle(self.best_configuration["best_model_or_ensemble"]["postprocessing_file"])
+            postprocessing_file = self.best_configuration["best_model_or_ensemble"]["postprocessing_file"]
+            warnings.warn(
+                f"unpickling postprocessing_file {postprocessing_file}: this path is read from "
+                "inference_information.json and is loaded with Python pickle without any allow list, "
+                "which gives whoever controls that file arbitrary code execution. Only proceed if the "
+                "inference_information.json is from a source you trust "
+                "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-8f32-8649-rv87).",
+                stacklevel=2,
+            )
+            pp_fns, pp_fn_kwargs = load_pickle(postprocessing_file)
             apply_postprocessing_to_folder(
                 folder_for_pp,
                 join(target_dir_base, "ensemble_predictions_postprocessed"),

tests/apps/nnunet/test_nnunetv2_runner_command.py

@@ -11,10 +11,17 @@
 
 from __future__ import annotations
 
+import os
+import sys
+import threading
+import types
 import unittest
+import warnings
 from unittest import mock
 
+from monai.apps.nnunet import nnunetv2_runner
 from monai.apps.nnunet.nnunetv2_runner import nnUNetV2Runner
+from monai.bundle import ConfigParser
 
 
 def _make_runner(export_validation_probabilities=False):
@@ -74,5 +81,114 @@ def test_validate_emits_bare_val_flag(self):
         self.assertNotIn("True", cmd)
 
 
+class TestTrainParallelCommand(unittest.TestCase):
+    def test_train_parallel_uses_argv_list_without_shell(self):
+        runner = _make_runner()
+        runner.dataset_name = "Dataset001_Test"
+        runner.nnunet_results = "/tmp/nnunet_results"
+
+        all_cmds = [
+            {
+                0: [
+                    (["python", "-m", "train", "--fold", "0"], {"CUDA_VISIBLE_DEVICES": "0"}),
+                    (["python", "-m", "train", "--fold", "1"], {"CUDA_VISIBLE_DEVICES": "0"}),
+                ],
+                1: [(["python", "-m", "train", "--fold", "2"], {"CUDA_VISIBLE_DEVICES": "1"})],
+            }
+        ]
+
+        with mock.patch.object(runner, "train_parallel_cmd", return_value=all_cmds):
+            with mock.patch("monai.apps.nnunet.nnunetv2_runner.subprocess.Popen") as popen:
+                popen.return_value.wait.return_value = None
+                runner.train_parallel()
+
+        self.assertEqual(popen.call_count, 3)
+        for call in popen.call_args_list:
+            self.assertIsInstance(call.args[0], list)
+            self.assertFalse(call.kwargs["shell"])
+
+    def test_commands_run_sequentially_per_device(self):
+        runner = _make_runner()
+        runner.dataset_name = "Dataset001_Test"
+        runner.nnunet_results = "/tmp/nnunet_results"
+
+        all_cmds = [
+            {0: [(["python", "-m", "train", "--fold", "0"], {}), (["python", "-m", "train", "--fold", "1"], {})]}
+        ]
+
+        events = []
+        lock = threading.Lock()
+
+        class _FakeProcess:
+            def __init__(self, cmd):
+                self.cmd = cmd
+
+            def wait(self):
+                with lock:
+                    events.append(("wait", self.cmd))
+                return 0
+
+        def _fake_popen(cmd, *args, **kwargs):
+            with lock:
+                events.append(("popen", cmd))
+            return _FakeProcess(cmd)
+
+        with mock.patch.object(runner, "train_parallel_cmd", return_value=all_cmds):
+            with mock.patch("monai.apps.nnunet.nnunetv2_runner.subprocess.Popen", side_effect=_fake_popen):
+                runner.train_parallel()
+
+        self.assertEqual(
+            events,
+            [
+                ("popen", ["python", "-m", "train", "--fold", "0"]),
+                ("wait", ["python", "-m", "train", "--fold", "0"]),
+                ("popen", ["python", "-m", "train", "--fold", "1"]),
+                ("wait", ["python", "-m", "train", "--fold", "1"]),
+            ],
+        )
+
+
+class TestPredictEnsemblePostprocessingWarnings(unittest.TestCase):
+    def test_postprocessing_pickle_warns_on_untrusted_file(self):
+        runner = _make_runner()
+        runner.dataset_name = "Dataset001_Test"
+        runner.nnunet_raw = "/tmp/nnunet_raw"
+        runner.nnunet_results = "/tmp/nnunet_results"
+        runner.best_configuration = {
+            "best_model_or_ensemble": {
+                "selected_model_or_models": [{"configuration": "3d_fullres"}],
+                "postprocessing_file": "/tmp/attacker_controlled_postprocessing.pkl",
+                "some_plans_file": "/tmp/plans.json",
+            }
+        }
+
+        ensemble_mod = types.ModuleType("nnunetv2.ensembling.ensemble")
+        ensemble_mod.ensemble_folders = mock.MagicMock()
+        pp_mod = types.ModuleType("nnunetv2.postprocessing.remove_connected_components")
+        pp_mod.apply_postprocessing_to_folder = mock.MagicMock()
+        fp_mod = types.ModuleType("nnunetv2.utilities.file_path_utilities")
+        fp_mod.get_output_folder = mock.MagicMock(return_value="/tmp/model_folder")
+
+        fake_modules = {
+            "nnunetv2.ensembling.ensemble": ensemble_mod,
+            "nnunetv2.postprocessing.remove_connected_components": pp_mod,
+            "nnunetv2.utilities.file_path_utilities": fp_mod,
+        }
+
+        load_pickle = mock.MagicMock(return_value=([], {}))
+        with mock.patch.dict(sys.modules, fake_modules):
+            with mock.patch.object(ConfigParser, "load_config_file", return_value=runner.best_configuration):
+                with mock.patch.object(nnunetv2_runner, "join", os.path.join):
+                    with mock.patch.object(nnunetv2_runner, "load_pickle", load_pickle):
+                        with warnings.catch_warnings(record=True) as caught:
+                            warnings.simplefilter("always")
+                            runner.predict_ensemble_postprocessing(
+                                run_predict=False, run_ensemble=False, run_postprocessing=True
+                            )
+
+        load_pickle.assert_called_once_with("/tmp/attacker_controlled_postprocessing.pkl")
+        self.assertTrue(any("unpickling postprocessing_file" in str(item.message) for item in caught))
+
+
 if __name__ == "__main__":
     unittest.main()