Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions changes/4173.bugfix.md
Original file line number Diff line number Diff line change
@@ -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`.
1 change: 1 addition & 0 deletions docs/user-guide/config.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions docs/user-guide/storage.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/zarr/core/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
58 changes: 45 additions & 13 deletions src/zarr/storage/_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down Expand Up @@ -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)


Expand All @@ -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
----------
Expand All @@ -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)
Expand All @@ -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.
Expand All @@ -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
-------
Expand All @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
"""
Expand Down
1 change: 1 addition & 0 deletions tests/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}},
}
]
)
Expand Down
49 changes: 45 additions & 4 deletions tests/test_store/test_local.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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
Loading