Add module for loading test data.

.devcontainer/devcontainer.json

@@ -14,6 +14,8 @@
 	"mounts": [
 		// Mount the local ~/.aws config to pass along AWS credentials for PBSS.
 		"source=${localEnv:HOME}/.aws,target=/home/bionemo/.aws,type=bind,consistency=cached",
+		"source=${localEnv:HOME}/.ngc,target=/home/bionemo/.ngc,type=bind,consistency=cached",
+		"source=${localEnv:HOME}/.cache,target=/home/bionemo/.cache,type=bind,consistency=cached",
 		"source=${localEnv:HOME}/.ssh,target=/home/bionemo/.ssh,readonly,type=bind,consistency=cached"
 	],
 	"containerEnv": {

.gitignore

@@ -7,10 +7,7 @@ docs/site/
 .gnupg
 
 # Lightning and project output files
-/test_data
 lightning_logs
-/models
-/data
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.vscode/settings.json

@@ -1,14 +1,23 @@
 {
     "cSpell.words": [
+        "adata",
         "allclose",
         "bionemo",
+        "boto",
+        "botocore",
+        "docstrings",
         "dtype",
         "NBVAL",
         "nemo",
+        "ngcsdk",
+        "pbss",
+        "platformdirs",
         "pretraining",
+        "pydantic",
         "rampup",
         "resamplers",
         "singlecell",
+        "tqdm",
         "uniref"
     ],
     "editor.rulers": [

docs/mkdocs.yml

@@ -9,6 +9,7 @@ theme:
     - content.code.annotate
     - navigation.tabs
     - navigation.sections
+    - navigation.indexes
     - navigation.instant
     - navigation.instant.prefetch
     - search.suggest
@@ -54,10 +55,14 @@ plugins:
   - mkdocs-jupyter
 
 markdown_extensions:
+  - pymdownx.details
+  - pymdownx.superfences
   - pymdownx.snippets:
       check_paths: true
       restrict_base_path: False
       base_path: "."
+  - def_list
+  - admonition
 
 copyright: |
   © Copyright 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

docs/scripts/gen_ref_pages.py

@@ -33,12 +33,21 @@
         parts = tuple(module_path.parts)
 
         if parts[-1] == "__init__":
-            parts = parts[:-1]
+            # parts = parts[:-1]
+            continue  # Don't generate ref pages for __init__.py
         elif parts[-1] == "__main__":
-            continue
+            continue  # Don't generate ref pages for __main__.py
 
         with mkdocs_gen_files.open(full_doc_path, "w") as fd:
             identifier = ".".join(parts)
             print("::: " + identifier, file=fd)
 
         mkdocs_gen_files.set_edit_path(full_doc_path, path.relative_to(root))
+
+    for path in sorted(src.rglob("*.md")):
+        doc_path = path.relative_to(src)
+        full_doc_path = Path("API_reference", doc_path)
+        with mkdocs_gen_files.open(full_doc_path, "w") as fd:
+            fd.write(path.read_text())
+        print(full_doc_path)
+        mkdocs_gen_files.set_edit_path(full_doc_path, path.relative_to(root))

scripts/download_artifacts.py

@@ -378,4 +378,5 @@ def main(models: Optional[List[str]], data: Optional[List[str]]) -> None:
     if not (args.models or args.data):
         logging.warning("No models or data were selected to download.")
     else:
+        logging.warning("This script is deprecated, use `bionemo.testing.data.load` instead.")
         main(models=args.models, data=args.data)

scripts/singlecell/geneformer/test_pretrain.py

@@ -27,6 +27,7 @@
 from bionemo.llm.model.biobert.transformer_specs import BiobertSpecOption
 from bionemo.llm.utils.datamodule_utils import parse_kwargs_to_arglist
 from bionemo.testing import megatron_parallel_state_utils
+from bionemo.testing.data.load import load
 
 
 # TODO(@jstjohn) use fixtures for pulling down data and checkpoints
@@ -39,7 +40,8 @@
     # From here, we want to get to the root of the repository: _before_ sub-packages/
     .parent.parent
 ).absolute()
-data_path: Path = bionemo2_root / "test_data/cellxgene_2023-12-15_small/processed_data"
+assert bionemo2_root != Path("/")
+data_path: Path = load("single_cell/testdata-20240506") / "cellxgene_2023-12-15_small" / "processed_data"
 
 
 def test_bionemo2_rootdir():

sub-packages/bionemo-esm2/tests/bionemo/esm2/model/test_model.py

@@ -36,6 +36,7 @@
 from bionemo.llm.model.biobert.model import MegatronBioBertModel
 from bionemo.llm.utils.weight_utils import nemo1_to_nemo2_biobert_key_mapping
 from bionemo.testing import megatron_parallel_state_utils
+from bionemo.testing.data.load import load
 
 
 bionemo2_root: Path = (
@@ -47,7 +48,7 @@
     .parent.parent
 ).absolute()
 assert bionemo2_root != Path("/")
-nemo1_checkpoint_path: Path = bionemo2_root / "models/protein/esm2nv/esm2nv_650M_converted.nemo"
+nemo1_checkpoint_path: Path = load("esm2/nv_650m:1.0")
 
 
 def reduce_hiddens(hiddens: Tensor, attention_mask: Tensor) -> Tensor:

sub-packages/bionemo-example_model/tests/bionemo/example_model/test_lightning_basic.py

@@ -19,6 +19,7 @@
 
 from bionemo.example_model import lightning_basic as lb
 from bionemo.testing import megatron_parallel_state_utils
+from bionemo.testing.data.load import BIONEMO_CACHE_DIR
 
 
 @pytest.mark.needs_gpu
@@ -33,5 +34,5 @@ def test_train_mnist_litautoencoder_with_megatron_strategy_single_gpu():
             enable_nemo_ckpt_io=False,
         )
         trainer = nl.Trainer(accelerator="gpu", devices=1, strategy=strategy, max_steps=10, num_nodes=1)
-        data_module = lb.MNISTDataModule()
+        data_module = lb.MNISTDataModule(str(BIONEMO_CACHE_DIR))
         trainer.fit(model, data_module)

sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py

@@ -40,6 +40,7 @@
 from bionemo.llm.model.biobert.model import BiobertSpecOption
 from bionemo.llm.utils.weight_utils import nemo1_to_nemo2_biobert_key_mapping
 from bionemo.testing import megatron_parallel_state_utils
+from bionemo.testing.data.load import load
 from bionemo.testing.utils import assert_matrix_correlation_above_value, assert_matrix_mape_below_value
 
 
@@ -54,11 +55,11 @@
     .parent.parent
 ).absolute()
 assert bionemo2_root != Path("/")
-nemo1_checkpoint_path: Path = bionemo2_root / "models/singlecell/geneformer/geneformer-qa.nemo"
-nemo1_release_checkpoint_path: Path = bionemo2_root / "models/singlecell/geneformer/geneformer-10M-240530.nemo"
-nemo_1_per_layer_outputs_path: Path = bionemo2_root / "test_data/nemo1-test-outputs-geneformer-qa.pt"
-nemo_1_expected_values_path: Path = bionemo2_root / "test_data/nemo1_geneformer_qa_test_golden_values.pt"
-data_path: Path = bionemo2_root / "test_data/cellxgene_2023-12-15_small/processed_data"
+nemo1_checkpoint_path: Path = load("geneformer/qa")
+nemo1_release_checkpoint_path: Path = load("geneformer/10M_240530")
+nemo_1_per_layer_outputs_path: Path = load("single_cell/nemo1-geneformer-per-layer-outputs")
+nemo_1_expected_values_path: Path = load("single_cell/nemo1-geneformer-golden-vals")
+data_path: Path = load("single_cell/testdata-20240506") / "cellxgene_2023-12-15_small" / "processed_data"
 
 
 CELLS_FOR_TEST: List[List[str]] = [

sub-packages/bionemo-scdl/tests/bionemo/scdl/conftest.py

@@ -18,6 +18,8 @@
 
 import pytest
 
+from bionemo.testing.data.load import load
+
 
 @pytest.fixture
 def test_directory() -> Path:
@@ -26,6 +28,4 @@ def test_directory() -> Path:
     Returns:
         A Path object that is the directory with test data.
     """
-    current_file = Path(__file__).resolve()
-    target_path = current_file.parents[5].absolute() / "test_data/scdl_data"
-    return target_path
+    return load("scdl/sample") / "scdl_data"

sub-packages/bionemo-testing/requirements.txt

@@ -1 +1,9 @@
+boto3
+email-validator
+platformdirs
+pooch
+pydantic>=2.7.0
 pytest
+pyyaml
+ngcsdk
+tqdm

sub-packages/bionemo-testing/src/bionemo/testing/data/README.md

@@ -0,0 +1,75 @@
+# BioNeMo test data management
+
+This library manages the downloading and caching of large or binary data files used in the documentation or test suite.
+These files should not be committed directly to the repo, and instead should be loaded at test-time when they are
+needed.
+
+We currently support two locations for test data or saved models:
+
+SwiftStack
+
+:   SwiftStack or `pbss` is an NVIDIA-internal, s3-compatible object store that allows for very large data and fast,
+    parallel read/writes. Most critically, `pbss` can be uploaded to without legal approvals for dataset redistribution.
+    These files will not be accessible by external collaborators.
+
+[NGC](https://catalog.ngc.nvidia.com/)
+
+:   NGC hosts containers, models, and resources, some of which require authentication and others that are generally
+    available. This library uses the model and resource types to save test data and reference model weights. These items
+    are accessible by external collaborators, but require legal approval before re-distributing test data.
+
+
+## Loading test or example data
+
+Test data are specified via yaml files in `sub-packages/bionemo-testing/src/bionemo/testing/data/resources`. As an
+example, in `esm2.yaml`:
+
+```yaml
+- tag: nv_650m:1.0
+  ngc: "nvidia/clara/esm2nv650m:1.0"
+  ngc_registry: model
+  pbss: "s3://bionemo-ci/models/esm2nv_650M_converted.nemo"
+  sha256: 1e38063cafa808306329428dd17ea6df78c9e5d6b3d2caf04237c555a1f131b7
+  owner: Farhad Ramezanghorbani <[email protected]>
+  description: >
+    A pretrained 650M parameter ESM2 model.
+    See https://ngc.nvidia.com/catalog/models/nvidia:clara:esm2nv650m.
+```
+
+To load these model weights during a test, use the [load][bionemo.testing.data.load.load] function with the filename and
+tag of the desired asset, which returns a path a the specified file:
+
+```python
+path_to_my_checkpoint = load("esm2/nv_650m:1.0")
+config = ESM2Config(nemo1_ckpt_path=path_to_my_checkpoint)
+```
+
+If this function is called without the data available on the local machine, it will be fetched from the default source
+(currently `pbss`.) Otherwise, it will return the cached directory. To download with NGC, pass `source="ngc"` to
+[load][bionemo.testing.data.load.load].
+
+## File unpacking and/or decompression
+
+All test artifacts are individual files. If a zip or tar archive is specified, it will be unpacked automatically, and
+the path to the directory will be returned via [load][bionemo.testing.data.load.load]. Compressed files ('gzip', 'bz2',
+or 'xz') are automatically decompressed before they are returned. The file's compression and/or archive format is
+determined based on the filename specified in the `pbss` URL.
+
+!!! note "Files in NGC resources"
+
+    NGC resources are folders, i.e., they may contain multiple files per resource.
+    [load][bionemo.testing.data.load.load] will _only_ download the filename matching the stem of the `pbss` url. The
+    same NGC resource can therefore be used to host multiple test assets that are used independently.
+
+
+## Adding new test assets
+
+To add new data, first ensure that the data is available from either NGC or `pbss`. Next, extend or create a new yaml
+file in `sub-packages/bionemo-testing/src/bionemo/testing/data/resources` with the required information. Owner emails
+must be provided for all assets. The description and `ngc` fields are currently optional. If the `sha256` is left
+unspecified, `pooch` will report the downloaded file's sha when loaded.
+
+!!! warning
+
+    SHAs should be provided for all files to ensure the download completes correctly, and to invalidate caches if the
+    files change.

sub-packages/bionemo-testing/src/bionemo/testing/data/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.