Lazy netcdf saves

lib/iris/fileformats/netcdf/_dask_locks.py

@@ -0,0 +1,131 @@
+# 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 containing code to create locks enabling dask workers to co-operate.
+
+This matter is complicated by needing different solutions for different dask scheduler
+types, i.e. local 'threads' scheduler, local 'processes' or distributed.
+
+In any case, a "iris.fileformats.netcdf.saver.Saver" object contains a netCDF4.Dataset
+targetting an output file, and creates a Saver.lock object to serialise write-accesses
+to the file from dask tasks : All dask-task file writes go via a
+"iris.fileformats.netcdf.saver.DeferredSaveWrapper" object, which also contains a link
+to the Saver.lock, and uses it to prevent workers from fouling each other.
+For each chunk written, the DeferredSaveWrapper acquires the common per-file lock;
+opens a Dataset on the file; performs a write to the relevant variable; closes the
+Dataset and then releases the lock.
+
+For a threaded scheduler, the Saver.lock is a simple threading.Lock(). The workers
+(threads) execute tasks which contain a DeferredSaveWrapper, as above. All of those
+contain the common lock, and this is simply **the same object** for all workers, since
+they share an address space.
+
+For a distributed scheduler, the Saver.lock is a `distributed.Lock()` which is
+identified with the output filepath. This is distributed to the workers by
+serialising the task function arguments, which will include the DeferredSaveWrapper.
+A worker behaves like a process, though it may execute on a remote machine. When a
+distributed.Lock is deserialised to reconstruct the worker task, this creates an object
+that communicates with the scheduler. These objects behave as a single common lock,
+as they all have the same string 'identity', so the scheduler implements inter-process
+communication so that the mutally exclude each other.
+
+It is also *conceivable* that multiple processes could write to the same file in
+parallel, if the operating system supports it. However, this also requires that the
+libnetcdf C library is built with parallel access option, which is not common.
+With the "ordinary" libnetcdf build, a process which attempts to rpen for writing a file
+which is _already_ open for writing simply raises an access error.
+In any case, Iris netcdf saver will not support this mode of operation, at present.
+
+We don't currently support a local "processes" type scheduler. If we did, the
+behaviour should be very similar to a distributed scheduler. It would need to use some
+other serialisable shared-lock solution in place of 'distributed.Lock', which requires
+a distributed scheduler to function.
+
+"""
+import threading
+
+import dask.array
+import dask.base
+import dask.multiprocessing
+import dask.threaded
+
+
+def dask_scheduler_is_distributed():
+ """Return whether a distributed.Client is active."""
+ # NOTE: this replicates logic in `dask.base.get_scheduler` : if a distributed client
+ # has been created + is still active, then the default scheduler will always be
+ # "distributed".
+ is_distributed = False
+ # NOTE: must still work when 'distributed' is not available.
+ try:
+ import distributed
+
+ client = distributed.get_client()
+ is_distributed = client is not None
+ except (ImportError, ValueError):
+ pass
+ return is_distributed
+
+
+def get_dask_array_scheduler_type():
+ """
+ Work out what type of scheduler an array.compute*() will use.
+
+ Returns one of 'distributed', 'threads' or 'processes'.
+ The return value is a valid argument for dask.config.set(scheduler=<type>).
+ This cannot distinguish between distributed local and remote clusters -- both of
+ those simply return 'distributed'.
+
+ NOTE: this takes account of how dask is *currently* configured. It will be wrong
+ if the config changes before the compute actually occurs.
+
+ """
+ if dask_scheduler_is_distributed():
+ result = "distributed"
+ else:
+ # Call 'get_scheduler', which respects the config settings, but pass an array
+ # so we default to the default scheduler for that type of object.
+ trial_dask_array = dask.array.zeros(1)
+ get_function = dask.base.get_scheduler(collections=[trial_dask_array])
+ # Detect the ones which we recognise.
+ if get_function == dask.threaded.get:
+ result = "threads"
+ elif get_function == dask.multiprocessing.get:
+ result = "processes"
+ else:
+ msg = f"Dask default scheduler for arrays is unrecognised : {get_function}"
+ raise ValueError(msg)
+
+ return result
+
+
+def get_worker_lock(identity: str):
+ """
+ Return a mutex Lock which can be shared by multiple Dask workers.
+
+ The type of Lock generated depends on the dask scheduler type, which must therefore
+ be set up before this is called.
+
+ """
+ scheduler_type = get_dask_array_scheduler_type()
+ if scheduler_type == "threads":
+ # N.B. the "identity" string is never used in this case, as the same actual
+ # lock object is used by all workers.
+ lock = threading.Lock()
+ elif scheduler_type == "distributed":
+ from dask.distributed import Lock as DistributedLock
+
+ lock = DistributedLock(identity)
+ else:
+ msg = (
+ "The configured dask array scheduler type is "
+ f'"{scheduler_type}", '
+ "which is not supported by the Iris netcdf saver."
+ )
+ raise ValueError(msg)
+
+ # NOTE: not supporting 'processes' scheduler, for now.
+ return lock

lib/iris/fileformats/netcdf/_thread_safe_nc.py

@@ -340,3 +340,41 @@ def __getstate__(self):
  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}>"

lib/iris/fileformats/netcdf/saver.py

@@ -23,6 +23,7 @@
 import warnings
 
 import cf_units
+import dask
 import dask.array as da
 import numpy as np
 import numpy.ma as ma
@@ -44,7 +45,7 @@
 from iris.coords import AncillaryVariable, AuxCoord, CellMeasure, DimCoord
 import iris.exceptions
 import iris.fileformats.cf
-from iris.fileformats.netcdf import _thread_safe_nc
+from iris.fileformats.netcdf import _dask_locks, _thread_safe_nc
 import iris.io
 import iris.util
 
@@ -499,7 +500,7 @@ def __setitem__(self, keys, arr):
 class Saver:
  """A manager for saving netcdf files."""
 
- def __init__(self, filename, netcdf_format):
+ def __init__(self, filename, netcdf_format, compute=True):
  """
  A manager for saving netcdf files.
 
@@ -512,6 +513,15 @@ def __init__(self, filename, netcdf_format):
  Underlying netCDF file format, one of 'NETCDF4', 'NETCDF4_CLASSIC',
  'NETCDF3_CLASSIC' or 'NETCDF3_64BIT'. Default is 'NETCDF4' format.
 
+ * compute (bool):
+ If True, the Saver performs normal 'synchronous' data writes, where data
+ is streamed directly into file variables during the save operation.
+ If False, the file is created as normal, but computation and streaming of
+ any lazy array content is instead deferred to :class:`dask.delayed` objects,
+ which are held in a list in the saver 'delayed_writes' property.
+ The relavant file variables are created empty, and the write can
+ subsequently be completed by computing the 'save.deferred_writes'.
+
  Returns:
  None.
 
@@ -548,18 +558,28 @@ def __init__(self, filename, netcdf_format):
  self._mesh_dims = {}
  #: A dictionary, mapping formula terms to owner cf variable name
  self._formula_terms_cache = {}
+ #: Target filepath
+ self.filepath = os.path.abspath(filename)
+ #: Whether lazy saving.
+ self.lazy_saves = not compute
+ #: A list of deferred writes for lazy saving : each is a (source, target) pair
+ self.deferred_writes = []
+ # N.B. the file-write-lock *type* actually depends on the dask scheduler type.
+ #: A per-file write lock to prevent dask attempting overlapping writes.
+ self.file_write_lock = _dask_locks.get_worker_lock(self.filepath)
  #: NetCDF dataset
+ self._dataset = None
  try:
  self._dataset = _thread_safe_nc.DatasetWrapper(
- filename, mode="w", format=netcdf_format
+ self.filepath, mode="w", format=netcdf_format
  )
  except RuntimeError:
- dir_name = os.path.dirname(filename)
+ dir_name = os.path.dirname(self.filepath)
  if not os.path.isdir(dir_name):
  msg = "No such file or directory: {}".format(dir_name)
  raise IOError(msg)
  if not os.access(dir_name, os.R_OK | os.W_OK):
- msg = "Permission denied: {}".format(filename)
+ msg = "Permission denied: {}".format(self.filepath)
  raise IOError(msg)
  else:
  raise
@@ -2444,8 +2464,7 @@ def _increment_name(self, varname):
 
  return "{}_{}".format(varname, num)
 
- @staticmethod
- def _lazy_stream_data(data, fill_value, fill_warn, cf_var):
+ def _lazy_stream_data(self, data, fill_value, fill_warn, cf_var):
  if hasattr(data, "shape") and data.shape == (1,) + cf_var.shape:
  # (Don't do this check for string data).
  # Reduce dimensionality where the data array has an extra dimension
@@ -2455,16 +2474,37 @@ def _lazy_stream_data(data, fill_value, fill_warn, cf_var):
  data = data.squeeze(axis=0)
 
  if is_lazy_data(data):
+ if self.lazy_saves:
+ # deferred lazy streaming
+ def store(data, cf_var, fill_value):
+ # Create a data-writeable object that we can stream into, which
+ # encapsulates the file to be opened + variable to be written.
+ writeable_var_wrapper = (
+ _thread_safe_nc.DeferredSaveWrapper(
+ self.filepath, cf_var, self.file_write_lock
+ )
+ )
+ # Add to the list of deferred writes, used in _deferred_save().
+ self.deferred_writes.append((data, writeable_var_wrapper))
+ # NOTE: in this case, no checking of fill-value violations so just
+ # return dummy values for this.
+ # TODO: just for now -- can probably make this work later
+ is_masked, contains_value = False, False
+ return is_masked, contains_value
 
- def store(data, cf_var, fill_value):
- # Store lazy data and check whether it is masked and contains
- # the fill value
- target = _FillValueMaskCheckAndStoreTarget(cf_var, fill_value)
- da.store([data], [target])
- return target.is_masked, target.contains_value
+ else:
+ # Immediate streaming store : check mask+fill as we go.
+ def store(data, cf_var, fill_value):
+ # Store lazy data and check whether it is masked and contains
+ # the fill value
+ target = _FillValueMaskCheckAndStoreTarget(
+ cf_var, fill_value
+ )
+ da.store([data], [target], lock=False)
+ return target.is_masked, target.contains_value
 
  else:
-
+ # Real data is always written directly, i.e. not via lazy save.
  def store(data, cf_var, fill_value):
  cf_var[:] = data
  is_masked = np.ma.is_masked(data)
@@ -2513,6 +2553,28 @@ def store(data, cf_var, fill_value):
  )
  warnings.warn(msg.format(cf_var.name, fill_value))
 
+ def _deferred_save(self):
+ """
+ Create a 'delayed' to trigger file completion for lazy saves.
+
+ This contains all the deferred writes, which complete the file by filling out
+ the data of variables initially created empty.
+
+ """
+ if self.deferred_writes:
+ # Create a single delayed da.store operation to complete the file.
+ sources, targets = zip(*self.deferred_writes)
+ result = da.store(sources, targets, compute=False, lock=False)
+ else:
+ # Return a delayed anyway, just for usage consistency.
+ @dask.delayed
+ def no_op():
+ return None
+
+ result = no_op()
+
+ return result
+
 
 def save(
  cube,
@@ -2530,6 +2592,7 @@ def save(
  least_significant_digit=None,
  packing=None,
  fill_value=None,
+ compute=True,
 ):
  """
  Save cube(s) to a netCDF file, given the cube and the filename.
@@ -2652,6 +2715,14 @@ def save(
  `:class:`iris.cube.CubeList`, or a single element, and each element of
  this argument will be applied to each cube separately.
 
+ * compute (bool):
+ When False, create the output file but defer writing any lazy array content to
+ its variables, such as (lazy) data and aux-coords points and bounds.
+ Instead return a class:`dask.delayed` which, when computed, will compute all
+ the lazy content and stream it to complete the file.
+ Several such data saves can be performed in parallel, by passing a list of them
+ into a :func:`dask.compute` call.
+
  Returns:
  None.
 
@@ -2752,7 +2823,9 @@ def is_valid_packspec(p):
  raise ValueError(msg)
 
  # Initialise Manager for saving
- with Saver(filename, netcdf_format) as sman:
+ # N.B. FOR NOW -- we are cheating and making all saves compute=False, as otherwise
+ # non-lazy saves do *not* work with the distributed scheduler.
+ with Saver(filename, netcdf_format, compute=False) as sman:
  # Iterate through the cubelist.
  for cube, packspec, fill_value in zip(cubes, packspecs, fill_values):
  sman.write(
@@ -2797,3 +2870,12 @@ def is_valid_packspec(p):
 
  # Add conventions attribute.
  sman.update_global_attributes(Conventions=conventions)
+
+ # For now, not using Saver(compute=True) as it doesn't work with distributed or
+ # process workers (only threaded).
+ result = sman._deferred_save()
+ if compute:
+ result = result.compute()
+ result = None
+
+ return result