Update Scheduler.rebalance return value when data is missing

distributed/client.py

@@ -587,7 +587,7 @@ def __init__(
         deserializers=None,
         extensions=DEFAULT_EXTENSIONS,
         direct_to_workers=None,
-        **kwargs
+        **kwargs,
     ):
         if timeout == no_default:
             timeout = dask.config.get("distributed.comm.timeouts.connect")
@@ -960,7 +960,7 @@ async def _start(self, timeout=no_default, **kwargs):
                 self.cluster = await LocalCluster(
                     loop=self.loop,
                     asynchronous=self._asynchronous,
-                    **self._startup_kwargs
+                    **self._startup_kwargs,
                 )
             except (OSError, socket.error) as e:
                 if e.errno != errno.EADDRINUSE:
@@ -970,7 +970,7 @@ async def _start(self, timeout=no_default, **kwargs):
                     scheduler_port=0,
                     loop=self.loop,
                     asynchronous=True,
-                    **self._startup_kwargs
+                    **self._startup_kwargs,
                 )
 
             # Wait for all workers to be ready
@@ -1422,7 +1422,7 @@ def submit(
         actor=False,
         actors=False,
         pure=None,
-        **kwargs
+        **kwargs,
     ):
 
         """ Submit a function application to the scheduler
@@ -1542,7 +1542,7 @@ def map(
         actor=False,
         actors=False,
         pure=None,
-        **kwargs
+        **kwargs,
     ):
         """ Map a function on a sequence of arguments
 
@@ -2538,7 +2538,7 @@ def get(
         priority=0,
         fifo_timeout="60s",
         actors=None,
-        **kwargs
+        **kwargs,
     ):
         """ Compute dask graph
 
@@ -2669,7 +2669,7 @@ def compute(
         fifo_timeout="60s",
         actors=None,
         traverse=True,
-        **kwargs
+        **kwargs,
     ):
         """ Compute dask collections on cluster
 
@@ -2817,7 +2817,7 @@ def persist(
         priority=0,
         fifo_timeout="60s",
         actors=None,
-        **kwargs
+        **kwargs,
     ):
         """ Persist dask collections on cluster
 
@@ -3013,6 +3013,10 @@ async def _rebalance(self, futures=None, workers=None):
         await _wait(futures)
         keys = list({tokey(f.key) for f in self.futures_of(futures)})
         result = await self.scheduler.rebalance(keys=keys, workers=workers)
+        if result["status"] == "missing-data":
+            raise ValueError(
+                f"During rebalance {len(result['keys'])} keys were found to be missing"
+            )
         assert result["status"] == "OK"
 
     def rebalance(self, futures=None, workers=None, **kwargs):
@@ -3023,7 +3027,7 @@ def rebalance(self, futures=None, workers=None, **kwargs):
         depending on keyword arguments.
 
         This operation is generally not well tested against normal operation of
-        the scheduler.  It it not recommended to use it while waiting on
+        the scheduler.  It is not recommended to use it while waiting on
         computations.
 
         Parameters
@@ -3085,7 +3089,7 @@ def replicate(self, futures, n=None, workers=None, branching_factor=2, **kwargs)
             n=n,
             workers=workers,
             branching_factor=branching_factor,
-            **kwargs
+            **kwargs,
         )
 
     def nthreads(self, workers=None, **kwargs):
@@ -3505,7 +3509,7 @@ def retire_workers(self, workers=None, close_workers=True, **kwargs):
             self.scheduler.retire_workers,
             workers=workers,
             close_workers=close_workers,
-            **kwargs
+            **kwargs,
         )
 
     def set_metadata(self, key, value):

distributed/scheduler.py

@@ -30,6 +30,7 @@
     second,
     compose,
     groupby,
+    concat,
 )
 from tornado.ioloop import IOLoop
 
@@ -3103,7 +3104,7 @@ async def rebalance(self, comm=None, keys=None, workers=None):
                 if not all(r["status"] == "OK" for r in result):
                     return {
                         "status": "missing-data",
-                        "keys": sum([r["keys"] for r in result if "keys" in r], []),
+                        "keys": tuple(concat(r["keys"].keys() for r in result)),
                     }
 
                 for sender, recipient, ts in msgs:

distributed/tests/test_client.py

@@ -2876,6 +2876,15 @@ def test_rebalance_unprepared(c, s, a, b):
     s.validate_state()
 
 
+@gen_cluster(client=True)
+async def test_rebalance_raises_missing_data(c, s, a, b):
+    with pytest.raises(ValueError, match=f"keys were found to be missing"):
+        futures = await c.scatter(range(100))
+        keys = [f.key for f in futures]
+        del futures
+        await c.rebalance(keys)
+
+
 @gen_cluster(client=True)
 def test_receive_lost_key(c, s, a, b):
     x = c.submit(inc, 1, workers=[a.address])
@@ -4864,8 +4873,8 @@ def test_bytes_keys(c, s, a, b):
 
 @gen_cluster(client=True)
 def test_unicode_ascii_keys(c, s, a, b):
-    uni_type = type(u"")
-    key = u"inc-123"
+    uni_type = type("")
+    key = "inc-123"
     future = c.submit(inc, 1, key=key)
     result = yield future
     assert type(future.key) is uni_type
@@ -4876,8 +4885,8 @@ def test_unicode_ascii_keys(c, s, a, b):
 
 @gen_cluster(client=True)
 def test_unicode_keys(c, s, a, b):
-    uni_type = type(u"")
-    key = u"inc-123\u03bc"
+    uni_type = type("")
+    key = "inc-123\u03bc"
     future = c.submit(inc, 1, key=key)
     result = yield future
     assert type(future.key) is uni_type
@@ -4889,8 +4898,8 @@ def test_unicode_keys(c, s, a, b):
     result2 = yield future2
     assert result2 == 3
 
-    future3 = yield c.scatter({u"data-123": 123})
-    result3 = yield future3[u"data-123"]
+    future3 = yield c.scatter({"data-123": 123})
+    result3 = yield future3["data-123"]
     assert result3 == 123