1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
|
from __future__ import annotations
from functools import partial
from typing import TYPE_CHECKING
import dask
import joblib
import pytest
from dask.base import normalize_token, tokenize
from packaging.version import Version
if Version(dask.__version__) < Version("2024.8.0"):
from dask.base import normalize_seq
else:
from dask.tokenize import normalize_seq
from filelock import FileLock
from scipy import sparse
import anndata as ad
from anndata.tests.helpers import subset_func # noqa: F401
if TYPE_CHECKING:
from collections.abc import Generator
from types import EllipsisType
@pytest.fixture
def backing_h5ad(tmp_path):
return tmp_path / "test.h5ad"
@pytest.fixture(
params=[("h5ad", None), ("zarr", 2), ("zarr", 3)], ids=["h5ad", "zarr2", "zarr3"]
)
def diskfmt(request):
if (fmt := request.param[0]) == "h5ad":
yield fmt
else:
with ad.settings.override(zarr_write_format=request.param[1]):
yield fmt
@pytest.fixture
def diskfmt2(diskfmt):
if diskfmt == "h5ad":
with ad.settings.override(zarr_write_format=2):
yield "zarr"
else:
yield "h5ad"
@pytest.fixture(
params=[
pytest.param((..., (slice(None), slice(None))), id="ellipsis"),
pytest.param(((...,), (slice(None), slice(None))), id="ellipsis_tuple"),
pytest.param(
((..., slice(0, 10)), (slice(None), slice(0, 10))), id="obs-ellipsis"
),
pytest.param(
((slice(0, 10), ...), (slice(0, 10), slice(None))), id="var-ellipsis"
),
pytest.param(
((slice(0, 10), slice(0, 10), ...), (slice(0, 10), slice(0, 10))),
id="obs-var-ellipsis",
),
pytest.param(
((..., slice(0, 10), slice(0, 10)), (slice(0, 10), slice(0, 10))),
id="ellipsis-obs-var",
),
pytest.param(
((slice(0, 10), ..., slice(0, 10)), (slice(0, 10), slice(0, 10))),
id="obs-ellipsis-var",
),
]
)
def ellipsis_index_with_equivalent(
request,
) -> tuple[tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]]:
return request.param
@pytest.fixture
def ellipsis_index(
ellipsis_index_with_equivalent: tuple[
tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]
],
) -> tuple[EllipsisType | slice, ...] | EllipsisType:
return ellipsis_index_with_equivalent[0]
@pytest.fixture
def equivalent_ellipsis_index(
ellipsis_index_with_equivalent: tuple[
tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]
],
) -> tuple[slice, slice]:
return ellipsis_index_with_equivalent[1]
@pytest.fixture(scope="session")
def local_cluster_addr(
tmp_path_factory: pytest.TempPathFactory, worker_id: str
) -> Generator[str, None, None]:
# Adapted from https://pytest-xdist.readthedocs.io/en/latest/how-to.html#making-session-scoped-fixtures-execute-only-once
import dask.distributed as dd
def make_cluster() -> dd.LocalCluster:
return dd.LocalCluster(n_workers=1, threads_per_worker=1)
if worker_id == "master":
with make_cluster() as cluster:
yield cluster.scheduler_address
return
# get the temp directory shared by all workers
root_tmp_dir = tmp_path_factory.getbasetemp().parent
fn = root_tmp_dir / "dask_scheduler_address.txt"
lock = FileLock(str(fn) + ".lock")
lock.acquire() # can’t use context manager, because we need to release the lock before yielding
address = fn.read_text() if fn.is_file() else None
if address:
lock.release()
yield address
return
with make_cluster() as cluster:
fn.write_text(cluster.scheduler_address)
lock.release()
yield cluster.scheduler_address
#####################
# Dask tokenization #
#####################
# TODO: Should we be exporting this?
# sparray classes don't have tokenize defined yet, see: https://github.com/dask/dask/issues/10375
def normalize_sparse_matrix(x, attrs):
return (
type(x).__name__,
normalize_seq(normalize_token(getattr(x, key)) for key in attrs),
)
for cls, attrs in [
(sparse.dia_array, ("data", "offsets", "shape")),
(sparse.bsr_array, ("data", "indices", "indptr", "blocksize", "shape")),
(sparse.coo_array, ("data", "row", "col", "shape")),
(sparse.csr_array, ("data", "indices", "indptr", "shape")),
(sparse.csc_array, ("data", "indices", "indptr", "shape")),
(sparse.lil_array, ("data", "rows", "shape")),
]:
normalize_token.register(cls, partial(normalize_sparse_matrix, attrs=attrs))
@normalize_token.register(sparse.dok_array)
def normalize_dok_matrix(x):
return type(x).__name__, normalize_token(sorted(x.items()))
@normalize_token.register(ad.AnnData)
def tokenize_anndata(adata: ad.AnnData):
res = []
if adata.X is not None:
res.append(tokenize(adata.X))
res.extend([tokenize(adata.obs), tokenize(adata.var)])
for attr in ["obsm", "varm", "obsp", "varp", "layers"]:
elem = getattr(adata, attr)
res.append(tokenize(list(dict(elem).items())))
res.append(joblib.hash(adata.uns))
if adata.raw is not None:
res.append(tokenize(adata.raw.to_adata()))
return tuple(res)
|
