Fsspec remfile with cache

src/nwb_benchmarks/benchmarks/network_tracking_remote_file_reading.py

@@ -7,10 +7,14 @@
     get_s3_url,
     network_activity_tracker,
     read_hdf5_fsspec_no_cache,
+    read_hdf5_fsspec_with_cache,
     read_hdf5_nwbfile_fsspec_no_cache,
+    read_hdf5_nwbfile_fsspec_with_cache,
     read_hdf5_nwbfile_remfile,
+    read_hdf5_nwbfile_remfile_with_cache,
     read_hdf5_nwbfile_ros3,
     read_hdf5_remfile,
+    read_hdf5_remfile_with_cache,
     read_hdf5_ros3,
 )
 
@@ -35,6 +39,16 @@ def track_network_activity_during_read(self, s3_url: str):
         return network_tracker.asv_network_statistics
 
 
+class FsspecWithCacheDirectFileReadBenchmark:
+    param_names = param_names
+    params = params
+
+    def track_network_activity_during_read(self, s3_url: str):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self.file, self.bytestream, self.tmpdir = read_hdf5_fsspec_with_cache(s3_url=s3_url)
+        return network_tracker.asv_network_statistics
+
+
 class RemfileDirectFileReadBenchmark:
     param_names = param_names
     params = params
@@ -45,6 +59,16 @@ def track_network_activity_during_read(self, s3_url: str):
         return network_tracker.asv_network_statistics
 
 
+class RemfileDirectFileReadBenchmarkWithCache:
+    param_names = param_names
+    params = params
+
+    def track_network_activity_during_read(self, s3_url: str):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self.file, self.bytestream, self.tmpdir = read_hdf5_remfile_with_cache(s3_url=s3_url)
+        return network_tracker.asv_network_statistics
+
+
 class Ros3DirectFileReadBenchmark:
     param_names = param_names
     params = params
@@ -66,6 +90,18 @@ def track_network_activity_during_read(self, s3_url: str):
         return network_tracker.asv_network_statistics
 
 
+class FsspecWithCacheNWBFileReadBenchmark:
+    param_names = param_names
+    params = params
+
+    def track_network_activity_during_read(self, s3_url: str):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_fsspec_with_cache(
+                s3_url=s3_url
+            )
+        return network_tracker.asv_network_statistics
+
+
 class RemfileNWBFileReadBenchmark:
     param_names = param_names
     params = params
@@ -76,6 +112,18 @@ def track_network_activity_during_read(self, s3_url: str):
         return network_tracker.asv_network_statistics
 
 
+class RemfileNWBFileReadBenchmarkWithCache:
+    param_names = param_names
+    params = params
+
+    def track_network_activity_during_read(self, s3_url: str):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_remfile_with_cache(
+                s3_url=s3_url
+            )
+        return network_tracker.asv_network_statistics
+
+
 class Ros3NWBFileReadBenchmark:
     param_names = param_names
     params = params

src/nwb_benchmarks/benchmarks/network_tracking_remote_slicing.py

@@ -7,9 +7,13 @@
     get_object_by_name,
     get_s3_url,
     network_activity_tracker,
+    read_hdf5_fsspec_with_cache,
     read_hdf5_nwbfile_fsspec_no_cache,
+    read_hdf5_nwbfile_fsspec_with_cache,
     read_hdf5_nwbfile_remfile,
+    read_hdf5_nwbfile_remfile_with_cache,
     read_hdf5_nwbfile_ros3,
+    read_hdf5_remfile_with_cache,
     robust_ros3_read,
 )
 
@@ -41,6 +45,23 @@ def track_network_activity_during_slice(self, s3_url: str, object_name: str, sli
         return network_tracker.asv_network_statistics
 
 
+class FsspecWithCacheContinuousSliceBenchmark:
+    param_names = param_names
+    params = params
+
+    def setup(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_fsspec_with_cache(
+            s3_url=s3_url
+        )
+        self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name=object_name)
+        self.data_to_slice = self.neurodata_object.data
+
+    def track_network_activity_during_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self._temp = self.data_to_slice[slice_range]
+        return network_tracker.asv_network_statistics
+
+
 class RemfileContinuousSliceBenchmark:
     param_names = param_names
     params = params
@@ -56,6 +77,23 @@ def track_network_activity_during_slice(self, s3_url: str, object_name: str, sli
         return network_tracker.asv_network_statistics
 
 
+class RemfileContinuousSliceBenchmarkWithCache:
+    param_names = param_names
+    params = params
+
+    def setup(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_remfile_with_cache(
+            s3_url=s3_url
+        )
+        self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name=object_name)
+        self.data_to_slice = self.neurodata_object.data
+
+    def track_network_activity_during_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        with network_activity_tracker(tshark_path=TSHARK_PATH) as network_tracker:
+            self._temp = self.data_to_slice[slice_range]
+        return network_tracker.asv_network_statistics
+
+
 class Ros3ContinuousSliceBenchmark:
     param_names = param_names
     params = params

src/nwb_benchmarks/benchmarks/time_remote_file_reading.py

@@ -3,10 +3,14 @@
 from nwb_benchmarks.core import (
     get_s3_url,
     read_hdf5_fsspec_no_cache,
+    read_hdf5_fsspec_with_cache,
     read_hdf5_nwbfile_fsspec_no_cache,
+    read_hdf5_nwbfile_fsspec_with_cache,
     read_hdf5_nwbfile_remfile,
+    read_hdf5_nwbfile_remfile_with_cache,
     read_hdf5_nwbfile_ros3,
     read_hdf5_remfile,
+    read_hdf5_remfile_with_cache,
     read_hdf5_ros3,
 )
 
@@ -38,13 +42,20 @@ class DirectFileReadBenchmark:
     def time_read_hdf5_fsspec_no_cache(self, s3_url: str):
         self.file, self.bytestream = read_hdf5_fsspec_no_cache(s3_url=s3_url)
 
+	def time_read_hdf5_fsspec_with_cache(self, s3_url: str):
+        self.file, self.bytestream, self.tmpdir = read_hdf5_fsspec_with_cache(s3_url=s3_url)
+
     def time_read_hdf5_remfile(self, s3_url: str):
         self.file, self.bytestream = read_hdf5_remfile(s3_url=s3_url)
 
+    def time_read_hdf5_remfile_with_cache(self, s3_url: str):
+        self.file, self.bytestream, self.tmpdir = read_hdf5_remfile_with_cache(s3_url=s3_url)
+
     def time_read_hdf5_ros3(self, s3_url: str):
         self.file, _ = read_hdf5_ros3(s3_url=s3_url, retry=False)
 
 
+
 class NWBFileReadBenchmark:
     """
     Time the read of the HDF5-backend files with `pynwb` using each streaming method.
@@ -60,8 +71,14 @@ class NWBFileReadBenchmark:
     def time_read_hdf5_nwbfile_fsspec_no_cache(self, s3_url: str):
         self.nwbfile, self.io, self.file, self.bytestream = read_hdf5_nwbfile_fsspec_no_cache(s3_url=s3_url)
 
+	def time_read_hdf5_nwbfile_fsspec_with_cache(self, s3_url: str):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_fsspec_with_cache(s3_url=s3_url)
+
     def time_read_hdf5_nwbfile_remfile(self, s3_url: str):
         self.nwbfile, self.io, self.file, self.bytestream = read_hdf5_nwbfile_remfile(s3_url=s3_url)
 
+    def time_read_hdf5_nwbfile_remfile_with_cache(self, s3_url: str):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_remfile_with_cache(s3_url=s3_url)
+
     def time_read_hdf5_nwbfile_ros3(self, s3_url: str):
         self.nwbfile, self.io, _ = read_hdf5_nwbfile_ros3(s3_url=s3_url, retry=False)

src/nwb_benchmarks/benchmarks/time_remote_slicing.py

@@ -6,9 +6,13 @@
 from nwb_benchmarks.core import (
     get_object_by_name,
     get_s3_url,
+    read_hdf5_fsspec_with_cache,
     read_hdf5_nwbfile_fsspec_no_cache,
+    read_hdf5_nwbfile_fsspec_with_cache,
     read_hdf5_nwbfile_remfile,
+    read_hdf5_nwbfile_remfile_with_cache,
     read_hdf5_nwbfile_ros3,
+    read_hdf5_remfile_with_cache,
 )
 
 # TODO: add the others
@@ -42,6 +46,24 @@ def time_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
         self._temp = self.data_to_slice[slice_range]
 
 
+class FsspecWithCacheContinuousSliceBenchmark:
+    rounds = 1
+    repeat = 3
+    param_names = param_names
+    params = params
+
+    def setup(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_fsspec_with_cache(
+            s3_url=s3_url
+        )
+        self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name=object_name)
+        self.data_to_slice = self.neurodata_object.data
+
+    def time_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        """Note: store as self._temp to avoid tracking garbage collection as well."""
+        self._temp = self.data_to_slice[slice_range]
+
+
 class RemfileContinuousSliceBenchmark:
     rounds = 1
     repeat = 3
@@ -58,6 +80,24 @@ def time_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
         self._temp = self.data_to_slice[slice_range]
 
 
+class RemfileContinuousSliceBenchmarkWithCache:
+    rounds = 1
+    repeat = 3
+    param_names = param_names
+    params = params
+
+    def setup(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        self.nwbfile, self.io, self.file, self.bytestream, self.tmpdir = read_hdf5_nwbfile_remfile_with_cache(
+            s3_url=s3_url
+        )
+        self.neurodata_object = get_object_by_name(nwbfile=self.nwbfile, object_name="ElectricalSeriesAp")
+        self.data_to_slice = self.neurodata_object.data
+
+    def time_slice(self, s3_url: str, object_name: str, slice_range: Tuple[slice]):
+        """Note: store as self._temp to avoid tracking garbage collection as well."""
+        self._temp = self.data_to_slice[slice_range]
+
+
 class Ros3ContinuousSliceBenchmark:
     rounds = 1
     repeat = 3

src/nwb_benchmarks/core/_streaming.py

@@ -1,3 +1,4 @@
+import tempfile
 import time
 import warnings
 from typing import Any, Callable, Tuple, Union
@@ -7,6 +8,7 @@
 import pynwb
 import remfile
 from fsspec.asyn import reset_lock
+from fsspec.implementations.cached import CachingFileSystem
 from fsspec.implementations.http import HTTPFile
 
 # Useful if running in verbose mode
@@ -26,6 +28,22 @@ def read_hdf5_fsspec_no_cache(
     file = h5py.File(name=byte_stream)
     return (file, byte_stream)
 
+def read_hdf5_fsspec_with_cache(
+    s3_url: str,
+) -> Tuple[h5py.File, HTTPFile]:
+    """Load the raw HDF5 file from an S3 URL using fsspec without a cache; does not formally read the NWB file."""
+    reset_lock()
+    fsspec.get_filesystem_class("https").clear_instance_cache()
+    filesystem = fsspec.filesystem("https")
+	tmpdir= tempfile.TemporaryDirectory()
+	fs = CachingFileSystem(
+    fs=fs,
+    cache_storage=tmpdir.name,  # Local folder for the cache
+		)
+    byte_stream = filesystem.open(path=s3_url, mode="rb")
+    file = h5py.File(name=byte_stream)
+    return (file, byte_stream, tmpdir)
+
 
 def read_hdf5_nwbfile_fsspec_no_cache(
     s3_url: str,
@@ -36,6 +54,20 @@ def read_hdf5_nwbfile_fsspec_no_cache(
     nwbfile = io.read()
     return (nwbfile, io, file, byte_stream)
 
+def read_hdf5_nwbfile_fsspec_with_cache(
+    s3_url: str,
+) -> Tuple[pynwb.NWBFile, pynwb.NWBHDF5IO, h5py.File, HTTPFile]:
+    """Read an HDF5 NWB file from an S3 URL using fsspec without a cache."""
+    (file, byte_stream) = read_hdf5_fsspec_no_cache(s3_url=s3_url)
+    io = pynwb.NWBHDF5IO(file=file, load_namespaces=True)
+    nwbfile = io.read()
+    tmpdir= tempfile.TemporaryDirectory()
+	fs = CachingFileSystem(
+    fs=fs,
+    cache_storage=tmpdir.name,  # Local folder for the cache
+		)
+    return (nwbfile, io, file, byte_stream, tmpdir)
+
 
 def read_hdf5_remfile(s3_url: str) -> Tuple[h5py.File, remfile.File]:
     """Load the raw HDF5 file from an S3 URL using remfile; does not formally read the NWB file."""
@@ -51,6 +83,21 @@ def read_hdf5_nwbfile_remfile(s3_url: str) -> Tuple[pynwb.NWBFile, pynwb.NWBHDF5
     nwbfile = io.read()
     return (nwbfile, io, file, byte_stream)
 
+def read_hdf5_remfile_with_cache(s3_url: str) -> Tuple[h5py.File, remfile.File]:
+    """Load the raw HDF5 file from an S3 URL using remfile; does not formally read the NWB file."""
+    tmpdir= tempfile.TemporaryDirectory()
+    disk_cache= remfile.DiskCache(tempdir.name)
+    byte_stream = remfile.File(url=s3_url,disk_cache=disk_cache)
+    file = h5py.File(name=byte_stream)
+    return (file, byte_stream, tmpdir)
+
+
+def read_hdf5_nwbfile_remfile_with_cache(s3_url: str) -> Tuple[pynwb.NWBFile, pynwb.NWBHDF5IO, h5py.File, remfile.File]:
+    """Read an HDF5 NWB file from an S3 URL using the ROS3 driver from h5py."""
+    (file, byte_stream, tmpdir) = read_hdf5_remfile_with_cache(s3_url=s3_url)
+    io = pynwb.NWBHDF5IO(file=file, load_namespaces=True)
+    nwbfile = io.read()
+    return (nwbfile, io, file, byte_stream, tmpdir)
 
 def robust_ros3_read(
     command: Callable,