Handle dask distributed progress (#614)

* distributed progress

* test local cluster

malariagen_data/anoph/base.py

@@ -32,6 +32,7 @@
     LoggingHelper,
     check_colab_location,
     check_types,
+    distributed_client,
     get_gcp_region,
     hash_params,
     init_filesystem,
@@ -174,9 +175,13 @@ def _dask_progress(self, desc=None, leave=False, **kwargs):  # pragma: no cover
         # Progress doesn't mix well with debug logging.
         show_progress = self._show_progress and not self._debug
         if show_progress:
-            return TqdmCallback(
-                desc=desc, leave=leave, tqdm_class=self._tqdm_class, **kwargs
-            )
+            if distributed_client():
+                # Cannot easily show progress, fall back to spinner.
+                return self._spinner(desc=desc)
+            else:
+                return TqdmCallback(
+                    desc=desc, leave=leave, tqdm_class=self._tqdm_class, **kwargs
+                )
         else:
             return nullcontext()
 

malariagen_data/util.py

@@ -1423,3 +1423,13 @@ def hash_columns(x):
             v = x[i, j]
             out[j] = out[j] * 33 + v
     return out
+
+
+def distributed_client():
+    from distributed import get_client
+
+    try:
+        client = get_client()
+    except ValueError:
+        client = None
+    return client