forked from SciTools/iris
-
Notifications
You must be signed in to change notification settings - Fork 0
/
_thread_safe_nc.py
380 lines (303 loc) · 13.6 KB
/
_thread_safe_nc.py
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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
# Copyright Iris contributors
#
# This file is part of Iris and is released under the LGPL license.
# See COPYING and COPYING.LESSER in the root of the repository for full
# licensing details.
"""
Module to ensure all calls to the netCDF4 library are thread-safe.
Intention is that no other Iris module should import the netCDF4 module.
"""
from abc import ABC
from threading import Lock
import typing
import netCDF4
import numpy as np
_GLOBAL_NETCDF4_LOCK = Lock()
# Doesn't need thread protection, but this allows all netCDF4 refs to be
# replaced with thread_safe refs.
default_fillvals = netCDF4.default_fillvals
class _ThreadSafeWrapper(ABC):
"""
Contains a netCDF4 class instance, ensuring wrapping all API calls within _GLOBAL_NETCDF4_LOCK.
Designed to 'gate keep' all the instance's API calls, but allowing the
same API as if working directly with the instance itself.
Using a contained object instead of inheritance, as we cannot successfully
subclass or monkeypatch netCDF4 classes, because they are only wrappers for
the C-layer.
"""
CONTAINED_CLASS = NotImplemented
# Allows easy type checking, avoiding difficulties with isinstance and mocking.
THREAD_SAFE_FLAG = True
@classmethod
def _from_existing(cls, instance):
"""Pass an existing instance to __init__, where it is contained."""
assert isinstance(instance, cls.CONTAINED_CLASS)
return cls(instance)
def __init__(self, *args, **kwargs):
"""Contain an existing instance, or generate a new one from arguments."""
if isinstance(args[0], self.CONTAINED_CLASS):
instance = args[0]
else:
with _GLOBAL_NETCDF4_LOCK:
instance = self.CONTAINED_CLASS(*args, **kwargs)
self._contained_instance = instance
def __getattr__(self, item):
if item == "_contained_instance":
# Special behaviour when accessing the _contained_instance itself.
return object.__getattribute__(self, item)
else:
with _GLOBAL_NETCDF4_LOCK:
return getattr(self._contained_instance, item)
def __setattr__(self, key, value):
if key == "_contained_instance":
# Special behaviour when accessing the _contained_instance itself.
object.__setattr__(self, key, value)
else:
with _GLOBAL_NETCDF4_LOCK:
return setattr(self._contained_instance, key, value)
def __getitem__(self, item):
with _GLOBAL_NETCDF4_LOCK:
return self._contained_instance.__getitem__(item)
def __setitem__(self, key, value):
with _GLOBAL_NETCDF4_LOCK:
return self._contained_instance.__setitem__(key, value)
class DimensionWrapper(_ThreadSafeWrapper):
"""
Accessor for a netCDF4.Dimension, always acquiring _GLOBAL_NETCDF4_LOCK.
All API calls should be identical to those for netCDF4.Dimension.
"""
CONTAINED_CLASS = netCDF4.Dimension
class VariableWrapper(_ThreadSafeWrapper):
"""
Accessor for a netCDF4.Variable, always acquiring _GLOBAL_NETCDF4_LOCK.
All API calls should be identical to those for netCDF4.Variable.
"""
CONTAINED_CLASS = netCDF4.Variable
def setncattr(self, *args, **kwargs) -> None:
"""
Calls netCDF4.Variable.setncattr within _GLOBAL_NETCDF4_LOCK.
Only defined explicitly in order to get some mocks to work.
"""
with _GLOBAL_NETCDF4_LOCK:
return self._contained_instance.setncattr(*args, **kwargs)
@property
def dimensions(self) -> typing.List[str]:
"""
Calls netCDF4.Variable.dimensions within _GLOBAL_NETCDF4_LOCK.
Only defined explicitly in order to get some mocks to work.
"""
with _GLOBAL_NETCDF4_LOCK:
# Return value is a list of strings so no need for
# DimensionWrapper, unlike self.get_dims().
return self._contained_instance.dimensions
# All Variable API that returns Dimension(s) is wrapped to instead return
# DimensionWrapper(s).
def get_dims(self, *args, **kwargs) -> typing.Tuple[DimensionWrapper]:
"""
Calls netCDF4.Variable.get_dims() within _GLOBAL_NETCDF4_LOCK, returning DimensionWrappers.
The original returned netCDF4.Dimensions are simply replaced with their
respective DimensionWrappers, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
dimensions_ = list(
self._contained_instance.get_dims(*args, **kwargs)
)
return tuple([DimensionWrapper._from_existing(d) for d in dimensions_])
class GroupWrapper(_ThreadSafeWrapper):
"""
Accessor for a netCDF4.Group, always acquiring _GLOBAL_NETCDF4_LOCK.
All API calls should be identical to those for netCDF4.Group.
"""
CONTAINED_CLASS = netCDF4.Group
# All Group API that returns Dimension(s) is wrapped to instead return
# DimensionWrapper(s).
@property
def dimensions(self) -> typing.Dict[str, DimensionWrapper]:
"""
Calls dimensions of netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning DimensionWrappers.
The original returned netCDF4.Dimensions are simply replaced with their
respective DimensionWrappers, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
dimensions_ = self._contained_instance.dimensions
return {
k: DimensionWrapper._from_existing(v)
for k, v in dimensions_.items()
}
def createDimension(self, *args, **kwargs) -> DimensionWrapper:
"""
Calls createDimension() from netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning DimensionWrapper.
The original returned netCDF4.Dimension is simply replaced with its
respective DimensionWrapper, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
new_dimension = self._contained_instance.createDimension(
*args, **kwargs
)
return DimensionWrapper._from_existing(new_dimension)
# All Group API that returns Variable(s) is wrapped to instead return
# VariableWrapper(s).
@property
def variables(self) -> typing.Dict[str, VariableWrapper]:
"""
Calls variables of netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning VariableWrappers.
The original returned netCDF4.Variables are simply replaced with their
respective VariableWrappers, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
variables_ = self._contained_instance.variables
return {
k: VariableWrapper._from_existing(v) for k, v in variables_.items()
}
def createVariable(self, *args, **kwargs) -> VariableWrapper:
"""
Calls createVariable() from netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning VariableWrapper.
The original returned netCDF4.Variable is simply replaced with its
respective VariableWrapper, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
new_variable = self._contained_instance.createVariable(
*args, **kwargs
)
return VariableWrapper._from_existing(new_variable)
def get_variables_by_attributes(
self, *args, **kwargs
) -> typing.List[VariableWrapper]:
"""
Calls get_variables_by_attributes() from netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning VariableWrappers.
The original returned netCDF4.Variables are simply replaced with their
respective VariableWrappers, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
variables_ = list(
self._contained_instance.get_variables_by_attributes(
*args, **kwargs
)
)
return [VariableWrapper._from_existing(v) for v in variables_]
# All Group API that returns Group(s) is wrapped to instead return
# GroupWrapper(s).
@property
def groups(self):
"""
Calls groups of netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning GroupWrappers.
The original returned netCDF4.Groups are simply replaced with their
respective GroupWrappers, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
groups_ = self._contained_instance.groups
return {k: GroupWrapper._from_existing(v) for k, v in groups_.items()}
@property
def parent(self):
"""
Calls parent of netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning a GroupWrapper.
The original returned netCDF4.Group is simply replaced with its
respective GroupWrapper, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
parent_ = self._contained_instance.parent
return GroupWrapper._from_existing(parent_)
def createGroup(self, *args, **kwargs):
"""
Calls createGroup() from netCDF4.Group/Dataset within _GLOBAL_NETCDF4_LOCK, returning GroupWrapper.
The original returned netCDF4.Group is simply replaced with its
respective GroupWrapper, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
new_group = self._contained_instance.createGroup(*args, **kwargs)
return GroupWrapper._from_existing(new_group)
class DatasetWrapper(GroupWrapper):
"""
Accessor for a netCDF4.Dataset, always acquiring _GLOBAL_NETCDF4_LOCK.
All API calls should be identical to those for netCDF4.Dataset.
"""
CONTAINED_CLASS = netCDF4.Dataset
@classmethod
def fromcdl(cls, *args, **kwargs):
"""
Calls netCDF4.Dataset.fromcdl() within _GLOBAL_NETCDF4_LOCK, returning a DatasetWrapper.
The original returned netCDF4.Dataset is simply replaced with its
respective DatasetWrapper, ensuring that downstream calls are
also performed within _GLOBAL_NETCDF4_LOCK.
"""
with _GLOBAL_NETCDF4_LOCK:
instance = cls.CONTAINED_CLASS.fromcdl(*args, **kwargs)
return cls._from_existing(instance)
class NetCDFDataProxy:
"""A reference to the data payload of a single NetCDF file variable."""
__slots__ = ("shape", "dtype", "path", "variable_name", "fill_value")
def __init__(self, shape, dtype, path, variable_name, fill_value):
self.shape = shape
self.dtype = dtype
self.path = path
self.variable_name = variable_name
self.fill_value = fill_value
@property
def ndim(self):
return len(self.shape)
def __getitem__(self, keys):
# Using a DatasetWrapper causes problems with invalid ID's and the
# netCDF4 library, presumably because __getitem__ gets called so many
# times by Dask. Use _GLOBAL_NETCDF4_LOCK directly instead.
with _GLOBAL_NETCDF4_LOCK:
dataset = netCDF4.Dataset(self.path)
try:
variable = dataset.variables[self.variable_name]
# Get the NetCDF variable data and slice.
var = variable[keys]
finally:
dataset.close()
return np.asanyarray(var)
def __repr__(self):
fmt = (
"<{self.__class__.__name__} shape={self.shape}"
" dtype={self.dtype!r} path={self.path!r}"
" variable_name={self.variable_name!r}>"
)
return fmt.format(self=self)
def __getstate__(self):
return {attr: getattr(self, attr) for attr in self.__slots__}
def __setstate__(self, state):
for key, value in state.items():
setattr(self, key, value)
class DeferredSaveWrapper:
"""
The "opposite" of a NetCDFDataProxy : An object mimicking the data access of a
netCDF4.Variable, but where the data is to be ***written to***.
It encapsulates the netcdf file and variable which are actually to be written to.
This opens the file each time, to enable writing the data chunk, then closes it.
TODO: could be improved with a caching scheme, but this just about works.
"""
def __init__(self, filepath, cf_var, file_write_lock):
self.path = filepath
self.varname = cf_var.name
self.lock = file_write_lock
def __setitem__(self, keys, array_data):
# Write to the variable.
# First acquire a file-specific lock for all workers writing to this file.
self.lock.acquire()
# Open the file for writing + write to the specific file variable.
# Exactly as above, in NetCDFDataProxy : a DatasetWrapper causes problems with
# invalid ID's and the netCDF4 library, for so-far unknown reasons.
# Instead, use _GLOBAL_NETCDF4_LOCK, and netCDF4 _directly_.
with _GLOBAL_NETCDF4_LOCK:
dataset = None
try:
dataset = netCDF4.Dataset(self.path, "r+")
var = dataset.variables[self.varname]
var[keys] = array_data
finally:
if dataset:
dataset.close()
self.lock.release()
def __repr__(self):
return f"<{self.__class__.__name__} path={self.path!r} var={self.varname!r}>"