speed up loading, and speed up tests

nowcasting_dataset/dataset/batch.py

@@ -5,6 +5,7 @@
 import os
 from pathlib import Path
 from typing import Optional, Union
+from concurrent import futures
 
 import xarray as xr
 from pydantic import BaseModel, Field
@@ -130,26 +131,44 @@ def save_netcdf(self, batch_i: int, path: Path):
             path: the path where it will be saved. This can be local or in the cloud.
 
         """
-        for data_source in self.data_sources:
-            if data_source is not None:
-                data_source.save_netcdf(batch_i=batch_i, path=path)
+
+        with futures.ThreadPoolExecutor() as executor:
+            # Submit tasks to the executor.
+            for data_source in self.data_sources:
+                if data_source is not None:
+                    _ = executor.submit(
+                        data_source.save_netcdf,
+                        batch_i=batch_i,
+                        path=path,
+                    )
 
     @staticmethod
     def load_netcdf(local_netcdf_path: Union[Path, str], batch_idx: int):
         """Load batch from netcdf file"""
         data_sources_names = Example.__fields__.keys()
 
-        # collect data sources
+        # set up futures executor
         batch_dict = {}
-        for data_source_name in data_sources_names:
-
-            local_netcdf_filename = os.path.join(
-                local_netcdf_path, data_source_name, f"{batch_idx}.nc"
-            )
-            if os.path.exists(local_netcdf_filename):
-                xr_dataset = xr.load_dataset(local_netcdf_filename)
-            else:
-                xr_dataset = None
+        with futures.ThreadPoolExecutor() as executor:
+            future_examples_per_source = []
+
+            # loop over data sources
+            for data_source_name in data_sources_names:
+
+                local_netcdf_filename = os.path.join(
+                    local_netcdf_path, data_source_name, f"{batch_idx}.nc"
+                )
+
+                # submit task
+                future_examples = executor.submit(
+                    xr.load_dataset,
+                    filename_or_obj=local_netcdf_filename,
+                )
+                future_examples_per_source.append([data_source_name, future_examples])
+
+        # Collect results from each thread.
+        for data_source_name, future_examples in future_examples_per_source:
+            xr_dataset = future_examples.result()
 
             batch_dict[data_source_name] = xr_dataset
 

tests/dataset/test_batch.py

@@ -10,21 +10,30 @@
 
 def test_model():
 
-    _ = Batch.fake()
+    con = Configuration()
+    con.process.batch_size = 4
+
+    _ = Batch.fake(configuration=con)
 
 
 def test_model_save_to_netcdf():
 
+    con = Configuration()
+    con.process.batch_size = 4
+
     with tempfile.TemporaryDirectory() as dirpath:
-        Batch.fake().save_netcdf(path=dirpath, batch_i=0)
+        Batch.fake(configuration=con).save_netcdf(path=dirpath, batch_i=0)
 
         assert os.path.exists(f"{dirpath}/satellite/0.nc")
 
 
 def test_model_load_from_netcdf():
 
+    con = Configuration()
+    con.process.batch_size = 4
+
     with tempfile.TemporaryDirectory() as dirpath:
-        Batch.fake().save_netcdf(path=dirpath, batch_i=0)
+        Batch.fake(configuration=con).save_netcdf(path=dirpath, batch_i=0)
 
         batch = Batch.load_netcdf(batch_idx=0, local_netcdf_path=dirpath)
 
@@ -33,7 +42,10 @@ def test_model_load_from_netcdf():
 
 def test_batch_to_batch_ml():
 
-    _ = BatchML.from_batch(batch=Batch.fake())
+    con = Configuration()
+    con.process.batch_size = 4
+
+    _ = BatchML.from_batch(batch=Batch.fake(configuration=con))
 
 
 def test_fake_dataset():