Adding NGC model file support to filesystem (#216)

* Adding NGC download cache option

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added ruff as a linting tool.
 - Ported utilities from Modulus Launch to main package.
 - EDM diffusion models and recipes for training and sampling.
+- NGC model registry download integration into package/filesystem.
 
 ### Changed
 

modulus/utils/filesystem.py

@@ -13,14 +13,18 @@
 # limitations under the License.
 
 import hashlib
+import json
 import logging
 import os
+import re
 import urllib.request
+import zipfile
 
 import fsspec
 import fsspec.implementations.cached
 import requests
 import s3fs
+from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
@@ -43,25 +47,105 @@ def _get_fs(path):
         return fsspec.filesystem("file")
 
 
-def _download_cached(path: str, recursive: bool = False) -> str:
+def _download_ngc_model_file(path: str, out_path: str, timeout: int = 300) -> str:
+    """Pulls files from model registry on NGC. Supports private registries when NGC
+    API key is set the the `NGC_API_KEY` environment variable. If download file is a zip
+    folder it will get unzipped.
+
+    Args:
+        path (str): NGC model file path of form:
+            `ngc://models/<org_id/team_id/model_id>@<version>/<path/in/repo>`
+        out_path (str): Output path to save file / folder as
+        timeout (int): Time out of requests, default 5 minutes
+
+    Raises:
+        ValueError: Invlaid url
+
+    Returns:
+        str: output file / folder path
+    """
+    # Strip ngc model url prefix
+    suffix = "ngc://models/"
+    # The regex check
+    pattern = re.compile(f"{suffix}[\w-]+/[\w-]+/[\w-]+@[A-Za-z0-9.]+/[\w/]+")
+    if not pattern.match(path):
+        raise ValueError(
+            "Invalid URL, should be of form ngc://models/<org_id/team_id/model_id>@<version>/<path/in/repo>"
+        )
+
+    path = path.replace(suffix, "")
+    (org, team, model_version, filename) = path.split("/", 4)
+    (model, version) = model_version.split("@", 1)
+    token = ""
+    # If API key environment variable
+    if "NGC_API_KEY" in os.environ:
+        try:
+            session = requests.Session()
+            session.auth = ("$oauthtoken", os.environ["NGC_API_KEY"])
+            headers = {"Accept": "application/json"}
+            authn_url = f"https://authn.nvidia.com/token?service=ngc&scope=group/ngc:{org}&group/ngc:{org}/{team}"
+            r = session.get(authn_url, headers=headers, timeout=5)
+            r.raise_for_status()
+            token = json.loads(r.content)["token"]
+        except requests.exceptions.RequestException:
+            logger.warning(
+                "Failed to get JWT using the API set in NGC_API_KEY environment variable"
+            )
+            raise  # Re-raise the exception
+
+    # Download file, apparently the URL for private registries is different than the public?
+    if len(token) > 0:
+        file_url = f"https://api.ngc.nvidia.com/v2/org/{org}/team/{team}/models/{model}/versions/{version}/files/{filename}"
+    else:
+        file_url = f"https://api.ngc.nvidia.com/v2/models/{org}/{team}/{model}/versions/{version}/files/{filename}"
+    local_url = f"{LOCAL_CACHE}/{filename}"
+    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
+    # Streaming here for larger files
+    with requests.get(file_url, headers=headers, stream=True, timeout=timeout) as r:
+        r.raise_for_status()
+        total_size_in_bytes = int(r.headers.get("content-length", 0))
+        chunk_size = 1024  # 1 kb
+        progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
+        progress_bar.set_description(f"Fetching {filename}")
+        with open(local_url, "wb") as f:
+            for chunk in r.iter_content(chunk_size=chunk_size):
+                progress_bar.update(len(chunk))
+                f.write(chunk)
+        progress_bar.close()
+
+    # Unzip contents if zip file (most model files are)
+    if zipfile.is_zipfile(local_url):
+        with zipfile.ZipFile(local_url, "r") as zip_ref:
+            zip_ref.extractall(out_path)
+        # Clean up zip
+        os.remove(local_url)
+    else:
+        os.rename(local_url, out_path)
+
+    return out_path
+
+
+def _download_cached(
+    path: str, recursive: bool = False, local_cache_path: str = LOCAL_CACHE
+) -> str:
     sha = hashlib.sha256(path.encode())
     filename = sha.hexdigest()
     try:
-        os.makedirs(LOCAL_CACHE, exist_ok=True)
+        os.makedirs(local_cache_path, exist_ok=True)
     except PermissionError as error:
         logger.error(
             "Failed to create cache folder, check permissions or set a cache"
-            + " location using the LOCAL_CACHE enviroment variable"
+            + " location using the LOCAL_CACHE environment variable"
         )
         raise error
     except OSError as error:
         logger.error(
             "Failed to create cache folder, set a cache"
-            + " location using the LOCAL_CACHE enviroment variable"
+            + " location using the LOCAL_CACHE environment variable"
         )
         raise error
 
-    cache_path = os.path.join(LOCAL_CACHE, filename)
+    cache_path = os.path.join(local_cache_path, filename)
 
     url = urllib.parse.urlparse(path)
 
@@ -71,6 +155,9 @@ def _download_cached(path: str, recursive: bool = False) -> str:
         if path.startswith("s3://"):
             fs = _get_fs(path)
             fs.get(path, cache_path, recursive=recursive)
+        elif path.startswith("ngc://models/"):
+            path = _download_ngc_model_file(path, cache_path)
+            return path
         elif url.scheme == "http":
             # urllib.request.urlretrieve(path, cache_path)
             # TODO: Check if this supports directory fetches
@@ -92,12 +179,23 @@ def _download_cached(path: str, recursive: bool = False) -> str:
 
 
 class Package:
-    """A package
-
-    Represents a potentially remote directory tree
+    """A generic file system abstraction. Can be used to represent local and remote
+    file systems. Remote files are automatically fetched and stored in the
+    $LOCAL_CACHE or $HOME/.cache/modulus folder. The `get` method can then be used
+    to access files present.
+
+    Presently one can use Package with the following directories:
+    - Package("/path/to/local/directory") = local file system
+    - Package("s3://bucket/path/to/directory") = object store file system
+    - Package("http://url/path/to/directory") = http file system
+    - Package("ngc://model/<org_id/team_id/model_id>@<version>") = ngc model file system
+
+    Args:
+        root (str): Root directory for file system
+        seperator (str, optional): directory seperator. Defaults to "/".
     """
 
-    def __init__(self, root: str, seperator: str):
+    def __init__(self, root: str, seperator: str = "/"):
         self.root = root
         self.seperator = seperator
 

test/utils/test_filesystem.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 import hashlib
+import os
 from pathlib import Path
 
+import pytest
+
 from modulus.utils import filesystem
 
 
@@ -53,3 +56,46 @@ def test_http_package():
 
     known_checksum = "e075b2836d03f7971f754354807dcdca51a7875c8297cb161557946736d1f7fc"
     assert calculate_checksum(path) == known_checksum
+
+
+@pytest.mark.skip("Skipping because slow, need better test solution")
+def test_ngc_model_file():
+    test_url = "ngc://models/nvidia/modulus/modulus_dlwp_cubesphere@v0.2"
+    package = filesystem.Package(test_url, seperator="/")
+    path = package.get("dlwp_cubesphere.zip")
+
+    path = Path(path)
+    folders = [f for f in path.iterdir()]
+    assert len(folders) == 1 and folders[0].name == "dlwp"
+
+    files = [f for f in folders[0].iterdir()]
+    assert len(files) == 11
+
+
+@pytest.mark.skipif(
+    "NGC_API_KEY" not in os.environ, reason="Skipping because no NGC API key"
+)
+def test_ngc_model_file_private():
+    test_url = "ngc://models/nvstaging/simnet/modulus_ci@v0.1"
+    package = filesystem.Package(test_url, seperator="/")
+    path = package.get("test.txt")
+
+    known_checksum = "d2a84f4b8b650937ec8f73cd8be2c74add5a911ba64df27458ed8229da804a26"
+    assert calculate_checksum(path) == known_checksum
+
+
+def test_ngc_model_file_invalid():
+    test_url = "ngc://models/nvidia/modulus/modulus_dlwp_cubesphere/v0.2"
+    package = filesystem.Package(test_url, seperator="/")
+    with pytest.raises(ValueError):
+        package.get("dlwp_cubesphere.zip")
+
+    test_url = "ngc://models/nvidia/modulus_dlwp_cubesphere@v0.2"
+    package = filesystem.Package(test_url, seperator="/")
+    with pytest.raises(ValueError):
+        package.get("dlwp_cubesphere.zip")
+
+    test_url = "ngc://models/nvidia/modulus/other/modulus_dlwp_cubesphere@v0.2"
+    package = filesystem.Package(test_url, seperator="/")
+    with pytest.raises(ValueError):
+        package.get("dlwp_cubesphere.zip")