diff --git a/monai/apps/auto3dseg/bundle_gen.py b/monai/apps/auto3dseg/bundle_gen.py index 680de6c99f..78d5dc3566 100644 --- a/monai/apps/auto3dseg/bundle_gen.py +++ b/monai/apps/auto3dseg/bundle_gen.py @@ -11,6 +11,7 @@ from __future__ import annotations +import ast import importlib import ntpath import os @@ -39,7 +40,9 @@ _run_cmd_torchrun, algo_to_json, ) +from monai.bundle.config_item import ConfigExpression from monai.bundle.config_parser import ConfigParser +from monai.bundle.reference_resolver import ReferenceResolver from monai.config import PathLike from monai.utils import ensure_tuple, look_up_option, run_cmd from monai.utils.enums import AlgoKeys @@ -97,6 +100,103 @@ def _confined_join(root: str, name: str) -> str: return path +def _resolve_path_value(value: Any, parser: ConfigParser, what: str) -> str: + """ + Resolve a config value that must be a plain path string, without evaluating it with ``eval``. + + The value can come from a work directory this process did not write, so only the shapes the bundle + templates use are accepted: a plain string, or a ``$`` expression made of string literals, ``+`` concatenation, ``str(...)`` of a reference, and references to other + entries, for example ``$@bundle_root + '/model_fold' + str(@fold)``. Calls other than ``str`` with one + positional argument, attribute access, subscripts to anything but the reference table, comprehensions + and ``_target_`` mappings are refused. + + Args: + value: the config value, as stored in the file. + parser: the parser that holds the other entries referenced with ``@``. + what: a description of the value, used in the error messages. + + Returns: + the resolved path string. + + Raises: + ValueError: if ``value`` is not a string, or is not one of the accepted shapes. + """ + if not isinstance(value, str): + raise ValueError(f"Invalid {what}: expected a string, got {type(value).__name__}.") + if value.startswith("@"): + resolved = _resolve_path_reference(value[1:], parser, what) + if not isinstance(resolved, str): + raise ValueError(f"Invalid {what}: reference {value!r} is not a plain string.") + return resolved + if not value.startswith(ConfigExpression.prefix): + return value + + refs = { + ref: _resolve_path_reference(ref, parser, what, allow_integer=True) + for ref in ReferenceResolver.match_refs_pattern(value) + } + expression = ReferenceResolver.update_refs_pattern(value, refs) + try: + node = ast.parse(expression[len(ConfigExpression.prefix) :], mode="eval").body + except SyntaxError as e: + raise ValueError(f"Invalid {what} expression: {value!r}.") from e + resolved = _path_expression(node, refs, what) + if not isinstance(resolved, str): + raise ValueError(f"Invalid {what}: the expression does not resolve to a string.") + return resolved + + +def _resolve_path_reference(ref: str, parser: ConfigParser, what: str, allow_integer: bool = False) -> str | int: + """Return the value of ``@ref``, refusing anything that would need evaluation itself. + + A string is returned as is. An integer is only returned when ``allow_integer`` is set, because the + templates use it through ``str()``, as in ``str(@fold)``; it is never concatenated directly. + """ + try: + ref_value = parser[ref] + except (KeyError, IndexError) as e: + raise ValueError(f"Invalid {what}: reference @{ref} does not exist.") from e + if isinstance(ref_value, str): + if ref_value.startswith(("@", "$")): + raise ValueError(f"Invalid {what}: reference @{ref} is not a plain value.") + return ref_value + if allow_integer and isinstance(ref_value, int) and not isinstance(ref_value, bool): + return ref_value + raise ValueError(f"Invalid {what}: reference @{ref} is not a plain value.") + + +def _path_expression(node: ast.AST, refs: dict[str, str | int], what: str) -> str | int: + """Evaluate the small expression grammar accepted by :func:`_resolve_path_value`.""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): + left = _path_expression(node.left, refs, what) + right = _path_expression(node.right, refs, what) + if isinstance(left, str) and isinstance(right, str): + return left + right + raise ValueError(f"Invalid {what}: only strings can be concatenated.") + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "str" + and len(node.args) == 1 + and not node.keywords + ): + return str(_path_expression(node.args[0], refs, what)) + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Name) + and node.value.id == ReferenceResolver._vars + and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str) + and node.slice.value in refs + ): + return refs[node.slice.value] + raise ValueError( + f"Invalid {what}: only string literals, + concatenation, str() of a reference and @references are allowed." + ) + + class BundleAlgo(Algo): """ An algorithm represented by a set of bundle configurations and scripts. @@ -360,11 +460,29 @@ def train( def get_score(self, *args, **kwargs): """ Returns validation scores of the model trained by the current Algo. + + The configuration is read from the algorithm output directory, which may belong to a work + directory this process did not create (``import_bundle_algo_history`` reads any work dir). + ``ckpt_path`` is therefore resolved as data rather than evaluated: string literals, + ``+`` concatenation and ``@references`` such as ``$@bundle_root + '/model_fold' + str(@fold)`` + work, while any other call, an attribute, a comprehension or a ``_target_`` mapping is refused. + + Args: + *args: unused, kept for compatibility with the ``Algo`` interface. + **kwargs: unused, kept for compatibility with the ``Algo`` interface. + + Returns: + the best average Dice score recorded in ``progress.yaml``. + + Raises: + ValueError: if ``ckpt_path`` in the saved configuration is not a plain path string. """ config_yaml = os.path.join(self.output_path, "configs", "hyper_parameters.yaml") - parser = ConfigParser() + parser = ConfigParser(globals=False) parser.read_config(config_yaml) - ckpt_path = parser.get_parsed_content("ckpt_path", default=self.output_path) + # resolved as data, never with eval(): the file can come from a work directory this process + # did not write, and a `$` expression there must not be able to call or import anything + ckpt_path = _resolve_path_value(parser.get("ckpt_path", default=self.output_path), parser, "ckpt_path") dict_file = ConfigParser.load_config_file(os.path.join(ckpt_path, "progress.yaml")) # dict_file: a list of scores saved in the form of dict in progress.yaml diff --git a/monai/apps/auto3dseg/utils.py b/monai/apps/auto3dseg/utils.py index 1946bb718a..ff8243b1ac 100644 --- a/monai/apps/auto3dseg/utils.py +++ b/monai/apps/auto3dseg/utils.py @@ -11,6 +11,7 @@ from __future__ import annotations +import logging import os from monai.apps.auto3dseg.bundle_gen import BundleAlgo @@ -59,8 +60,10 @@ def import_bundle_algo_history( if best_metric is None: try: best_metric = algo.get_score() - except Exception: - pass + except Exception as e: + # a missing score is normal for an untrained algo; anything else means the directory + # could not be read as one, which is worth surfacing rather than swallowing + logging.warning(f"Could not read the validation score of {name!r} in {output_folder}: {e}") is_trained = best_metric is not None diff --git a/tests/auto3dseg/test_algo_score_config.py b/tests/auto3dseg/test_algo_score_config.py new file mode 100644 index 0000000000..cdff2858b8 --- /dev/null +++ b/tests/auto3dseg/test_algo_score_config.py @@ -0,0 +1,156 @@ +# 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 json +import os +import tempfile +import unittest + +from monai.apps.auto3dseg.bundle_gen import BundleAlgo +from monai.apps.auto3dseg.utils import import_bundle_algo_history + + +class TestAlgoScoreConfig(unittest.TestCase): + """Regression tests for the score of an algorithm output directory. + + ``import_bundle_algo_history`` reads any work directory, so the directory and the config inside it + can come from someone else. Reading the saved score must not evaluate code from that config. + """ + + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.root = self.tempdir.name + self.marker = os.path.join(self.root, "PWNED") + self.algo_dir = os.path.join(self.root, "algo1") + os.makedirs(os.path.join(self.algo_dir, "configs")) + + def tearDown(self): + self.tempdir.cleanup() + + def _write_config(self, **values) -> None: + with open(os.path.join(self.algo_dir, "configs", "hyper_parameters.yaml"), "w") as f: + for key, value in values.items(): + f.write(f"{key}: {json.dumps(value)}\n") + + def _write_progress(self, ckpt_dir: str, score: float) -> None: + os.makedirs(ckpt_dir, exist_ok=True) + with open(os.path.join(ckpt_dir, "progress.yaml"), "w") as f: + f.write(f"- best_avg_dice_score: {score}\n") + + def _algo(self) -> BundleAlgo: + algo = BundleAlgo(template_path=self.root) + algo.output_path = self.algo_dir + return algo + + def _assert_refused(self, **values) -> None: + self._write_config(**values) + with self.assertRaises(ValueError): + self._algo().get_score() + self.assertFalse(os.path.exists(self.marker), "the ckpt_path payload was executed") + + def test_plain_ckpt_path_still_resolves(self): + ckpt_dir = os.path.join(self.algo_dir, "ckpt") + self._write_progress(ckpt_dir, 0.75) + self._write_config(ckpt_path=ckpt_dir) + + self.assertAlmostEqual(self._algo().get_score(), 0.75) + + def test_string_concatenation_still_resolves(self): + ckpt_dir = os.path.join(self.algo_dir, "model_fold_0") + self._write_progress(ckpt_dir, 0.42) + self._write_config(ckpt_path=f"${self.algo_dir!r} + '/model_fold_0'") + + self.assertAlmostEqual(self._algo().get_score(), 0.42) + + def test_reference_expression_still_resolves(self): + """The form used by the shipped templates, ``$@bundle_root + '/model'``, keeps working.""" + ckpt_dir = os.path.join(self.algo_dir, "model") + self._write_progress(ckpt_dir, 0.33) + self._write_config(bundle_root=self.algo_dir, ckpt_path="$@bundle_root + '/model'") + + self.assertAlmostEqual(self._algo().get_score(), 0.33) + + def test_integer_reference_through_str_still_resolves(self): + """The ``dints`` and ``swinunetr`` templates use ``$@bundle_root + '/model_fold' + str(@fold)``.""" + ckpt_dir = os.path.join(self.algo_dir, "model_fold0") + self._write_progress(ckpt_dir, 0.61) + self._write_config(bundle_root=self.algo_dir, fold=0, ckpt_path="$@bundle_root + '/model_fold' + str(@fold)") + + self.assertAlmostEqual(self._algo().get_score(), 0.61) + + def test_integer_reference_cannot_be_concatenated(self): + """An integer reference is only usable through str(); concatenating it directly is refused.""" + self._assert_refused(bundle_root=self.algo_dir, fold=0, ckpt_path="$@bundle_root + @fold") + + def test_missing_reference_is_refused(self): + """A reference to an id that is not in the file raises ValueError with the documented message.""" + self._write_config(ckpt_path="$@no_such_entry + '/model'") + with self.assertRaisesRegex(ValueError, "does not exist"): + self._algo().get_score() + + def test_unknown_reference_key_is_refused(self): + """A subscript that does not come from a matched @reference raises ValueError, not KeyError.""" + self._assert_refused(ckpt_path="$__local_refs['missing']") + + def test_whole_value_reference_still_resolves(self): + ckpt_dir = os.path.join(self.algo_dir, "ckpt") + self._write_progress(ckpt_dir, 0.5) + self._write_config(ckpt=ckpt_dir, ckpt_path="@ckpt") + + self.assertAlmostEqual(self._algo().get_score(), 0.5) + + def test_expression_in_ckpt_path_does_not_run(self): + """A ``$`` expression in ``ckpt_path`` cannot import or call anything.""" + payload = f"$__import__('pathlib').Path({self.marker!r}).write_text('pwned') or {self.algo_dir!r}" + + self._assert_refused(ckpt_path=payload) + + def test_module_names_are_not_in_scope(self): + self._assert_refused(ckpt_path="$str(torch.tensor(1))") + + def test_subclasses_escape_is_refused(self): + """The usual ``__subclasses__()`` sandbox escape has no way in: no attribute or call is accepted.""" + payload = ( + f"$[c for c in ().__class__.__base__.__subclasses__() if c.__name__ == 'catch_warnings'][0]()" + f"._module.__builtins__['__import__']('pathlib').Path({self.marker!r}).write_text('pwned')" + ) + + self._assert_refused(ckpt_path=payload) + + def test_target_mapping_is_refused(self): + """A ``_target_`` mapping is not instantiated: ``get_score`` never asks for a component.""" + self._assert_refused(ckpt_path={"_target_": "subprocess.Popen", "args": [["touch", self.marker]]}) + + def test_history_import_does_not_run_a_config_payload(self): + """End to end: a work directory with a malicious config is imported without executing it.""" + payload = f"$__import__('pathlib').Path({self.marker!r}).write_text('pwned') or {self.algo_dir!r}" + self._write_config(ckpt_path=payload) + with open(os.path.join(self.algo_dir, "algo_object.json"), "w") as f: + json.dump( + { + "_target_": "monai.apps.auto3dseg.bundle_gen.BundleAlgo", + "_state_": {"template_path": self.root, "output_path": self.algo_dir}, + }, + f, + ) + + with self.assertLogs(level="WARNING"): + history = import_bundle_algo_history(self.root, only_trained=False) + + self.assertFalse(os.path.exists(self.marker), "the work directory payload was evaluated") + self.assertEqual(len(history), 1) + self.assertFalse(history[0]["is_trained"]) + + +if __name__ == "__main__": + unittest.main()