Merge pull request #175 from NeurodataWithoutBorders/lindi_download

Author: rly

src/nwb_benchmarks/benchmarks/network_tracking_create_json_from_remote.py

@@ -0,0 +1,38 @@
+import os
+
+from asv_runner.benchmarks.mark import skip_benchmark
+
+from nwb_benchmarks import TSHARK_PATH
+from nwb_benchmarks.core import (
+    BaseBenchmark,
+    create_lindi_reference_file_system,
+    network_activity_tracker,
+)
+
+from .params_remote_file_reading import hdf5_params
+
+
+class LindiCreateJSONFromRemoteFileBenchmark(BaseBenchmark):
+    """
+    Track the network activity during the creation of a LINDI JSON file for a remote NWB HDF5 file using lindi.
+    """
+
+    params = hdf5_params
+
+    def setup(self, params: dict[str, str]):
+        https_url = params["https_url"]
+        self.lindi_file = os.path.basename(https_url) + ".nwb.lindi.json"
+        self.teardown(params)
+
+    def teardown(self, params: dict[str, str]):
+        if os.path.exists(self.lindi_file):
+            os.remove(self.lindi_file)
+
+    # TODO This benchmark takes a long time to index all of the chunks for these files! Do not run until ready
+    @skip_benchmark
+    def track_network_create_lindi_json(self, params: dict[str, str]):
+        """Read a remote HDF5 file to create a LINDI JSON file."""
+        https_url = params["https_url"]
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            create_lindi_reference_file_system(https_url=https_url, outfile_path=self.lindi_file)
+        return network_tracker.asv_network_statistics

src/nwb_benchmarks/benchmarks/network_tracking_remote_file_reading.py

@@ -5,16 +5,14 @@
 network activity tracker.
 """
 
-import os
 import shutil
 
-from asv_runner.benchmarks.mark import skip_benchmark, skip_benchmark_if
+from asv_runner.benchmarks.mark import skip_benchmark_if
 
 from nwb_benchmarks import TSHARK_PATH
 from nwb_benchmarks.core import (
     BaseBenchmark,
-    create_lindi_reference_file_system,
-    download_file,
+    download_asset_if_not_exists,
     network_activity_tracker,
     read_hdf5_h5py_fsspec_https_no_cache,
     read_hdf5_h5py_fsspec_https_with_cache,
@@ -39,7 +37,6 @@
 
 from .params_remote_file_reading import (
     hdf5_params,
-    lindi_hdf5_params,
     lindi_remote_rfs_params,
     zarr_params,
 )
@@ -339,38 +336,13 @@ def track_network_read_hdf5_pynwb_remfile_preloaded_with_cache(self, params: dic
         return network_tracker.asv_network_statistics
 
 
-class LindiCreateLocalJSONFileBenchmark(BaseBenchmark):
-    """
-    Track the network activity during read of remote HDF5 files and the creation of a LINDI JSON file using lindi.
-    """
-
-    params = lindi_hdf5_params
-
-    def setup(self, params: dict[str, str]):
-        https_url = params["https_url"]
-        self.lindi_file = os.path.basename(https_url) + ".nwb.lindi.json"
-        self.teardown(params)
-
-    def teardown(self, params: dict[str, str]):
-        if os.path.exists(self.lindi_file):
-            os.remove(self.lindi_file)
-
-    # TODO This benchmark takes a long time to index all of the chunks for these files! Do not run until ready
-    @skip_benchmark
-    def track_network_read_create_lindi_json(self, params: dict[str, str]):
-        """Read a remote HDF5 file to create a LINDI JSON file."""
-        https_url = params["https_url"]
-        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
-            create_lindi_reference_file_system(https_url=https_url, outfile_path=self.lindi_file)
-        return network_tracker.asv_network_statistics
-
-
 class LindiLocalJSONFileReadBenchmark(BaseBenchmark):
     """
     Track the network activity during read of remote HDF5 files by reading the local LINDI JSON files with lindi and
     h5py or pynwb.
 
-    This downloads the already created remote LINDI JSON files during setup.
+    This downloads the remote LINDI JSON file during setup if it does not already exist in the persistent download
+    directory.
 
     Note: in all cases, store the in-memory objects to be consistent with timing benchmarks.
     """
@@ -380,17 +352,13 @@ class LindiLocalJSONFileReadBenchmark(BaseBenchmark):
     def setup(self, params: dict[str, str]):
         """Download the LINDI JSON file."""
         https_url = params["https_url"]
-        self.lindi_file = os.path.basename(https_url) + ".lindi.json"
-        self.teardown(params)
-        download_file(url=https_url, local_path=self.lindi_file)
+        self.lindi_file = download_asset_if_not_exists(https_url=https_url)
 
     def teardown(self, params: dict[str, str]):
         if hasattr(self, "io"):
             self.io.close()
         if hasattr(self, "client"):
             self.client.close()
-        if os.path.exists(self.lindi_file):
-            os.remove(self.lindi_file)
 
     @skip_benchmark_if(TSHARK_PATH is None)
     def track_network_read_lindi_h5py(self, params: dict[str, str]):

src/nwb_benchmarks/benchmarks/network_tracking_remote_slicing.py

@@ -14,13 +14,13 @@
 from nwb_benchmarks import TSHARK_PATH
 from nwb_benchmarks.core import (
     BaseBenchmark,
+    download_read_hdf5_pynwb_lindi,
     get_object_by_name,
     network_activity_tracker,
     read_hdf5_pynwb_fsspec_https_no_cache,
     read_hdf5_pynwb_fsspec_https_with_cache,
     read_hdf5_pynwb_fsspec_s3_no_cache,
     read_hdf5_pynwb_fsspec_s3_with_cache,
-    read_hdf5_pynwb_lindi,
     read_hdf5_pynwb_remfile_no_cache,
     read_hdf5_pynwb_remfile_with_cache,
     read_hdf5_pynwb_ros3,
@@ -336,21 +336,41 @@ class LindiLocalJSONContinuousSliceBenchmark(ContinuousSliceBenchmark):
     Time the read of a continuous data slice from remote HDF5 NWB files by reading the local LINDI JSON files with
     lindi and pynwb.
 
-    This downloads the already created remote LINDI JSON files during setup.
+    This downloads the remote LINDI JSON file during setup if it does not already exist in the persistent download
+    directory.
+    """
+
+    params = lindi_remote_rfs_params
+
+    def setup(self, params: dict[str, str | Tuple[slice]]):
+        https_url = params["https_url"]
+        object_name = params["object_name"]
+
+        self.nwbfile, self.io, self.client = download_read_hdf5_pynwb_lindi(https_url=https_url)
+        self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name=object_name)
+        self.data_to_slice = self.neurodata_object.data
+
+
+class LindiLocalJSONPreloadedContinuousSliceBenchmark(ContinuousSliceBenchmark):
+    """
+    Time the read of a continuous data slice from remote HDF5 NWB files by reading the local LINDI JSON files with
+    lindi and pynwb after preloading the data into any caches.
 
-    This should be about the same as reading a data slice from an NWB file that is instantiated with a remote LINDI JSON
-    file because, in that case, the first thing that LINDI does is download the remote file to a temporary directory.
+    This downloads the remote LINDI JSON file during setup if it does not already exist in the persistent download
+    directory.
     """
 
     params = lindi_remote_rfs_params
 
     def setup(self, params: dict[str, str | Tuple[slice]]):
         https_url = params["https_url"]
         object_name = params["object_name"]
+        slice_range = params["slice_range"]
 
-        self.nwbfile, self.io, self.client = read_hdf5_pynwb_lindi(rfs=https_url)
+        self.nwbfile, self.io, self.client = download_read_hdf5_pynwb_lindi(https_url=https_url)
         self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name=object_name)
         self.data_to_slice = self.neurodata_object.data
+        self._temp = self.data_to_slice[slice_range]
 
 
 class ZarrPyNWBS3ContinuousSliceBenchmark(ContinuousSliceBenchmark):

src/nwb_benchmarks/benchmarks/params_remote_file_reading.py

@@ -36,10 +36,6 @@
     # ),
 )
 
-# Parameters for LINDI when HDF5 files are remote without using an existing LINDI JSON reference file system on
-# the remote server (i.e., we create the LINDI JSON file for these in these tests)
-lindi_hdf5_params = hdf5_params
-
 # Parameters for LINDI pointing to a remote LINDI reference file system JSON file. I.e., here we do not
 # to create the JSON but can load it directly from the remote store
 lindi_remote_rfs_params = (
@@ -48,20 +44,23 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
     ),
     dict(
         name="OphysTestCase",
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
     ),
     dict(
         name="IcephysTestCase",
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
     ),
     # dict(

src/nwb_benchmarks/benchmarks/params_remote_slicing.py

@@ -154,6 +154,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="ElectricalSeries",
         slice_range=(slice(0, 262_144), slice(0, 384)),  # 12 chunks
@@ -163,6 +164,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="ElectricalSeries",
         slice_range=(slice(0, 262_144 * 2), slice(0, 384)),  # 24 chunks
@@ -172,6 +174,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="ElectricalSeries",
         slice_range=(slice(0, 262_144 * 3), slice(0, 384)),  # 36 chunks
@@ -181,6 +184,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="ElectricalSeries",
         slice_range=(slice(0, 262_144 * 4), slice(0, 384)),  # 48 chunks
@@ -190,6 +194,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-npI3/sub-npI3_behavior+ecephys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="ElectricalSeries",
         slice_range=(slice(0, 262_144 * 5), slice(0, 384)),  # 60 chunks
@@ -198,7 +203,8 @@
         name="OphysTestCase1",
         https_url=get_https_url(
             dandiset_id="213889",
-            dandi_path="sub-R6_ses-20200206T210000_behavior+ophys/sub-R6_ses-20200206T210000_behavior+ophys.lindi.json",
+            dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="TwoPhotonSeries",
         slice_range=(slice(0, 20), slice(0, 796), slice(0, 512)),  # 1 chunk
@@ -207,7 +213,8 @@
         name="OphysTestCase2",
         https_url=get_https_url(
             dandiset_id="213889",
-            dandi_path="sub-R6_ses-20200206T210000_behavior+ophys/sub-R6_ses-20200206T210000_behavior+ophys.lindi.json",
+            dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="TwoPhotonSeries",
         slice_range=(slice(0, 20 * 2), slice(0, 796), slice(0, 512)),  # 2 chunks
@@ -216,7 +223,8 @@
         name="OphysTestCase3",
         https_url=get_https_url(
             dandiset_id="213889",
-            dandi_path="sub-R6_ses-20200206T210000_behavior+ophys/sub-R6_ses-20200206T210000_behavior+ophys.lindi.json",
+            dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="TwoPhotonSeries",
         slice_range=(slice(0, 20 * 3), slice(0, 796), slice(0, 512)),  # 3 chunks
@@ -225,7 +233,8 @@
         name="OphysTestCase4",
         https_url=get_https_url(
             dandiset_id="213889",
-            dandi_path="sub-R6_ses-20200206T210000_behavior+ophys/sub-R6_ses-20200206T210000_behavior+ophys.lindi.json",
+            dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="TwoPhotonSeries",
         slice_range=(slice(0, 20 * 4), slice(0, 796), slice(0, 512)),  # 4 chunks
@@ -234,7 +243,8 @@
         name="OphysTestCase5",
         https_url=get_https_url(
             dandiset_id="213889",
-            dandi_path="sub-R6_ses-20200206T210000_behavior+ophys/sub-R6_ses-20200206T210000_behavior+ophys.lindi.json",
+            dandi_path="sub-R6/sub-R6_behavior+ophys.nwb.lindi.json",
+            follow_redirects=False,
         ),
         object_name="TwoPhotonSeries",
         slice_range=(slice(0, 20 * 5), slice(0, 796), slice(0, 512)),  # 5 chunks
@@ -244,6 +254,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
         object_name="data_00002_AD0",
         slice_range=(slice(0, 81_920),),  # 10 chunks
@@ -253,6 +264,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
         object_name="data_00002_AD0",
         slice_range=(slice(0, 81_920 * 2),),  # 20 chunks
@@ -262,6 +274,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
         object_name="data_00002_AD0",
         slice_range=(slice(0, 81_920 * 3),),  # 30 chunks
@@ -271,6 +284,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
         object_name="data_00002_AD0",
         slice_range=(slice(0, 81_920 * 4),),  # 40 chunks
@@ -280,6 +294,7 @@
         https_url=get_https_url(
             dandiset_id="213889",
             dandi_path="sub-1214579789_ses-1214621812_icephys/sub-1214579789_ses-1214621812_icephys.lindi.json",
+            follow_redirects=False,
         ),
         object_name="data_00002_AD0",
         slice_range=(slice(0, 81_920 * 5),),  # 50 chunks