diff --git a/distributed/cli/dask_scheduler.py b/distributed/cli/dask_scheduler.py index 7302027666..dabeea5dbc 100755 --- a/distributed/cli/dask_scheduler.py +++ b/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", diff --git a/distributed/cli/dask_worker.py b/distributed/cli/dask_worker.py index 9c89569876..c32d8338d8 100755 --- a/distributed/cli/dask_worker.py +++ b/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( diff --git a/distributed/cli/tests/test_dask_scheduler.py b/distributed/cli/tests/test_dask_scheduler.py index dc5b7b14a4..ced9f33904 100644 --- a/distributed/cli/tests/test_dask_scheduler.py +++ b/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 = {} diff --git a/distributed/client.py b/distributed/client.py index 5f14cb4911..8ae44daad4 100644 --- a/distributed/client.py +++ b/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 diff --git a/distributed/comm/__init__.py b/distributed/comm/__init__.py index 0b9801c5cd..ce5d7ac3ba 100644 --- a/distributed/comm/__init__.py +++ b/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 diff --git a/distributed/comm/addressing.py b/distributed/comm/addressing.py index f5e1660da1..9f506a0214 100644 --- a/distributed/comm/addressing.py +++ b/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. """ diff --git a/distributed/comm/core.py b/distributed/comm/core.py index be454bde03..1f83a3c139 100644 --- a/distributed/comm/core.py +++ b/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() diff --git a/distributed/comm/tcp.py b/distributed/comm/tcp.py index 1b85ede232..7a91f3b5ac 100644 --- a/distributed/comm/tcp.py +++ b/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 ) diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index 306e875ece..096d861a36 100644 --- a/distributed/comm/tests/test_comms.py +++ b/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): diff --git a/distributed/comm/tests/test_uds.py b/distributed/comm/tests/test_uds.py new file mode 100644 index 0000000000..a2c5b29e3a --- /dev/null +++ b/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 diff --git a/distributed/comm/uds.py b/distributed/comm/uds.py new file mode 100644 index 0000000000..c25e12edbc --- /dev/null +++ b/distributed/comm/uds.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import logging +import os +import socket +from typing import ClassVar + +import tornado.netutil as netutil +from tornado.tcpclient import TCPClient +from tornado.tcpserver import TCPServer + +import dask + +from distributed.comm.registry import Backend +from distributed.comm.tcp import ( + MAX_BUFFER_SIZE, + TCP, + TCPConnector, + TCPListener, +) +from distributed.utils import ( + get_uds_path, +) + +logger = logging.getLogger(__name__) + + +class UnixSocketResolver(netutil.Resolver): + """A resolver for Unix Domain Sockets. This is used by tornado to lookup hostnames. For UDS, this should always return socket type and pathname (instead of a real DNS lookup).""" + + async def resolve( + self, host: str, port: int, family: socket.AddressFamily = socket.AF_UNSPEC + ) -> list[tuple[int, str]]: + return [(socket.AF_UNIX, host)] + + +class UDSListener(TCPListener): + """A Listener for Unix Domain Sockets, based on the TCPListener class. Ensures the address is an absolute path instead of a hostname:port string.""" + + prefix = "unix://" + comm_class = TCP + + def __init__( + self, + address, + *args, + **kwargs, + ): + path = get_uds_path(address) + if ":" not in path: + path = f"{path}:0" + super().__init__(path, *args, **kwargs) # fake port 0 + + def get_host_port(self): + """ + The listening address as a tuple. Port is always 0. + """ + self._check_started() + + if self.bound_address is None: + self.bound_address = self.tcp_server._sockets[0].getsockname() + return (self.bound_address, 0) # fake port 0 + + async def _handle_stream(self, stream, address): + """ + We override the super class's _handle_stream to pass in 0 as a fake port (it will be ignored anyway). + """ + return await super()._handle_stream(stream, (self.bound_address, 0)) + + async def start(self): + self.tcp_server = TCPServer(max_buffer_size=MAX_BUFFER_SIZE, **self.server_args) + self.tcp_server.handle_stream = self._handle_stream + # When shuffling data between workers, there can + # really be O(cluster size) connection requests + # on a single worker socket, make sure the backlog + # is large enough not to lose any. + backlog = int(dask.config.get("distributed.comm.socket-backlog")) + socket = netutil.bind_unix_socket( + self.ip, # self.ip is actually the path to the socket + mode=0o600, + backlog=backlog, + ) + self.tcp_server.add_socket(socket) + self.bound_address = self.ip # ip is path to unix socket + + def stop(self): + super().stop() + if os.path.exists(self.bound_address): + try: + os.remove(self.bound_address) + except OSError as e: + logger.debug( + f"Attempted removal of socket at {self.bound_address} failed with error: {e}" + ) + + +class UDSConnector(TCPConnector): + client: ClassVar[TCPClient] = TCPClient(resolver=UnixSocketResolver()) + + prefix = "unix://" + comm_class = TCP + + +class UDSBackend(Backend): + """A Backend for Unix Domain Sockets. It overrides the TCP class's functions for parsing addresses, since UDS does not require port numbers.""" + + _connector_class = UDSConnector + _listener_class = UDSListener + + def get_connector(self): + return self._connector_class() + + def get_listener(self, loc, handle_comm, deserialize, **connection_args): + return self._listener_class(loc, handle_comm, deserialize, **connection_args) + + def get_address_host(self, loc): + path = loc.split("unix://")[-1] + if os.path.isabs(path): + return path + else: + # something like `unix://127.0.0.1:0` was passed in + # this happens when a cluster sets protocl to 'unix', but doesn't explicitly override the default host and port + # in this case, return a new uds socket path + return get_uds_path(path) + + def resolve_address(self, loc): + return loc + + def get_local_address_for(self, loc): + return loc diff --git a/distributed/deploy/local.py b/distributed/deploy/local.py index e090f6649e..a090009222 100644 --- a/distributed/deploy/local.py +++ b/distributed/deploy/local.py @@ -60,6 +60,8 @@ class LocalCluster(SpecCluster): 'localhost:8787' or '0.0.0.0:8787'. Defaults to ':8787'. Set to ``None`` to disable the dashboard. Use ':0' for a random port. + Set to an absolute filesystem path to listen on a Unix Domain Socket at that path. + Set to 'unix://' to listen on a Unix socket with a random name, in the location configured by dask.config['temporary-directory'] or a default temporary directory. When specifying only a port like ':8787', the dashboard will bind to the given interface from the ``host`` parameter. If ``host`` is empty, binding will occur on all interfaces '0.0.0.0'. To avoid firewall issues when deploying locally, set ``host`` to 'localhost'. diff --git a/distributed/node.py b/distributed/node.py index 9f9a731f51..500223ee96 100644 --- a/distributed/node.py +++ b/distributed/node.py @@ -1,12 +1,14 @@ from __future__ import annotations import logging +import os import ssl import warnings import weakref from contextlib import suppress import tlz +import tornado.netutil as netutil from tornado.httpserver import HTTPServer import dask @@ -160,20 +162,34 @@ def start_http_server( change_port = False retries_left = 3 - while True: - try: - if not change_port: - self.http_server.listen(**http_address) - else: - self.http_server.listen(**tlz.merge(http_address, {"port": 0})) - break - except Exception: - change_port = True - retries_left = retries_left - 1 - if retries_left < 1: - raise - - bound_addresses = get_tcp_server_addresses(self.http_server) + + if os.path.isabs(http_address["address"]): # unix socket + try: # remove any old sockets + os.remove(http_address["address"]) + except OSError: + pass + dashboard_socket = netutil.bind_unix_socket( + http_address["address"], mode=0o600 + ) + self.http_server.add_socket(dashboard_socket) + bound_addresses = [(f"unix://{http_address['address']}", 0)] + else: + while True: + try: + if not change_port: + self.http_server.listen(**http_address) + else: + self.http_server.listen( + **tlz.merge(http_address, {"port": 0}) + ) + break + except Exception: + change_port = True + retries_left = retries_left - 1 + if retries_left < 1: + raise + + bound_addresses = get_tcp_server_addresses(self.http_server) # If more than one address is configured we just use the first here # Socket addresses representation: https://docs.python.org/3/library/socket.html#socket-families diff --git a/distributed/objects.py b/distributed/objects.py index 53b7b91565..eddc386fbe 100644 --- a/distributed/objects.py +++ b/distributed/objects.py @@ -7,8 +7,6 @@ from dask.widgets import get_environment, get_template -from distributed.utils import format_dashboard_link - class HasWhat(dict): """A dictionary of all workers and which keys that worker has.""" @@ -35,7 +33,7 @@ def _format_dashboard_address(server): if "host" in server else urlparse(server["address"]).hostname ) - return format_dashboard_link(host, server["services"]["dashboard"]) + return host, server["services"]["dashboard"] except KeyError: return None diff --git a/distributed/scheduler.py b/distributed/scheduler.py index 8f4ba87da3..85c2a9df16 100644 --- a/distributed/scheduler.py +++ b/distributed/scheduler.py @@ -12,6 +12,7 @@ import os import pickle import random +import socket import textwrap import uuid import warnings @@ -22,6 +23,7 @@ Callable, Collection, Container, + Generator, Hashable, Iterable, Iterator, @@ -94,6 +96,7 @@ unparse_host_port, ) from distributed.comm.addressing import addresses_from_user_args +from distributed.compatibility import WINDOWS from distributed.core import ( ErrorMessage, OKMessage, @@ -4279,13 +4282,24 @@ def _repr_html_(self) -> str: def identity(self, n_workers: int = -1) -> dict[str, Any]: """Basic information about ourselves and our cluster""" + + def _map_services_to_uri(services: dict) -> Generator[tuple[str, int | str]]: + """the services dict contains k/v-pairs where v is either an int port number, of a unix socket url""" + for k, v in services.items(): + s = list(v._sockets.values())[0] + if not WINDOWS and s.family is socket.AF_UNIX: + yield k, f"unix://{s.getsockname()}" + else: + yield k, v.port + if n_workers == -1: n_workers = len(self.workers) + d = { "type": type(self).__name__, "id": str(self.id), "address": self.address, - "services": {key: v.port for (key, v) in self.services.items()}, + "services": {k: v for k, v in _map_services_to_uri(self.services)}, "started": self.time_started, "n_workers": len(self.workers), "total_threads": self.total_nthreads, diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py index 1ebeafe299..5810145e70 100644 --- a/distributed/tests/test_utils.py +++ b/distributed/tests/test_utils.py @@ -8,6 +8,7 @@ import multiprocessing import os import queue +import shutil import socket import sys import traceback @@ -42,6 +43,7 @@ get_ip_interface, get_mp_context, get_traceback, + get_uds_path, is_kernel, is_valid_xml, iscoroutinefunction, @@ -630,6 +632,12 @@ def test_format_dashboard_link(): assert "host" in format_dashboard_link("host", 1234) assert "1234" in format_dashboard_link("host", 1234) + if not WINDOWS: + assert ( + format_dashboard_link("localhost", "/tmp/dashboard.sock") + == "http+unix:///tmp/dashboard.sock/status" + ) + try: os.environ["host"] = "hello" assert "hello" not in format_dashboard_link("host", 1234) @@ -1086,3 +1094,60 @@ def test_tuple_comparable_eq(obj1, obj2, expected): def test_tuple_comparable_error(): with pytest.raises(ValueError): TupleComparable("string") + + +@pytest.mark.parametrize( + "env_var", + [ + "XDG_RUNTIME_DIR", + "DASK_TEMPORARY_DIRECTORY", + None, + ], +) +@pytest.mark.parametrize( + "arg,expectation", + [ + ["", None], + ["localhost", None], + ["unix://", None], + ["unix:///dask.sock", "/dask.sock"], + ["/dask.sock", "/dask.sock"], + ], +) +@pytest.mark.skipif(WINDOWS, reason="No unix sockets on Windows") +def test_get_uds_path(request, tmp_path, env_var, arg, expectation, monkeypatch): + if env_var != "XDG_RUNTIME_DIR": + # on Ubuntu this env var is always set, so we need to unset it to test the other options + monkeypatch.setenv("XDG_RUNTIME_DIR", "") + + if env_var: + base_path = str(tmp_path) + monkeypatch.setenv(env_var, base_path) + elif MACOS: + # on MACOS we don't use tempfile.gettempdir + base_path = f"/tmp/dask-{os.environ.get('USER')}" + else: + # a path will be created using tempfile.gettempdir + base_path = f"/tmp/pytest/{request.node.name}" + import tempfile + + monkeypatch.setattr(tempfile, "gettempdir", lambda: base_path) + + if env_var == "DASK_TEMPORARY_DIRECTORY": + # hack: setting this env var doesn't have effect within this test because the dask config is already loaded at this point + monkeypatch.setitem(dask.config.config, "temporary-directory", base_path) + + arg = arg.replace("", base_path) + expectation = expectation.replace("", base_path) if expectation else "" + + try: + result = get_uds_path(arg) + if os.path.isabs(str(expectation)): + assert result == expectation + else: + # we don't know what the socket filename is, but it should live under result_base_path + assert result.startswith(f"{base_path}/dask_run") + assert result.endswith(".sock") + assert os.path.exists(os.path.dirname(result)) + finally: + shutil.rmtree(base_path, ignore_errors=True) diff --git a/distributed/utils.py b/distributed/utils.py index 8fc57cd379..2779094fe7 100644 --- a/distributed/utils.py +++ b/distributed/utils.py @@ -52,7 +52,7 @@ import tblib.pickling_support from tornado import escape -from distributed.compatibility import asyncio_run +from distributed.compatibility import MACOS, asyncio_run from distributed.config import get_loop_factory try: @@ -1267,12 +1267,37 @@ def warn_on_duration(duration: str | float | timedelta, msg: str) -> Generator[N warnings.warn(msg.format(duration=diff), stacklevel=2) -def format_dashboard_link(host, port): - template = dask.config.get("distributed.dashboard.link") - if dask.config.get("distributed.scheduler.dashboard.tls.cert"): - scheme = "https" - else: - scheme = "http" +def format_dashboard_link(process_address: str, port_or_uri: str | int) -> str: + """Return a formatted link to the dashboard based on: + `process_address`: the address of the process (e.g. Scheduler) that is calling this. + `port_or_uri`: value of an entry in a Client, Cluster, or Scheduler's `services` list. Can be either a port, in which case the hostname is assumed to be the same as the hostname for `process_address`. Or can be an absolute path for a service that listens on a Unix Domain Socket. + """ + if isinstance(port_or_uri, str) and ( + port_or_uri.startswith("unix://") or os.path.isabs(port_or_uri) + ): # unix socket + host = port_or_uri[7:] if port_or_uri.startswith("unix://") else port_or_uri + port = None + scheme = "http+unix" + template = "{scheme}://{host}/status" + else: # no unix socket + port = port_or_uri + try: + protocol, rest = process_address.split("://") + except ValueError: # no protocol prefix given + protocol, rest = "", process_address + + if protocol == "inproc": + host = "localhost" + else: + host = rest.split(":")[0] + + if dask.config.get("distributed.scheduler.dashboard.tls.cert"): + scheme = "https" + else: + scheme = "http" + + template = dask.config.get("distributed.dashboard.link") + return template.format( **toolz.merge(os.environ, dict(scheme=scheme, host=host, port=port)) ) @@ -1522,8 +1547,15 @@ def clean_dashboard_address(addrs: AnyType, default_listen_ip: str = "") -> list [{'address': '', 'port': 8787}, {'address': '', 'port': 8887}] >>> clean_dashboard_address(":8787,:8887") [{'address': '', 'port': 8787}, {'address': '', 'port': 8887}] + >>> clean_dashboard_address("/tmp/dashboard.sock") + [{'address': '/tmp/dashboard.sock', 'port': 0}] + >>> clean_dashboard_address("unix://") + [{'address': '/tmp/.sock', 'port': 0}] """ + if isinstance(addrs, str) and (addrs.startswith("unix://") or os.path.isabs(addrs)): + return [{"address": get_uds_path(addrs), "port": 0}] + if default_listen_ip == "0.0.0.0": default_listen_ip = "" # for IPV6 @@ -1894,3 +1926,38 @@ def url_escape(url, *args, **kwargs): Escape a URL path segment. Cache results for better performance. """ return escape.url_escape(url, *args, **kwargs) + + +def get_uds_path(address: str | None = None) -> str: + """ + Take an address for a Unix Domain Socket and return an absolute path. + If address is already an absolute path (possibly prefixed by unix://), return it. Underlying directories will not be created. + In all other cases, generate a random filename in 'dask_run' under one of these locations: + 1. $XDG_RUNTIME_DIR (set; will be created) + 2. dask.config["temporary-directory"] (if set; will be created) + 3. tempfile.gettempdir() (will be created) + """ + addr = str(address) + if addr.startswith("unix://"): + addr = addr[7:] + if os.path.isabs(addr): + return addr + else: + xdg_dir = os.environ.get("XDG_RUNTIME_DIR", False) + if xdg_dir: + base_path = f"{xdg_dir}" + else: + base_path = dask.config.get("temporary-directory") + if not base_path: + if MACOS: + # MacOS throws `OSError: AF_UNIX path too long` with tempfile.gettempdir(), so use a hardcoded path under /tmp/ + base_path = f"/tmp/dask-{os.environ.get('USER')}" + else: + base_path = tempfile.gettempdir() + + base_path = f"{base_path}/dask_run" + os.makedirs(base_path, mode=0o700, exist_ok=True) + + import secrets # use token_hex to generate a random filename + + return os.path.join(base_path, f"{secrets.token_hex()}.sock") diff --git a/docs/source/communications.rst b/docs/source/communications.rst index b9406e6b3c..3298536f25 100644 --- a/docs/source/communications.rst +++ b/docs/source/communications.rst @@ -39,6 +39,8 @@ source tree: communication between endpoints as long as they are situated in the same process. +* ``unix`` is a transport using Unix Domain Sockets. Unix sockets bypass the system's network layer and are therefore more efficient than TCP sockets. Moreover, they do not use ports but instead listen at filesystem paths, and can be secured using normal file permissions. This makes them a safe choice in multi-user environments, without needing TLS. (Not available on Windows.) + Some URIs may be valid for listening but not for connecting. For example, the URI ``tcp://`` will listen on all IPv4 and IPv6 addresses and on an arbitrary port, but you cannot connect to that address.