From 718323515e1b37350aa78c0652ad423c072bb21e Mon Sep 17 00:00:00 2001 From: Aljes Date: Tue, 15 Sep 2026 16:14:26 +0100 Subject: [PATCH] Batch Zarr attribute writes, which were wedging split in CI pytest-timeout caught the hang I could not locate before. The main thread was stuck in zarr's sync-over-async bridge: split_store -> subset_h5ad -> subset_axis_group -> copy_attrs -> zarr Attributes.__setitem__ -> Group.update_attributes -> zarr.core.sync.sync Every `attrs[k] = v` on a Zarr group persists the whole metadata document through that bridge, so copy_attrs was doing one round trip per attribute and set_encoding two per element. `split` repeats all of it once per output group, which is why only that command showed it; the thread dump had 288 threads and four idle asyncio workers. Attributes are now written in a single `put()` on Zarr, falling back to assignment on HDF5, which has no batch API. That removes most of the round trips rather than just making them less likely to collide. The earlier Rich live-display fix stands on its own merits -- nesting two live displays on one console is wrong regardless -- but it was not the cause. Ten consecutive runs of the split and zarr tests: no hangs, no failures. Co-Authored-By: Claude Opus 5 --- src/adata/elements/spec.py | 24 +++++++++++++++++++++--- src/adata/storage/__init__.py | 23 +++++++++++++++++++++-- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/src/adata/elements/spec.py b/src/adata/elements/spec.py index 541dd1a..1466717 100644 --- a/src/adata/elements/spec.py +++ b/src/adata/elements/spec.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import Any, Dict, Optional, Tuple ANNDATA = "anndata" RAW = "raw" @@ -95,6 +95,25 @@ def encoding_type(obj: Any) -> Optional[str]: return encoding_of(obj)[0] +def set_attrs(obj: Any, values: Dict[str, Any]) -> None: + """Set several attributes in as few writes as the backend allows. + + On Zarr each assignment persists the whole metadata document through its + sync-over-async bridge, so setting attributes one at a time costs a round + trip each. HDF5 has no such batch API and falls back to assignment. + """ + if not values: + return + + put = getattr(obj.attrs, "put", None) + if callable(put): + put({**dict(obj.attrs), **values}) + return + + for key, value in values.items(): + obj.attrs[key] = value + + def set_encoding(obj: Any, enc_type: str, version: Optional[str] = None) -> None: """Stamp ``encoding-type``/``encoding-version`` on a group or dataset. @@ -103,5 +122,4 @@ def set_encoding(obj: Any, enc_type: str, version: Optional[str] = None) -> None """ if version is None: version = CURRENT_VERSION[enc_type] - obj.attrs["encoding-type"] = enc_type - obj.attrs["encoding-version"] = version + set_attrs(obj, {"encoding-type": enc_type, "encoding-version": version}) diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 96d351d..58e7bad 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -229,8 +229,27 @@ def _normalize_attr_value(value: Any, target_backend: str) -> Any: def copy_attrs(src_attrs: Any, dst_attrs: Any, *, target_backend: str) -> None: - for k, v in src_attrs.items(): - dst_attrs[k] = _normalize_attr_value(v, target_backend) + """Copy attributes across, normalising values for the target backend. + + Written in one go on Zarr. Each `attrs[k] = v` there persists the whole + metadata document through zarr's sync-over-async bridge, so setting them + one at a time means a separate round trip per attribute -- slow, and the + place a `split` over many groups was observed to wedge in CI. + """ + normalized = { + str(k): _normalize_attr_value(v, target_backend) + for k, v in src_attrs.items() + } + if not normalized: + return + + put = getattr(dst_attrs, "put", None) + if target_backend == "zarr" and callable(put): + put({**dict(dst_attrs), **normalized}) + return + + for k, v in normalized.items(): + dst_attrs[k] = v def dataset_create_kwargs(