Keep argument boundaries when building Auto3DSeg commands

monai/apps/auto3dseg/auto_runner.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import os
+import shlex
 import shutil
 import warnings
 from copy import deepcopy
@@ -804,22 +805,21 @@ def _train_algo_in_nni(self, history: list[dict[str, Any]]) -> None:
 
             # pyrefly: ignore [redundant-cast]
             max_trial = min(self.hpo_tasks, cast(int, default_nni_config["maxTrialNumber"]))
-            cmd = "nnictl create --config " + nni_config_filename + " --port 8088"
+            cmd = ["nnictl", "create", "--config", nni_config_filename, "--port", "8088"]
 
             if mode_dry_run:
-                logger.info(f"AutoRunner HPO is in dry-run mode. Please manually launch: {cmd}")
+                logger.info(f"AutoRunner HPO is in dry-run mode. Please manually launch: {shlex.join(cmd)}")
                 continue
 
-            run_cmd(cmd.split(), check=True)
+            run_cmd(cmd, check=True)
 
             n_trainings = len(import_bundle_algo_history(self.work_dir, only_trained=True))
             while n_trainings - last_total_tasks < max_trial:
                 sleep(1)
                 n_trainings = len(import_bundle_algo_history(self.work_dir, only_trained=True))
 
-            cmd = "nnictl stop --all"
             # pyrefly: ignore [bad-argument-type]
-            run_cmd(cmd.split(), check=True)
+            run_cmd(["nnictl", "stop", "--all"], check=True)
             logger.info(f"NNI completes HPO on {name}")
             last_total_tasks = n_trainings
 

monai/apps/auto3dseg/bundle_gen.py

@@ -15,6 +15,7 @@
 import ntpath
 import os
 import re
+import shlex
 import shutil
 import subprocess
 import sys
@@ -284,19 +285,19 @@ def _create_cmd(self, train_params: None | dict = None) -> tuple[str, str]:
 
             return (
                 _prepare_cmd_bcprun(
-                    f"{train_py} run",
+                    f"{shlex.quote(train_py)} run",
                     cmd_prefix=f"{self.device_setting['CMD_PREFIX']}",
                     config_file=config_files,
                     **params,
                 ),
                 "",
             )
         elif int(self.device_setting["n_devices"]) > 1:
-            return _prepare_cmd_torchrun(f"{train_py} run", config_file=config_files, **params), ""
+            return _prepare_cmd_torchrun(f"{shlex.quote(train_py)} run", config_file=config_files, **params), ""
         else:
             return (
                 _prepare_cmd_default(
-                    f"{train_py} run",
+                    f"{shlex.quote(train_py)} run",
                     cmd_prefix=f"{self.device_setting['CMD_PREFIX']}",
                     config_file=config_files,
                     **params,
@@ -332,7 +333,7 @@ def _run_cmd(self, cmd: str, devices_info: str = "") -> subprocess.CompletedProc
                 cmd, nnodes=1, nproc_per_node=self.device_setting["n_devices"], env=ps_environ, check=True
             )
         else:
-            return run_cmd(cmd.split(), run_cmd_verbose=True, env=ps_environ, check=True)
+            return run_cmd(shlex.split(cmd), run_cmd_verbose=True, env=ps_environ, check=True)
 
     def train(
         self, train_params: None | dict = None, device_setting: None | dict = None
@@ -456,10 +457,54 @@ def load_state_dict(self, state: dict) -> None:
         Args:
             state: A dictionary containing the state to restore.
         """
+        allowed = set(self.state_dict())
+        unknown = sorted(set(state) - allowed)
+        if unknown:
+            raise ValueError(
+                f"Unexpected keys in {type(self).__name__} state: {unknown}. Expected a subset of {sorted(allowed)}."
+            )
+        if "device_setting" in state:
+            state = dict(state)
+            # merge into the current settings so fields missing from the saved state keep their defaults
+            state["device_setting"] = {**self.device_setting, **_validate_device_setting(state["device_setting"])}
         for key, value in state.items():
             setattr(self, key, value)
 
 
+def _validate_device_setting(device_setting: Any) -> dict:
+    """
+    Check a restored ``device_setting`` dictionary before it is used to build training commands.
+
+    Args:
+        device_setting: the value restored from a saved state.
+
+    Returns:
+        a copy of ``device_setting`` with the numeric fields converted to ``int``.
+
+    Raises:
+        ValueError: if the value is not a dictionary, has unknown keys, or has malformed fields.
+    """
+    known = {"CUDA_VISIBLE_DEVICES", "n_devices", "NUM_NODES", "MN_START_METHOD", "CMD_PREFIX"}
+    if not isinstance(device_setting, dict):
+        raise ValueError(f"device_setting must be a dict, got {type(device_setting).__name__}.")
+    unknown = sorted(set(device_setting) - known)
+    if unknown:
+        raise ValueError(f"Unexpected keys in device_setting: {unknown}.")
+    checked = dict(device_setting)
+    for key in ("n_devices", "NUM_NODES"):
+        if key in checked:
+            value = checked[key]
+            if isinstance(value, bool) or not re.fullmatch(r"\d+", str(value)):
+                raise ValueError(f"device_setting[{key!r}] must be a non-negative integer, got {value!r}.")
+            checked[key] = int(value)
+    if "CUDA_VISIBLE_DEVICES" in checked and not re.fullmatch(r"[\w,-]*", str(checked["CUDA_VISIBLE_DEVICES"])):
+        raise ValueError(f"device_setting['CUDA_VISIBLE_DEVICES'] is malformed: {checked['CUDA_VISIBLE_DEVICES']!r}.")
+    for key in ("MN_START_METHOD", "CMD_PREFIX"):
+        if key in checked and not isinstance(checked[key], str):
+            raise ValueError(f"device_setting[{key!r}] must be a string, got {type(checked[key]).__name__}.")
+    return checked
+
+
 # path to download the algo_templates
 default_algo_zip = (
     f"https://github.com/Project-MONAI/research-contributions/releases/download/algo_templates/{ALGO_HASH}.tar.gz"

monai/apps/auto3dseg/ensemble_builder.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import os
+import shlex
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
@@ -628,16 +629,19 @@ def _create_cmd(self) -> None:
             return
 
         # define base cmd for subprocess
-        base_cmd = f"monai.apps.auto3dseg EnsembleRunner ensemble \
-                --data_src_cfg_name {self.data_src_cfg_name} \
-                --work_dir {self.work_dir} \
-                --num_fold {self.num_fold} \
-                --ensemble_method_name {self.ensemble_method_name} \
-                --mgpu True"
+        # every value is shell-quoted so it stays a single argument when the command is split
+        base_cmd = (
+            "monai.apps.auto3dseg EnsembleRunner ensemble"
+            f" --data_src_cfg_name {shlex.quote(str(self.data_src_cfg_name))}"
+            f" --work_dir {shlex.quote(str(self.work_dir))}"
+            f" --num_fold {shlex.quote(str(self.num_fold))}"
+            f" --ensemble_method_name {shlex.quote(str(self.ensemble_method_name))}"
+            " --mgpu True"
+        )
 
         if self.kwargs and isinstance(self.kwargs, Mapping):
             for k, v in self.kwargs.items():
-                base_cmd += f" --{k}={v}"
+                base_cmd += " " + shlex.quote(f"--{k}={v}")
         # define env for subprocess
         ps_environ = os.environ.copy()
         ps_environ["CUDA_VISIBLE_DEVICES"] = str(self.device_setting["CUDA_VISIBLE_DEVICES"])

monai/auto3dseg/utils.py

@@ -15,6 +15,7 @@
 import logging
 import os
 import pickle
+import shlex
 import subprocess
 import sys
 import warnings
@@ -585,14 +586,20 @@ def list_to_python_fire_arg_str(args: list) -> str:
 
 
 def check_and_set_optional_args(params: dict) -> str:
-    """convert `params` into '--key_1=value_1 --key_2=value_2 ...'"""
+    """
+    Convert `params` into ' --key_1=value_1 --key_2=value_2 ...'.
+
+    Each ``--key=value`` option is shell-quoted as a whole, so a value is always a single argument when the
+    command is split with ``shlex.split`` (as ``_run_cmd_torchrun`` and ``BundleAlgo._run_cmd`` do), even if
+    it contains whitespace or quotes.
+    """
     cmd_mod_opt = ""
     for k, v in params.items():
         if isinstance(v, dict):
             raise ValueError("Nested dict is not supported.")
         elif isinstance(v, list):
             v = list_to_python_fire_arg_str(v)
-        cmd_mod_opt += f" --{k}={v}"
+        cmd_mod_opt += " " + shlex.quote(f"--{k}={v}")
     return cmd_mod_opt
 
 
@@ -685,7 +692,7 @@ def _run_cmd_torchrun(cmd: str, **kwargs: Any) -> subprocess.CompletedProcess:
     """
     params = kwargs.copy()
 
-    cmd_list = cmd.split()
+    cmd_list = shlex.split(cmd)
 
     # append arguments to the command list
     torchrun_list = ["torchrun"]

tests/auto3dseg/test_command_args.py

@@ -0,0 +1,137 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import shlex
+import unittest
+from unittest import mock
+
+from monai.apps.auto3dseg.bundle_gen import BundleAlgo
+from monai.apps.auto3dseg.ensemble_builder import EnsembleRunner
+from monai.auto3dseg.utils import _prepare_cmd_default, _run_cmd_torchrun, check_and_set_optional_args
+
+
+def _device_setting(n_devices: int = 1) -> dict:
+    return {
+        "CUDA_VISIBLE_DEVICES": ",".join(str(i) for i in range(n_devices)),
+        "n_devices": n_devices,
+        "NUM_NODES": 1,
+        "MN_START_METHOD": "bcprun",
+        "CMD_PREFIX": "",
+    }
+
+
+class TestOptionalArgs(unittest.TestCase):
+    def test_plain_values_keep_their_tokens(self):
+        """Values without special characters produce the same arguments as before."""
+        cmd = _prepare_cmd_default("train.py run", num_epochs=2, lr=0.1, config_file=["a.yaml", "b.yaml"])
+        self.assertEqual(
+            shlex.split(cmd),
+            ["python", "train.py", "run", "--num_epochs=2", "--lr=0.1", "--config_file='a.yaml,b.yaml'"],
+        )
+
+    def test_value_with_whitespace_is_one_argument(self):
+        for value in ("/data/my dir", "/tmp/x\t--work_dir\t/other", "a b --c=d"):
+            with self.subTest(value=value):
+                args = shlex.split(check_and_set_optional_args({"work_dir": value}))
+                self.assertEqual(args, [f"--work_dir={value}"])
+
+    def test_torchrun_keeps_argument_boundaries(self):
+        cmd = "train.py run" + check_and_set_optional_args({"output_dir": "/tmp/a b"})
+        with mock.patch("monai.auto3dseg.utils.run_cmd") as run:
+            _run_cmd_torchrun(cmd, nnodes=1, nproc_per_node=2)
+        self.assertEqual(
+            run.call_args.args[0],
+            ["torchrun", "--nnodes", "1", "--nproc_per_node", "2", "train.py", "run", "--output_dir=/tmp/a b"],
+        )
+
+
+class TestEnsembleRunnerCommand(unittest.TestCase):
+    def _runner(self, work_dir: str, **kwargs) -> EnsembleRunner:
+        runner = EnsembleRunner.__new__(EnsembleRunner)
+        runner.data_src_cfg_name = "input.yaml"
+        runner.work_dir = work_dir
+        runner.num_fold = 5
+        runner.ensemble_method_name = "AlgoEnsembleBestByFold"
+        runner.kwargs = kwargs
+        runner.device_setting = _device_setting(n_devices=2)
+        return runner
+
+    def test_values_stay_single_arguments(self):
+        work_dir = "/tmp/my work\t--work_dir\t/elsewhere"
+        runner = self._runner(work_dir, sigmoid="True --output_dir=/elsewhere")
+        with mock.patch("monai.apps.auto3dseg.ensemble_builder._run_cmd_torchrun") as run:
+            runner._create_cmd()
+        args = shlex.split(run.call_args.args[0])
+        self.assertEqual(args[args.index("--work_dir") + 1], work_dir)
+        self.assertEqual(args.count("--work_dir"), 1)
+        self.assertIn("--sigmoid=True --output_dir=/elsewhere", args)
+        self.assertFalse(any(a.startswith("--output_dir") for a in args))
+
+
+class TestBundleAlgoCommand(unittest.TestCase):
+    def test_output_path_with_whitespace_is_one_argument(self):
+        algo = BundleAlgo(template_path=None)
+        algo.output_path = "/tmp/my algo"
+        algo.device_setting = _device_setting(n_devices=1)
+        with mock.patch("monai.apps.auto3dseg.bundle_gen.run_cmd") as run:
+            algo.train({"num_epochs": 1})
+        args = run.call_args.args[0]
+        self.assertEqual(args[:3], ["python", os.path.join("/tmp/my algo", "scripts", "train.py"), "run"])
+        self.assertIn("--num_epochs=1", args)
+
+
+class TestBundleAlgoLoadState(unittest.TestCase):
+    def test_round_trip(self):
+        algo = BundleAlgo(template_path="/templates")
+        algo.name = "segresnet_0"
+        algo.device_setting = _device_setting(n_devices=2)
+        state = algo.state_dict()
+        restored = BundleAlgo()
+        restored.load_state_dict(state)
+        self.assertEqual(restored.state_dict(), state)
+
+    def test_unknown_key_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, "Unexpected keys"):
+            BundleAlgo().load_state_dict({"train": "x"})
+
+    def test_malformed_device_setting_is_rejected(self):
+        for setting in (
+            {"n_devices": "2 --master_port 1"},
+            {"NUM_NODES": "1 2"},
+            {"CUDA_VISIBLE_DEVICES": "0 --x"},
+            {"CMD_PREFIX": ["python"]},
+            {"OTHER": 1},
+            "not a dict",
+        ):
+            with self.subTest(setting=setting):
+                with self.assertRaises(ValueError):
+                    BundleAlgo().load_state_dict({"device_setting": setting})
+
+    def test_partial_device_setting_keeps_defaults(self):
+        algo = BundleAlgo()
+        defaults = dict(algo.device_setting)
+        algo.load_state_dict({"device_setting": {"CUDA_VISIBLE_DEVICES": "0", "n_devices": 1, "NUM_NODES": 1}})
+        self.assertEqual(algo.device_setting["CMD_PREFIX"], defaults["CMD_PREFIX"])
+        self.assertEqual(algo.device_setting["MN_START_METHOD"], defaults["MN_START_METHOD"])
+        self.assertEqual(algo.device_setting["CUDA_VISIBLE_DEVICES"], "0")
+
+    def test_numeric_strings_are_converted(self):
+        algo = BundleAlgo()
+        algo.load_state_dict({"device_setting": {"n_devices": "2", "NUM_NODES": "1", "CUDA_VISIBLE_DEVICES": "0,1"}})
+        self.assertEqual(algo.device_setting["n_devices"], 2)
+        self.assertEqual(algo.device_setting["NUM_NODES"], 1)
+
+
+if __name__ == "__main__":
+    unittest.main()