Unix socket support

distributed/cli/dask_scheduler.py

@@ -37,7 +37,7 @@
     help="Preferred network interface like 'eth0' or 'ib0'",
 )
 @click.option(
-    "--protocol", type=str, default=None, help="Protocol like tcp, tls, or ucx"
+    "--protocol", type=str, default=None, help="Protocol like tcp, tls, unix, or ucx"
 )
 @click.option(
     "--tls-ca-file",

distributed/cli/dask_worker.py

@@ -116,7 +116,7 @@
     "--interface", type=str, default=None, help="Network interface like 'eth0' or 'ib0'"
 )
 @click.option(
-    "--protocol", type=str, default=None, help="Protocol like tcp, tls, or ucx"
+    "--protocol", type=str, default=None, help="Protocol like tcp, tls, unix, or ucx"
 )
 @click.option("--nthreads", type=int, default=0, help="Number of threads per process.")
 @click.option(

distributed/cli/tests/test_dask_scheduler.py

@@ -4,6 +4,7 @@
 import os
 import shutil
 import signal
+import socket
 import subprocess
 import sys
 import tempfile
@@ -353,6 +354,46 @@ def test_dashboard_port_zero(loop):
             assert get_dashboard_port(c) > 0
 
 
+@pytest.mark.skipif(WINDOWS, reason="POSIX only")
+def test_dashboard_unix_socket(loop):
+    pytest.importorskip("bokeh")
+    port = open_port()
+
+    def _connect_unix_socket(path):
+        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        try:
+            s.connect(path)
+        except Exception:
+            return False
+        else:
+            return True
+        finally:
+            s.close()
+
+    with popen(
+        [
+            sys.executable,
+            "-m",
+            "dask",
+            "scheduler",
+            "--host",
+            f"127.0.0.1:{port}",
+            "--dashboard-address",
+            "unix://",
+        ],
+    ):
+        with Client(f"tcp://127.0.0.1:{port}", loop=loop) as c:
+            proto, rest = c.dashboard_link.split("://")
+            assert proto == "http+unix"
+            assert rest.endswith(".sock/status")
+            assert ":" not in rest  # no port
+
+            path = rest.split("/status")[0]
+            assert os.path.isabs(path)
+            assert os.path.exists(path)
+            assert _connect_unix_socket(path)
+
+
 PRELOAD_TEXT = """
 _scheduler_info = {}
 

distributed/client.py

@@ -1337,16 +1337,10 @@ def dashboard_link(self):
             scheduler, info = self._get_scheduler_info(n_workers=0)
             if scheduler is None:
                 return None
-            else:
-                protocol, rest = scheduler.address.split("://")
-
-            port = info["services"]["dashboard"]
-            if protocol == "inproc":
-                host = "localhost"
-            else:
-                host = rest.split(":")[0]
 
-            return format_dashboard_link(host, port)
+            return format_dashboard_link(
+                scheduler.address, info["services"]["dashboard"]
+            )
 
     def _get_scheduler_info(self, n_workers):
         from distributed.scheduler import Scheduler

distributed/comm/__init__.py

@@ -14,13 +14,16 @@
 from distributed.comm.core import Comm, CommClosedError, connect, listen
 from distributed.comm.registry import backends
 from distributed.comm.utils import get_tcp_server_address, get_tcp_server_addresses
+from distributed.compatibility import WINDOWS
 
 
 def _register_transports():
-    from distributed.comm import inproc, tcp, ws
+    from distributed.comm import inproc, tcp, uds, ws
 
     backends["tcp"] = tcp.TCPBackend()
     backends["tls"] = tcp.TLSBackend()
+    if not WINDOWS:
+        backends["unix"] = uds.UDSBackend()
 
     try:
         # If `distributed-ucxx` is installed, it takes over the protocol="ucx" support

distributed/comm/addressing.py

@@ -14,6 +14,8 @@ def parse_address(addr: str, strict: bool = False) -> tuple[str, str]:
 
     >>> parse_address('tcp://127.0.0.1')
     ('tcp', '127.0.0.1')
+    >>> parse_address('unix:///tmp/socket')
+    ('unix', '/tmp/socket')
 
     If strict is set to true the address must have a scheme.
     """

distributed/comm/core.py

@@ -13,7 +13,11 @@
 from dask.utils import parse_timedelta
 
 from distributed.comm import registry
-from distributed.comm.addressing import get_address_host, parse_address, resolve_address
+from distributed.comm.addressing import (
+    get_address_host,
+    parse_address,
+    resolve_address,
+)
 from distributed.metrics import time
 from distributed.protocol.compression import get_compression_settings
 from distributed.protocol.pickle import HIGHEST_PROTOCOL
@@ -321,6 +325,7 @@ async def connect(
     scheme, loc = parse_address(addr)
     backend = registry.get_backend(scheme)
     connector = backend.get_connector()
+
     comm = None
 
     start = time()

distributed/comm/tcp.py

@@ -37,6 +37,7 @@
     get_tcp_server_address,
     to_frames,
 )
+from distributed.compatibility import WINDOWS
 from distributed.protocol.utils import host_array, pack_frames_prelude, unpack_frames
 from distributed.system import MEMORY_LIMIT
 from distributed.utils import ensure_ip, ensure_memoryview, get_ip, nbytes
@@ -59,7 +60,7 @@ def set_tcp_timeout(comm):
     """
     Set kernel-level TCP timeout on the stream.
     """
-    if comm.closed():
+    if comm.closed() or (not WINDOWS and comm.socket.family is socket.AF_UNIX):
         return
 
     timeout = dask.config.get("distributed.comm.timeouts.tcp")
@@ -124,7 +125,10 @@ def get_stream_address(comm):
     if comm.closed():
         raise CommClosedError()
 
-    return unparse_host_port(*comm.socket.getsockname()[:2])
+    if not WINDOWS and comm.socket.family is socket.AF_UNIX:
+        return comm.socket.getsockname() or comm.socket.getpeername()
+    else:
+        return unparse_host_port(*comm.socket.getsockname()[:2])
 
 
 def convert_stream_closed_error(obj, exc):
@@ -567,6 +571,7 @@ async def connect(self, address, deserialize=True, **connection_args):
             raise FatalCommClosedError() from err
 
         local_address = self.prefix + get_stream_address(stream)
+
         comm = self.comm_class(
             stream, local_address, self.prefix + address, deserialize
         )

distributed/comm/tests/test_comms.py

@@ -25,7 +25,7 @@
 )
 from distributed.comm.registry import backends, get_backend
 from distributed.comm.tcp import get_stream_address
-from distributed.compatibility import asyncio_run
+from distributed.compatibility import WINDOWS, asyncio_run
 from distributed.config import get_loop_factory
 from distributed.metrics import time
 from distributed.protocol import Serialized, deserialize, serialize, to_serialize
@@ -55,6 +55,17 @@ def tcp(monkeypatch, request):
     return tcp
 
 
+@pytest.fixture(params=["tornado"])
+def uds(monkeypatch, request):
+    """Set the TCP backend to either tornado or asyncio"""
+    if request.param == "tornado":
+        import distributed.comm.uds as uds
+    else:
+        raise NotImplementedError()
+    monkeypatch.setitem(backends, "uds", uds.UDSBackend())
+    return uds
+
+
 ca_file = get_cert("tls-ca-cert.pem")
 
 # The Subject field of our test certs
@@ -171,6 +182,8 @@ def test_get_address_host(tcp):
 
     assert f("tcp://127.0.0.1:123") == "127.0.0.1"
     assert f("inproc://%s/%d/123" % (get_ip(), os.getpid())) == get_ip()
+    if not WINDOWS:
+        assert f("unix:///tmp/dask.sock") == "/tmp/dask.sock"
 
 
 def test_resolve_address(tcp):
@@ -192,6 +205,9 @@ def test_resolve_address(tcp):
     assert f("tcp://localhost:456") == "tcp://127.0.0.1:456"
     assert f("tls://localhost:456") == "tls://127.0.0.1:456"
 
+    if not WINDOWS:
+        assert f("unix:///tmp/dask.sock") == "unix:///tmp/dask.sock"
+
 
 def test_get_local_address_for(tcp):
     f = get_local_address_for
@@ -277,6 +293,53 @@ async def client_communicate(key, delay=0):
     assert set(l) == {1234} | set(range(N))
 
 
+@pytest.mark.skipif(WINDOWS, reason="No unix sockets on Windows")
+@gen_test()
+async def test_uds_specific(uds):
+    """
+    Test concrete UDS API.
+    """
+
+    async def handle_comm(comm):
+        assert comm.peer_address == (f"unix://{host}:0")
+        assert comm.extra_info == {}
+        msg = await comm.read()
+        msg["op"] = "pong"
+        await comm.write(msg)
+        await comm.close()
+
+    listener = await uds.UDSListener("localhost", handle_comm)
+    host, port = listener.get_host_port()
+
+    assert host.endswith(".sock")
+    assert port == 0  # we fake port 0 when using UDS
+
+    l = []
+
+    async def client_communicate(key, delay=0):
+        comm = await connect(listener.contact_address)
+        assert comm.peer_address == f"unix://{host}:0"
+        assert comm.extra_info == {}
+        await comm.write({"op": "ping", "data": key})
+        if delay:
+            await asyncio.sleep(delay)
+        msg = await comm.read()
+        assert msg == {"op": "pong", "data": key}
+        l.append(key)
+        await comm.close()
+
+    await client_communicate(key=1234)
+
+    # Many clients at once
+    N = 100
+    futures = [client_communicate(key=i, delay=0.05) for i in range(N)]
+    await asyncio.gather(*futures)
+    assert set(l) == {1234} | set(range(N))
+
+    listener.stop()
+    assert not os.path.exists(host)  # assert socket deleted
+
+
 @pytest.mark.parametrize("sni", [None, "localhost"])
 @gen_test()
 async def test_tls_specific(tcp, sni):

distributed/comm/tests/test_uds.py

@@ -0,0 +1,23 @@
+import pytest
+
+from distributed.compatibility import WINDOWS
+
+from distributed.comm.addressing import parse_address, unparse_address
+from distributed.comm.registry import backends, get_backend
+from distributed.comm.uds import UDSBackend
+
+
+@pytest.mark.skipif(WINDOWS, reason="No unix sockets on Windows")
+def test_registered():
+    assert "unix" in backends
+    backend = get_backend("unix")
+    assert isinstance(backend, UDSBackend)
+
+
+@pytest.mark.skipif(WINDOWS, reason="No unix sockets on Windows")
+def test_parse_uds_address():
+    addr = "unix:///tmp/dask-test.sock"
+    scheme, loc = parse_address(addr)
+    assert scheme == "unix"
+    assert loc == "/tmp/dask-test.sock"
+    assert unparse_address(scheme, loc) == addr