Refactor zarr data copying to use store.set instead of fsspec operations (#194)

* Initial replacement of fsspec interactions with store.set

* Move _get_fs_and_path to progress tracker since it's the only use

* Implement retries

Author: mosegontar

src/reformatters/common/retry.py

@@ -0,0 +1,15 @@
+import time
+from collections.abc import Callable
+from typing import Any
+
+
+def retry(func: Callable[[], Any], max_attempts: int = 6) -> Any:
+    """Simple retry utility that sleeps for a short time between attempts."""
+    for attempt in range(max_attempts):
+        try:
+            return func()
+        except Exception:
+            if attempt == max_attempts - 1:  # Last attempt failed
+                raise
+            time.sleep(attempt)
+            continue

src/reformatters/common/update_progress_tracker.py

@@ -3,18 +3,33 @@
 import threading
 from collections.abc import Sequence
 
+import fsspec  # type: ignore
 import zarr.storage
 
 from reformatters.common.config_models import BaseInternalAttrs, DataVar
 from reformatters.common.fsspec import fsspec_apply
 from reformatters.common.logging import get_logger
-from reformatters.common.zarr import _get_fs_and_path
 
 log = get_logger(__name__)
 
 PROCESSED_VARIABLES_KEY = "processed_variables"
 
 
+def _get_fs_and_path(
+    store: zarr.abc.store.Store,
+) -> tuple[fsspec.AbstractFileSystem, str]:
+    """Gross work around to allow us to make other store types quack like FsspecStore."""
+    fs = getattr(store, "fs", None)
+    if not isinstance(fs, fsspec.AbstractFileSystem):
+        raise ValueError(
+            "primary_store must have an fs that is an instance of fsspec.AbstractFileSystem"
+        )
+    path = getattr(store, "path", None)
+    if not isinstance(path, str):
+        raise ValueError("primary_store must have a path attribute that is a string")
+    return fs, path
+
+
 class UpdateProgressTracker:
     """
     Tracks which variables have been processed within a time slice of a job.

src/reformatters/common/zarr.py

@@ -4,14 +4,13 @@
 from typing import Literal
 from uuid import uuid4
 
-import fsspec  # type: ignore
 import xarray as xr
 import zarr
 from fsspec.implementations.local import LocalFileSystem  # type: ignore
 
 from reformatters.common.config import Config
-from reformatters.common.fsspec import fsspec_apply
 from reformatters.common.logging import get_logger
+from reformatters.common.retry import retry
 
 logger = get_logger(__name__)
 
@@ -113,15 +112,11 @@ def copy_data_var(
         f"Copying data var chunks to primary store ({primary_store}) for {relative_dir}."
     )
 
-    fs, path = _get_fs_and_path(primary_store)
-    fs.auto_mkdir = True
-
-    # We want to support local and s3fs filesystems. fsspec local filesystem is sync,
-    # but our s3fs from zarr.storage.FsspecStore is async and here we work around it.
-    # The AsyncFileSystem wrapper on LocalFilesystem raises NotImplementedError when _put is called.
-    source = f"{tmp_store / relative_dir}/"
-    dest = f"{path}/{relative_dir}"
-    fsspec_apply(fs, "put", source, dest, recursive=True, auto_mkdir=True)
+    for file in tmp_store.glob(f"{relative_dir}**/*"):
+        if not file.is_file():
+            continue
+        key = str(file.relative_to(tmp_store))
+        sync_to_store(primary_store, key, file.read_bytes())
 
     if track_progress_callback is not None:
         track_progress_callback()
@@ -151,25 +146,18 @@ def copy_zarr_metadata(
     metadata_files.append(tmp_store / "zarr.json")
     metadata_files.extend(tmp_store.glob("*/zarr.json"))
 
-    fs, path = _get_fs_and_path(primary_store)
-
-    # This could be partially parallelized BUT make sure to write the coords before the metadata.
     for file in metadata_files:
-        relative = file.relative_to(tmp_store)
-        dest = f"{path}/{relative}"
-        fsspec_apply(fs, "put_file", file, dest)
-
-
-def _get_fs_and_path(
-    store: zarr.abc.store.Store,
-) -> tuple[fsspec.AbstractFileSystem, str]:
-    """Gross work around to allow us to make other store types quack like FsspecStore."""
-    fs = getattr(store, "fs", None)
-    if not isinstance(fs, fsspec.AbstractFileSystem):
-        raise ValueError(
-            "primary_store must have an fs that is an instance of fsspec.AbstractFileSystem"
-        )
-    path = getattr(store, "path", None)
-    if not isinstance(path, str):
-        raise ValueError("primary_store must have a path attribute that is a string")
-    return fs, path
+        relative_path = str(file.relative_to(tmp_store))
+        sync_to_store(primary_store, relative_path, file.read_bytes())
+
+
+def sync_to_store(store: zarr.abc.store.Store, key: str, data: bytes) -> None:
+    retry(
+        lambda: zarr.core.sync.sync(
+            store.set(
+                key,
+                zarr.core.buffer.default_buffer_prototype().buffer.from_bytes(data),
+            )
+        ),
+        max_attempts=6,
+    )

tests/common/test_retry.py

@@ -0,0 +1,23 @@
+from unittest.mock import Mock
+
+import pytest
+
+from reformatters.common.retry import retry
+
+
+def test_retry_succeeds_on_first_attempt() -> None:
+    mock_func = Mock(return_value="success")
+    result = retry(mock_func)
+    assert result == "success"
+
+
+def test_retry_succeeds_after_failures() -> None:
+    mock_func = Mock(side_effect=[ValueError("fail"), "success"])
+    result = retry(mock_func, max_attempts=3)
+    assert result == "success"
+
+
+def test_retry_fails_after_max_attempts() -> None:
+    mock_func = Mock(side_effect=ValueError("persistent failure"))
+    with pytest.raises(ValueError, match="persistent failure"):
+        retry(mock_func, max_attempts=2)