O(1) rebalance

distributed/client.py

@@ -3059,11 +3059,14 @@ def upload_file(self, filename, **kwargs):
         )
 
     async def _rebalance(self, futures=None, workers=None):
-        await _wait(futures)
-        keys = list({stringify(f.key) for f in self.futures_of(futures)})
+        if futures is not None:
+            await _wait(futures)
+            keys = list({stringify(f.key) for f in self.futures_of(futures)})
+        else:
+            keys = None
         result = await self.scheduler.rebalance(keys=keys, workers=workers)
         if result["status"] == "missing-data":
-            raise ValueError(
+            raise KeyError(
                 f"During rebalance {len(result['keys'])} keys were found to be missing"
             )
         assert result["status"] == "OK"

distributed/distributed-schema.yaml

@@ -395,7 +395,7 @@ properties:
             description: >-
               Settings for memory management
             properties:
-              recent_to_old_time:
+              recent-to-old-time:
                 type: string
                 description: >-
                   When there is an increase in process memory (as observed by the
@@ -404,6 +404,49 @@ properties:
                   non-time-sensitive heuristics. This should be set to be longer than
                   the duration of most dask tasks.
 
+              rebalance:
+                type: object
+                description: >-
+                  Settings for memory rebalance operations
+                properties:
+                  measure:
+                    enum:
+                      - process
+                      - optimistic
+                      - managed
+                      - managed_in_memory
+                    description: >-
+                      Which of the properties of distributed.scheduler.MemoryState
+                      should be used for measuring worker memory usage
+                  sender-min:
+                    type: number
+                    minimum: 0
+                    maximum: 1
+                    description: >-
+                      Fraction of worker process memory at which we stop potentially
+                      receiving data from other workers. Ignored when max_memory is not
+                      set.
+
+                  recipient-max:
+                    type: number
+                    minimum: 0
+                    maximum: 1
+                    description: >-
+                      Fraction of worker process memory at which we start potentially
+                      transferring data to other workers.
+
+                  sender-recipient-gap:
+                    type: number
+                    minimum: 0
+                    maximum: 1
+                    description: >-
+                      Fraction of worker process memory, around the cluster mean, where
+                      a worker is neither a sender nor a recipient of data during a
+                      rebalance operation. E.g. if the mean cluster occupation is 50%,
+                      sender-recipient-gap=0.1 means that only nodes above 55% will
+                      donate data and only nodes below 45% will receive them. This helps
+                      avoid data from bouncing around the cluster repeatedly.
+
               target:
                 oneOf:
                   - {type: number, minimum: 0, maximum: 1}

distributed/distributed.yaml

@@ -88,7 +88,43 @@ distributed:
       # system) that is not accounted for by the dask keys stored on the worker, ignore
       # it for this long before considering it in non-critical memory measures.
       # This should be set to be longer than the duration of most dask tasks.
-      recent_to_old_time: 30s
+      recent-to-old-time: 30s
+
+      rebalance:
+        # Memory measure to rebalance upon. Possible choices are:
+        # process
+        #     Total process memory, as measured by the OS.
+        # optimistic
+        #     Managed by dask (instantaneous) + unmanaged (without any increases
+        #     happened in the last <distributed.worker.memory.recent-to-old-time>).
+        #     Recommended for use on CPython with large (2MiB+) numpy-based data chunks.
+        # managed_in_memory
+        #     Only consider the data allocated by dask in RAM. Recommended if RAM is not
+        #     released in a timely fashion back to the OS after the Python objects are
+        #     dereferenced, but remains available for reuse by PyMalloc.
+        #
+        #     If this is your problem on Linux, you should alternatively consider
+        #     setting the MALLOC_TRIM_THRESHOLD_ environment variable (note the final
+        #     underscore) to a low value; refer to the mallopt man page and to the
+        #     comments about M_TRIM_THRESHOLD on
+        #     https://sourceware.org/git/?p=glibc.git;a=blob;f=malloc/malloc.c
+        # managed
+        #     Only consider data allocated by dask, including that spilled to disk.
+        #     Recommended if disk occupation of the spill file is an issue.
+        measure: optimistic
+        # Fraction of worker process memory at which we start potentially sending
+        # data to other workers. Ignored when max_memory is not set.
+        sender-min: 0.30
+        # Fraction of worker process memory at which we stop potentially accepting
+        # data from other workers. Ignored when max_memory is not set.
+        recipient-max: 0.60
+        # Fraction of worker process memory, around the cluster mean, where a worker is
+        # neither a sender nor a recipient of data during a rebalance operation. E.g. if
+        # the mean cluster occupation is 50%, sender-recipient-gap=0.10 means that only
+        # nodes above 55% will donate data and only nodes below 45% will receive them.
+        # This helps avoid data from bouncing around the cluster repeatedly.
+        # Ignored when max_memory is not set.
+        sender-recipient-gap: 0.10
 
       # Fractions of worker process memory at which we take action to avoid memory
       # blowup. Set any of the values to False to turn off the behavior entirely.