adds nbval test for ngc download

NVIDIA · Nov 8, 2024 · d4a236f · d4a236f
1 parent 43a6f42
commit d4a236f
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 0 deletions.
diff --git a/sub-packages/bionemo-core/pyproject.toml b/sub-packages/bionemo-core/pyproject.toml
@@ -19,6 +19,7 @@ dependencies = [
     "torch>=2.2.1",
     'boto3',
     'lightning>=2.2.1',
+    'nest_asyncio',
     'ngcsdk',
     'pooch',
     'pydantic>=2.7.0',

diff --git a/sub-packages/bionemo-core/src/bionemo/core/data/load.py b/sub-packages/bionemo-core/src/bionemo/core/data/load.py
@@ -23,6 +23,7 @@
 from typing import Literal, Optional, Sequence, TextIO
 
 import boto3
+import nest_asyncio
 import ngcsdk
 import pooch
 from botocore.config import Config
@@ -87,6 +88,8 @@ class NGCDownloader:
     def __call__(self, url: str, output_file: str | Path, _: pooch.Pooch) -> None:
         """Download a file from NGC."""
         client = default_ngc_client()
+        client.configure()
+        nest_asyncio.apply()
 
         download_fns = {
             "model": client.registry.model.download_version,