Lazy import dask.distributed to reduce import time of xarray (#7172)

* Lazy import testing and tutorial

* Lazy import distributed to avoid a costly import

* Revert changes to __init__

* Explain why we lazy import

* Add release note

* dask.distritubed.lock now supports blocking argument

Co-authored-by: Deepak Cherian <[email protected]>

Author: hmaarrfk

doc/whats-new.rst

@@ -61,8 +61,7 @@ Internal Changes
 ~~~~~~~~~~~~~~~~
 - Doctests fail on any warnings (:pull:`7166`)
   By `Maximilian Roos <https://github.com/max-sixty>`_.
-
-
+- Improve import time by lazy loading ``dask.distributed`` (:pull: `7172`).
 - Explicitly specify ``longdouble=False`` in :py:func:`cftime.date2num` when
   encoding times to preserve existing behavior and prevent future errors when it
   is eventually set to ``True`` by default in cftime (:pull:`7171`).  By

xarray/backends/locks.py

@@ -11,11 +11,6 @@
     # no need to worry about serializing the lock
     SerializableLock = threading.Lock  # type: ignore
 
-try:
-    from dask.distributed import Lock as DistributedLock
-except ImportError:
-    DistributedLock = None  # type: ignore
-
 
 # Locks used by multiple backends.
 # Neither HDF5 nor the netCDF-C library are thread-safe.
@@ -41,14 +36,6 @@ def _get_multiprocessing_lock(key):
     return multiprocessing.Lock()
 
 
-_LOCK_MAKERS = {
-    None: _get_threaded_lock,
-    "threaded": _get_threaded_lock,
-    "multiprocessing": _get_multiprocessing_lock,
-    "distributed": DistributedLock,
-}
-
-
 def _get_lock_maker(scheduler=None):
     """Returns an appropriate function for creating resource locks.
 
@@ -61,7 +48,23 @@ def _get_lock_maker(scheduler=None):
     --------
     dask.utils.get_scheduler_lock
     """
-    return _LOCK_MAKERS[scheduler]
+
+    if scheduler is None:
+        return _get_threaded_lock
+    elif scheduler == "threaded":
+        return _get_threaded_lock
+    elif scheduler == "multiprocessing":
+        return _get_multiprocessing_lock
+    elif scheduler == "distributed":
+        # Lazy import distributed since it is can add a significant
+        # amount of time to import
+        try:
+            from dask.distributed import Lock as DistributedLock
+        except ImportError:
+            DistributedLock = None  # type: ignore
+        return DistributedLock
+    else:
+        raise KeyError(scheduler)
 
 
 def _get_scheduler(get=None, collection=None) -> str | None:
@@ -128,15 +131,12 @@ def acquire(lock, blocking=True):
     if blocking:
         # no arguments needed
         return lock.acquire()
-    elif DistributedLock is not None and isinstance(lock, DistributedLock):
-        # distributed.Lock doesn't support the blocking argument yet:
-        # https://github.com/dask/distributed/pull/2412
-        return lock.acquire(timeout=0)
     else:
         # "blocking" keyword argument not supported for:
         # - threading.Lock on Python 2.
         # - dask.SerializableLock with dask v1.0.0 or earlier.
         # - multiprocessing.Lock calls the argument "block" instead.
+        # - dask.distributed.Lock uses the blocking argument as the first one
         return lock.acquire(blocking)