fix(store): relocate LocalStore temporary files to configurable tmp_dir

`LocalStore` now writes temporary files to a dedicated `tmp_dir` rather than
directly inside the store path. Previously, `_atomic_write` wrote temp files
ending in `.partial` inside the store which triggered `ZarrUserWarning` during
concurrent reading and listing of the store.

The location defaults to the system temporary directory (via
`tempfile.gettempdir()`), and is overridable via the `tmp_dir` argument to
`LocalStore` or globally via the `store.local.tmp_dir` config option (env var
`ZARR_STORE__LOCAL__TMP_DIR`). The temporary directory must be on the same
filesystem as the store, otherwise `_atomic_write` will raise `OSError` with
`EXDEV`. The pre-existing try/except now also has a specific catch for `EXDEV`
and re-raises, informing the user to change the `tmp_dir` value to a location on
the same filesystem rather than emitting a raw `OSError`. Documented this new
option in config.md and storage.md. Fixes #4161.

Author: omar-merhebi

changes/4173.bugfix.md

@@ -0,0 +1,12 @@
+`LocalStore` now writes the temporary files used by atomic writes to a
+configurable location rather than alongside each chunk inside the store.
+Previously the temporary `.partial` files were placed in the store directory and
+would surface when listing or concurrently reading a store, producing spurious
+`ZarrUserWarning`s.
+
+The temporary location defaults to the system temporary directory and can be
+overridden per store via the `tmp_dir` argument to `LocalStore`, or globally via
+the `store.local.tmp_dir` config option (environment variable
+`ZARR_STORE__LOCAL__TMP_DIR`). It must be on the same filesystem as the store.
+If it is not, a write fails with a clear error telling you to point `tmp_dir` at
+the store's filesystem, instead of a raw cross-device `OSError`.

docs/user-guide/config.md

@@ -44,6 +44,7 @@ Configuration options include the following:
 - Selections of implementations of codecs, codec pipelines and buffers
 - Enabling GPU support with `zarr.config.enable_gpu()`. See [GPU support](gpu.md) for more.
 - Control request merging when reading multiple chunks from the same shard with `array.sharding_coalesce_max_gap_bytes` and `array.sharding_coalesce_max_bytes`. Reads of nearby chunks are coalesced into a single request to the store when separated by at most `sharding_coalesce_max_gap_bytes` and the resulting merged read is no larger than `sharding_coalesce_max_bytes`.
+- Set the temporary write location for `LocalStore` writes with `store.local.tmp_dir`.
 
 For selecting custom implementations of codecs, pipelines, buffers and ndbuffers,
 first register the implementations in the registry and then select them in the config.

docs/user-guide/storage.md

@@ -113,6 +113,11 @@ group = zarr.open_group(store=store, mode='r')
 print(group)
 ```
 
+By default, `LocalStore` writes the temporary files used by atomic writes to the system temporary
+directory. Set a different location with the `tmp_dir` argument or globally via the
+`store.local.tmp_dir` config option (`ZARR_STORE__LOCAL__TMP_DIR`). The temporary directory
+should be on the same filesystem as the store, or writes may fail with a cross-device error.
+
 ### Zip Store
 
 The [`zarr.storage.ZipStore`][] stores the contents of a Zarr hierarchy in a single

src/zarr/core/config.py

@@ -152,6 +152,7 @@ def enable_gpu(self) -> ConfigSet:
             },
             "buffer": "zarr.buffer.cpu.Buffer",
             "ndbuffer": "zarr.buffer.cpu.NDBuffer",
+            "store": {"local": {"tmp_dir": None}},
         }
     ],
     deprecations=deprecations,

src/zarr/storage/_local.py

@@ -2,12 +2,14 @@
 
 import asyncio
 import contextlib
+import errno
 import io
 import os
 import shutil
 import sys
 import uuid
 from pathlib import Path
+from tempfile import gettempdir
 from typing import TYPE_CHECKING, BinaryIO, Literal, Self
 
 from zarr.abc.store import (
@@ -20,6 +22,7 @@
 from zarr.core.buffer import Buffer
 from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import AccessModeLiteral, concurrent_map
+from zarr.core.config import config as zarr_config
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable, Iterator
@@ -62,26 +65,36 @@ def _safe_move(src: Path, dst: Path) -> None:
 def _atomic_write(
     path: Path,
     mode: Literal["r+b", "wb"],
+    tmp_dir: Path,
     exclusive: bool = False,
 ) -> Iterator[BinaryIO]:
-    tmp_path = path.with_suffix(f".{uuid.uuid4().hex}.partial")
+    tmp_path = tmp_dir / f"{uuid.uuid4().hex}.partial"
     try:
         with tmp_path.open(mode) as f:
             yield f
         if exclusive:
             _safe_move(tmp_path, path)
         else:
             tmp_path.replace(path)
-    except Exception:
+    except Exception as e:
         tmp_path.unlink(missing_ok=True)
+        if isinstance(e, OSError) and e.errno == errno.EXDEV:
+            msg = (
+                f"Cannot finalize atomic write {path}: tmp dir {tmp_dir} and the "
+                "store location are not on the same filesystem. Set tmp "
+                "location to a location on the same filesystem as the store "
+                "using store.local.tmp_dir config or ZARR_STORE__LOCAL__TMP_DIR."
+            )
+            raise OSError(errno.EXDEV, msg) from e
         raise
 
 
-def _put(path: Path, value: Buffer, exclusive: bool = False) -> int:
+def _put(path: Path, value: Buffer, tmp_dir: Path, exclusive: bool = False) -> int:
     path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_dir.mkdir(parents=True, exist_ok=True)
     # write takes any object supporting the buffer protocol
     view = value.as_buffer_like()
-    with _atomic_write(path, "wb", exclusive=exclusive) as f:
+    with _atomic_write(path, "wb", tmp_dir, exclusive=exclusive) as f:
         return f.write(view)
 
 
@@ -95,6 +108,9 @@ class LocalStore(Store):
         Directory to use as root of store.
     read_only : bool
         Whether the store is read-only
+    tmp_dir : str or Path, optional
+        Where to write the store's temporary files during atomic write.
+        `None` defaults to value of `tempfile.gettempdir()`.
 
     Attributes
     ----------
@@ -110,7 +126,9 @@ class LocalStore(Store):
 
     root: Path
 
-    def __init__(self, root: Path | str, *, read_only: bool = False) -> None:
+    def __init__(
+        self, root: Path | str, *, read_only: bool = False, tmp_dir: Path | str | None = None
+    ) -> None:
         super().__init__(read_only=read_only)
         if isinstance(root, str):
             root = Path(root)
@@ -119,17 +137,28 @@ def __init__(self, root: Path | str, *, read_only: bool = False) -> None:
                 f"'root' must be a string or Path instance. Got an instance of {type(root)} instead."
             )
         self.root = root
+        self._tmp_dir = tmp_dir
+
+    def _resolve_tmp_dir(self) -> Path:
+        value = self._tmp_dir
+        if value is None:
+            value = zarr_config.get("store.local.tmp_dir", None)
+        if value is None:
+            value = gettempdir()
+        return Path(value)
 
     def with_read_only(self, read_only: bool = False) -> Self:
         # docstring inherited
-        return type(self)(
-            root=self.root,
-            read_only=read_only,
-        )
+        return type(self)(root=self.root, read_only=read_only, tmp_dir=self._tmp_dir)
 
     @classmethod
     async def open(
-        cls, root: Path | str, *, read_only: bool = False, mode: AccessModeLiteral | None = None
+        cls,
+        root: Path | str,
+        *,
+        read_only: bool = False,
+        mode: AccessModeLiteral | None = None,
+        tmp_dir: Path | str | None = None,
     ) -> Self:
         """
         Create and open the store.
@@ -144,6 +173,9 @@ async def open(
             Mode in which to create the store. This only affects opening the store,
             and the final read-only state of the store is controlled through the
             read_only parameter.
+        tmp_dir : str or Path, optional
+            Directory for the temporary files used by atomic writes. Must be on
+            the same filesystem as root. Defaults to system temporary directory.
 
         Returns
         -------
@@ -156,7 +188,7 @@ async def open(
             read_only_creation = mode in ["r", "r+"]
         else:
             read_only_creation = read_only
-        store = cls(root, read_only=read_only_creation)
+        store = cls(root, read_only=read_only_creation, tmp_dir=tmp_dir)
         await store._open()
 
         # Set read_only state
@@ -226,7 +258,7 @@ def set_sync(self, key: str, value: Buffer) -> None:
                 f"Got an instance of {type(value)} instead."
             )
         path = self.root / key
-        _put(path, value)
+        _put(path, value, self._resolve_tmp_dir())
 
     def delete_sync(self, key: str) -> None:
         self._ensure_open_sync()
@@ -290,7 +322,7 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
                 f"LocalStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
             )
         path = self.root / key
-        await asyncio.to_thread(_put, path, value, exclusive=exclusive)
+        await asyncio.to_thread(_put, path, value, self._resolve_tmp_dir(), exclusive=exclusive)
 
     async def delete(self, key: str) -> None:
         """

tests/test_config.py

@@ -103,6 +103,7 @@ def test_config_defaults_set() -> None:
                 },
                 "buffer": "zarr.buffer.cpu.Buffer",
                 "ndbuffer": "zarr.buffer.cpu.NDBuffer",
+                "store": {"local": {"tmp_dir": None}},
             }
         ]
     )

tests/test_store/test_local.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
+import errno
+import os
 import pathlib
 import re
+from tempfile import gettempdir
 
 import numpy as np
 import pytest
@@ -164,7 +167,7 @@ async def test_move(
 @pytest.mark.parametrize("exclusive", [True, False])
 def test_atomic_write_successful(tmp_path: pathlib.Path, exclusive: bool) -> None:
     path = tmp_path / "data"
-    with _atomic_write(path, "wb", exclusive=exclusive) as f:
+    with _atomic_write(path, "wb", tmp_path, exclusive=exclusive) as f:
         f.write(b"abc")
     assert path.read_bytes() == b"abc"
     assert list(path.parent.iterdir()) == [path]  # no temp files
@@ -174,7 +177,7 @@ def test_atomic_write_successful(tmp_path: pathlib.Path, exclusive: bool) -> Non
 def test_atomic_write_incomplete(tmp_path: pathlib.Path, exclusive: bool) -> None:
     path = tmp_path / "data"
     with pytest.raises(RuntimeError):  # noqa: PT012
-        with _atomic_write(path, "wb", exclusive=exclusive) as f:
+        with _atomic_write(path, "wb", tmp_path, exclusive=exclusive) as f:
             f.write(b"a")
             raise RuntimeError
     assert not path.exists()
@@ -186,7 +189,7 @@ def test_atomic_write_non_exclusive_preexisting(tmp_path: pathlib.Path) -> None:
     with path.open("wb") as f:
         f.write(b"xyz")
     assert path.read_bytes() == b"xyz"
-    with _atomic_write(path, "wb", exclusive=False) as f:
+    with _atomic_write(path, "wb", tmp_path, exclusive=False) as f:
         f.write(b"abc")
     assert path.read_bytes() == b"abc"
     assert list(path.parent.iterdir()) == [path]  # no temp files
@@ -198,7 +201,45 @@ def test_atomic_write_exclusive_preexisting(tmp_path: pathlib.Path) -> None:
         f.write(b"xyz")
     assert path.read_bytes() == b"xyz"
     with pytest.raises(FileExistsError):
-        with _atomic_write(path, "wb", exclusive=True) as f:
+        with _atomic_write(path, "wb", tmp_path, exclusive=True) as f:
             f.write(b"abc")
     assert path.read_bytes() == b"xyz"
     assert list(path.parent.iterdir()) == [path]  # no temp files
+
+
+def test_tmp_dir_arg(tmp_path: pathlib.Path) -> None:
+    store = LocalStore(tmp_path, tmp_dir=tmp_path / "scratch")
+    assert store._resolve_tmp_dir() == tmp_path / "scratch"
+
+
+def test_tmp_dir_from_config(tmp_path: pathlib.Path) -> None:
+    with zarr.config.set({"store.local.tmp_dir": str(tmp_path / "cfg")}):
+        store = LocalStore(tmp_path)
+        assert store._resolve_tmp_dir() == tmp_path / "cfg"
+
+
+def test_tmp_dir_default(tmp_path: pathlib.Path) -> None:
+    store = LocalStore(tmp_path)
+    assert store._resolve_tmp_dir() == pathlib.Path(gettempdir())
+
+
+@pytest.mark.parametrize("exclusive", [True, False])
+def test_atomic_write_cross_device_raises(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch, exclusive: bool
+) -> None:
+    def _raise_exdev(*args: object, **kwargs: object) -> None:
+        raise OSError(errno.EXDEV, "Invalid cross-device link")
+
+    if exclusive:
+        monkeypatch.setattr("zarr.storage._local._safe_move", _raise_exdev)
+    else:
+        monkeypatch.setattr(os, "replace", _raise_exdev)
+
+    path = tmp_path / "data"
+    with pytest.raises(OSError, match="same filesystem") as excinfo:
+        with _atomic_write(path, "wb", tmp_path, exclusive=exclusive) as f:
+            f.write(b"abc")
+
+    assert excinfo.value.errno == errno.EXDEV  # errno preserved
+    assert not path.exists()  # target never got created
+    assert list(tmp_path.iterdir()) == []  # tmp cleans up