Repository: Project-MONAI/MONAI

Issue:
Add NVIDIA NGC source support to the bundle.download
**Is your feature request related to a problem? Please describe.**
Currently, we have the `monai.bundle.download()` API to automatically download the bundle from MONAI model-zoo or direct URL. Would be nice to add the NGC source support to this function.
And also use a system environment to control the default source, for example changing the default arg to:
```py
source: str = os.environ.get("BUNDLE_DOWNLOAD_SRC", "github")
```
If `source` is not `github` or `ngc` and no direct URL provided, raise `ValueError`.
Better to include this feature in MONAI 1.1 in case we need it soon.




Discussion:
It seems the MONAI Bundle needs to be public on NGC before we implement it in MONAI.
Hi @mingxin-zheng ,

We can implement a draft PR referring to (or call the API directly):
https://github.com/Project-MONAI/MONAI/blob/dev/monai/apps/mmars/mmars.py#L103
Then when the NGC bundles are published, we can test the PR immediately and merge to MONAI 1.1 ASAP.

Thanks.

Update: I'm starting to do this feature now.

Resolving patch:
diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -25,6 +25,7 @@
 import torch
 from torch.cuda import is_available
 
+from monai.apps.mmars.mmars import _get_all_ngc_models
 from monai.apps.utils import _basename, download_url, extractall, get_logger
 from monai.bundle.config_item import ConfigComponent
 from monai.bundle.config_parser import ConfigParser
@@ -42,6 +43,9 @@
 
 logger = get_logger(module_name=__name__)
 
+# set BUNDLE_DOWNLOAD_SRC="ngc" to use NGC source in default for bundle download
+download_source = os.environ.get("BUNDLE_DOWNLOAD_SRC", "github")
+
 
 def _update_args(args: Optional[Union[str, Dict]] = None, ignore_none: bool = True, **kwargs) -> Dict:
     """
@@ -130,9 +134,11 @@ def _get_git_release_url(repo_owner: str, repo_name: str, tag_name: str, filenam
     return f"https://github.com/{repo_owner}/{repo_name}/releases/download/{tag_name}/{filename}"
 
 
+def _get_ngc_bundle_url(model_name: str, version: str):
+    return f"https://api.ngc.nvidia.com/v2/models/nvidia/monaitoolkit/{model_name}/versions/{version}/zip"
+
+
 def _download_from_github(repo: str, download_path: Path, filename: str, progress: bool = True):
-    if len(repo.split("/")) != 3:
-        raise ValueError("if source is `github`, repo should be in the form of `repo_owner/repo_name/release_tag`.")
     repo_owner, repo_name, tag_name = repo.split("/")
     if ".zip" not in filename:
         filename += ".zip"
@@ -142,6 +148,45 @@ def _download_from_github(repo: str, download_path: Path, filename: str, progres
     extractall(filepath=filepath, output_dir=download_path, has_base=True)
 
 
+def _add_ngc_prefix(name: str, prefix: str = "monai_"):
+    if name.startswith(prefix):
+        return name
+    return f"{prefix}{name}"
+
+
+def _remove_ngc_prefix(name: str, prefix: str = "monai_"):
+    if name.startswith(prefix):
+        return name[len(prefix) :]
+    return name
+
+
+def _download_from_ngc(download_path: Path, filename: str, version: str, remove_prefix: Optional[str], progress: bool):
+    # ensure prefix is contained
+    filename = _add_ngc_prefix(filename)
+    url = _get_ngc_bundle_url(model_name=filename, version=version)
+    filepath = download_path / f"{filename}_v{version}.zip"
+    if remove_prefix:
+        filename = _remove_ngc_prefix(filename)
+    extract_path = download_path / f"{filename}"
+    download_url(url=url, filepath=filepath, hash_val=None, progress=progress)
+    extractall(filepath=filepath, output_dir=extract_path, has_base=True)
+
+
+def _get_latest_bundle_version(source: str, name: str, repo: str):
+    if source == "ngc":
+        name = _add_ngc_prefix(name)
+        model_dict = _get_all_ngc_models(name)
+        for v in model_dict.values():
+            if v["name"] == name:
+                return v["latest"]
+        return None
+    elif source == "github":
+        repo_owner, repo_name, tag_name = repo.split("/")
+        return get_bundle_versions(name, repo=os.path.join(repo_owner, repo_name), tag=tag_name)["latest_version"]
+    else:
+        raise ValueError(f"To get the latest bundle version, source should be 'github' or 'ngc', got {source}.")
+
+
 def _process_bundle_dir(bundle_dir: Optional[PathLike] = None):
     if bundle_dir is None:
         get_dir, has_home = optional_import("torch.hub", name="get_dir")
@@ -156,9 +201,10 @@ def download(
     name: Optional[str] = None,
     version: Optional[str] = None,
     bundle_dir: Optional[PathLike] = None,
-    source: str = "github",
-    repo: str = "Project-MONAI/model-zoo/hosting_storage_v1",
+    source: str = download_source,
+    repo: Optional[str] = None,
     url: Optional[str] = None,
+    remove_prefix: Optional[str] = "monai_",
     progress: bool = True,
     args_file: Optional[str] = None,
 ):
@@ -175,9 +221,12 @@ def download(
         # Execute this module as a CLI entry, and download bundle from the model-zoo repo:
         python -m monai.bundle download --name <bundle_name> --version "0.1.0" --bundle_dir "./"
 
-        # Execute this module as a CLI entry, and download bundle:
+        # Execute this module as a CLI entry, and download bundle from specified github repo:
         python -m monai.bundle download --name <bundle_name> --source "github" --repo "repo_owner/repo_name/release_tag"
 
+        # Execute this module as a CLI entry, and download bundle from ngc with latest version:
+        python -m monai.bundle download --name <bundle_name> --source "ngc" --bundle_dir "./"
+
         # Execute this module as a CLI entry, and download bundle via URL:
         python -m monai.bundle download --name <bundle_name> --url <url>
 
@@ -190,18 +239,27 @@ def download(
 
     Args:
         name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`.
-            for example: "spleen_ct_segmentation", "prostate_mri_anatomy" in the model-zoo:
+            for example:
+            "spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo:
             https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1.
-        version: version name of the target bundle to download, like: "0.1.0".
+            "monai_brats_mri_segmentation" in ngc:
+            https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.
+        version: version name of the target bundle to download, like: "0.1.0". If `None`, will download
+            the latest version.
         bundle_dir: target directory to store the downloaded data.
             Default is `bundle` subfolder under `torch.hub.get_dir()`.
         source: storage location name. This argument is used when `url` is `None`.
-            "github" is currently the only supported value.
-        repo: repo name. This argument is used when `url` is `None`.
-            If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
+            In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
+            it should be "ngc" or "github".
+        repo: repo name. This argument is used when `url` is `None` and `source` is "github".
+            If used, it should be in the form of "repo_owner/repo_name/release_tag".
         url: url to download the data. If not `None`, data will be downloaded directly
             and `source` will not be checked.
             If `name` is `None`, filename is determined by `monai.apps.utils._basename(url)`.
+        remove_prefix: This argument is used when `source` is "ngc". Currently, all ngc bundles
+            have the ``monai_`` prefix, which is not existing in their model zoo contrasts. In order to
+            maintain the consistency between these two sources, remove prefix is necessary.
+            Therefore, if specified, downloaded folder name will remove the prefix.
         progress: whether to display a progress bar.
         args_file: a JSON or YAML file to provide default values for all the args in this function.
             so that the command line inputs can be simplified.
@@ -215,17 +273,20 @@ def download(
         source=source,
         repo=repo,
         url=url,
+        remove_prefix=remove_prefix,
         progress=progress,
     )
 
     _log_input_summary(tag="download", args=_args)
-    source_, repo_, progress_, name_, version_, bundle_dir_, url_ = _pop_args(
-        _args, "source", "repo", "progress", name=None, version=None, bundle_dir=None, url=None
+    source_, progress_, remove_prefix_, repo_, name_, version_, bundle_dir_, url_ = _pop_args(
+        _args, "source", "progress", remove_prefix=None, repo=None, name=None, version=None, bundle_dir=None, url=None
     )
 
     bundle_dir_ = _process_bundle_dir(bundle_dir_)
-    if name_ is not None and version_ is not None:
-        name_ = "_v".join([name_, version_])
+    if repo_ is None:
+        repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"
+    if len(repo_.split("/")) != 3:
+        raise ValueError("repo should be in the form of `repo_owner/repo_name/release_tag`.")
 
     if url_ is not None:
         if name_ is not None:
@@ -234,14 +295,27 @@ def download(
             filepath = bundle_dir_ / f"{_basename(url_)}"
         download_url(url=url_, filepath=filepath, hash_val=None, progress=progress_)
         extractall(filepath=filepath, output_dir=bundle_dir_, has_base=True)
-    elif source_ == "github":
-        if name_ is None:
-            raise ValueError(f"To download from source: Github, `name` must be provided, got {name_}.")
-        _download_from_github(repo=repo_, download_path=bundle_dir_, filename=name_, progress=progress_)
     else:
-        raise NotImplementedError(
-            f"Currently only download from provided URL in `url` or Github is implemented, got source: {source_}."
-        )
+        if name_ is None:
+            raise ValueError(f"To download from source: {source_}, `name` must be provided.")
+        if version_ is None:
+            version_ = _get_latest_bundle_version(source=source_, name=name_, repo=repo_)
+        if source_ == "github":
+            if version_ is not None:
+                name_ = "_v".join([name_, version_])
+            _download_from_github(repo=repo_, download_path=bundle_dir_, filename=name_, progress=progress_)
+        elif source_ == "ngc":
+            _download_from_ngc(
+                download_path=bundle_dir_,
+                filename=name_,
+                version=version_,
+                remove_prefix=remove_prefix_,
+                progress=progress_,
+            )
+        else:
+            raise NotImplementedError(
+                f"Currently only download from `url`, source 'github' or 'ngc' are implemented, got source: {source_}."
+            )
 
 
 def load(
@@ -250,8 +324,8 @@ def load(
     model_file: Optional[str] = None,
     load_ts_module: bool = False,
     bundle_dir: Optional[PathLike] = None,
-    source: str = "github",
-    repo: str = "Project-MONAI/model-zoo/hosting_storage_v1",
+    source: str = download_source,
+    repo: Optional[str] = None,
     progress: bool = True,
     device: Optional[str] = None,
     key_in_ckpt: Optional[str] = None,
@@ -263,18 +337,25 @@ def load(
     Load model weights or TorchScript module of a bundle.
 
     Args:
-        name: bundle name, for example: "spleen_ct_segmentation", "prostate_mri_anatomy" in the model-zoo:
+        name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`.
+            for example:
+            "spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo:
             https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1.
-        version: version name of the target bundle to download, like: "0.1.0".
+            "monai_brats_mri_segmentation" in ngc:
+            https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.
+        version: version name of the target bundle to download, like: "0.1.0". If `None`, will download
+            the latest version.
         model_file: the relative path of the model weights or TorchScript module within bundle.
             If `None`, "models/model.pt" or "models/model.ts" will be used.
         load_ts_module: a flag to specify if loading the TorchScript module.
         bundle_dir: directory the weights/TorchScript module will be loaded from.
             Default is `bundle` subfolder under `torch.hub.get_dir()`.
         source: storage location name. This argument is used when `model_file` is not existing locally and need to be
-            downloaded first. "github" is currently the only supported value.
-        repo: repo name. This argument is used when `model_file` is not existing locally and need to be
-            downloaded first. If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
+            downloaded first.
+            In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
+            it should be "ngc" or "github".
+        repo: repo name. This argument is used when `url` is `None` and `source` is "github".
+            If used, it should be in the form of "repo_owner/repo_name/release_tag".
         progress: whether to display a progress bar when downloading.
         device: target device of returned weights or module, if `None`, prefer to "cuda" if existing.
         key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model
@@ -421,7 +502,7 @@ def get_bundle_versions(
 
     bundles_info = _get_all_bundles_info(repo=repo, tag=tag, auth_token=auth_token)
     if bundle_name not in bundles_info:
-        raise ValueError(f"bundle: {bundle_name} is not existing.")
+        raise ValueError(f"bundle: {bundle_name} is not existing in repo: {repo}.")
     bundle_info = bundles_info[bundle_name]
     all_versions = sorted(bundle_info.keys())
 


Test patch:
diff --git a/tests/ngc_mmar_loading.py b/tests/ngc_bundle_download.py
similarity index 68%
rename from tests/ngc_mmar_loading.py
rename to tests/ngc_bundle_download.py
--- a/tests/ngc_mmar_loading.py
+++ b/tests/ngc_bundle_download.py
@@ -11,14 +11,50 @@
 
 import os
 import sys
+import tempfile
 import unittest
 
 import torch
 from parameterized import parameterized
 
+from monai.apps import check_hash
 from monai.apps.mmars import MODEL_DESC, load_from_mmar
+from monai.bundle import download
 from monai.config import print_debug_info
 from monai.networks.utils import copy_model_state
+from tests.utils import skip_if_downloading_fails, skip_if_quick, skip_if_windows
+
+TEST_CASE_NGC_1 = [
+    "spleen_ct_segmentation",
+    "0.3.7",
+    None,
+    "monai_spleen_ct_segmentation",
+    "models/model.pt",
+    "b418a2dc8672ce2fd98dc255036e7a3d",
+]
+TEST_CASE_NGC_2 = [
+    "monai_spleen_ct_segmentation",
+    "0.3.7",
+    "monai_",
+    "spleen_ct_segmentation",
+    "models/model.pt",
+    "b418a2dc8672ce2fd98dc255036e7a3d",
+]
+
+
+@skip_if_windows
+class TestNgcBundleDownload(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_NGC_1, TEST_CASE_NGC_2])
+    @skip_if_quick
+    def test_ngc_download_bundle(self, bundle_name, version, remove_prefix, download_name, file_path, hash_val):
+        with skip_if_downloading_fails():
+            with tempfile.TemporaryDirectory() as tempdir:
+                download(
+                    name=bundle_name, source="ngc", version=version, bundle_dir=tempdir, remove_prefix=remove_prefix
+                )
+                full_file_path = os.path.join(tempdir, download_name, file_path)
+                self.assertTrue(os.path.exists(full_file_path))
+                self.assertTrue(check_hash(filepath=full_file_path, val=hash_val))
 
 
 @unittest.skip("deprecating mmar tests")
diff --git a/tests/test_bundle_download.py b/tests/test_bundle_download.py
--- a/tests/test_bundle_download.py
+++ b/tests/test_bundle_download.py
@@ -31,18 +31,16 @@
 
 TEST_CASE_1 = ["test_bundle", None]
 
-TEST_CASE_2 = ["test_bundle_v0.1.1", None]
+TEST_CASE_2 = ["test_bundle", "0.1.1"]
 
-TEST_CASE_3 = ["test_bundle", "0.1.1"]
-
-TEST_CASE_4 = [
+TEST_CASE_3 = [
     ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"],
     "test_bundle",
     "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/test_bundle.zip",
     "a131d39a0af717af32d19e565b434928",
 ]
 
-TEST_CASE_5 = [
+TEST_CASE_4 = [
     ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"],
     "test_bundle",
     "Project-MONAI/MONAI-extra-test-data/0.8.1",
@@ -50,7 +48,7 @@
     "model.pt",
 ]
 
-TEST_CASE_6 = [
+TEST_CASE_5 = [
     ["test_output.pt", "test_input.pt"],
     "test_bundle",
     "0.1.1",
@@ -62,9 +60,9 @@
 
 @skip_if_windows
 class TestDownload(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
     @skip_if_quick
-    def test_download_bundle(self, bundle_name, version):
+    def test_github_download_bundle(self, bundle_name, version):
         bundle_files = ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"]
         repo = "Project-MONAI/MONAI-extra-test-data/0.8.1"
         hash_val = "a131d39a0af717af32d19e565b434928"
@@ -72,7 +70,7 @@ def test_download_bundle(self, bundle_name, version):
             # download a whole bundle from github releases
             with tempfile.TemporaryDirectory() as tempdir:
                 cmd = ["coverage", "run", "-m", "monai.bundle", "download", "--name", bundle_name, "--source", "github"]
-                cmd += ["--bundle_dir", tempdir, "--repo", repo, "--progress", "False"]
+                cmd += ["--bundle_dir", tempdir, "--repo", repo]
                 if version is not None:
                     cmd += ["--version", version]
                 command_line_tests(cmd)
@@ -82,7 +80,7 @@ def test_download_bundle(self, bundle_name, version):
                     if file == "network.json":
                         self.assertTrue(check_hash(filepath=file_path, val=hash_val))
 
-    @parameterized.expand([TEST_CASE_4])
+    @parameterized.expand([TEST_CASE_3])
     @skip_if_quick
     def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val):
         with skip_if_downloading_fails():
@@ -103,7 +101,7 @@ def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val):
 
 
 class TestLoad(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_5])
+    @parameterized.expand([TEST_CASE_4])
     @skip_if_quick
     def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file):
         with skip_if_downloading_fails():
@@ -150,7 +148,7 @@ def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file)
                 output_2 = model_2.forward(input_tensor)
                 assert_allclose(output_2, expected_output, atol=1e-4, rtol=1e-4, type_test=False)
 
-    @parameterized.expand([TEST_CASE_6])
+    @parameterized.expand([TEST_CASE_5])
     @skip_if_quick
     @SkipIfBeforePyTorchVersion((1, 7, 1))
     def test_load_ts_module(self, bundle_files, bundle_name, version, repo, device, model_file):