From fd157b8b9d3043d59a976c843e030cbd6335116c Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 9 Jul 2024 17:42:25 +0800
Subject: [PATCH 01/21] support download from private registry

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 114 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 110 insertions(+), 4 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 598d938cbd..b98bdcd51d 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -12,10 +12,12 @@
 from __future__ import annotations
 
 import ast
+import io
 import json
 import os
 import re
 import warnings
+import zipfile
 from collections.abc import Mapping, Sequence
 from pathlib import Path
 from pydoc import locate
@@ -171,6 +173,13 @@ def _get_ngc_bundle_url(model_name: str, version: str) -> str:
     return f"https://api.ngc.nvidia.com/v2/models/nvidia/monaitoolkit/{model_name.lower()}/versions/{version}/zip"
 
 
+def _get_nvstaging_bundle_url(model_name: str, version: str, org: str, team: str) -> str:
+    team_str = ""
+    if team != "no-team":
+        team_str = f"team/{team}"
+    return f"https://api.ngc.nvidia.com/v2/org/{org}/{team_str}/models/{model_name.lower()}/versions/{version}/zip"
+
+
 def _get_monaihosting_bundle_url(model_name: str, version: str) -> str:
     monaihosting_root_path = "https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting"
     return f"{monaihosting_root_path}/{model_name.lower()}/versions/{version}/files/{model_name}_v{version}.zip"
@@ -219,6 +228,44 @@ def _download_from_ngc(
     extractall(filepath=filepath, output_dir=extract_path, has_base=True)
 
 
+def _download_from_nvstaging(
+    download_path: Path,
+    filename: str,
+    version: str,
+    remove_prefix: str | None,
+    org: str,
+    team: str,
+    headers: str | None = None,
+) -> None:
+    # ensure prefix is contained
+    filename = _add_ngc_prefix(filename)
+    request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, org=org, team=team)
+    print(request_url)
+    response = requests_get(request_url, headers=headers)
+    if remove_prefix:
+        filename = _remove_ngc_prefix(filename, prefix=remove_prefix)
+    extract_path = download_path / f"{filename}"
+    with zipfile.ZipFile(io.BytesIO(response.content)) as z:
+        z.extractall(extract_path)
+        logger.info(f"Writing into directory: {extract_path}.")
+
+
+def _get_ngc_token(api_key, retry=0):
+    """Try to connect to NGC."""
+    url = "https://authn.nvidia.com/token?service=ngc"
+    headers = {"Accept": "application/json", "Authorization": "ApiKey " + api_key}
+    if has_requests:
+        response = requests_get(url, headers=headers)
+        if not response.ok:
+            if retry < 3:
+                logger.info(f"Retrying {retry} time(s) to GET {url}.")
+                return _get_ngc_token(url, headers, retry + 1)
+            raise RuntimeError("NGC API response is not ok. Failed to get token.")
+        else:
+            token = response.json()["token"]
+        return token
+
+
 def _get_latest_bundle_version_monaihosting(name):
     url = "https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting"
     full_url = f"{url}/{name.lower()}"
@@ -232,7 +279,23 @@ def _get_latest_bundle_version_monaihosting(name):
     return model_info["model"]["latestVersionIdStr"]
 
 
-def _get_latest_bundle_version(source: str, name: str, repo: str) -> dict[str, list[str] | str] | Any | None:
+def _get_latest_bundle_version_private_registry(name, org, team, headers={}):
+    team_str = ""
+    if team != "no-team":
+        team_str = f"team/{team}"
+    url = f"https://api.ngc.nvidia.com/v2/org/{org}/{team_str}/models"
+    full_url = f"{url}/{name.lower()}"
+    requests_get, has_requests = optional_import("requests", name="get")
+    if has_requests:
+        resp = requests_get(full_url, headers=headers)
+        resp.raise_for_status()
+    else:
+        raise ValueError("NGC API requires requests package.  Please install it.")
+    model_info = json.loads(resp.text)
+    return model_info["model"]["latestVersionIdStr"]
+
+
+def _get_latest_bundle_version(source: str, name: str, repo: str, **kwargs) -> dict[str, list[str] | str] | Any | None:
     if source == "ngc":
         name = _add_ngc_prefix(name)
         model_dict = _get_all_ngc_models(name)
@@ -242,6 +305,11 @@ def _get_latest_bundle_version(source: str, name: str, repo: str) -> dict[str, l
         return None
     elif source == "monaihosting":
         return _get_latest_bundle_version_monaihosting(name)
+    elif source == "nvstaging":
+        org = kwargs.pop("org", "nvstaging")
+        team = kwargs.pop("team", "monaitoolkit")
+        headers = kwargs.pop("headers", {})
+        return _get_latest_bundle_version_private_registry(name, org, team, headers)
     elif source == "github":
         repo_owner, repo_name, tag_name = repo.split("/")
         return get_bundle_versions(name, repo=f"{repo_owner}/{repo_name}", tag=tag_name)["latest_version"]
@@ -279,6 +347,9 @@ def download(
     remove_prefix: str | None = "monai_",
     progress: bool = True,
     args_file: str | None = None,
+    org: str | None = None,
+    team: str | None = None,
+    api_key: str | None = None,
 ) -> None:
     """
     download bundle from the specified source or url. The bundle should be a zip file and it
@@ -354,11 +425,23 @@ def download(
         url=url,
         remove_prefix=remove_prefix,
         progress=progress,
+        org=org,
+        team=team,
     )
 
     _log_input_summary(tag="download", args=_args)
-    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
+    source_, progress_, remove_prefix_, repo_, name_, version_, bundle_dir_, url_, org_, team_ = _pop_args(
+        _args,
+        "source",
+        "progress",
+        remove_prefix=None,
+        repo=None,
+        name=None,
+        version=None,
+        bundle_dir=None,
+        url=None,
+        org=None,
+        team=None,
     )
 
     bundle_dir_ = _process_bundle_dir(bundle_dir_)
@@ -376,10 +459,23 @@ def download(
         download_url(url=url_, filepath=filepath, hash_val=None, progress=progress_)
         extractall(filepath=filepath, output_dir=bundle_dir_, has_base=True)
     else:
+        headers = {}
         if name_ is None:
             raise ValueError(f"To download from source: {source_}, `name` must be provided.")
+        if source == "nvstaging":
+            api_key = os.getenv("NGC_API_KEY", api_key)
+            org_ = "nvstaging" if org_ is None else org_
+            team_ = "monaitoolkit" if team_ is None else team_
+            if api_key is None:
+                raise ValueError("API key is required for nvstaging source.")
+            else:
+                token = _get_ngc_token(api_key)
+                headers = {"Authorization": f"Bearer {token}"}
+
         if version_ is None:
-            version_ = _get_latest_bundle_version(source=source_, name=name_, repo=repo_)
+            version_ = _get_latest_bundle_version(
+                source=source_, name=name_, repo=repo_, org=org_, team=team_, headers=headers
+            )
         if source_ == "github":
             if version_ is not None:
                 name_ = "_v".join([name_, version_])
@@ -394,6 +490,16 @@ def download(
                 remove_prefix=remove_prefix_,
                 progress=progress_,
             )
+        elif source_ == "nvstaging":
+            _download_from_nvstaging(
+                download_path=bundle_dir_,
+                filename=name_,
+                version=version_,
+                remove_prefix=remove_prefix_,
+                org=org_,
+                team=team_,
+                headers=headers,
+            )
         elif source_ == "huggingface_hub":
             extract_path = os.path.join(bundle_dir_, name_)
             huggingface_hub.snapshot_download(repo_id=repo_, revision=version_, local_dir=extract_path)

From b184692c5c909fe3c60ae59387fca2995f4433cc Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 17:26:30 +0800
Subject: [PATCH 02/21] support zip download

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index b98bdcd51d..78df4d0866 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -240,12 +240,20 @@ def _download_from_nvstaging(
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
     request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, org=org, team=team)
-    print(request_url)
-    response = requests_get(request_url, headers=headers)
+    if has_requests:
+        response = requests_get(request_url, headers=headers)
+        response.raise_for_status()
+    else:
+        raise ValueError("NGC API requires requests package.  Please install it.")
+
+    zip_path = download_path / f"{filename}_v{version}.zip"
+    with open(zip_path, "wb") as f:
+        f.write(response.content)
+    logger.info(f"Downloading: {zip_path}.")
     if remove_prefix:
         filename = _remove_ngc_prefix(filename, prefix=remove_prefix)
     extract_path = download_path / f"{filename}"
-    with zipfile.ZipFile(io.BytesIO(response.content)) as z:
+    with zipfile.ZipFile(zip_path, "r") as z:
         z.extractall(extract_path)
         logger.info(f"Writing into directory: {extract_path}.")
 
@@ -309,6 +317,7 @@ def _get_latest_bundle_version(source: str, name: str, repo: str, **kwargs) -> d
         org = kwargs.pop("org", "nvstaging")
         team = kwargs.pop("team", "monaitoolkit")
         headers = kwargs.pop("headers", {})
+        name = _add_ngc_prefix(name)
         return _get_latest_bundle_version_private_registry(name, org, team, headers)
     elif source == "github":
         repo_owner, repo_name, tag_name = repo.split("/")

From a1ea5572f46fc9b04bebcecea556bfb3d5afb53a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 16 Jul 2024 09:27:14 +0000
Subject: [PATCH 03/21] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 monai/bundle/scripts.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 78df4d0866..4b4c71b269 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -12,7 +12,6 @@
 from __future__ import annotations
 
 import ast
-import io
 import json
 import os
 import re

From 35bdfce36c7905dcc58518416eb197536bc589e8 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 19:55:23 +0800
Subject: [PATCH 04/21] add tests

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 .github/workflows/blossom-ci.yml    |  1 +
 .github/workflows/conda.yml         |  2 ++
 .github/workflows/cron.yml          |  6 ++++++
 .github/workflows/docker.yml        |  2 ++
 .github/workflows/integration.yml   |  1 +
 .github/workflows/pythonapp-min.yml |  6 ++++++
 .github/workflows/setupapp.yml      |  4 ++++
 tests/test_bundle_download.py       | 19 ++++++++++++++++++-
 8 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml
index bf507bab3b..cdb134e6ad 100644
--- a/.github/workflows/blossom-ci.yml
+++ b/.github/workflows/blossom-ci.yml
@@ -93,6 +93,7 @@ jobs:
         run: blossom-ci
         env:
           OPERATION: 'START-CI-JOB'
+          NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
           CI_SERVER: ${{ secrets.CI_SERVER }}
           REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
index 8e2807111e..a12d79cea1 100644
--- a/.github/workflows/conda.yml
+++ b/.github/workflows/conda.yml
@@ -60,6 +60,8 @@ jobs:
         conda deactivate
     - name: Test env (CPU ${{ runner.os }})
       shell: bash -el {0}
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         conda activate monai
         $(pwd)/runtests.sh --build --unittests
diff --git a/.github/workflows/cron.yml b/.github/workflows/cron.yml
index 0f9e6cd480..8658508fdf 100644
--- a/.github/workflows/cron.yml
+++ b/.github/workflows/cron.yml
@@ -50,6 +50,8 @@ jobs:
         python -m pip install -r requirements-dev.txt
         python -m pip list
     - name: Run tests report coverage
+      env:
+          NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         export LAUNCH_DELAY=$[ $RANDOM % 16 * 60 ]
         echo "Sleep $LAUNCH_DELAY"
@@ -94,6 +96,8 @@ jobs:
         python -m pip install -r requirements-dev.txt
         python -m pip list
     - name: Run tests report coverage
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         export LAUNCH_DELAY=$[ $RANDOM % 16 * 60 ]
         echo "Sleep $LAUNCH_DELAY"
@@ -196,6 +200,8 @@ jobs:
     - name: Run tests report coverage
       # The docker image process has done the compilation.
       # BUILD_MONAI=1 is necessary for triggering the USE_COMPILED flag.
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         cd /opt/monai
         nvidia-smi
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index 65716f86f9..5581bc5392 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -90,6 +90,8 @@ jobs:
     runs-on: [self-hosted, linux, X64, docker]
     steps:
     - name: Import
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 CUDA_VISIBLE_DEVICES=  # cpu-only
         python -c 'import monai; monai.config.print_debug_info()'
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
index c82530a551..9810ae8bb1 100644
--- a/.github/workflows/integration.yml
+++ b/.github/workflows/integration.yml
@@ -68,6 +68,7 @@ jobs:
       shell: bash
       env:
         BUILD_MONAI: 1
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: ./runtests.sh --build --net
 
     - name: Add reaction
diff --git a/.github/workflows/pythonapp-min.yml b/.github/workflows/pythonapp-min.yml
index 02d8f5058e..504517102a 100644
--- a/.github/workflows/pythonapp-min.yml
+++ b/.github/workflows/pythonapp-min.yml
@@ -60,6 +60,8 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (CPU ${{ runner.os }})
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
@@ -106,6 +108,8 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (CPU ${{ runner.os }})
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
@@ -155,6 +159,8 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (pytorch ${{ matrix.pytorch-version }})
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
diff --git a/.github/workflows/setupapp.yml b/.github/workflows/setupapp.yml
index a76635e224..6a4b018725 100644
--- a/.github/workflows/setupapp.yml
+++ b/.github/workflows/setupapp.yml
@@ -49,6 +49,8 @@ jobs:
         python -m pip install --upgrade torch torchvision
         python -m pip install -r requirements-dev.txt
     - name: Run unit tests report coverage
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -m pip list
         git config --global --add safe.directory /__w/MONAI/MONAI
@@ -104,6 +106,8 @@ jobs:
         python -m pip install --upgrade pip wheel
         python -m pip install -r requirements-dev.txt
     - name: Run quick tests CPU ubuntu
+      env:
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -m pip list
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
diff --git a/tests/test_bundle_download.py b/tests/test_bundle_download.py
index 89fbe5e8b2..c527982741 100644
--- a/tests/test_bundle_download.py
+++ b/tests/test_bundle_download.py
@@ -56,7 +56,7 @@
 TEST_CASE_5 = [
     ["models/model.pt", "models/model.ts", "configs/train.json"],
     "brats_mri_segmentation",
-    "https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting/brats_mri_segmentation/versions/0.3.9/files/brats_mri_segmentation_v0.3.9.zip",
+    "https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting/brats_mri_segmentation/versions/0.4.0/files/brats_mri_segmentation_v0.4.0.zip",
 ]
 
 TEST_CASE_6 = [["models/model.pt", "configs/train.json"], "renalStructures_CECT_segmentation", "0.1.0"]
@@ -173,6 +173,23 @@ def test_monaihosting_url_download_bundle(self, bundle_files, bundle_name, url):
                     file_path = os.path.join(tempdir, bundle_name, file)
                     self.assertTrue(os.path.exists(file_path))
 
+    @parameterized.expand([TEST_CASE_5])
+    @skip_if_quick
+    def test_nvstaging_source_download_bundle(self, bundle_files, bundle_name, _url):
+        with skip_if_downloading_fails():
+            # download a single file from url, also use `args_file`
+            with tempfile.TemporaryDirectory() as tempdir:
+                def_args = {"name": bundle_name, "bundle_dir": tempdir}
+                def_args_file = os.path.join(tempdir, "def_args.json")
+                parser = ConfigParser()
+                parser.export_config_file(config=def_args, filepath=def_args_file)
+                cmd = ["coverage", "run", "-m", "monai.bundle", "download", "--args_file", def_args_file]
+                cmd += ["--progress", "False", "--source", "nvstaging"]
+                command_line_tests(cmd)
+                for file in bundle_files:
+                    file_path = os.path.join(tempdir, bundle_name, file)
+                    self.assertTrue(os.path.exists(file_path))
+
     @parameterized.expand([TEST_CASE_6])
     @skip_if_quick
     def test_monaihosting_source_download_bundle(self, bundle_files, bundle_name, version):

From b4aff816a19a0255b1338ffd3dfc2ef8624b9976 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 19:59:23 +0800
Subject: [PATCH 05/21] fix pre-commit error

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 .github/workflows/pythonapp-min.yml | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/pythonapp-min.yml b/.github/workflows/pythonapp-min.yml
index 504517102a..cf43996d84 100644
--- a/.github/workflows/pythonapp-min.yml
+++ b/.github/workflows/pythonapp-min.yml
@@ -60,8 +60,6 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (CPU ${{ runner.os }})
-      env:
-        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
@@ -69,6 +67,7 @@ jobs:
       shell: bash
       env:
         QUICKTEST: True
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
 
   min-dep-py3:  # min dependencies installed tests for different python
     runs-on: ubuntu-latest
@@ -108,14 +107,13 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (CPU ${{ runner.os }})
-      env:
-        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
         ./runtests.sh --min
       env:
         QUICKTEST: True
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
 
   min-dep-pytorch:  # min dependencies installed tests for different pytorch
     runs-on: ubuntu-latest
@@ -159,11 +157,10 @@ jobs:
         BUILD_MONAI=0 python setup.py develop  # no compile of extensions
       shell: bash
     - name: Run quick tests (pytorch ${{ matrix.pytorch-version }})
-      env:
-        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
         python -c "import monai; monai.config.print_config()"
         ./runtests.sh --min
       env:
         QUICKTEST: True
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}

From e438623f4dc6f0059cc1af77230b732e278ec14b Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:09:39 +0800
Subject: [PATCH 06/21] fix pre-commit issue

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 .github/workflows/docker.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index 5581bc5392..d9345a06d4 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -90,8 +90,6 @@ jobs:
     runs-on: [self-hosted, linux, X64, docker]
     steps:
     - name: Import
-      env:
-        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
       run: |
         export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 CUDA_VISIBLE_DEVICES=  # cpu-only
         python -c 'import monai; monai.config.print_debug_info()'
@@ -102,3 +100,4 @@ jobs:
       shell: bash
       env:
         QUICKTEST: True
+        NGC_API_KEY: ${{ secrets.NGC_API_KEY }}

From 46b675f5713cd7b70984d905e288a1b7ce20545b Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:17:58 +0800
Subject: [PATCH 07/21] fix mypy

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 4b4c71b269..7240a4d869 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -234,7 +234,7 @@ def _download_from_nvstaging(
     remove_prefix: str | None,
     org: str,
     team: str,
-    headers: str | None = None,
+    headers: dict = {},
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
@@ -266,7 +266,7 @@ def _get_ngc_token(api_key, retry=0):
         if not response.ok:
             if retry < 3:
                 logger.info(f"Retrying {retry} time(s) to GET {url}.")
-                return _get_ngc_token(url, headers, retry + 1)
+                return _get_ngc_token(url, retry + 1)
             raise RuntimeError("NGC API response is not ok. Failed to get token.")
         else:
             token = response.json()["token"]
@@ -302,7 +302,7 @@ def _get_latest_bundle_version_private_registry(name, org, team, headers={}):
     return model_info["model"]["latestVersionIdStr"]
 
 
-def _get_latest_bundle_version(source: str, name: str, repo: str, **kwargs) -> dict[str, list[str] | str] | Any | None:
+def _get_latest_bundle_version(source: str, name: str, repo: str, **kwargs: Any) -> dict[str, list[str] | str] | Any | None:
     if source == "ngc":
         name = _add_ngc_prefix(name)
         model_dict = _get_all_ngc_models(name)

From ae95599b94687015470b9d00faacdd1672ccd31a Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:26:29 +0800
Subject: [PATCH 08/21] fix flake8

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 7240a4d869..3a0486724f 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -228,13 +228,7 @@ def _download_from_ngc(
 
 
 def _download_from_nvstaging(
-    download_path: Path,
-    filename: str,
-    version: str,
-    remove_prefix: str | None,
-    org: str,
-    team: str,
-    headers: dict = {},
+    download_path: Path, filename: str, version: str, remove_prefix: str | None, org: str, team: str, headers: dict = {}
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
@@ -302,7 +296,9 @@ def _get_latest_bundle_version_private_registry(name, org, team, headers={}):
     return model_info["model"]["latestVersionIdStr"]
 
 
-def _get_latest_bundle_version(source: str, name: str, repo: str, **kwargs: Any) -> dict[str, list[str] | str] | Any | None:
+def _get_latest_bundle_version(
+    source: str, name: str, repo: str, **kwargs: Any
+) -> dict[str, list[str] | str] | Any | None:
     if source == "ngc":
         name = _add_ngc_prefix(name)
         model_dict = _get_all_ngc_models(name)

From b9f1548e18e5f7452644f74580f9b1fdffd17e56 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:41:02 +0800
Subject: [PATCH 09/21] fix flake8

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 3a0486724f..16199ab061 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -228,12 +228,13 @@ def _download_from_ngc(
 
 
 def _download_from_nvstaging(
-    download_path: Path, filename: str, version: str, remove_prefix: str | None, org: str, team: str, headers: dict = {}
+    download_path: Path, filename: str, version: str, remove_prefix: str | None, org: str, team: str, headers: dict | None = None
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
     request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, org=org, team=team)
     if has_requests:
+        headers = {} if headers is None else headers
         response = requests_get(request_url, headers=headers)
         response.raise_for_status()
     else:
@@ -280,7 +281,7 @@ def _get_latest_bundle_version_monaihosting(name):
     return model_info["model"]["latestVersionIdStr"]
 
 
-def _get_latest_bundle_version_private_registry(name, org, team, headers={}):
+def _get_latest_bundle_version_private_registry(name, org, team, headers=None):
     team_str = ""
     if team != "no-team":
         team_str = f"team/{team}"
@@ -288,6 +289,7 @@ def _get_latest_bundle_version_private_registry(name, org, team, headers={}):
     full_url = f"{url}/{name.lower()}"
     requests_get, has_requests = optional_import("requests", name="get")
     if has_requests:
+        headers = {} if headers is None else headers
         resp = requests_get(full_url, headers=headers)
         resp.raise_for_status()
     else:

From ec3b62c374db4eb396d70165f653902b72e94eca Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:50:39 +0800
Subject: [PATCH 10/21] fix flake8

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 16199ab061..51a4918b22 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -228,7 +228,13 @@ def _download_from_ngc(
 
 
 def _download_from_nvstaging(
-    download_path: Path, filename: str, version: str, remove_prefix: str | None, org: str, team: str, headers: dict | None = None
+    download_path: Path,
+    filename: str,
+    version: str,
+    remove_prefix: str | None,
+    org: str,
+    team: str,
+    headers: dict | None = None,
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)

From a9bd84717a76273fe45c7b075de303183265cba0 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 11:53:05 +0800
Subject: [PATCH 11/21] use env instead of arg

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 46 +++++++++++++++--------------------------
 1 file changed, 17 insertions(+), 29 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 51a4918b22..f6f32526b9 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -172,11 +172,8 @@ def _get_ngc_bundle_url(model_name: str, version: str) -> str:
     return f"https://api.ngc.nvidia.com/v2/models/nvidia/monaitoolkit/{model_name.lower()}/versions/{version}/zip"
 
 
-def _get_nvstaging_bundle_url(model_name: str, version: str, org: str, team: str) -> str:
-    team_str = ""
-    if team != "no-team":
-        team_str = f"team/{team}"
-    return f"https://api.ngc.nvidia.com/v2/org/{org}/{team_str}/models/{model_name.lower()}/versions/{version}/zip"
+def _get_nvstaging_bundle_url(model_name: str, version: str, repo: str) -> str:
+    return f"https://api.ngc.nvidia.com/v2/{repo}/models/{model_name.lower()}/versions/{version}/zip"
 
 
 def _get_monaihosting_bundle_url(model_name: str, version: str) -> str:
@@ -232,13 +229,12 @@ def _download_from_nvstaging(
     filename: str,
     version: str,
     remove_prefix: str | None,
-    org: str,
-    team: str,
+    repo: str | None,
     headers: dict | None = None,
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
-    request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, org=org, team=team)
+    request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, repo=repo)
     if has_requests:
         headers = {} if headers is None else headers
         response = requests_get(request_url, headers=headers)
@@ -287,11 +283,8 @@ def _get_latest_bundle_version_monaihosting(name):
     return model_info["model"]["latestVersionIdStr"]
 
 
-def _get_latest_bundle_version_private_registry(name, org, team, headers=None):
-    team_str = ""
-    if team != "no-team":
-        team_str = f"team/{team}"
-    url = f"https://api.ngc.nvidia.com/v2/org/{org}/{team_str}/models"
+def _get_latest_bundle_version_private_registry(name, repo, headers=None):
+    url = f"https://api.ngc.nvidia.com/v2/{repo}/models"
     full_url = f"{url}/{name.lower()}"
     requests_get, has_requests = optional_import("requests", name="get")
     if has_requests:
@@ -317,11 +310,9 @@ def _get_latest_bundle_version(
     elif source == "monaihosting":
         return _get_latest_bundle_version_monaihosting(name)
     elif source == "nvstaging":
-        org = kwargs.pop("org", "nvstaging")
-        team = kwargs.pop("team", "monaitoolkit")
         headers = kwargs.pop("headers", {})
         name = _add_ngc_prefix(name)
-        return _get_latest_bundle_version_private_registry(name, org, team, headers)
+        return _get_latest_bundle_version_private_registry(name, repo, headers)
     elif source == "github":
         repo_owner, repo_name, tag_name = repo.split("/")
         return get_bundle_versions(name, repo=f"{repo_owner}/{repo_name}", tag=tag_name)["latest_version"]
@@ -359,8 +350,6 @@ def download(
     remove_prefix: str | None = "monai_",
     progress: bool = True,
     args_file: str | None = None,
-    org: str | None = None,
-    team: str | None = None,
     api_key: str | None = None,
 ) -> None:
     """
@@ -437,12 +426,10 @@ def download(
         url=url,
         remove_prefix=remove_prefix,
         progress=progress,
-        org=org,
-        team=team,
     )
 
     _log_input_summary(tag="download", args=_args)
-    source_, progress_, remove_prefix_, repo_, name_, version_, bundle_dir_, url_, org_, team_ = _pop_args(
+    source_, progress_, remove_prefix_, repo_, name_, version_, bundle_dir_, url_ = _pop_args(
         _args,
         "source",
         "progress",
@@ -452,14 +439,18 @@ def download(
         version=None,
         bundle_dir=None,
         url=None,
-        org=None,
-        team=None,
     )
 
     bundle_dir_ = _process_bundle_dir(bundle_dir_)
     if repo_ is None:
+        org_ = os.getenv("NGC_ORG", None)
+        team_ = os.getenv("NGC_TEAM", None)
+        if org_ is not None:
+            repo_ = f"org/{org_}/team/{team_}" if team_ is not None else f"org/{org_}"
         repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"
-    if len(repo_.split("/")) != 3 and source_ != "huggingface_hub":
+    if len(repo_.split("/")) not in (2, 4) and source_ == "nvstaging":
+        raise ValueError("repo should be in the form of `org/org_name/team/team_name` or `org/org_name`.")
+    if len(repo_.split("/")) != 3 and source_ == "github":
         raise ValueError("repo should be in the form of `repo_owner/repo_name/release_tag`.")
     elif len(repo_.split("/")) != 2 and source_ == "huggingface_hub":
         raise ValueError("Hugging Face Hub repo should be in the form of `repo_owner/repo_name`")
@@ -476,8 +467,6 @@ def download(
             raise ValueError(f"To download from source: {source_}, `name` must be provided.")
         if source == "nvstaging":
             api_key = os.getenv("NGC_API_KEY", api_key)
-            org_ = "nvstaging" if org_ is None else org_
-            team_ = "monaitoolkit" if team_ is None else team_
             if api_key is None:
                 raise ValueError("API key is required for nvstaging source.")
             else:
@@ -486,7 +475,7 @@ def download(
 
         if version_ is None:
             version_ = _get_latest_bundle_version(
-                source=source_, name=name_, repo=repo_, org=org_, team=team_, headers=headers
+                source=source_, name=name_, repo=repo_, headers=headers
             )
         if source_ == "github":
             if version_ is not None:
@@ -508,8 +497,7 @@ def download(
                 filename=name_,
                 version=version_,
                 remove_prefix=remove_prefix_,
-                org=org_,
-                team=team_,
+                repo=repo_,
                 headers=headers,
             )
         elif source_ == "huggingface_hub":

From 37681c5a7d6c3ec79053ebd5e20a34dab6e1aeb6 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 11:58:00 +0800
Subject: [PATCH 12/21] remove api_key

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index f6f32526b9..bf995986cb 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -350,7 +350,6 @@ def download(
     remove_prefix: str | None = "monai_",
     progress: bool = True,
     args_file: str | None = None,
-    api_key: str | None = None,
 ) -> None:
     """
     download bundle from the specified source or url. The bundle should be a zip file and it
@@ -400,10 +399,11 @@ def download(
             Default is `bundle` subfolder under `torch.hub.get_dir()`.
         source: storage location name. This argument is used when `url` is `None`.
             In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
-            it should be "ngc", "monaihosting", "github", or "huggingface_hub".
+            it should be "ngc", "monaihosting", "github", "nvstaging", or "huggingface_hub".
         repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
             If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
             If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
+            If `source` is "nvstaging", it should be in the form of "org/org_name" or "org/org_name/team/team_name".
         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)`.
@@ -466,7 +466,7 @@ def download(
         if name_ is None:
             raise ValueError(f"To download from source: {source_}, `name` must be provided.")
         if source == "nvstaging":
-            api_key = os.getenv("NGC_API_KEY", api_key)
+            api_key = os.getenv("NGC_API_KEY", None)
             if api_key is None:
                 raise ValueError("API key is required for nvstaging source.")
             else:

From d7fd69d44f6e1aae20bc013e0a64bf121874869f Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 12:02:55 +0800
Subject: [PATCH 13/21] update ci env

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 .github/workflows/conda.yml         | 2 ++
 .github/workflows/cron.yml          | 6 ++++++
 .github/workflows/docker.yml        | 2 ++
 .github/workflows/integration.yml   | 2 ++
 .github/workflows/pythonapp-min.yml | 6 ++++++
 .github/workflows/setupapp.yml      | 4 ++++
 monai/bundle/scripts.py             | 3 +++
 7 files changed, 25 insertions(+)

diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
index a12d79cea1..394685acd3 100644
--- a/.github/workflows/conda.yml
+++ b/.github/workflows/conda.yml
@@ -62,6 +62,8 @@ jobs:
       shell: bash -el {0}
       env:
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         conda activate monai
         $(pwd)/runtests.sh --build --unittests
diff --git a/.github/workflows/cron.yml b/.github/workflows/cron.yml
index 8658508fdf..cc113b0446 100644
--- a/.github/workflows/cron.yml
+++ b/.github/workflows/cron.yml
@@ -52,6 +52,8 @@ jobs:
     - name: Run tests report coverage
       env:
           NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+          NGC_ORG: ${{ secrets.NGC_ORG }}
+          NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         export LAUNCH_DELAY=$[ $RANDOM % 16 * 60 ]
         echo "Sleep $LAUNCH_DELAY"
@@ -98,6 +100,8 @@ jobs:
     - name: Run tests report coverage
       env:
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         export LAUNCH_DELAY=$[ $RANDOM % 16 * 60 ]
         echo "Sleep $LAUNCH_DELAY"
@@ -202,6 +206,8 @@ jobs:
       # BUILD_MONAI=1 is necessary for triggering the USE_COMPILED flag.
       env:
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         cd /opt/monai
         nvidia-smi
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index d9345a06d4..17ffe4cf90 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -101,3 +101,5 @@ jobs:
       env:
         QUICKTEST: True
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
index 9810ae8bb1..5be2ebb86c 100644
--- a/.github/workflows/integration.yml
+++ b/.github/workflows/integration.yml
@@ -69,6 +69,8 @@ jobs:
       env:
         BUILD_MONAI: 1
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: ./runtests.sh --build --net
 
     - name: Add reaction
diff --git a/.github/workflows/pythonapp-min.yml b/.github/workflows/pythonapp-min.yml
index cf43996d84..b0d37937e9 100644
--- a/.github/workflows/pythonapp-min.yml
+++ b/.github/workflows/pythonapp-min.yml
@@ -68,6 +68,8 @@ jobs:
       env:
         QUICKTEST: True
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
 
   min-dep-py3:  # min dependencies installed tests for different python
     runs-on: ubuntu-latest
@@ -114,6 +116,8 @@ jobs:
       env:
         QUICKTEST: True
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
 
   min-dep-pytorch:  # min dependencies installed tests for different pytorch
     runs-on: ubuntu-latest
@@ -164,3 +168,5 @@ jobs:
       env:
         QUICKTEST: True
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
diff --git a/.github/workflows/setupapp.yml b/.github/workflows/setupapp.yml
index 6a4b018725..7e01f55cd9 100644
--- a/.github/workflows/setupapp.yml
+++ b/.github/workflows/setupapp.yml
@@ -51,6 +51,8 @@ jobs:
     - name: Run unit tests report coverage
       env:
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         python -m pip list
         git config --global --add safe.directory /__w/MONAI/MONAI
@@ -108,6 +110,8 @@ jobs:
     - name: Run quick tests CPU ubuntu
       env:
         NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+        NGC_ORG: ${{ secrets.NGC_ORG }}
+        NGC_TEAM: ${{ secrets.NGC_TEAM }}
       run: |
         python -m pip list
         python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))'
diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index bf995986cb..5631e5d9ba 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -379,6 +379,9 @@ def download(
         # Execute this module as a CLI entry, and download bundle via URL:
         python -m monai.bundle download --name <bundle_name> --url <url>
 
+        # Execute this module as a CLI entry, and download bundle from nvstaging with latest version:
+        python -m monai.bundle download --name <bundle_name> --source "nvstaging" --bundle_dir "./" --repo "org/org_name"
+
         # 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.
         # The content of the JSON / YAML file is a dictionary. For example:

From e08fa1a68292c08386093ebb010dea347d12b332 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 12:06:11 +0800
Subject: [PATCH 14/21] update doc string

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 5631e5d9ba..8463383bb4 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -403,10 +403,12 @@ def download(
         source: storage location name. This argument is used when `url` is `None`.
             In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
             it should be "ngc", "monaihosting", "github", "nvstaging", or "huggingface_hub".
+            If source is "nvstaging", you need specify the NGC_API_KEY in the environment variable.
         repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
             If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
             If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
             If `source` is "nvstaging", it should be in the form of "org/org_name" or "org/org_name/team/team_name".
+                Or you can specify the environment variable NGC_ORG and NGC_TEAM.
         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)`.

From ed131276e16bebb462e1eeab28b1d6b99d8e9559 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 12:06:43 +0800
Subject: [PATCH 15/21] fix flake8

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 8463383bb4..fee7a3b34c 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -435,15 +435,7 @@ def download(
 
     _log_input_summary(tag="download", args=_args)
     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,
+        _args, "source", "progress", remove_prefix=None, repo=None, name=None, version=None, bundle_dir=None, url=None
     )
 
     bundle_dir_ = _process_bundle_dir(bundle_dir_)
@@ -479,9 +471,7 @@ def download(
                 headers = {"Authorization": f"Bearer {token}"}
 
         if version_ is None:
-            version_ = _get_latest_bundle_version(
-                source=source_, name=name_, repo=repo_, headers=headers
-            )
+            version_ = _get_latest_bundle_version(source=source_, name=name_, repo=repo_, headers=headers)
         if source_ == "github":
             if version_ is not None:
                 name_ = "_v".join([name_, version_])

From 1eb15cfed50eff4170cdc23a936462af6d90a9d0 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 13:17:43 +0800
Subject: [PATCH 16/21] fix ci

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 .github/workflows/blossom-ci.yml | 2 ++
 monai/bundle/scripts.py          | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml
index cdb134e6ad..736d4bfdf6 100644
--- a/.github/workflows/blossom-ci.yml
+++ b/.github/workflows/blossom-ci.yml
@@ -94,6 +94,8 @@ jobs:
         env:
           OPERATION: 'START-CI-JOB'
           NGC_API_KEY: ${{ secrets.NGC_API_KEY }}
+          NGC_ORG: ${{ secrets.NGC_ORG }}
+          NGC_TEAM: ${{ secrets.NGC_TEAM }}
           CI_SERVER: ${{ secrets.CI_SERVER }}
           REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index fee7a3b34c..0d20eadd4b 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -229,7 +229,7 @@ def _download_from_nvstaging(
     filename: str,
     version: str,
     remove_prefix: str | None,
-    repo: str | None,
+    repo: str,
     headers: dict | None = None,
 ) -> None:
     # ensure prefix is contained

From 49a6244bde34b765841dce8de990d03683bcc0b1 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 13:31:09 +0800
Subject: [PATCH 17/21] fix format

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 0d20eadd4b..36b7bd67c6 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -225,12 +225,7 @@ def _download_from_ngc(
 
 
 def _download_from_nvstaging(
-    download_path: Path,
-    filename: str,
-    version: str,
-    remove_prefix: str | None,
-    repo: str,
-    headers: dict | None = None,
+    download_path: Path, filename: str, version: str, remove_prefix: str | None, repo: str, headers: dict | None = None
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
@@ -407,8 +402,8 @@ def download(
         repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
             If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
             If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
-            If `source` is "nvstaging", it should be in the form of "org/org_name" or "org/org_name/team/team_name".
-                Or you can specify the environment variable NGC_ORG and NGC_TEAM.
+            If `source` is "nvstaging", it should be in the form of "org/org_name" or "org/org_name/team/team_name",
+            or you can specify the environment variable NGC_ORG and NGC_TEAM.
         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)`.

From 57b5d03e3c08a21a83921fec33022f0eadbe9050 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 16:17:57 +0800
Subject: [PATCH 18/21] address comments

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py       | 28 ++++++++++++++--------------
 tests/test_bundle_download.py |  4 ++--
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 36b7bd67c6..d79e3865dd 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -172,7 +172,7 @@ def _get_ngc_bundle_url(model_name: str, version: str) -> str:
     return f"https://api.ngc.nvidia.com/v2/models/nvidia/monaitoolkit/{model_name.lower()}/versions/{version}/zip"
 
 
-def _get_nvstaging_bundle_url(model_name: str, version: str, repo: str) -> str:
+def _get_ngc_private_bundle_url(model_name: str, version: str, repo: str) -> str:
     return f"https://api.ngc.nvidia.com/v2/{repo}/models/{model_name.lower()}/versions/{version}/zip"
 
 
@@ -224,12 +224,12 @@ def _download_from_ngc(
     extractall(filepath=filepath, output_dir=extract_path, has_base=True)
 
 
-def _download_from_nvstaging(
+def _download_from_ngc_private(
     download_path: Path, filename: str, version: str, remove_prefix: str | None, repo: str, headers: dict | None = None
 ) -> None:
     # ensure prefix is contained
     filename = _add_ngc_prefix(filename)
-    request_url = _get_nvstaging_bundle_url(model_name=filename, version=version, repo=repo)
+    request_url = _get_ngc_private_bundle_url(model_name=filename, version=version, repo=repo)
     if has_requests:
         headers = {} if headers is None else headers
         response = requests_get(request_url, headers=headers)
@@ -304,7 +304,7 @@ def _get_latest_bundle_version(
         return None
     elif source == "monaihosting":
         return _get_latest_bundle_version_monaihosting(name)
-    elif source == "nvstaging":
+    elif source == "ngc_private":
         headers = kwargs.pop("headers", {})
         name = _add_ngc_prefix(name)
         return _get_latest_bundle_version_private_registry(name, repo, headers)
@@ -374,8 +374,8 @@ def download(
         # Execute this module as a CLI entry, and download bundle via URL:
         python -m monai.bundle download --name <bundle_name> --url <url>
 
-        # Execute this module as a CLI entry, and download bundle from nvstaging with latest version:
-        python -m monai.bundle download --name <bundle_name> --source "nvstaging" --bundle_dir "./" --repo "org/org_name"
+        # Execute this module as a CLI entry, and download bundle from ngc_private with latest version:
+        python -m monai.bundle download --name <bundle_name> --source "ngc_private" --bundle_dir "./" --repo "org/org_name"
 
         # 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.
@@ -397,12 +397,12 @@ def download(
             Default is `bundle` subfolder under `torch.hub.get_dir()`.
         source: storage location name. This argument is used when `url` is `None`.
             In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
-            it should be "ngc", "monaihosting", "github", "nvstaging", or "huggingface_hub".
-            If source is "nvstaging", you need specify the NGC_API_KEY in the environment variable.
+            it should be "ngc", "monaihosting", "github", "ngc_private", or "huggingface_hub".
+            If source is "ngc_private", you need specify the NGC_API_KEY in the environment variable.
         repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
             If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
             If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
-            If `source` is "nvstaging", it should be in the form of "org/org_name" or "org/org_name/team/team_name",
+            If `source` is "ngc_private", it should be in the form of "org/org_name" or "org/org_name/team/team_name",
             or you can specify the environment variable NGC_ORG and NGC_TEAM.
         url: url to download the data. If not `None`, data will be downloaded directly
             and `source` will not be checked.
@@ -440,7 +440,7 @@ def download(
         if org_ is not None:
             repo_ = f"org/{org_}/team/{team_}" if team_ is not None else f"org/{org_}"
         repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"
-    if len(repo_.split("/")) not in (2, 4) and source_ == "nvstaging":
+    if len(repo_.split("/")) not in (2, 4) and source_ == "ngc_private":
         raise ValueError("repo should be in the form of `org/org_name/team/team_name` or `org/org_name`.")
     if len(repo_.split("/")) != 3 and source_ == "github":
         raise ValueError("repo should be in the form of `repo_owner/repo_name/release_tag`.")
@@ -457,10 +457,10 @@ def download(
         headers = {}
         if name_ is None:
             raise ValueError(f"To download from source: {source_}, `name` must be provided.")
-        if source == "nvstaging":
+        if source == "ngc_private":
             api_key = os.getenv("NGC_API_KEY", None)
             if api_key is None:
-                raise ValueError("API key is required for nvstaging source.")
+                raise ValueError("API key is required for ngc_private source.")
             else:
                 token = _get_ngc_token(api_key)
                 headers = {"Authorization": f"Bearer {token}"}
@@ -481,8 +481,8 @@ def download(
                 remove_prefix=remove_prefix_,
                 progress=progress_,
             )
-        elif source_ == "nvstaging":
-            _download_from_nvstaging(
+        elif source_ == "ngc_private":
+            _download_from_ngc_private(
                 download_path=bundle_dir_,
                 filename=name_,
                 version=version_,
diff --git a/tests/test_bundle_download.py b/tests/test_bundle_download.py
index c527982741..fe7caf5c17 100644
--- a/tests/test_bundle_download.py
+++ b/tests/test_bundle_download.py
@@ -175,7 +175,7 @@ def test_monaihosting_url_download_bundle(self, bundle_files, bundle_name, url):
 
     @parameterized.expand([TEST_CASE_5])
     @skip_if_quick
-    def test_nvstaging_source_download_bundle(self, bundle_files, bundle_name, _url):
+    def test_ngc_private_source_download_bundle(self, bundle_files, bundle_name, _url):
         with skip_if_downloading_fails():
             # download a single file from url, also use `args_file`
             with tempfile.TemporaryDirectory() as tempdir:
@@ -184,7 +184,7 @@ def test_nvstaging_source_download_bundle(self, bundle_files, bundle_name, _url)
                 parser = ConfigParser()
                 parser.export_config_file(config=def_args, filepath=def_args_file)
                 cmd = ["coverage", "run", "-m", "monai.bundle", "download", "--args_file", def_args_file]
-                cmd += ["--progress", "False", "--source", "nvstaging"]
+                cmd += ["--progress", "False", "--source", "ngc_private"]
                 command_line_tests(cmd)
                 for file in bundle_files:
                     file_path = os.path.join(tempdir, bundle_name, file)

From 082933b1dda155e1c72f76c92b3333a969ce8534 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 16:59:26 +0800
Subject: [PATCH 19/21] minor fix

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index d79e3865dd..04f82dbb8f 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -439,13 +439,14 @@ def download(
         team_ = os.getenv("NGC_TEAM", None)
         if org_ is not None:
             repo_ = f"org/{org_}/team/{team_}" if team_ is not None else f"org/{org_}"
-        repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"
+        else:
+            repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"
     if len(repo_.split("/")) not in (2, 4) and source_ == "ngc_private":
-        raise ValueError("repo should be in the form of `org/org_name/team/team_name` or `org/org_name`.")
+        raise ValueError(f"repo should be in the form of `org/org_name/team/team_name` or `org/org_name`, got {repo_}.")
     if len(repo_.split("/")) != 3 and source_ == "github":
-        raise ValueError("repo should be in the form of `repo_owner/repo_name/release_tag`.")
+        raise ValueError(f"repo should be in the form of `repo_owner/repo_name/release_tag`, got {repo_}.")
     elif len(repo_.split("/")) != 2 and source_ == "huggingface_hub":
-        raise ValueError("Hugging Face Hub repo should be in the form of `repo_owner/repo_name`")
+        raise ValueError(f"Hugging Face Hub repo should be in the form of `repo_owner/repo_name`, got {repo_}.")
     if url_ is not None:
         if name_ is not None:
             filepath = bundle_dir_ / f"{name_}.zip"

From cb18482b6c3bfb3be5396c4ccd2a8c9e39ecae58 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 17:17:05 +0800
Subject: [PATCH 20/21] address comments

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index 04f82dbb8f..ab746a14dd 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -235,7 +235,7 @@ def _download_from_ngc_private(
         response = requests_get(request_url, headers=headers)
         response.raise_for_status()
     else:
-        raise ValueError("NGC API requires requests package.  Please install it.")
+        raise ValueError("NGC API requires requests package. Please install it.")
 
     zip_path = download_path / f"{filename}_v{version}.zip"
     with open(zip_path, "wb") as f:
@@ -273,7 +273,7 @@ def _get_latest_bundle_version_monaihosting(name):
         resp = requests_get(full_url)
         resp.raise_for_status()
     else:
-        raise ValueError("NGC API requires requests package.  Please install it.")
+        raise ValueError("NGC API requires requests package. Please install it.")
     model_info = json.loads(resp.text)
     return model_info["model"]["latestVersionIdStr"]
 
@@ -287,7 +287,7 @@ def _get_latest_bundle_version_private_registry(name, repo, headers=None):
         resp = requests_get(full_url, headers=headers)
         resp.raise_for_status()
     else:
-        raise ValueError("NGC API requires requests package.  Please install it.")
+        raise ValueError("NGC API requires requests package. Please install it.")
     model_info = json.loads(resp.text)
     return model_info["model"]["latestVersionIdStr"]
 

From ac013dd5b6ba9e065e3038ebc1fd33619182012c Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 17 Jul 2024 17:31:23 +0800
Subject: [PATCH 21/21] add comments

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
---
 monai/bundle/scripts.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
index ab746a14dd..975b6cbdbd 100644
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -256,6 +256,7 @@ def _get_ngc_token(api_key, retry=0):
     if has_requests:
         response = requests_get(url, headers=headers)
         if not response.ok:
+            # retry 3 times, if failed, raise an error.
             if retry < 3:
                 logger.info(f"Retrying {retry} time(s) to GET {url}.")
                 return _get_ngc_token(url, retry + 1)