Repository: NVIDIA/earth2studio (language: python)

Issue:
🐛[BUG]: LaggedEnsemble overwrite input data 
### Version

main

### On which installation method(s) does this occur?

Source

### Describe the issue

In `earth2studio.perturbations.lagged`, the `LaggedEnsemble` perturbation strategy overwrite the passed data array, 

```python
x[i] = fetch_data(
                source=self.source,
                time=coords["time"] + lag,
                variable=coords["variable"],
                lead_time=coords["lead_time"],
                device=x.device,
            )[0]
```

This can lead to bugs if the input array is reused in another perturbation strategy. We propose the following fix:

```python
y = torch.clone(x)
        for i, lag in enumerate(self.lags):
            y[i] = fetch_data(
                source=self.source,
                time=coords["time"] + lag,
                variable=coords["variable"],
                lead_time=coords["lead_time"],
                device=y.device,
            )[0]
```


Discussion / hints:

Fix patch:
diff --git a/docs/userguide/about/install.md b/docs/userguide/about/install.md
index 7882a344..0542ab89 100644
--- a/docs/userguide/about/install.md
+++ b/docs/userguide/about/install.md
@@ -174,3 +174,5 @@ The import ones are:
 
 - `EARTH2STUDIO_CACHE`: The location of the cache used for Earth2Studio. This is a file
 path where things like models and cached data from data sources will be stored.
+- `EARTH2STUDIO_PACKAGE_TIMEOUT`: The max number of seconds for a download operation of
+a model package file from a remote store such as NGC, Huggingface or S3.
diff --git a/docs/userguide/developer/documentation.md b/docs/userguide/developer/documentation.md
index 22f95662..8688f1be 100644
--- a/docs/userguide/developer/documentation.md
+++ b/docs/userguide/developer/documentation.md
@@ -130,3 +130,10 @@ make docs-dev FILENAME=<example filename .py>
 ```
 
 Build files will always be in `docs/_build/html`.
+Since the docs are static, Python can be used to host them [locally](http://localhost:8000):
+
+```bash
+cd docs/_build/html
+
+python -m http.server 8000
+```
diff --git a/docs/userguide/developer/overview.md b/docs/userguide/developer/overview.md
index 8eb3b28e..a6221eb3 100644
--- a/docs/userguide/developer/overview.md
+++ b/docs/userguide/developer/overview.md
@@ -1,3 +1,7 @@
+<!-- markdownlint-disable MD025 -->
+
+(developer_overview)=
+
 # Overview
 
 This guide assumes you have looked at the repository and are familiar with the majority
diff --git a/docs/userguide/index.md b/docs/userguide/index.md
index 16047421..12954bdf 100644
--- a/docs/userguide/index.md
+++ b/docs/userguide/index.md
@@ -64,6 +64,7 @@ run(["2024-01-01"], 10, model, ds, io)
 - [Overview](developer/overview)
 - [Style](developer/style)
 - [Documentation](developer/documentation)
+- [Testing](developer/testing)
 
 ## Support
 
@@ -112,6 +113,7 @@ advanced/lexicon
 developer/overview
 developer/style
 developer/documentation
+developer/testing
 ```
 
 ```{toctree}
diff --git a/earth2studio/models/auto.py b/earth2studio/models/auto.py
index b1b85af0..7697526e 100644
--- a/earth2studio/models/auto.py
+++ b/earth2studio/models/auto.py
@@ -20,6 +20,7 @@
 import warnings
 from typing import Any
 
+import aiohttp
 import fsspec
 import s3fs
 from fsspec.callbacks import Callback, TqdmCallback
@@ -192,15 +193,26 @@ def __init__(
                 self.root = f"https://api.ngc.nvidia.com/v2/models/{org}/{team}/{model}/versions/{version}/files/"
             else:
                 self.root = f"https://api.ngc.nvidia.com/v2/models/{org}/{model}/versions/{version}/files/"
-            self.fs = HTTPFileSystem(block_size=2**10)
+            self.fs = HTTPFileSystem(
+                block_size=Package.default_blocksize(),
+                client_kwargs={
+                    "timeout": aiohttp.ClientTimeout(total=Package.default_timeout())
+                },
+            )
         elif root.startswith("hf://"):
-            self.fs = HfFileSystem(target_options={"default_block_size": 2**20})
+            # https://github.com/huggingface/huggingface_hub/blob/v0.23.4/src/huggingface_hub/hf_file_system.py#L816
+            if "HF_HUB_DOWNLOAD_TIMEOUT" not in os.environ:
+                os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = str(Package.default_timeout())
+            self.fs = HfFileSystem(
+                target_options={"default_block_size": Package.default_blocksize()}
+            )
         elif root.startswith("s3://"):
             self.fs = s3fs.S3FileSystem(
                 anon=True,
                 client_kwargs={},
-                target_options={"default_block_size": 2**20},
+                default_block_size=Package.default_blocksize(),
             )
+            self.fs.read_timeout = Package.default_timeout()
         else:
             protocol = split_protocol(root)[0]
             self.fs = fsspec.filesystem(protocol)
@@ -226,6 +238,34 @@ def default_cache(cls, path: str = "") -> str:
         default_cache = os.environ.get("EARTH2STUDIO_CACHE", default_cache)
         return os.path.join(default_cache, path)
 
+    @classmethod
+    def default_timeout(cls) -> int:
+        """Default remote store timeout in seconds
+
+        Returns
+        -------
+        int
+            Time out in seconds
+        """
+        default_timeout = 300
+        try:
+            timeout = os.environ.get("EARTH2STUDIO_PACKAGE_TIMEOUT", default_timeout)
+            default_timeout = int(timeout)
+        except ValueError:
+            pass
+        return default_timeout
+
+    @classmethod
+    def default_blocksize(cls) -> int:
+        """Default remote store block size
+
+        Returns
+        -------
+        int
+            Download block size in bytes
+        """
+        return 2**20
+
     @property
     def cache(self) -> str:
         """Cache path"""
diff --git a/earth2studio/models/batch.py b/earth2studio/models/batch.py
index 309b2079..58ba188d 100644
--- a/earth2studio/models/batch.py
+++ b/earth2studio/models/batch.py
@@ -267,12 +267,9 @@ def _compress_batch(
             If model's input_coords do not contain the batch dimension
         """
 
-        if (
-            next(iter(model.input_coords)) != "batch"
-            and next(iter(model.output_coords(model.input_coords))) != "batch"
-        ):
+        if next(iter(model.input_coords)) != "batch":
             raise ValueError(
-                "Model coordinate systems not compatible with batch processing"
+                "Model input coordinate systems not compatible with batch processing"
             )
 
         flatten_coords: CoordSystem
diff --git a/earth2studio/perturbation/lagged.py b/earth2studio/perturbation/lagged.py
index c7123e84..e87df325 100644
--- a/earth2studio/perturbation/lagged.py
+++ b/earth2studio/perturbation/lagged.py
@@ -89,14 +89,14 @@ def __call__(
                 f"Warning! The number of lags, {len(self.lags)}, does not match "
                 f"the number of ensemble members requested, {len(coords['ensemble'])}."
             )
-
+        y = torch.clone(x)
         for i, lag in enumerate(self.lags):
-            x[i] = fetch_data(
+            y[i] = fetch_data(
                 source=self.source,
                 time=coords["time"] + lag,
                 variable=coords["variable"],
                 lead_time=coords["lead_time"],
-                device=x.device,
+                device=y.device,
             )[0]
 
-        return x, coords
+        return y, coords


Test patch:
diff --git a/docs/userguide/developer/testing.md b/docs/userguide/developer/testing.md
index e69de29b..abab66a0 100644
--- a/docs/userguide/developer/testing.md
+++ b/docs/userguide/developer/testing.md
@@ -0,0 +1,114 @@
+# Testing Guide
+
+Testing is a core part of Earth2Studio development and we have strict requirements for
+unit testing any new feature.
+Presently we expect new functionality to have over 90% coverage within reason.
+The core tool used for testing in Earth2Studio is [PyTest](https://docs.pytest.org/en/8.2.x/)
+and all unit tests can be found inside the [test](https://github.com/NVIDIA/earth2studio/tree/main/test)
+folder.
+Ensure that developer dependencies are install as outline in the [developer overview](#developer_overview)
+section.
+
+## PyTest
+
+Earth2Studio has several configurations for testing listed below.
+Many unit tests, particularly ones for data sources, have [timeout limits](https://github.com/pytest-dev/pytest-timeout)
+and are allowed to fail using the [xfail](https://docs.pytest.org/en/6.2.x/skipping.html#xfail-mark-test-functions-as-expected-to-fail)
+decorator.
+It is normal for some tests to fail, particularly for machines with slower internet
+connections.
+Users should always look at the test summary to under to understand which tests produced
+an error but did not fail.
+
+### Quick Unit Test Suite
+
+This will run the quickest test suite which skips many tests that are longer running.
+
+```bash
+pytest  test/
+```
+
+### Standard Unit Test Suite
+
+To run the full unit test suite, add the `--slow` flag.
+Compared to the quick test suite, this one will run much longer
+
+```bash
+pytest --slow test/
+```
+
+### Complete Unit Test Suite
+
+The complete unit tests ran by the CI pipeline is ran with the `--ci-cache`.
+This option enables tests that load and test full model archiectures with their model
+package which can be quite expensive.
+These tests will look for model files inside a specific folder unique to the CI pipeline
+on Nvidia systems or whatever is set in the `EARTH2STUDIO_CACHE` environment variable.
+See [configuration](#configuration_userguide) section for details on this environment
+variable.
+
+```bash
+ pytest --slow --ci-cache test/
+```
+
+:::{warning}
+The full model cache is rather large (order of 10Gb) which will be downloaded running
+this command.
+The CI pipeline will already have all model files downloaded on the runner machine to
+skip this step.
+Additional information on this can be found below.
+:::
+
+#### CI Model Package Cache
+
+In the CI pipeline a NFS mounted folder is used as a pre-cached store of all
+pre-trained model checkpoint files (automodels) so each pipeline does not need to rerun
+the download.
+The cache folder is managed internally by the developers.
+This should reflect a local cache if a user was to fetch and load all possible models.
+The point of truth for generating this cache of model checkpoints **should** always be
+`test/models/test_auto_models.py`.
+This test all listed auto-models and create a `./cache` folder with the following
+command:
+
+```bash
+pytest --model-download -s test/models/test_auto_models.py
+
+ls -l ./cache
+```
+
+## Coverage
+
+Coverage is calculated using the [coverage.py](https://coverage.readthedocs.io/en/7.5.3/)
+package configured in the [pyproject.toml](https://github.com/NVIDIA/earth2studio/blob/main/pyproject.toml).
+Presently the CI pipeline will fail if unit test coverage drops below 90%.
+Running coverage with pytest just requires the following:
+
+```bash
+# Run tests
+coverage pytest --slow test/
+
+# Compile reports
+coverage combine
+```
+
+The CI pipeline uses the commands defined in the [Makefile](https://github.com/NVIDIA/earth2studio/blob/main/Makefile):
+
+```bash
+make pytest
+
+make coverage
+```
+
+## Contributing Tests
+
+We expect all contributors to also provide unit tests for their feature.
+Tests are always evolving based on new errors discovered and features added.
+Right now the best way to contribute unit tests for a new feature is to modify the tests
+for a feature of a similar type.
+E.g. for a new data source, copy and modify the tests for say GFS or ARCO.
+For additonal infomation about the tests, reach out with an issue and the developers can
+provide more information.
+
+If you are contributing a model, be sure to make the developers aware of this to
+properly integrate it into the CI pipeline.
diff --git a/test/conftest.py b/test/conftest.py
index e7d771d5..bbf36861 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -20,7 +20,7 @@
 import pytest
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def model_cache_context():
     class EnvContextManager:
         def __init__(self, **kwargs):
diff --git a/test/models/px/test_dlwp.py b/test/models/px/test_dlwp.py
index 4ca5a8d0..6de49491 100644
--- a/test/models/px/test_dlwp.py
+++ b/test/models/px/test_dlwp.py
@@ -238,16 +238,23 @@ def test_dlwp_exceptions(dc, dlwp_phoo_cs_transform, device):
         p(x, coords)
 
 
+@pytest.fixture(scope="module")
+def model(model_cache_context) -> DLWP:
+    # Test only on cuda device
+    with model_cache_context():
+        package = DLWP.load_default_package()
+        p = DLWP.load_model(package)
+        return p
+
+
 @pytest.mark.ci_cache
 @pytest.mark.timeout(360)
 @pytest.mark.parametrize("device", ["cuda:0"])
-def test_dlwp_package(device, model_cache_context):
+def test_dlwp_package(device, model):
     torch.cuda.empty_cache()
     time = np.array([np.datetime64("1993-04-05T00:00")])
-    with model_cache_context():
-        with torch.device(device):
-            package = DLWP.load_default_package()
-            p = DLWP.load_model(package).to(device)
+    # Test the cached model package DLWP
+    p = model.to(device)
 
     dc = p.input_coords.copy()
     del dc["batch"]
diff --git a/test/models/px/test_fcn.py b/test/models/px/test_fcn.py
index e8610ad2..a77de6e2 100644
--- a/test/models/px/test_fcn.py
+++ b/test/models/px/test_fcn.py
@@ -157,17 +157,23 @@ def test_fcn_exceptions(dc, device):
         p(x, coords)
 
 
+@pytest.fixture(scope="module")
+def model(model_cache_context) -> FCN:
+    # Test only on cuda device
+    with model_cache_context():
+        package = FCN.load_default_package()
+        p = FCN.load_model(package)
+        return p
+
+
 @pytest.mark.ci_cache
 @pytest.mark.timeout(15)
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
-def test_fcn_package(device, model_cache_context):
+def test_fcn_package(model, device):
     torch.cuda.empty_cache()
     time = np.array([np.datetime64("1993-04-05T00:00")])
     # Test the cached model package FCN
-    # Test only on cuda device
-    with model_cache_context():
-        package = FCN.load_default_package()
-        p = FCN.load_model(package).to(device)
+    p = model.to(device)
 
     dc = p.input_coords.copy()
     del dc["batch"]
diff --git a/test/models/px/test_sfno.py b/test/models/px/test_sfno.py
index 53144865..018c8e60 100644
--- a/test/models/px/test_sfno.py
+++ b/test/models/px/test_sfno.py
@@ -157,17 +157,23 @@ def test_sfno_exceptions(dc, device):
         p(x, coords)
 
 
+@pytest.fixture(scope="module")
+def model(model_cache_context) -> SFNO:
+    # Test only on cuda device
+    with model_cache_context():
+        package = SFNO.load_default_package()
+        p = SFNO.load_model(package)
+        return p
+
+
 @pytest.mark.ci_cache
 @pytest.mark.timeout(360)
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
-def test_sfno_package(device, model_cache_context):
+def test_sfno_package(device, model):
     torch.cuda.empty_cache()
     time = np.array([np.datetime64("1993-04-05T00:00")])
     # Test the cached model package SFNO
-    # Test only on cuda device
-    with model_cache_context():
-        package = SFNO.load_default_package()
-        p = SFNO.load_model(package).to(device)
+    p = model.to(device)
 
     # Create "domain coords"
     dc = {k: p.input_coords[k] for k in ["lat", "lon"]}
diff --git a/test/models/test_auto_models.py b/test/models/test_auto_models.py
index e77c383c..94979ce1 100644
--- a/test/models/test_auto_models.py
+++ b/test/models/test_auto_models.py
@@ -82,12 +82,23 @@ def cache_folder(tmp_path_factory):
 
 @pytest.mark.parametrize(
     "url,file",
-    [(None, "test.txt"), ("hf://NickGeneva/earth_ai", "README.md")],
+    [
+        (None, "test.txt"),
+        ("hf://NickGeneva/earth_ai", "README.md"),
+        (
+            "ngc://models/nvidia/modulus/sfno_73ch_small@0.1.0",
+            "sfno_73ch_small/metadata.json",
+        ),
+        ("s3://noaa-swpc-pds", "text/3-day-geomag-forecast.txt"),
+    ],
 )
 def test_package(url, file, cache_folder, model_cache_context):
     if url is None:
         url = "file://" / cache_folder
-    with model_cache_context(EARTH2STUDIO_CACHE=str(cache_folder.resolve())):
+    with model_cache_context(
+        EARTH2STUDIO_CACHE=str(cache_folder.resolve()),
+        EARTH2STUDIO_PACKAGE_TIMEOUT="30",
+    ):
         package = Package(str(url))
         file_path = package.resolve(file)
         assert Path(file_path).is_file()