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,8 @@
 import numpy as np
 
 from .. import backends, conventions, Dataset
-from .common import ArrayWriter, GLOBAL_LOCK
+from .common import (ArrayWriter, get_scheduler, get_scheduler_lock,
+                     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 +67,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
@@ -131,6 +132,32 @@ def _protect_dataset_variables_inplace(dataset, cache):
             variable.data = data
 
 
+def _get_lock(engine, scheduler, format, path_or_file):
+    """ Get the lock(s) that apply to a particular scheduler/engine/format"""
+
+    locks = []
+    SchedulerLock = get_scheduler_lock(scheduler)
+    if format in ['NETCDF4', None] and engine in ['h5netcdf', 'netcdf4']:
+        locks.append(HDF5_LOCK)
+
+    try:
+        # per file lock
+        # Dask locks take a name argument (e.g. filename)
+        locks.append(SchedulerLock(path_or_file))
+    except TypeError:
+        # threading/multiprocessing lock
+        locks.append(SchedulerLock())
+
+    # When we have more than one lock, use the CombinedLock wrapper class
+    lock = CombinedLock(locks) if len(locks) > 1 else locks[0]
+
+    # Question: Should we be dropping one of these two locks when they are they
+    # are basically the same. For instance, when using netcdf4 and dask is not
+    # installed, locks will be [threading.Lock(), threading.Lock()]
+
+    return lock
+
+
 def open_dataset(filename_or_obj, group=None, decode_cf=True,
                  mask_and_scale=True, decode_times=True, autoclose=False,
                  concat_characters=True, decode_coords=True, engine=None,
@@ -622,8 +649,20 @@ 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
 
+    # handle scheduler specific logic
+    scheduler = get_scheduler()
+    if (dataset.chunks and scheduler in ['distributed', 'multiprocessing'] and
+            engine != 'netcdf4'):
+        raise NotImplementedError("Writing netCDF files with the %s backend "
+                                  "is not currently supported with dask's %s "
+                                  "scheduler" % (engine, scheduler))
+    lock = _get_lock(engine, scheduler, format, path_or_file)
+    autoclose = (dataset.chunks and
+                 scheduler in ['distributed', 'multiprocessing'])
+
     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,54 @@
 NONE_VAR_NAME = '__values__'
 
 
-# dask.utils.SerializableLock if available, otherwise just a threading.Lock
-GLOBAL_LOCK = Lock()
+def get_scheduler(get=None, collection=None):
+    """ Determine the dask scheduler that is being used.
+
+    None is returned if not dask scheduler is active.
+
+    See also
+    --------
+    dask.utils.effective_get
+    """
+    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):
+    """ Get the appropriate lock for a certain situation based onthe dask
+       scheduler used.
+
+    See Also
+    --------
+    dask.utils.get_scheduler_lock
+    """
+
+    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 +126,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 +167,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 +252,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 +262,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 +273,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 +319,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 +456,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):