5052 Fix bundle.load API (#5053)

Fixes #5052 .

### Description
This PR updated the `bundle.load` API to fix the wrong path issue.

### Status
**Ready**

### 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`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/` folder.

Signed-off-by: Nic Ma <nma@nvidia.com>

Author: Nic-Ma

monai/bundle/scripts.py

@@ -153,6 +153,7 @@ def _process_bundle_dir(bundle_dir: Optional[PathLike] = None):
 
 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",
@@ -170,11 +171,14 @@ def download(
 
     .. code-block:: bash
 
+        # 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:
-        python -m monai.bundle download --name "bundle_name" --source "github" --repo "repo_owner/repo_name/release_tag"
+        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 via URL:
-        python -m monai.bundle download --name "bundle_name" --url <url>
+        python -m monai.bundle download --name <bundle_name> --url <url>
 
         # Set default args of `run` in a JSON / YAML file, help to record and simplify the command line.
         # Other args still can override the default args at runtime.
@@ -185,6 +189,9 @@ 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:
+            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".
         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`.
@@ -200,19 +207,28 @@ def download(
 
     """
     _args = _update_args(
-        args=args_file, name=name, bundle_dir=bundle_dir, source=source, repo=repo, url=url, progress=progress
+        args=args_file,
+        name=name,
+        version=version,
+        bundle_dir=bundle_dir,
+        source=source,
+        repo=repo,
+        url=url,
+        progress=progress,
     )
 
     _log_input_summary(tag="download", args=_args)
-    source_, repo_, progress_, name_, bundle_dir_, url_ = _pop_args(
-        _args, "source", "repo", "progress", name=None, bundle_dir=None, url=None
+    source_, repo_, progress_, name_, version_, bundle_dir_, url_ = _pop_args(
+        _args, "source", "repo", "progress", 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 url_ is not None:
-        if name is not None:
-            filepath = bundle_dir_ / f"{name}.zip"
+        if name_ is not None:
+            filepath = bundle_dir_ / f"{name_}.zip"
         else:
             filepath = bundle_dir_ / f"{_basename(url_)}"
         download_url(url=url_, filepath=filepath, hash_val=None, progress=progress_)
@@ -229,6 +245,7 @@ def download(
 
 def load(
     name: str,
+    version: Optional[str] = None,
     model_file: Optional[str] = None,
     load_ts_module: bool = False,
     bundle_dir: Optional[PathLike] = None,
@@ -245,7 +262,9 @@ def load(
     Load model weights or TorchScript module of a bundle.
 
     Args:
-        name: bundle name.
+        name: bundle name, for example: "spleen_ct_segmentation", "prostate_mri_anatomy" in the 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".
         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.
@@ -280,7 +299,7 @@ def load(
         model_file = os.path.join("models", "model.ts" if load_ts_module is True else "model.pt")
     full_path = os.path.join(bundle_dir_, name, model_file)
     if not os.path.exists(full_path):
-        download(name=name, bundle_dir=bundle_dir_, source=source, repo=repo, progress=progress)
+        download(name=name, version=version, bundle_dir=bundle_dir_, source=source, repo=repo, progress=progress)
 
     if device is None:
         device = "cuda:0" if is_available() else "cpu"

tests/test_bundle_download.py

@@ -28,31 +28,31 @@
     skip_if_windows,
 )
 
-TEST_CASE_1 = [
-    ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"],
-    "test_bundle",
-    "Project-MONAI/MONAI-extra-test-data/0.8.1",
-    "a131d39a0af717af32d19e565b434928",
-]
+TEST_CASE_1 = ["test_bundle", None]
+
+TEST_CASE_2 = ["test_bundle_v0.1.1", None]
 
-TEST_CASE_2 = [
+TEST_CASE_3 = ["test_bundle", "0.1.1"]
+
+TEST_CASE_4 = [
     ["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_3 = [
+TEST_CASE_5 = [
     ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"],
     "test_bundle",
     "Project-MONAI/MONAI-extra-test-data/0.8.1",
     "cuda" if torch.cuda.is_available() else "cpu",
     "model.pt",
 ]
 
-TEST_CASE_4 = [
+TEST_CASE_6 = [
     ["test_output.pt", "test_input.pt"],
     "test_bundle",
+    "0.1.1",
     "Project-MONAI/MONAI-extra-test-data/0.8.1",
     "cuda" if torch.cuda.is_available() else "cpu",
     "model.ts",
@@ -61,22 +61,27 @@
 
 @skip_if_windows
 class TestDownload(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
     @skip_if_quick
-    def test_download_bundle(self, bundle_files, bundle_name, repo, hash_val):
+    def test_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"
         with skip_if_downloading_fails():
             # 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"]
+                if version is not None:
+                    cmd += ["--version", version]
                 command_line_tests(cmd)
                 for file in bundle_files:
-                    file_path = os.path.join(tempdir, bundle_name, file)
+                    file_path = os.path.join(tempdir, "test_bundle", file)
                     self.assertTrue(os.path.exists(file_path))
                     if file == "network.json":
                         self.assertTrue(check_hash(filepath=file_path, val=hash_val))
 
-    @parameterized.expand([TEST_CASE_2])
+    @parameterized.expand([TEST_CASE_4])
     @skip_if_quick
     def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val):
         with skip_if_downloading_fails():
@@ -97,7 +102,7 @@ def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val):
 
 
 class TestLoad(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_3])
+    @parameterized.expand([TEST_CASE_5])
     @skip_if_quick
     def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file):
         with skip_if_downloading_fails():
@@ -144,16 +149,17 @@ def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file)
                 output_2 = model_2.forward(input_tensor)
                 torch.testing.assert_allclose(output_2, expected_output)
 
-    @parameterized.expand([TEST_CASE_4])
+    @parameterized.expand([TEST_CASE_6])
     @skip_if_quick
     @SkipIfBeforePyTorchVersion((1, 7, 1))
-    def test_load_ts_module(self, bundle_files, bundle_name, repo, device, model_file):
+    def test_load_ts_module(self, bundle_files, bundle_name, version, repo, device, model_file):
         with skip_if_downloading_fails():
             # load ts module
             with tempfile.TemporaryDirectory() as tempdir:
                 # load ts module
                 model_ts, metadata, extra_file_dict = load(
                     name=bundle_name,
+                    version=version,
                     model_file=model_file,
                     load_ts_module=True,
                     bundle_dir=tempdir,