From 702af6a05a2d10e8e2847aa182773aeb85278564 Mon Sep 17 00:00:00 2001 From: Xavier Roynard Date: Mon, 3 Aug 2026 13:12:39 +0000 Subject: [PATCH 1/2] feat(zarr): make save_to_disk write path fsspec-aware (#485) --- src/plaid/storage/zarr/writer.py | 87 +++++++++- tests/storage/test_zarr_fsspec_write.py | 210 ++++++++++++++++++++++++ 2 files changed, 290 insertions(+), 7 deletions(-) create mode 100644 tests/storage/test_zarr_fsspec_write.py diff --git a/src/plaid/storage/zarr/writer.py b/src/plaid/storage/zarr/writer.py index 004dc5ac..b85dfa59 100644 --- a/src/plaid/storage/zarr/writer.py +++ b/src/plaid/storage/zarr/writer.py @@ -17,6 +17,7 @@ from pathlib import Path from typing import Any, Callable, Generator, Optional, Union +import fsspec import numpy as np import yaml import zarr @@ -30,6 +31,69 @@ from ...infos import Infos +def _is_local_target(target: Union[str, Path]) -> bool: + """Return whether ``target`` resolves to the local filesystem. + + A plain path (absolute or relative) or a ``file://`` URL is local; any other + fsspec protocol (``memory://``, ``s3://``, ``gs://``, ``http://``, ...) is + remote. + + Args: + target (Union[str, Path]): Output folder path or fsspec URL. + + Returns: + bool: ``True`` if the target lives on the local filesystem. + + Raises: + ImportError: If ``target`` uses a protocol whose fsspec handler is not + installed (e.g. ``s3://`` without ``s3fs``). The original fsspec + message is preserved so the user knows which extra to install. + """ + fs, _ = fsspec.core.url_to_fs(str(target)) + protocols = fs.protocol + if isinstance(protocols, str): + protocols = (protocols,) + return any(p in ("file", "local") for p in protocols) + + +def _join_target(base: Union[str, Path], *parts: str) -> Union[Path, str]: + """Join path components, preserving fsspec URLs. + + ``pathlib.Path`` collapses the ``//`` in a URL (``memory://root`` becomes + ``memory:/root``), silently corrupting the target. This helper keeps local + targets as ``Path`` (unchanged behavior) and joins remote URLs with plain + ``/`` so the protocol separator survives. + + Args: + base (Union[str, Path]): Base path or fsspec URL. + *parts (str): Additional path components to append. + + Returns: + Union[Path, str]: A ``Path`` for local targets, a ``str`` URL otherwise. + """ + if _is_local_target(base): + return Path(base).joinpath(*parts) + return "/".join([str(base).rstrip("/"), *parts]) + + +def _open_split_group(split_target: str, mode: str) -> Any: + """Open (or create) the Zarr group for a split, local or remote. + + ``zarr.open_group`` natively resolves an fsspec URL to a ``FsspecStore``, so + the same call works for local paths and remote URLs. This wrapper exists so + the parallel worker (which runs in a separate process and cannot share an + open handle) can reopen the exact same store from a plain string. + + Args: + split_target (str): Local path or fsspec URL of the split root. + mode (str): Zarr open mode (``"w"``, ``"a"``, ...). + + Returns: + Any: An open Zarr group. + """ + return zarr.open_group(split_target, mode=mode) + + def _auto_chunks(shape: tuple[int, ...], target_n: int) -> tuple[int, ...]: """Computes automatic chunk sizes for Zarr arrays based on shape and target size. @@ -134,9 +198,12 @@ def _zarr_worker_batch_job(args) -> int: # pragma: no cover """ split_root_path, gen_func, var_features_keys, batch, start_index = args - # split_root = zarr.open_group(split_root_path, mode="a") - store = zarr.storage.LocalStore(split_root_path) - split_root = zarr.group(store=store) + # Reopen the split group from a plain string so it works for both local + # paths and fsspec URLs (memory://, s3://, ...). ``zarr.open_group`` builds + # a ``LocalStore`` or ``FsspecStore`` from the target automatically; the + # previous hardcoded ``zarr.storage.LocalStore`` silently wrote a remote URL + # to a literal local directory instead of the intended remote target. + split_root = _open_split_group(split_root_path, mode="a") sample_counter = start_index written = 0 @@ -185,15 +252,21 @@ def generate_datasetdict_to_disk( None: This function does not return a value; it writes the dataset directly to disk. """ - output_folder = Path(output_folder) / "data" - output_folder.mkdir(exist_ok=True, parents=True) + # Build the ``data`` subfolder in a way that works for local paths and + # fsspec URLs alike. For local targets we keep the original ``Path`` + mkdir + # behavior; for remote targets we join with ``/`` (``Path`` would collapse + # the ``://`` separator) and skip the local ``mkdir`` — fsspec stores create + # keys lazily on write and have no notion of an empty directory. + data_target = _join_target(output_folder, "data") + if _is_local_target(output_folder): + Path(data_target).mkdir(exist_ok=True, parents=True) var_features_keys = list(variable_schema.keys()) gen_kwargs_ = gen_kwargs or {sn: {} for sn in generators.keys()} for split_name, gen_func in generators.items(): - split_root_path = str(output_folder / split_name) - _ = zarr.open_group(split_root_path, mode="w") # create/overwrite + split_root_path = str(_join_target(data_target, split_name)) + _ = _open_split_group(split_root_path, mode="w") # create/overwrite batch_ids_list = gen_kwargs_.get(split_name, {}).get("shards_ids", []) total_samples = ( diff --git a/tests/storage/test_zarr_fsspec_write.py b/tests/storage/test_zarr_fsspec_write.py new file mode 100644 index 00000000..451cf034 --- /dev/null +++ b/tests/storage/test_zarr_fsspec_write.py @@ -0,0 +1,210 @@ +"""Tests for the fsspec-aware Zarr write path (issue #485). + +These tests exercise ``plaid.storage.zarr.writer.generate_datasetdict_to_disk`` +against a remote (non-local) fsspec target using the in-memory filesystem +(``memory://``). They verify that: + +- local and remote targets are correctly discriminated; +- URL joining preserves the ``://`` protocol separator (``pathlib.Path`` would + collapse it); +- samples are written at the intended remote location in both sequential and + parallel (``num_proc > 1``) modes; +- no literal ``memory:`` directory is created on the local filesystem (the + previous ``LocalStore`` implementation wrote remote URLs to a local folder). +""" + +import os + +import fsspec +import numpy as np +import pytest +import zarr + +from plaid.containers.sample import Sample +from plaid.storage.zarr.writer import ( + _is_local_target, + _join_target, + _open_split_group, + generate_datasetdict_to_disk, +) + +# --------------------------------------------------------------------------- +# Module-level (picklable) helpers required by the parallel writer. +# --------------------------------------------------------------------------- + +_GLOBAL_VALUES = {0: 10.0, 1: 20.0, 2: 30.0, 3: 40.0} + + +def _make_sample(value: float) -> Sample: + """Build a minimal sample carrying a single global scalar.""" + sample = Sample() + sample.add_global("myglobal", float(value)) + return sample + + +class _MemorySampleGenerator: + """Picklable generator matching the ``gen_func(shards_ids)`` contract. + + Defined at module level so ``multiprocessing`` can pickle it when + ``num_proc > 1``. + """ + + def __call__(self, shards_ids=None): + if shards_ids is None: + shards_ids = [[]] + for shard in shards_ids: + for sample_id in shard: + yield _make_sample(_GLOBAL_VALUES[sample_id]) + + +_VARIABLE_SCHEMA = { + "Global/myglobal": {"dtype": "float64", "ndim": 0}, + "Global/myglobal_times": {"dtype": "float64", "ndim": 0}, + "Global": {"dtype": "float64", "ndim": 1}, + "Global_times": {"dtype": "float64", "ndim": 1}, +} + + +# --------------------------------------------------------------------------- +# Unit tests for the target helpers. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("target", "expected_local"), + [ + ("/abs/local/path", True), + ("relative/local/path", True), + ("file:///abs/local/path", True), + ("memory://root", False), + ], +) +def test_is_local_target(target, expected_local): + assert _is_local_target(target) is expected_local + + +def test_is_local_target_missing_backend_raises_clear_error(): + """An unavailable fsspec protocol surfaces the install hint from fsspec.""" + with pytest.raises(ImportError, match="s3fs"): + _is_local_target("s3://bucket/key") + + +def test_join_target_local_returns_path(tmp_path): + joined = _join_target(tmp_path, "data", "train") + assert joined == tmp_path / "data" / "train" + + +def test_join_target_remote_preserves_protocol(): + """``pathlib.Path`` would collapse ``memory://root`` to ``memory:/root``.""" + joined = _join_target("memory://root", "data", "train") + assert joined == "memory://root/data/train" + # a trailing slash on the base must not produce a doubled separator + assert _join_target("memory://root/", "data") == "memory://root/data" + + +# --------------------------------------------------------------------------- +# End-to-end write to a remote (memory://) target. +# --------------------------------------------------------------------------- + + +def _clear_memory_fs(): + """Reset the shared in-memory filesystem between test runs.""" + fs = fsspec.filesystem("memory") + fs.store.clear() + if hasattr(fs, "pseudo_dirs"): + fs.pseudo_dirs.clear() + fs.pseudo_dirs.append("") + + +def _read_split_groups(output_url: str, split_name: str): + data_target = _join_target(_join_target(output_url, "data"), split_name) + return zarr.open_group(str(data_target), mode="r") + + +@pytest.mark.parametrize("num_proc", [1, 2]) +def test_generate_datasetdict_to_disk_parallel_still_works_locally(num_proc, tmp_path): + """The store-selection change must not regress the local parallel path. + + ``memory://`` is per-process, so a ``num_proc > 1`` write cannot be read + back from the parent process; the parallel reopen is therefore validated on + a shared local target instead. + """ + output_folder = tmp_path / f"dataset_np{num_proc}" + + if num_proc == 1: + shards_ids = [[0, 1]] + else: + # one shard per worker to exercise the parallel reopen path + shards_ids = [[0], [1]] + + generate_datasetdict_to_disk( + output_folder=output_folder, + generators={"train": _MemorySampleGenerator()}, + variable_schema=_VARIABLE_SCHEMA, + gen_kwargs={"train": {"shards_ids": shards_ids}}, + num_proc=num_proc, + verbose=False, + ) + + group = zarr.open_group(str(output_folder / "data" / "train"), mode="r") + assert sorted(group.group_keys()) == [ + "sample_000000000", + "sample_000000001", + ] + + +def test_generate_datasetdict_to_disk_writes_to_memory_fs(): + """Sequential write to a remote (memory://) fsspec target. + + A parallel (``num_proc > 1``) equivalent is intentionally omitted here: the + ``memory://`` filesystem is per-process, so samples written by worker + processes are not visible from the parent. The parallel store-selection path + is covered by ``test_generate_datasetdict_to_disk_parallel_still_works_locally`` + and ``test_open_split_group_roundtrip_on_memory_fs``. + """ + _clear_memory_fs() + output_url = "memory://dataset_seq" + + generate_datasetdict_to_disk( + output_folder=output_url, + generators={"train": _MemorySampleGenerator()}, + variable_schema=_VARIABLE_SCHEMA, + gen_kwargs={"train": {"shards_ids": [[0, 1]]}}, + num_proc=1, + verbose=False, + ) + + # samples landed at the intended remote location + group = _read_split_groups(output_url, "train") + sample_groups = sorted(group.group_keys()) + assert sample_groups == ["sample_000000000", "sample_000000001"] + + # each sample carries the flattened global feature and is readable back + for sample_name in sample_groups: + array_keys = list(group[sample_name].array_keys()) + assert "Global__myglobal" in array_keys + + # the remote write must NOT create a literal local "memory:" directory + assert not os.path.exists("memory:") + assert not os.path.exists("memory:/") + + +def test_open_split_group_roundtrip_on_memory_fs(): + """Create (mode 'w') then reopen (mode 'a') a split group on memory://. + + This mirrors the sequential-create / parallel-worker-reopen handshake used + by ``generate_datasetdict_to_disk``. + """ + _clear_memory_fs() + target = str(_join_target(_join_target("memory://reopen", "data"), "train")) + + created = _open_split_group(target, mode="w") + created.create_group("sample_000000000").create_array("x", data=np.arange(3)) + + reopened = _open_split_group(target, mode="a") + reopened.create_group("sample_000000001").create_array("x", data=np.arange(3) + 10) + + assert sorted(reopened.group_keys()) == [ + "sample_000000000", + "sample_000000001", + ] From 5b1991675d0d6c86e7904e17ddab75b24614d41c Mon Sep 17 00:00:00 2001 From: Xavier Roynard Date: Tue, 4 Aug 2026 11:49:25 +0000 Subject: [PATCH 2/2] docs(zarr): document fsspec-aware save_to_disk and update changelog --- CHANGELOG.md | 1 + docs/source/tutorials/storage.md | 23 ++++++++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c1a6772..04f25fd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- (storage/zarr) make the Zarr backend write path fsspec-aware: `save_to_disk` can now write to any fsspec target (`memory://`, `s3://`, `gs://`, …) in addition to local paths, in both sequential and parallel (`num_proc > 1`) modes. Remote targets require the matching fsspec backend to be installed (e.g. `s3fs` for `s3://`); a missing backend surfaces fsspec's own install hint. No new hard dependency is added (`fsspec` is already pulled in by `zarr`). Scope is limited to the zarr write path; the read path is tracked separately. - (storage/writer) add optional `sample_callback` to `save_to_disk`, invoked once per sample right after it is written to disk as `sample_callback(split_name, index, sample_path)`. This lets callers process samples one by one once written instead of waiting for the whole dataset. Currently supported for the `cgns` backend, including parallel writing (`num_proc > 1`), where the callback runs inside the worker processes and must be picklable and process-safe. ### Fixed diff --git a/docs/source/tutorials/storage.md b/docs/source/tutorials/storage.md index 276bffb7..5a7f624f 100644 --- a/docs/source/tutorials/storage.md +++ b/docs/source/tutorials/storage.md @@ -10,7 +10,28 @@ End‑to‑end workflows for creating, saving, and loading PLAID datasets with t - **`sample_constructor`** is a simple function that takes a single identifier (of any type) and returns a PLAID `Sample`. The identifier can be an integer, a file path, a string, a tuple — anything that makes sense for your data. - **`ids`** is a dictionary mapping split names to **sliceable sequences** of identifiers — anything with `__getitem__` and `__len__` (list, tuple, numpy array, …). PLAID handles iteration, generator creation, and parallel sharding internally. -- **`save_to_disk`** writes a dataset locally; **`push_to_hub`** uploads it to Hugging Face Hub. +- **`save_to_disk`** writes a dataset to `output_folder`; **`push_to_hub`** uploads it to Hugging Face Hub. + +!!! info "Remote (fsspec) write targets — zarr backend" + With the **zarr** backend, `output_folder` may be a plain local path **or any + [fsspec](https://filesystem-spec.readthedocs.io/) URL** (`memory://`, + `s3://bucket/prefix`, `gs://…`, …), in both sequential and parallel + (`num_proc > 1`) modes: + + ```python + save_to_disk( + output_folder="s3://my-bucket/datasets/shapenetcar", + sample_constructor=sample_constructor, + ids=ids, + backend="zarr", + num_proc=N_PROC, + ) + ``` + + Remote targets require the matching fsspec backend to be installed + (`s3fs` for `s3://`, `gcsfs` for `gs://`, …); a missing backend raises an + `ImportError` carrying fsspec's own install hint. The other backends + (`cgns`, `hf_datasets`) currently expect a local `output_folder`. - **`init_from_disk`** / **`download_from_hub`** / **`init_streaming_from_hub`** load datasets back into PLAID. - Backend converters turn raw backend samples into PLAID `Sample` objects.