Improve performance of qml.data.load() when partially loading a dataset

[brownj85]

Context:
Link to shortcut story

Loading individual attributes of datasets took much longer than loading a whole dataset. This is because the fsspec library was mapping the HDF5 reads directly to HTTP requests, which only loaded a few KB each.

Description of the Change:
open_hdf5_s3() now opens the remote dataset in read-buffered mode, which reads data in 8MB chunks into a memory-mapped cache. This results in much fewer requests and faster loading.

Benefits:
Acceptable performance for partial loading of large datasets. The download throughput for partial loading is now comparable to downloading the whole dataset (about 15-20% less mb/s).

Possible Drawbacks:
None

Related GitHub Issues:

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (5f246a9) 99.64% compared to head (0d0e7bf) 99.63%.
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #4674      +/-   ##
==========================================
- Coverage   99.64%   99.63%   -0.01%     
==========================================
  Files         377      377              
  Lines       33999    33735     -264     
==========================================
- Hits        33878    33613     -265     
- Misses        121      122       +1     

Files	Coverage Δ
pennylane/data/base/hdf5.py	100.00% <100.00%> (ø)
pennylane/data/data_manager/__init__.py	98.52% <100.00%> (-0.03%)	⬇️

... and 42 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[timmysilv]

nice speed-up! to be clear, we used to only use (disk) caches when users requested a specific path, but now we always use one (in-memory), right? is there any loss of functionality by removing the option of specifying a cache path?

[anthayes92]

Thanks @brownj85, looks good overall! Just wondering how was the speedup verified?

[brownj85]

Thanks @brownj85, looks good overall! Just wondering how was the speedup verified

Just by my own testing. I'm assuming @timmysilv and @obliviateandsurrender tested it as well

[brownj85]

nice speed-up! to be clear, we used to only use (disk) caches when users requested a specific path, but now we always use one (in-memory), right? is there any loss of functionality by removing the option of specifying a cache path?

Pretty much - I used blockcache because I thought it did what mmap does, but it actually just stores the response data after the fact and doesn't do any read buffering. mmap is in-memory but it can commit blocks to a temporary file as needed. No loss of functionality, the cache_dir was just there to make the block cache work

[timmysilv]

sounds good. I didn't do any testing personally so I'm a bit curious as to what your own testing entailed. that said, I trust your judgement here and I'm happy with this change!

[DSGuala]

Tested with H2 and H2O. The performance improved significantly on downloading individual attributes 💪