fix distributed writes

doc/whats-new.rst

@@ -148,6 +148,8 @@ Bug fixes
 - Compatibility fixes to plotting module for Numpy 1.14 and Pandas 0.22
   (:issue:`1813`).
   By `Joe Hamman <https://github.com/jhamman>`_.
+- Fixed to_netcdf when using dask distributed (:issue:`1464`).
+  By `Joe Hamman <https://github.com/jhamman>`_..
 - Bug fix in encoding coordinates with ``{'_FillValue': None}`` in netCDF
   metadata (:issue:`1865`).
   By `Chris Roth <https://github.com/czr137>`_.

xarray/backends/api.py

@@ -10,7 +10,7 @@
 import numpy as np
 
 from .. import backends, conventions, Dataset
-from .common import ArrayWriter, GLOBAL_LOCK
+from .common import ArrayWriter, HDF5_LOCK, CombinedLock
 from ..core import indexing
 from ..core.combine import auto_combine
 from ..core.utils import close_on_error, is_remote_uri
@@ -66,9 +66,9 @@ def _default_lock(filename, engine):
             else:
                 # TODO: identify netcdf3 files and don't use the global lock
                 # for them
-                lock = GLOBAL_LOCK
+                lock = HDF5_LOCK
         elif engine in {'h5netcdf', 'pynio'}:
-            lock = GLOBAL_LOCK
+            lock = HDF5_LOCK
         else:
             lock = False
     return lock
@@ -622,8 +622,18 @@ def to_netcdf(dataset, path_or_file=None, mode='w', format=None, group=None,
     # if a writer is provided, store asynchronously
     sync = writer is None
 
+    # TODO Move this logic outside of this function
+    from .common import get_scheduler, get_scheduler_lock
+    scheduler = get_scheduler()
+    # I think we want to include the filename here to support concurrent writes
+    # using save_mfdataset
+    file_lock = get_scheduler_lock(scheduler)(path_or_file)
+    lock = CombinedLock([HDF5_LOCK, file_lock])
+    autoclose = scheduler == 'distributed'
+
     target = path_or_file if path_or_file is not None else BytesIO()
-    store = store_open(target, mode, format, group, writer)
+    store = store_open(target, mode, format, group, writer,
+                       autoclose=autoclose, lock=lock)
 
     if unlimited_dims is None:
         unlimited_dims = dataset.encoding.get('unlimited_dims', None)

xarray/backends/common.py

@@ -8,17 +8,20 @@
 import contextlib
 from collections import Mapping, OrderedDict
 import warnings
+import multiprocessing
+import threading
 
 from ..conventions import cf_encoder
 from ..core import indexing
 from ..core.utils import FrozenOrderedDict, NdimSizeLenMixin
 from ..core.pycompat import iteritems, dask_array_type
 
+# Import default lock
 try:
-    from dask.utils import SerializableLock as Lock
+    from dask.utils import SerializableLock
+    HDF5_LOCK = SerializableLock()
 except ImportError:
-    from threading import Lock
-
+    HDF5_LOCK = threading.Lock()
 
 # Create a logger object, but don't add any handlers. Leave that to user code.
 logger = logging.getLogger(__name__)
@@ -27,8 +30,38 @@
 NONE_VAR_NAME = '__values__'
 
 
-# dask.utils.SerializableLock if available, otherwise just a threading.Lock
-GLOBAL_LOCK = Lock()
+def get_scheduler(get=None, collection=None):
+    try:
+        from dask.utils import effective_get
+        actual_get = effective_get(get, collection)
+        try:
+            from dask.distributed import Client
+            if isinstance(actual_get.__self__, Client):
+                return 'distributed'
+        except (ImportError, AttributeError):
+            try:
+                import dask.multiprocessing
+                if actual_get == dask.multiprocessing.get:
+                    return 'multiprocessing'
+                else:
+                    return 'threaded'
+            except ImportError:
+                return 'threaded'
+    except ImportError:
+        return None
+
+
+def get_scheduler_lock(scheduler):
+    if scheduler == 'distributed':
+        from dask.distributed import Lock
+        return Lock
+    elif scheduler == 'multiprocessing':
+        return multiprocessing.Lock
+    elif scheduler == 'threaded':
+        from dask.utils import SerializableLock
+        return SerializableLock
+    else:
+        return threading.Lock
 
 
 def _encode_variable_name(name):
@@ -77,6 +110,39 @@ def robust_getitem(array, key, catch=Exception, max_retries=6,
             time.sleep(1e-3 * next_delay)
 
 
+class CombinedLock(object):
+    """A combination of multiple locks.
+
+    Like a locked door, a CombinedLock is locked if any of its constituent
+    locks are locked.
+    """
+
+    def __init__(self, locks):
+        self.locks = locks
+
+    def acquire(self, *args):
+        return all(lock.acquire(*args) for lock in self.locks)
+
+    def release(self, *args):
+        for lock in self.locks:
+            lock.release(*args)
+
+    def __enter__(self):
+        for lock in self.locks:
+            lock.__enter__()
+
+    def __exit__(self, *args):
+        for lock in self.locks:
+            lock.__exit__(*args)
+
+    @property
+    def locked(self):
+        return any(lock.locked for lock in self.locks)
+
+    def __repr__(self):
+        return "CombinedLock(%s)" % [repr(lock) for lock in self.locks]
+
+
 class BackendArray(NdimSizeLenMixin, indexing.ExplicitlyIndexed):
 
     def __array__(self, dtype=None):
@@ -85,7 +151,9 @@ def __array__(self, dtype=None):
 
 
 class AbstractDataStore(Mapping):
-    _autoclose = False
+    _autoclose = None
+    _ds = None
+    _isopen = False
 
     def __iter__(self):
         return iter(self.variables)
@@ -168,7 +236,7 @@ def __exit__(self, exception_type, exception_value, traceback):
 
 
 class ArrayWriter(object):
-    def __init__(self, lock=GLOBAL_LOCK):
+    def __init__(self, lock=HDF5_LOCK):
         self.sources = []
         self.targets = []
         self.lock = lock
@@ -178,11 +246,7 @@ def add(self, source, target):
             self.sources.append(source)
             self.targets.append(target)
         else:
-            try:
-                target[...] = source
-            except TypeError:
-                # workaround for GH: scipy/scipy#6880
-                target[:] = source
+            target[...] = source
 
     def sync(self):
         if self.sources:
@@ -193,9 +257,9 @@ def sync(self):
 
 
 class AbstractWritableDataStore(AbstractDataStore):
-    def __init__(self, writer=None):
+    def __init__(self, writer=None, lock=HDF5_LOCK):
         if writer is None:
-            writer = ArrayWriter()
+            writer = ArrayWriter(lock=lock)
         self.writer = writer
 
     def encode(self, variables, attributes):
@@ -239,6 +303,9 @@ def set_variable(self, k, v):  # pragma: no cover
         raise NotImplementedError
 
     def sync(self):
+        if self._isopen and self._autoclose:
+            # datastore will be reopened during write
+            self.close()
         self.writer.sync()
 
     def store_dataset(self, dataset):
@@ -373,27 +440,41 @@ class DataStorePickleMixin(object):
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        del state['ds']
+        del state['_ds']
+        del state['_isopen']
         if self._mode == 'w':
             # file has already been created, don't override when restoring
             state['_mode'] = 'a'
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
-        self.ds = self._opener(mode=self._mode)
+        self._ds = None
+        self._isopen = False
+
+    @property
+    def ds(self):
+        if self._ds is not None and self._isopen:
+            return self._ds
+        ds = self._opener(mode=self._mode)
+        self._isopen = True
+        return ds
 
     @contextlib.contextmanager
-    def ensure_open(self, autoclose):
+    def ensure_open(self, autoclose=None):
         """
         Helper function to make sure datasets are closed and opened
         at appropriate times to avoid too many open file errors.
 
         Use requires `autoclose=True` argument to `open_mfdataset`.
         """
-        if self._autoclose and not self._isopen:
+
+        if autoclose is None:
+            autoclose = self._autoclose
+
+        if not self._isopen:
             try:
-                self.ds = self._opener()
+                self._ds = self._opener()
                 self._isopen = True
                 yield
             finally:

xarray/backends/h5netcdf_.py

@@ -10,7 +10,8 @@
 from ..core.utils import FrozenOrderedDict, close_on_error
 from ..core.pycompat import iteritems, bytes_type, unicode_type, OrderedDict
 
-from .common import WritableCFDataStore, DataStorePickleMixin, find_root
+from .common import (WritableCFDataStore, DataStorePickleMixin, find_root,
+                     HDF5_LOCK)
 from .netCDF4_ import (_nc4_group, _encode_nc4_variable, _get_datatype,
                        _extract_nc4_variable_encoding, BaseNetCDF4Array)
 
@@ -64,12 +65,12 @@ class H5NetCDFStore(WritableCFDataStore, DataStorePickleMixin):
     """
 
     def __init__(self, filename, mode='r', format=None, group=None,
-                 writer=None, autoclose=False):
+                 writer=None, autoclose=False, lock=HDF5_LOCK):
         if format not in [None, 'NETCDF4']:
             raise ValueError('invalid format for h5netcdf backend')
         opener = functools.partial(_open_h5netcdf_group, filename, mode=mode,
                                    group=group)
-        self.ds = opener()
+        self._ds = opener()
         if autoclose:
             raise NotImplementedError('autoclose=True is not implemented '
                                       'for the h5netcdf backend pending '
@@ -81,7 +82,7 @@ def __init__(self, filename, mode='r', format=None, group=None,
         self._opener = opener
         self._filename = filename
         self._mode = mode
-        super(H5NetCDFStore, self).__init__(writer)
+        super(H5NetCDFStore, self).__init__(writer, lock=lock)
 
     def open_store_variable(self, name, var):
         with self.ensure_open(autoclose=False):
@@ -173,7 +174,10 @@ def prepare_variable(self, name, variable, check_encoding=False,
 
         for k, v in iteritems(attrs):
             nc4_var.setncattr(k, v)
-        return nc4_var, variable.data
+
+        target = H5NetCDFArrayWrapper(name, self)
+
+        return target, variable.data
 
     def sync(self):
         with self.ensure_open(autoclose=True):

xarray/backends/netCDF4_.py

@@ -16,7 +16,7 @@
 from ..core.pycompat import iteritems, basestring, OrderedDict, PY3, suppress
 
 from .common import (WritableCFDataStore, robust_getitem, BackendArray,
-                     DataStorePickleMixin, find_root)
+                     DataStorePickleMixin, find_root, HDF5_LOCK)
 from .netcdf3 import (encode_nc3_attr_value, encode_nc3_variable)
 
 # This lookup table maps from dtype.byteorder to a readable endian
@@ -43,6 +43,11 @@ def __init__(self, variable_name, datastore):
             dtype = np.dtype('O')
         self.dtype = dtype
 
+    def __setitem__(self, key, value):
+        with self.datastore.ensure_open(autoclose=True):
+            data = self.get_array()
+            data[key] = value
+
     def get_array(self):
         self.datastore.assert_open()
         return self.datastore.ds.variables[self.variable_name]
@@ -231,14 +236,14 @@ class NetCDF4DataStore(WritableCFDataStore, DataStorePickleMixin):
     """
 
     def __init__(self, netcdf4_dataset, mode='r', writer=None, opener=None,
-                 autoclose=False):
+                 autoclose=False, lock=HDF5_LOCK):
 
         if autoclose and opener is None:
             raise ValueError('autoclose requires an opener')
 
         _disable_auto_decode_group(netcdf4_dataset)
 
-        self.ds = netcdf4_dataset
+        self._ds = netcdf4_dataset
         self._autoclose = autoclose
         self._isopen = True
         self.format = self.ds.data_model
@@ -249,12 +254,12 @@ def __init__(self, netcdf4_dataset, mode='r', writer=None, opener=None,
             self._opener = functools.partial(opener, mode=self._mode)
         else:
             self._opener = opener
-        super(NetCDF4DataStore, self).__init__(writer)
+        super(NetCDF4DataStore, self).__init__(writer, lock=lock)
 
     @classmethod
     def open(cls, filename, mode='r', format='NETCDF4', group=None,
              writer=None, clobber=True, diskless=False, persist=False,
-             autoclose=False):
+             autoclose=False, lock=HDF5_LOCK):
         import netCDF4 as nc4
         if (len(filename) == 88 and
                 LooseVersion(nc4.__version__) < "1.3.1"):
@@ -274,7 +279,7 @@ def open(cls, filename, mode='r', format='NETCDF4', group=None,
                                    format=format)
         ds = opener()
         return cls(ds, mode=mode, writer=writer, opener=opener,
-                   autoclose=autoclose)
+                   autoclose=autoclose, lock=lock)
 
     def open_store_variable(self, name, var):
         with self.ensure_open(autoclose=False):
@@ -398,7 +403,9 @@ def prepare_variable(self, name, variable, check_encoding=False,
             # OrderedDict as the input to setncatts
             nc4_var.setncattr(k, v)
 
-        return nc4_var, variable.data
+        target = NetCDF4ArrayWrapper(name, self)
+
+        return target, variable.data
 
     def sync(self):
         with self.ensure_open(autoclose=True):