Repository: dask/distributed

Issue:
Bokeh is validating properties
I'm noticing that my laptop is hot even with an idle cluster.  I brought up the profile-server chart to look at what was going on and it looks like Bokeh is pretty active validating properties, which it shouldn't be doing.  I looked at the line profile and found that for example, this line was triggering validation

https://github.com/dask/distributed/blob/5f01fe6010a0a3621956920d0260ac5f0111f17d/distributed/dashboard/components/scheduler.py#L2408-L2413

despite the containing method being wrapped with `@without_property_validation`

cc @jsignell @bryevdv 

Discussion:
fjetter: +1 if there are no other suggestion in the issue ticket

Patch:
diff --git a/distributed/comm/utils.py b/distributed/comm/utils.py
index 5301265c..0ce4f8f8 100644
--- a/distributed/comm/utils.py
+++ b/distributed/comm/utils.py
@@ -19,19 +19,18 @@ if isinstance(OFFLOAD_THRESHOLD, str):
 
 
 async def to_frames(
-    msg, serializers=None, on_error="message", context=None, allow_offload=True
+    msg,
+    allow_offload=True,
+    **kwargs,
 ):
     """
     Serialize a message into a list of Distributed protocol frames.
+    Any kwargs are forwarded to protocol.dumps().
     """
 
     def _to_frames():
         try:
-            return list(
-                protocol.dumps(
-                    msg, serializers=serializers, on_error=on_error, context=context
-                )
-            )
+            return list(protocol.dumps(msg, **kwargs))
         except Exception as e:
             logger.info("Unserializable Message: %s", msg)
             logger.exception(e)
diff --git a/distributed/comm/ws.py b/distributed/comm/ws.py
index be767951..8a64f209 100644
--- a/distributed/comm/ws.py
+++ b/distributed/comm/ws.py
@@ -12,6 +12,8 @@ from tornado.httpserver import HTTPServer
 from tornado.iostream import StreamClosedError
 from tornado.websocket import WebSocketClosedError, WebSocketHandler, websocket_connect
 
+import dask
+
 from ..utils import ensure_bytes, nbytes
 from .addressing import parse_host_port, unparse_host_port
 from .core import Comm, CommClosedError, Connector, FatalCommClosedError, Listener
@@ -22,6 +24,11 @@ from .utils import ensure_concrete_host, from_frames, get_tcp_server_address, to
 logger = logging.getLogger(__name__)
 
 
+BIG_BYTES_SHARD_SIZE = dask.utils.parse_bytes(
+    dask.config.get("distributed.comm.websockets.shard")
+)
+
+
 class WSHandler(WebSocketHandler):
     def __init__(
         self,
@@ -106,6 +113,7 @@ class WSHandlerComm(Comm):
                 "recipient": self.remote_info,
                 **self.handshake_options,
             },
+            frame_split_size=BIG_BYTES_SHARD_SIZE,
         )
         n = struct.pack("Q", len(frames))
         try:
diff --git a/distributed/dashboard/components/__init__.py b/distributed/dashboard/components/__init__.py
index f0798417..9cba35f3 100644
--- a/distributed/dashboard/components/__init__.py
+++ b/distributed/dashboard/components/__init__.py
@@ -79,6 +79,7 @@ def add_periodic_callback(doc, component, interval):
     _attach(doc, component)
 
 
+@without_property_validation
 def update(ref):
     comp = ref()
     if comp is not None:
diff --git a/distributed/distributed-schema.yaml b/distributed/distributed-schema.yaml
index 440a39fb..f0beee81 100644
--- a/distributed/distributed-schema.yaml
+++ b/distributed/distributed-schema.yaml
@@ -759,6 +759,20 @@ properties:
                       Alternatively, the key can be appended to the cert file
                       above, and this field left blank
 
+          websockets:
+            type: object
+            properties:
+              shard:
+                type:
+                - string
+                description: |
+                  The maximum size of a websocket frame to send through a comm.
+
+                  This is somewhat duplicative of distributed.comm.shard, but websockets
+                  often have much smaller maximum message sizes than othe protocols, so
+                  this attribute is used to set a smaller default shard size and to
+                  allow separate control of websocket message sharding.
+
       diagnostics:
         type: object
         properties:
diff --git a/distributed/distributed.yaml b/distributed/distributed.yaml
index 28ca8cf7..86e2e548 100644
--- a/distributed/distributed.yaml
+++ b/distributed/distributed.yaml
@@ -196,6 +196,9 @@ distributed:
         key: null
         cert: null
 
+    websockets:
+      shard: 8MiB
+
   diagnostics:
     nvml: True
 
diff --git a/distributed/protocol/core.py b/distributed/protocol/core.py
index c4bd909e..1be2d761 100644
--- a/distributed/protocol/core.py
+++ b/distributed/protocol/core.py
@@ -16,7 +16,9 @@ from .utils import msgpack_opts
 logger = logging.getLogger(__name__)
 
 
-def dumps(msg, serializers=None, on_error="message", context=None) -> list:
+def dumps(
+    msg, serializers=None, on_error="message", context=None, frame_split_size=None
+) -> list:
     """Transform Python message to bytestream suitable for communication
 
     Developer Notes
@@ -53,7 +55,11 @@ def dumps(msg, serializers=None, on_error="message", context=None) -> list:
                     sub_header, sub_frames = obj.header, obj.frames
                 else:
                     sub_header, sub_frames = serialize_and_split(
-                        obj, serializers=serializers, on_error=on_error, context=context
+                        obj,
+                        serializers=serializers,
+                        on_error=on_error,
+                        context=context,
+                        size=frame_split_size,
                     )
                     _inplace_compress_frames(sub_header, sub_frames)
                 sub_header["num-sub-frames"] = len(sub_frames)
diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py
index c88d36a8..51815677 100644
--- a/distributed/protocol/serialize.py
+++ b/distributed/protocol/serialize.py
@@ -406,7 +406,9 @@ def deserialize(header, frames, deserializers=None):
     return loads(header, frames)
 
 
-def serialize_and_split(x, serializers=None, on_error="message", context=None):
+def serialize_and_split(
+    x, serializers=None, on_error="message", context=None, size=None
+):
     """Serialize and split compressable frames
 
     This function is a drop-in replacement of `serialize()` that calls `serialize()`
@@ -428,7 +430,7 @@ def serialize_and_split(x, serializers=None, on_error="message", context=None):
         frames, header.get("compression") or [None] * len(frames)
     ):
         if compression is None:  # default behavior
-            sub_frames = frame_split_size(frame)
+            sub_frames = frame_split_size(frame, n=size)
             num_sub_frames.append(len(sub_frames))
             offsets.append(len(out_frames))
             out_frames.extend(sub_frames)
diff --git a/distributed/protocol/utils.py b/distributed/protocol/utils.py
index 3f5a2f8f..cf4f1815 100644
--- a/distributed/protocol/utils.py
+++ b/distributed/protocol/utils.py
@@ -25,6 +25,7 @@ def frame_split_size(frame, n=BIG_BYTES_SHARD_SIZE) -> list:
     >>> frame_split_size([b'12345', b'678'], n=3)  # doctest: +SKIP
     [b'123', b'45', b'678']
     """
+    n = n or BIG_BYTES_SHARD_SIZE
     frame = memoryview(frame)
 
     if frame.nbytes <= n:
diff --git a/setup.cfg b/setup.cfg
index 2ef82cb0..c8646ecc 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -53,8 +53,9 @@ markers =
 
 # pytest-timeout settings
 # 'thread' kills off the whole test suite. 'signal' only kills the offending test.
-# However, 'signal' does not work on Windows. The CI script modifies this config file on
-# the fly on Linux and MacOS.
+# However, 'signal' doesn't work on Windows (due to lack of SIGALRM) and doesn't work on
+# the MacOS GitHub CI (although it's been reported to work on MacBooks).
+# The CI script modifies this config file on the fly on Linux.
 timeout_method = thread
 # This should not be reduced; Windows CI has been observed to be occasionally
 # exceptionally slow.


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 8301b671..84e6d8b7 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -96,8 +96,8 @@ jobs:
 
       - name: Reconfigure pytest-timeout
         shell: bash -l {0}
-        # No SIGALRM available on Windows
-        if: ${{ matrix.os != 'windows-latest' }}
+        # No SIGALRM available on Windows. On MacOS, it kills the whole test suite.
+        if: ${{ matrix.os == 'ubuntu-latest' }}
         run: sed -i.bak 's/timeout_method = thread/timeout_method = signal/' setup.cfg
 
       - name: Test
diff --git a/distributed/protocol/tests/test_serialize.py b/distributed/protocol/tests/test_serialize.py
index d946b014..36359304 100644
--- a/distributed/protocol/tests/test_serialize.py
+++ b/distributed/protocol/tests/test_serialize.py
@@ -11,6 +11,8 @@ try:
 except ImportError:
     np = None
 
+import dask
+
 from distributed import Nanny, wait
 from distributed.comm.utils import from_frames, to_frames
 from distributed.protocol import (
@@ -442,6 +444,21 @@ def test_compression_numpy_list():
     assert header["compression"] == [False, False]
 
 
+@gen_test()
+async def test_frame_split():
+    data = b"1234abcd" * (2 ** 20)  # 8 MiB
+    assert dask.sizeof.sizeof(data) == dask.utils.parse_bytes("8MiB")
+
+    size = dask.utils.parse_bytes("3MiB")
+    split_frames = await to_frames({"x": to_serialize(data)}, frame_split_size=size)
+    print(split_frames)
+    assert len(split_frames) == 3 + 2  # Three splits and two headers
+
+    size = dask.utils.parse_bytes("5MiB")
+    split_frames = await to_frames({"x": to_serialize(data)}, frame_split_size=size)
+    assert len(split_frames) == 2 + 2  # Two splits and two headers
+
+
 @pytest.mark.parametrize(
     "data,is_serializable",
     [
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 353781e5..9527709a 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -22,7 +22,7 @@ from dask.utils import apply, parse_timedelta, stringify
 
 from distributed import Client, Nanny, Worker, fire_and_forget, wait
 from distributed.comm import Comm
-from distributed.compatibility import LINUX, WINDOWS
+from distributed.compatibility import LINUX, MACOS, WINDOWS
 from distributed.core import ConnectionPool, Status, connect, rpc
 from distributed.metrics import time
 from distributed.protocol.pickle import dumps
@@ -3102,6 +3102,7 @@ async def test_transition_counter(c, s, a, b):
     assert s.transition_counter > 1
 
 
+@pytest.mark.skipif(MACOS and sys.version_info < (3, 9), reason="GH#5056")
 @pytest.mark.slow
 @gen_cluster(
     client=True,
@@ -3110,8 +3111,8 @@ async def test_transition_counter(c, s, a, b):
     timeout=60,
 )
 async def test_worker_heartbeat_after_cancel(c, s, *workers):
-    """This test is intended to ensure that after cancelation of a graph, the
-    worker heartbeat is always successful. The hearbeat may not be successful if
+    """This test is intended to ensure that after cancellation of a graph, the
+    worker heartbeat is always successful. The heartbeat may not be successful if
     the worker and scheduler state drift and the scheduler doesn't handle
     unknown information gracefully. One example would be a released/cancelled
     computation where the worker returns metrics about duration, type, etc. and