Check nnUNet plans network names against allowed modules

docs/source/apps.rst

@@ -297,3 +297,4 @@ FastMRIReader
 .. autofunction:: monai.apps.nnunet.convert_nnunet_to_monai_bundle
 .. autofunction:: monai.apps.nnunet.convert_monai_bundle_to_nnunet
 .. autofunction:: monai.apps.nnunet.get_network_from_nnunet_plans
+.. autodata:: monai.apps.nnunet.DEFAULT_NNUNET_NETWORK_MODULES

monai/apps/nnunet/__init__.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 from .nnunet_bundle import (
+    DEFAULT_NNUNET_NETWORK_MODULES,
     ModelnnUNetWrapper,
     convert_monai_bundle_to_nnunet,
     convert_nnunet_to_monai_bundle,

monai/apps/nnunet/nnunet_bundle.py

@@ -12,7 +12,9 @@
 
 import os
 import shutil
+from collections.abc import Sequence
 from pathlib import Path
+from pydoc import locate
 from typing import Any
 
 import numpy as np
@@ -32,8 +34,60 @@
     "convert_nnunet_to_monai_bundle",
     "convert_monai_bundle_to_nnunet",
     "ModelnnUNetWrapper",
+    "DEFAULT_NNUNET_NETWORK_MODULES",
 ]
 
+#: Modules that nnUNet plans may reference for the network class and its imported keyword arguments.
+DEFAULT_NNUNET_NETWORK_MODULES: tuple[str, ...] = ("dynamic_network_architectures", "torch.nn")
+
+
+def _in_allowed_modules(dotted_name: str, allowed_modules: Sequence[str]) -> bool:
+    return any(dotted_name == m or dotted_name.startswith(m + ".") for m in allowed_modules)
+
+
+def _check_nnunet_network_spec(
+    class_name: Any, init_kwargs: Any, kwargs_requiring_import: Any, allowed_modules: Sequence[str] | None = None
+) -> None:
+    """
+    Check the network specification read from an nnUNet ``plans.json`` before nnUNet resolves and builds it.
+
+    nnUNet imports ``network_class_name`` and every keyword argument listed in ``_kw_requires_import`` by dotted
+    name and calls them, so each of these names must be under one of ``allowed_modules`` and, if it can be
+    imported, must refer to a ``torch.nn.Module`` subclass. The module prefix is checked before anything is
+    imported.
+
+    Args:
+        class_name: the ``network_class_name`` from the plans.
+        init_kwargs: the ``arch_kwargs`` from the plans.
+        kwargs_requiring_import: the ``_kw_requires_import`` names from the plans.
+        allowed_modules: dotted module prefixes that may be referenced. Defaults to
+            ``DEFAULT_NNUNET_NETWORK_MODULES``. Add the module of a custom architecture here to load it.
+
+    Raises:
+        ValueError: if a name is outside ``allowed_modules`` or does not resolve to a ``torch.nn.Module`` subclass.
+    """
+    allowed = tuple(DEFAULT_NNUNET_NETWORK_MODULES if allowed_modules is None else allowed_modules)
+    hint = "Pass the module in `allowed_network_modules` if this architecture is trusted."
+
+    def _check(name: Any, what: str) -> None:
+        if not isinstance(name, str) or not _in_allowed_modules(name, allowed):
+            raise ValueError(f"Refusing to load {what} {name!r} from nnUNet plans: not in {list(allowed)}. {hint}")
+        resolved = locate(name)
+        # an unresolved name is left to nnUNet, which only searches inside dynamic_network_architectures for it
+        if resolved is not None and not (isinstance(resolved, type) and issubclass(resolved, torch.nn.Module)):
+            raise ValueError(f"Refusing to load {what} {name!r} from nnUNet plans: not a torch.nn.Module subclass.")
+
+    _check(class_name, "network class")
+    init_kwargs = init_kwargs or {}
+    if not isinstance(init_kwargs, dict):
+        raise ValueError(
+            f"Invalid network arguments in nnUNet plans: expected a dict, got {type(init_kwargs).__name__}."
+        )
+    for key in kwargs_requiring_import or ():
+        value = init_kwargs.get(key)
+        if value is not None:
+            _check(value, f"network argument {key!r}")
+
 
 def get_nnunet_trainer(
     dataset_name_or_id: str | int,
@@ -152,6 +206,9 @@ class ModelnnUNetWrapper(torch.nn.Module):
         The folder path where the model and related files are stored.
     model_name : str, optional
         The name of the model file, by default "model.pt".
+    allowed_network_modules : Sequence[str], optional
+        Module prefixes that the plans may reference for the network class and imported network arguments,
+        by default ``DEFAULT_NNUNET_NETWORK_MODULES``.
 
     Attributes
     ----------
@@ -166,7 +223,13 @@ class ModelnnUNetWrapper(torch.nn.Module):
     restoring network architecture, and setting up the predictor for inference.
     """
 
-    def __init__(self, predictor: object, model_folder: str | Path, model_name: str = "model.pt"):  # type: ignore
+    def __init__(
+        self,
+        predictor: object,
+        model_folder: str | Path,
+        model_name: str = "model.pt",
+        allowed_network_modules: Sequence[str] | None = None,
+    ):  # type: ignore
         super().__init__()
         self.predictor = predictor
 
@@ -216,6 +279,12 @@ def __init__(self, predictor: object, model_folder: str | Path, model_name: str
                 f"Unable to locate trainer class {trainer_name} in nnunetv2.training.nnUNetTrainer. "
                 f"Please place it there (in any .py file)!"
             )
+        _check_nnunet_network_spec(
+            configuration_manager.network_arch_class_name,
+            configuration_manager.network_arch_init_kwargs,
+            configuration_manager.network_arch_init_kwargs_req_import,
+            allowed_network_modules,
+        )
         network = trainer_class.build_network_architecture(
             configuration_manager.network_arch_class_name,
             configuration_manager.network_arch_init_kwargs,
@@ -294,7 +363,9 @@ def forward(self, x: MetaTensor) -> MetaTensor:
         return MetaTensor(out_tensor, meta=x.meta)
 
 
-def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "model.pt") -> ModelnnUNetWrapper:
+def get_nnunet_monai_predictor(
+    model_folder: str | Path, model_name: str = "model.pt", allowed_network_modules: Sequence[str] | None = None
+) -> ModelnnUNetWrapper:
     """
     Initializes and returns a `nnUNetMONAIModelWrapper` containing the corresponding `nnUNetPredictor`.
     The model folder should contain the following files, created during training:
@@ -325,6 +396,9 @@ def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "mode
         The folder where the model is stored.
     model_name : str, optional
         The name of the model file, by default "model.pt".
+    allowed_network_modules : Sequence[str], optional
+        Module prefixes that ``plans.json`` may reference for the network class and imported network arguments,
+        by default ``DEFAULT_NNUNET_NETWORK_MODULES``. Extend it to load a trusted custom architecture.
 
     Returns
     -------
@@ -344,7 +418,7 @@ def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "mode
         allow_tqdm=True,
     )
     # initializes the network architecture, loads the checkpoint
-    wrapper = ModelnnUNetWrapper(predictor, model_folder, model_name)
+    wrapper = ModelnnUNetWrapper(predictor, model_folder, model_name, allowed_network_modules)
     return wrapper
 
 
@@ -428,6 +502,7 @@ def get_network_from_nnunet_plans(
     configuration: str,
     model_ckpt: str | None = None,
     model_key_in_ckpt: str = "model",
+    allowed_network_modules: Sequence[str] | None = None,
 ) -> torch.nn.Module | Any:
     """
     Load and initialize a nnUNet network based on nnUNet plans and configuration.
@@ -444,6 +519,9 @@ def get_network_from_nnunet_plans(
         Path to the model checkpoint file. If None, the network is returned without loading weights (default is None).
     model_key_in_ckpt : str, optional
         The key in the checkpoint file that contains the model state dictionary (default is "model").
+    allowed_network_modules : Sequence[str], optional
+        Module prefixes that ``plans_file`` may reference for the network class and imported network arguments,
+        by default ``DEFAULT_NNUNET_NETWORK_MODULES``. Extend it to load a trusted custom architecture.
 
     Returns
     -------
@@ -465,6 +543,12 @@ def get_network_from_nnunet_plans(
 
     enable_deep_supervision = True
 
+    _check_nnunet_network_spec(
+        configuration_manager.network_arch_class_name,
+        configuration_manager.network_arch_init_kwargs,
+        configuration_manager.network_arch_init_kwargs_req_import,
+        allowed_network_modules,
+    )
     network = get_network_from_plans(
         configuration_manager.network_arch_class_name,
         configuration_manager.network_arch_init_kwargs,