AutoRunner enhancements: local folder and algorithm subsets (#5623)

AutoRunner class improvements

- adds auto-runner option **templates_path_or_url** to optionally accept
a local templates folder location. The default (None) is the same: to
download the release zip e.g.
AutoRunner(templates_path_or_url='/my/local/algorithm_templates')

- enhances **algos** input option to define a subset of algorithm names
to run (e.g. AutoRunner(algos=['segresnet','dints'])

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: myron <amyronenko@nvidia.com>

Author: myron

monai/apps/auto3dseg/auto_runner.py

@@ -14,7 +14,7 @@
 import subprocess
 from copy import deepcopy
 from time import sleep
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import torch
@@ -64,7 +64,11 @@ class AutoRunner:
         work_dir: working directory to save the intermediate and final results.
         input: the configuration dictionary or the file path to the configuration in form of YAML.
             The configuration should contain datalist, dataroot, modality, multigpu, and class_names info.
-        algos: optionally specify a list of algorithms to use.
+        algos: optionally specify algorithms to use.  If a dictionary, must be in the form
+            {"algname": dict(_target_="algname.scripts.algo.AlgnameAlgo", template_path="algname"), ...}
+            If a list or a string, defines a subset of names of the algorithms to use, e.g. 'segresnet' or
+            ['segresnet', 'dints'] out of the full set of algorithm templates provided by templates_path_or_url.
+            Defaults to None, to use all available algorithms.
         analyze: on/off switch to run DataAnalyzer and generate a datastats report. Defaults to None, to automatically
             decide based on cache, and run data analysis only if we have not completed this step yet.
         algo_gen: on/off switch to run AlgoGen and generate templated BundleAlgos. Defaults to None, to automatically
@@ -79,6 +83,8 @@ class AutoRunner:
             datasets.
         not_use_cache: if the value is True, it will ignore all cached results in data analysis,
             algorithm generation, or training, and start the pipeline from scratch.
+        templates_path_or_url: the folder with the algorithm templates or a url. If None provided, the default template
+            zip url will be downloaded and extracted into the work_dir.
         kwargs: image writing parameters for the ensemble inference. The kwargs format follows the SaveImage
             transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage.
 
@@ -106,6 +112,27 @@ class AutoRunner:
             runner = AutoRunner(work_dir=work_dir, input=input)
             runner.run()
 
+        - User can specify a subset of algorithms to use and run AutoRunner:
+
+        .. code-block:: python
+
+            work_dir = "./work_dir"
+            input = "path_to_yaml_data_cfg"
+            algos = ["segresnet", "dints"]
+            runner = AutoRunner(work_dir=work_dir, input=input, algos=algos)
+            runner.run()
+
+        - User can specify a a local folder with algorithms templates and run AutoRunner:
+
+        .. code-block:: python
+
+            work_dir = "./work_dir"
+            input = "path_to_yaml_data_cfg"
+            algos = "segresnet"
+            templates_path_or_url = "./local_path_to/algorithm_templates"
+            runner = AutoRunner(work_dir=work_dir, input=input, algos=algos, templates_path_or_url=templates_path_or_url)
+            runner.run()
+
         - User can specify training parameters by:
 
         .. code-block:: python
@@ -181,14 +208,15 @@ def __init__(
         self,
         work_dir: str = "./work_dir",
         input: Union[Dict[str, Any], str, None] = None,
-        algos: Optional[Union[Tuple, List]] = None,
+        algos: Optional[Union[Dict, List, str]] = None,
         analyze: Optional[bool] = None,
         algo_gen: Optional[bool] = None,
         train: Optional[bool] = None,
         hpo: bool = False,
         hpo_backend: str = "nni",
         ensemble: bool = True,
         not_use_cache: bool = False,
+        templates_path_or_url: Optional[str] = None,
         **kwargs,
     ):
 
@@ -198,6 +226,7 @@ def __init__(
         self.work_dir = os.path.abspath(work_dir)
         self.data_src_cfg_name = os.path.join(self.work_dir, "input.yaml")
         self.algos = algos
+        self.templates_path_or_url = templates_path_or_url
 
         if input is None and os.path.isfile(self.data_src_cfg_name):
             input = self.data_src_cfg_name
@@ -558,6 +587,7 @@ def run(self):
             bundle_generator = BundleGen(
                 algos=self.algos,
                 algo_path=self.work_dir,
+                templates_path_or_url=self.templates_path_or_url,
                 data_stats_filename=self.datastats_filename,
                 data_src_cfg_name=self.data_src_cfg_name,
             )

monai/apps/auto3dseg/bundle_gen.py

@@ -14,10 +14,13 @@
 import shutil
 import subprocess
 import sys
+import time
+import warnings
 from copy import deepcopy
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Any, Dict, List, Mapping
+from typing import Any, Dict, List, Mapping, Optional, Union
+from urllib.parse import urlparse
 
 import torch
 
@@ -284,15 +287,75 @@ def get_output_path(self):
 }
 
 
+def _download_algos_url(url: str, at_path: str):
+    """
+    Downloads the algorithm templates release archive, and extracts it into a parent directory of the at_path folder.
+    Returns a dictionary of the algorithm templates.
+    """
+    at_path = os.path.abspath(at_path)
+    zip_download_dir = TemporaryDirectory()
+    algo_compressed_file = os.path.join(zip_download_dir.name, "algo_templates.tar.gz")
+
+    download_attempts = 3
+    for i in range(download_attempts):
+        try:
+            download_and_extract(url=url, filepath=algo_compressed_file, output_dir=os.path.dirname(at_path))
+        except Exception as e:
+            msg = f"Download and extract of {url} failed, attempt {i+1}/{download_attempts}."
+            if i < download_attempts - 1:
+                warnings.warn(msg)
+                time.sleep(i)
+            else:
+                zip_download_dir.cleanup()
+                raise ValueError(msg) from e
+        else:
+            break
+
+    zip_download_dir.cleanup()
+
+    algos_all = deepcopy(default_algos)
+    for name in algos_all:
+        algos_all[name]["template_path"] = os.path.join(at_path, algos_all[name]["template_path"])
+
+    return algos_all
+
+
+def _copy_algos_folder(folder, at_path):
+    """
+    Copies the algorithm templates folder to at_path.
+    Returns a dictionary of of algorithm templates.
+    """
+    folder = os.path.abspath(folder)
+    at_path = os.path.abspath(at_path)
+
+    if folder != at_path:
+        if os.path.exists(at_path):
+            shutil.rmtree(at_path)
+        shutil.copytree(folder, at_path)
+
+    algos_all = {}
+    for name in os.listdir(at_path):
+        if os.path.exists(os.path.join(folder, name, "scripts", "algo.py")):
+            algos_all[name] = dict(
+                _target_=f"{name}.scripts.algo.{name.capitalize()}Algo", template_path=os.path.join(at_path, name)
+            )
+    if len(algos_all) == 0:
+        raise ValueError(f"Unable to find any algos in {folder}")
+
+    return algos_all
+
+
 class BundleGen(AlgoGen):
     """
     This class generates a set of bundles according to the cross-validation folds, each of them can run independently.
 
     Args:
         algo_path: the directory path to save the algorithm templates. Default is the current working dir.
-        algos: if dictionary, it outlines the algorithm to use. if None, automatically download the zip file
-            from the default link. if string, it represents the download link.
-            The current default options are released at:
+        algos: If dictionary, it outlines the algorithm to use. If a list or a string, defines a subset of names of
+            the algorithms to use, e.g. ('segresnet', 'dints') out of the full set of algorithm templates provided
+            by templates_path_or_url. Defaults to None - to use all available algorithms.
+        templates_path_or_url: the folder with the algorithm templates or a url. If None provided, the default template
+            zip url will be downloaded and extracted into the algo_path. The current default options are released at:
             https://github.com/Project-MONAI/research-contributions/tree/main/auto3dseg
         data_stats_filename: the path to the data stats file (generated by DataAnalyzer)
         data_src_cfg_name: the path to the data source config YAML file. The config will be in a form of
@@ -303,22 +366,39 @@ class BundleGen(AlgoGen):
         python -m monai.apps.auto3dseg BundleGen generate --data_stats_filename="../algorithms/data_stats.yaml"
     """
 
-    def __init__(self, algo_path: str = ".", algos=None, data_stats_filename=None, data_src_cfg_name=None):
-        self.algos: Any = []
-
-        if algos is None or isinstance(algos, str):
-            # trigger the download process
-            zip_download_dir = TemporaryDirectory()
-            algo_compressed_file = os.path.join(zip_download_dir.name, "algo_templates.tar.gz")
-            download_and_extract(default_algo_zip if algos is None else algos, algo_compressed_file, algo_path)
-            zip_download_dir.cleanup()
-            sys.path.insert(0, os.path.join(algo_path, "algorithm_templates"))
-            algos = deepcopy(default_algos)
-            for name in algos:
-                algos[name]["template_path"] = os.path.join(
-                    algo_path, "algorithm_templates", algos[name]["template_path"]
-                )
+    def __init__(
+        self,
+        algo_path: str = ".",
+        algos: Optional[Union[Dict, List, str]] = None,
+        templates_path_or_url: Optional[str] = None,
+        data_stats_filename: Optional[str] = None,
+        data_src_cfg_name: Optional[str] = None,
+    ):
+
+        if algos is None or isinstance(algos, (list, tuple, str)):
+
+            if templates_path_or_url is None:
+                templates_path_or_url = default_algo_zip
+
+            at_path = os.path.join(os.path.abspath(algo_path), "algorithm_templates")
+
+            if os.path.isdir(templates_path_or_url):
+                # if a local folder, copy if necessary
+                algos_all = _copy_algos_folder(folder=templates_path_or_url, at_path=at_path)
+            elif urlparse(templates_path_or_url).scheme in ("http", "https"):
+                # if url, trigger the download and extract process
+                algos_all = _download_algos_url(url=templates_path_or_url, at_path=at_path)
+            else:
+                raise ValueError(f"{self.__class__} received invalid templates_path_or_url: {templates_path_or_url}")
+
+            if algos is not None:
+                algos = {k: v for k, v in algos_all.items() if k in ensure_tuple(algos)}  # keep only provided
+                if len(algos) == 0:
+                    raise ValueError(f"Unable to find provided algos in {algos_all}")
+            else:
+                algos = algos_all
 
+        self.algos: Any = []
         if isinstance(algos, dict):
             for algo_name, algo_params in algos.items():
 
@@ -327,7 +407,9 @@ def __init__(self, algo_path: str = ".", algos=None, data_stats_filename=None, d
                     sys.path.append(template_path)
 
                 try:
-                    self.algos.append(ConfigParser(algo_params).get_parsed_content())
+                    onealgo = ConfigParser(algo_params).get_parsed_content()
+                    onealgo.name = algo_name
+                    self.algos.append(onealgo)
                 except RuntimeError as e:
                     msg = """Please make sure the folder structure of an Algo Template follows
                         [algo_name]
@@ -339,9 +421,8 @@ def __init__(self, algo_path: str = ".", algos=None, data_stats_filename=None, d
                             └── validate.py
                     """
                     raise RuntimeError(msg) from e
-                self.algos[-1].name = algo_name
         else:
-            self.algos = ensure_tuple(algos)
+            raise ValueError("Unexpected error algos is not a dict")
 
         self.data_stats_filename = data_stats_filename
         self.data_src_cfg_filename = data_src_cfg_name