🚀 (zarr) make save_to_disk write path fsspec-aware (#485)

[xroynard]

Summary

Makes the Zarr backend write path fsspec-aware so datasets can be written to any fsspec target (memory://, s3://, gs://, ...), not only the local filesystem. Scope is limited to sub-issue #485 (write side, zarr backend); the read path and the layer-1 metadata plumbing are tracked separately (#486, #487, #489).

Closes #485

Problem

In plaid/storage/zarr/writer.py::generate_datasetdict_to_disk, the write path was pinned to local storage in two places:

• Parallel worker used a hardcoded zarr.storage.LocalStore(split_root_path). When given a remote URL, this silently wrote the data to a literal local directory named after the URL (e.g. a memory: folder) instead of the intended remote target.
• The data folder was built with Path(output_folder) / "data" + mkdir(...). pathlib.Path collapses the :// separator (memory://root → memory:/root), corrupting remote targets, and mkdir is meaningless on object stores.

Changes

• _is_local_target(target) — discriminates local paths / file:// URLs from remote fsspec protocols via fsspec.core.url_to_fs. Missing backends surface fsspec's own install hint (e.g. Install s3fs to access S3).
• _join_target(base, *parts) — joins components while preserving the URL protocol separator for remote targets, keeping Path semantics for local ones.
• _open_split_group(target, mode) — thin wrapper over zarr.open_group, which already resolves an fsspec URL to a FsspecStore. Used by both the sequential create and the parallel-worker reopen, replacing the hardcoded LocalStore.
• The data/ folder is now joined URL-safely and mkdir is only invoked for local targets.

zarr.open_group already accepts fsspec URLs, and zarr pulls in fsspec (already imported at module top in the sibling reader.py), so no new dependency is added.

Tests

New tests/storage/test_zarr_fsspec_write.py (11 cases):

• unit coverage of _is_local_target (local/relative/file:///memory://, plus the missing-backend ImportError) and _join_target (local Path vs remote URL, trailing-slash handling);
• sequential end-to-end write to a memory:// target, asserting samples land at the intended remote location, the flattened global feature is readable back, and no literal memory: directory is created locally;
• _open_split_group create-then-reopen round-trip on memory:// (mirrors the sequential-create / parallel-worker-reopen handshake);
• local parallel (num_proc=2) write, guarding against a regression in the store-selection change.

Note: the parallel path is validated against a local target because memory:// is per-process — worker writes are not visible from the parent — so a remote num_proc>1 assertion would be a harness artifact, not a code check.

Existing storage tests (test_storage.py, test_zarr_init.py) pass unchanged. ruff check / ruff format --check clean.

Docs & changelog

• Docs: docs/source/tutorials/storage.md — added an admonition documenting that, with the zarr backend, output_folder accepts any fsspec URL (with an s3:// example and the backend-install caveat).
• Changelog: new entry under [Unreleased] / Added.

---

Checklist

• Typing enforced
• Documentation updated
• Changelog updated
• Tests and Example updates
• Coverage should be 100%

🔗 Related issues

Closes #485

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[xroynard]

How to test

Automated (the 11 new cases + regression):

cd src/plaid  # repo root
uv sync
uv run pytest tests/storage/test_zarr_fsspec_write.py -v
# and the existing storage suite, to confirm no regression:
uv run pytest tests/storage/test_storage.py tests/storage/test_zarr_init.py

What each part covers:

• helper units — _is_local_target (local / relative / file:// → local, memory:// → remote, and s3:// without s3fs raising the fsspec install hint) and _join_target (local Path vs remote URL, trailing-slash handling);
• sequential remote write — test_generate_datasetdict_to_disk_writes_to_memory_fs writes to memory:// and asserts the samples land at the intended remote location, the flattened global feature reads back, and no literal memory: directory is created locally (the bug the previous LocalStore had);
• create-then-reopen round-trip on memory:// (test_open_split_group_roundtrip_on_memory_fs) — mirrors the sequential-create / parallel-worker-reopen handshake;
• local parallel write (num_proc=2) — guards against a regression in the store-selection change.

Manual sanity check (no cloud creds needed), in a Python REPL:

import fsspec, zarr
from plaid.storage.zarr.writer import _is_local_target, _join_target

# discrimination
assert _is_local_target("/tmp/x") and _is_local_target("file:///tmp/x")
assert not _is_local_target("memory://root")

# URL join no longer collapses "://"
assert _join_target("memory://root", "data", "train") == "memory://root/data/train"

# end-to-end: write a DatasetDict to memory:// and read it back
# (build a small `generators`/`variable_schema` as in the sequential test, then)
# generate_datasetdict_to_disk(output_folder="memory://demo", ...)
group = zarr.open_group("memory://demo/data/train", mode="r")
print(sorted(group.group_keys()))   # sample_000000000, ...
import os; assert not os.path.exists("memory:")  # no local leak

Note: I don't assert the parallel path against memory:// — the in-memory FS is per-process, so worker writes aren't visible from the parent. num_proc>1 is therefore validated on a local target, and the remote reopen logic is covered separately by the round-trip test.

Testing a real cloud target (optional, outside CI): install the backend (uv pip install s3fs / gcsfs), set credentials, then pass output_folder="s3://your-bucket/prefix" to generate_datasetdict_to_disk. A missing backend surfaces fsspec's own install hint (e.g. Install s3fs to access S3).

[xroynard]

Hi @tmolcard — this lands the native remote-write direction we discussed in #478 and tracked in #479: save_to_disk can now write straight to an fsspec target (s3://, gs://, memory://, …) instead of the write→upload→delete dance your S3 pipeline currently does via the post-write callback. For now it's the zarr backend only (this PR is scoped to #485).

Since your read-from-S3 → build PLAID in memory → write-back-to-S3 pipeline is basically the motivating use case, your feedback would be really valuable if you have a moment. A no-creds sanity check on the branch:

uv run pytest tests/storage/test_zarr_fsspec_write.py -v

And the real thing, if you can point it at a bucket (uv pip install s3fs first):

save_to_disk(output_folder="s3://your-bucket/prefix", sample_constructor=..., ids=..., backend="zarr", num_proc=N)

Two things I'd genuinely like your read on:

1. Does a zarr remote target work for your pipeline, or is CGNS-on-S3 the one you actually need? (CGNS remote write is heavier — it usually needs a local buffer — so knowing whether zarr is acceptable for you helps prioritise.)
2. If zarr works, does writing straight to s3:// let you drop the callback-based upload-and-delete entirely, or do you still need per-sample hooks for another reason?

[casenave]

remove (issue #485)

[casenave]

Did your check that this works, or at least the local writes are still working as before ?