Fix compatibility with distributed >= 2022.5.1 and traitlets >= 5.2.0…

…, and raise the lower bound of required versions (#573)

* fix: distributed(>=2022.5.2) compatibility with dask-gateway (client)

* fix: traitlets(>=5.2.2) compatibility with dask-gateway-server

* fix: require recent click and dask version to avoid issues

* Relax lower bound of requirements for dask-gateway the client.

dask-gateway-server/dask_gateway_server/app.py

@@ -13,7 +13,7 @@
 from .backends import Backend
 from .routes import default_routes
 from .traitlets import Application, Type
-from .utils import AccessLogger, LogFormatter, classname, normalize_address, run_main
+from .utils import AccessLogger, classname, normalize_address, run_main
 
 
 class GenerateConfig(Application):
@@ -145,8 +145,6 @@ def _default_address(self):
     def _validate_address(self, proposal):
         return normalize_address(proposal.value)
 
-    _log_formatter_cls = LogFormatter
-
     classes = List([Backend, Authenticator])
 
     @catch_config_error

dask-gateway-server/dask_gateway_server/backends/inprocess.py

@@ -69,7 +69,7 @@ async def do_start_cluster(self, cluster):
     async def do_stop_cluster(self, cluster):
         scheduler = self.schedulers.pop(cluster.name)
 
-        await scheduler.close(fast=True)
+        await scheduler.close()
         scheduler.stop()
 
         workdir = cluster.state.get("workdir")

dask-gateway-server/dask_gateway_server/backends/kubernetes/controller.py

@@ -20,7 +20,6 @@
 from ...utils import (
     AccessLogger,
     FrozenAttrDict,
-    LogFormatter,
     RateLimiter,
     TaskPool,
     normalize_address,
@@ -284,8 +283,6 @@ def _validate_proxy_prefix(self, proposal):
         config=True,
     )
 
-    _log_formatter_cls = LogFormatter
-
     # Fail if the config file errors
     raise_config_file_errors = True
 

dask-gateway-server/dask_gateway_server/traitlets.py

@@ -3,7 +3,22 @@
 from traitlets import Unicode
 from traitlets.config import Application
 
-# Override default values for logging
+# We replace the class of the default formatter used via a configuration change.
+#
+# References:
+#
+# - Traitlets' Application.logging_config defaults:
+#   https://github.com/ipython/traitlets/blob/e2c731ef72dd41d4be527d4d93dd87ccc409830d/traitlets/config/application.py#L229-L256
+# - Python official schema for Application.logging_config:
+#   https://docs.python.org/3/library/logging.config.html#logging-config-dictschema
+#
+Application.logging_config = {
+    "formatters": {
+        "console": {
+            "class": "dask_gateway_server.utils.LogFormatter",
+        },
+    },
+}
 Application.log_level.default_value = "INFO"
 Application.log_format.default_value = (
     "%(log_color)s[%(levelname)1.1s %(asctime)s.%(msecs).03d "

dask-gateway-server/dask_gateway_server/utils.py

@@ -175,10 +175,11 @@ def classname(cls):
 
 
 class LogFormatter(ColoredFormatter):
-    def __init__(self, fmt=None, datefmt=None):
+    def __init__(self, fmt=None, datefmt=None, style=None):
         super().__init__(
             fmt=fmt,
             datefmt=datefmt,
+            style=style,
             reset=False,
             log_colors={
                 "DEBUG": "blue",

dask-gateway-server/requirements.txt

@@ -3,4 +3,4 @@
 aiohttp
 colorlog
 cryptography
-traitlets
+traitlets>=5.2.2.post1

dask-gateway/requirements.txt

@@ -1,8 +1,8 @@
 # These are dependencies for installing the dask-gateway package.
 #
 aiohttp
-click
-dask>=2.2.0
-distributed>=2021.01.1
+click>=8.1.3
+dask>=2022.1.0
+distributed>=2022.1.0,!=2022.5.1,!=2022.5.2
 pyyaml
 tornado

tests/test_db_backend.py

@@ -469,7 +469,7 @@ async def test_cluster_fails_after_connect():
             async with gateway.new_cluster() as cluster:
                 # Kill scheduler
                 scheduler = g.gateway.backend.schedulers[cluster.name]
-                await scheduler.close(fast=True)
+                await scheduler.close()
                 scheduler.stop()
 
                 # Gateway notices and cleans up cluster in time