Reader for Hugging Face's SafeTensor format

[TomNicholas]

We should be able to write a new VirtualiZarr reader for the Hugging Face SafeTensors format.

Hugging Face safetensors is an interesting example - it's uncompressed so basically just like reading netCDF3, having no internal chunking. But it also puts all the metadata at the start of the file, making it a bit like cloud-optimized HDF5. See also huggingface/safetensors#527 (comment)

Originally posted by @TomNicholas in #218

cc the Pangeo ML people: @weiji14 @negin513 @maxrjones

[TomNicholas]

The format specification seems very straightforward.

If it really is that simple then I think this reader could potentially be implemented without even using the safetensors library, instead just using fsspec and parsing the bytes ourselves. Our unit tests should still probably check correctness against the safetensors library itself though.

[TomNicholas]

There is an interesting issue on safetensors about "multi-part uploads". Apparently it's not officially supported but nevertheless widespread. This suggests desire for the model weights to be chunked and/or version-controlled, which the virtual zarr approach could obviously help with.

[TomNicholas]

Creating persisted virtual references looks like:
(data engineer code)

# reading into in-memory ChunkManifest objects
vds = open_virtual_dataset(<huggingface/safetensors/file.st>, reader='safetensors')

# serializing references to some kind of persistent storage format
vds.virtualize.to_kerchunk("refs.json")
# vds.virtualize.to_icechunk(icechunkstore)

Reading data from persisted virtual references:
(Data user code)

# icechunk

# open the icechunk repo from url...
# get the store object from the repo object
ds = xr.open_zarr(repo.store)

[TomNicholas]

@nenb a rough plan:

1. Write an virtualizarr reader for safetensors format
   • Resist temptation to support wide range of features
2. Dtype support?
3. Download some example safetensors and try it out by creating and serializing the virtual references to icechunk
4. Look at a bunch of real safetensors data and see if any has blockers of features that won't or can't be implemented by virtualizarr / your reader,
5. Do some performance comparison to show the value of reading (from object storage)

How to write a reader:

• Given a path to safetensors file
• remote files would require some way to scan through and get bytes (e.g. fsspec), but you should prototype locally
• Scan through it to find metadata (inc. chunk locations byte ranges, dtypes, array shapes, variable names)
• Create some virtualizarr objects with this info in them:
  • One ManifestArray per array in the file
    • Metadata goes in the zarr.core.metadata.ArrayV3Metadata class
    • Chunk references go in the ChunkManifest class

[TomNicholas]

Similarity to netCDF3:

• IIUC safetensors are not chunked or compressed
• which means that it's analogous to netCDF3
• So you have to choose an arbitrary chunking scheme up front (at virtualizarr read time) to go in your manifest
• But luckily zarr can read partial bits of uncompressed chunks (see v3 stores: Implement efficient get/set_partial_values zarr-python#1106) which means you probably just want to choose the whole array in each safetensors file to be one big chunk