fsspec-aware save_to_disk: stream samples to the target during writing (avoid full local materialization before remote push)

[xroynard]

Problem

The current remote flow is save_to_disk(local_folder) then push_to_hub(local_folder). This uploads only after the whole dataset has been fully materialized on local disk, so peak local disk usage = size of the entire dataset. At TB scale (the use case raised in #478) this is the blocker: the machine may simply not have room to stage the full dataset before a single byte is uploaded.

Note this isn't a "missing S3 door in push_to_hub" problem — it's a timing problem: the upload happens post-materialization instead of being interleaved with per-sample writing.

Proposed direction

Make save_to_disk fsspec-aware so each sample is written directly to the target (local or remote) as it is produced. Then:

• uploads become incremental (per sample),
• peak local disk drops to ~1 sample instead of the whole dataset,
• push_to_hub on a pre-written folder becomes just one special case rather than the only path to remote.

Scope

• In scope: file-based backends (cgns, zarr). They already write sample by sample in the writer loop (for sample in gen_func(): sample.save_to_dir(...)), so redirecting those writes through an fsspec AbstractFileSystem is a natural, low-cost change — consistent with @tmolcard's point that per-sample upload for cgns/zarr "isn't very complicated."
• Out of scope: hf_datasets backend. The Arrow table is materialized locally regardless, so streaming writes to a remote target don't apply there.

Relation to #478

The sample_callback hook added in #478 is effectively a workaround for this limitation (upload-then-delete each sample by hand because the library writes locally only). An fsspec-aware save_to_disk would make the hook unnecessary for this use case (it may retain value for other per-sample side effects). #478 can land as the minimal immediate fix; this issue tracks the more direct solution.

Open questions

• fsspec handle management under num_proc > 1 (independent handles per worker, process-safe).
• cgns/zarr writes that need a seekable handle may still need a small per-sample local temp buffer, uploaded then discarded — still bounded peak disk (~1 sample), unlike today.
• optional deps (s3fs/gcsfs/…) gated behind a [remote] extra with a clear error when the protocol handler is missing.

[xroynard]

Storage API: function inventory across layers

Before discussing the write-side gap, here is the full map of the save/load surface, since the same operation is renamed at each layer.

The storage module has 3 layers over 3 backends (cgns, hf_datasets, zarr, registered in storage/registry.py), plus a transverse metadata layer.

Layer 1 — public API (storage/reader.py + storage/writer.py, re-exported by plaid top-level). Only 5 verbs, backend-agnostic (dispatch via infos.storage_backend):

Function	Dir.	Target
save_to_disk	write	local disk (default backend="hf_datasets")
push_to_hub	write	HF Hub (uploads an already-written folder)
init_from_disk	read	local disk (auto-detects backend)
download_from_hub	read	HF Hub → disk (materialized)
init_streaming_from_hub	read	HF Hub (lazy/iterable)

Layer 2 — backend Protocol (storage/backend_api.py), names differ from layer 1: generate_to_disk ↔ save_to_disk, push_local_to_hub ↔ push_to_hub, init_from_disk, download_from_hub, init_datasetdict_streaming_from_hub ↔ init_streaming_from_hub, configure_dataset_card.

Layer 3 — per-backend impls (storage/{cgns,hf_datasets,zarr}/), renamed again: generate_datasetdict_to_disk, push_local_datasetdict_to_hub, init_datasetdict_from_disk, download_datasetdict_from_hub, init_datasetdict_streaming_from_hub.

Transverse — metadata (storage/common/), handling infos / problem_definitions / constants / schemas independently of the data backend: load_*_from_disk / load_*_from_hub / save_*_to_disk / push_*_to_hub.

So the apparent proliferation of names is just per-layer renaming + metadata handled separately — there are really only 5 verbs.

[xroynard]

Why incremental/remote writing doesn't exist yet: the read/write asymmetry

The key structural point behind this issue: reading has 3 doors, writing has only 2.

	local	hub (materialized)	hub (lazy/incremental)
read	init_from_disk	download_from_hub	init_streaming_from_hub ✅
write	save_to_disk	push_to_hub	— (missing) ❌

There is no write-side equivalent of init_streaming_from_hub. Concretely:

• save_to_disk writes everything locally first — Path(output_folder), save_metadata_to_disk, save_infos_to_disk, then delegates data to backend.generate_to_disk.
• push_to_hub only uploads an already-complete folder.

Hence peak local disk = the entire dataset before a single byte is uploaded. This is a timing problem (upload post-materialization instead of interleaved with per-sample writing), not a missing S3 door in push_to_hub. The sample_callback hook is currently the manual workaround for exactly this missing door.

[xroynard]

Feasibility per backend

The fsspec-aware write path is not uniformly reachable across the 3 backends.

zarr — closest to ready. The reader already streams via fsspec: mapper = fsspec.get_mapper(base_url + path) then zarr.open(mapper). The writer, however, writes to zarr.storage.LocalStore only, and iterates write_sample per sample — so per-sample write already exists, it's just pinned to a local store. Two concrete blockers on the read path today, both worth fixing while we're here:

• the streaming base URL is hardcoded to https://huggingface.co/datasets/{repo_id}/resolve/main/data/..., so it can't point at S3/GCS/local even though the underlying fsspec.get_mapper could handle any protocol;
• the store passed to the writer is a concrete LocalStore rather than an fsspec-derived store.
  Making the writer accept an fsspec URL/mapper (any of file://, s3://, gs://, http://) is the natural first target.

hf_datasets — the "hub" here is exclusively HuggingFace, not generic fsspec: common/reader.py uses hf_hub_download / snapshot_download with repo_type="dataset" and a repo_id. Incremental push would go through the datasets library / Hub API, a different mechanism than an fsspec mapper. This is the backend where "fsspec-aware" is least directly applicable.

cgns — file-per-sample layout, no streaming reader today; would need a store abstraction built from scratch.

Suggested scope: land fsspec-aware save + load on the zarr backend first (unblock any fsspec protocol on both sides, remove the hardcoded HF URL), treat hf_datasets and cgns as follow-ups.

[xroynard]

Not redundant with init_streaming_from_hub

To close the "isn't this already covered?" question: no.

• init_streaming_from_hub is read-side, HF-only — it lazily iterates an existing, already-published dataset. Its base URL is hardcoded to huggingface.co/.../resolve/main/....
• What this issue asks for is write-side, fsspec-generic — incrementally producing a dataset to an arbitrary target (file://, s3://, gs://, HF) without full local materialization.

They sit on opposite ends of the table above (lazy-read vs incremental-write) and target different destinations (HF-only vs any fsspec protocol). Building the fsspec write path also lets us generalize init_streaming_from_hub's hardcoded URL into a proper fsspec read path — so the two are complementary, not overlapping.

[xroynard]

Mirroring the fsspec-aware plan onto the read side (init_from_disk)

The comments above cover the write-side gap and why this isn't redundant with init_streaming_from_hub. This one is strictly about the read counterpart: can init_from_disk become fsspec-aware the same way, and how hard is each piece? Ordering the work by increasing difficulty, verified against main:

① Zarr data — easy (do first).
zarr/reader.py::init_datasetdict_from_disk hardcodes the store:

zarr.open(zarr.storage.LocalStore(local_path / sn), mode="r")

zarr.open() takes an fsspec URL directly in read mode (checked: zarr.open("memory://…", mode="r") round-trips). So this is just selecting the store from the path/storage_options instead of pinning LocalStore. ⚠️ Trap worth noting: LocalStore("s3://…") does not raise — it treats the string as a literal local path and reads the wrong place. So it's a correctness fix, not an exception to guard.

② Split discovery — easy, but needs the same treatment.
Both backends discover splits with local_path.iterdir() (is_dir() filtering). That call is local-only; on a remote target it has to go through fs.ls(...) on the fsspec filesystem derived from the URL. Small, but it's a second hardcoded-local spot in the same function, easy to miss.

③ Constants (common/reader.py::load_constants_from_disk) — hard.
Backed by np.memmap over data.mmap. np.memmap needs a real seekable local file — passing an fsspec handle fails with UnsupportedOperation: fileno (checked). _materialize_memmaps already exists to copy memmaps into plain arrays, but the source still has to be a local file. So remote constants can't be read in place: the bytes must land on local disk first.

④ CGNS data — hardest.
cgns/reader.py::init_datasetdict_from_disk builds CGNSDataset(path), whose __getitem__ calls Sample.load_from_dir → CGM.load(str(meshes_dir / "mesh_*.cgns")) (pyCGNS reads the CGNS/HDF5 file through a native C file handle — no fsspec object accepted). No streaming reader exists here either. Two options, both non-trivial: (a) materialize remote sample dirs into a local temp dir then reuse the local reader, or (b) rewrite Sample.load around fsspec. This is a design choice, not a config flag.

Suggested read-side order: ① + ② on zarr together (small, high value, testable on memory://), then ③ constants and ④ cgns as separate follow-ups since both require staging to local disk regardless. This mirrors the write-side "zarr first" scoping already proposed above, so read and write stay symmetric backend by backend.