def test_register_backend_entrypoint():
# Code adapted from pandas backend entry point testing
# https://github.com/pandas-dev/pandas/blob/2470690b9f0826a8feb426927694fa3500c3e8d2/pandas/tests/plotting/test_backend.py#L50-L76
dist = pkg_resources.get_distribution("distributed")
if dist.module_path not in distributed.__file__:
# We are running from a non-installed distributed, and this test is invalid
pytest.skip("Testing a non-installed distributed")
mod = types.ModuleType("dask_udp")
mod.UDPBackend = lambda: 1
sys.modules[mod.__name__] = mod
entry_point_name = "distributed.comm.backends"
backends_entry_map = pkg_resources.get_entry_map("distributed")
if entry_point_name not in backends_entry_map:
backends_entry_map[entry_point_name] = dict()
backends_entry_map[entry_point_name]["udp"] = pkg_resources.EntryPoint(
"udp", mod.__name__, attrs=["UDPBackend"], dist=dist)
# The require is disabled here since particularly unit tests may install
# dirty or dev versions which are conflicting with backend entrypoints if
# they are demanding for exact, stable versions. This should not fail the
# test
result = get_backend("udp", require=False)
assert result == 1
def resolve_address(addr: str) -> str:
"""
Apply scheme-specific address resolution to *addr*, replacing
all symbolic references with concrete location specifiers.
In practice, this can mean hostnames are resolved to IP addresses.
>>> resolve_address('tcp://localhost:8786')
'tcp://127.0.0.1:8786'
"""
scheme, loc = parse_address(addr)
backend = registry.get_backend(scheme)
return unparse_address(scheme, backend.resolve_address(loc))
async def test_inproc_handshakes_concurrently():
async def handle_comm():
pass
async with listen("inproc://", handle_comm) as listener:
addr = listener.listen_address
scheme, loc = parse_address(addr)
connector = get_backend(scheme).get_connector()
comm1 = await connector.connect(loc)
comm2 = await connector.connect(loc)
await comm1.close()
await comm2.close()
def get_address_host(addr: str) -> str:
"""
Return a hostname / IP address identifying the machine this address
is located on.
In contrast to get_address_host_port(), this function should always
succeed for well-formed addresses.
>>> get_address_host('tcp://1.2.3.4:80')
'1.2.3.4'
"""
scheme, loc = parse_address(addr)
backend = registry.get_backend(scheme)
return backend.get_address_host(loc)
async def test_inproc_continues_listening_after_handshake_error():
async def handle_comm():
pass
async with listen("inproc://", handle_comm) as listener:
addr = listener.listen_address
scheme, loc = parse_address(addr)
connector = get_backend(scheme).get_connector()
comm = await connector.connect(loc)
await comm.close()
comm = await connector.connect(loc)
await comm.close()
def get_local_address_for(addr: str) -> str:
"""
Get a local listening address suitable for reaching *addr*.
For instance, trying to reach an external TCP address will return
a local TCP address that's routable to that external address.
>>> get_local_address_for('tcp://8.8.8.8:1234')
'tcp://192.168.1.68'
>>> get_local_address_for('tcp://127.0.0.1:1234')
'tcp://127.0.0.1'
"""
scheme, loc = parse_address(addr)
backend = registry.get_backend(scheme)
return unparse_address(scheme, backend.get_local_address_for(loc))
def get_address_host_port(addr: str, strict: bool = False) -> tuple[str, int]:
"""
Get a (host, port) tuple out of the given address.
For definition of strict check parse_address
ValueError is raised if the address scheme doesn't allow extracting
the requested information.
>>> get_address_host_port('tcp://1.2.3.4:80')
('1.2.3.4', 80)
>>> get_address_host_port('tcp://[::1]:80')
('::1', 80)
"""
scheme, loc = parse_address(addr, strict=strict)
backend = registry.get_backend(scheme)
try:
return backend.get_address_host_port(loc)
except NotImplementedError:
raise ValueError(
f"don't know how to extract host and port for address {addr!r}"
)
def listen(addr, handle_comm, deserialize=True, **kwargs):
"""
Create a listener object with the given parameters. When its ``start()``
method is called, the listener will listen on the given address
(a URI such as ``tcp://0.0.0.0``) and call *handle_comm* with a
``Comm`` object for each incoming connection.
*handle_comm* can be a regular function or a coroutine.
"""
try:
scheme, loc = parse_address(addr, strict=True)
except ValueError:
if kwargs.get("ssl_context"):
addr = "tls://" + addr
else:
addr = "tcp://" + addr
scheme, loc = parse_address(addr, strict=True)
backend = registry.get_backend(scheme)
return backend.get_listener(loc, handle_comm, deserialize, **kwargs)
def test_register_backend_entrypoint():
# Code adapted from pandas backend entry point testing
# https://github.com/pandas-dev/pandas/blob/2470690b9f0826a8feb426927694fa3500c3e8d2/pandas/tests/plotting/test_backend.py#L50-L76
dist = pkg_resources.get_distribution("distributed")
if dist.module_path not in distributed.__file__:
# We are running from a non-installed distributed, and this test is invalid
pytest.skip("Testing a non-installed distributed")
mod = types.ModuleType("dask_udp")
mod.UDPBackend = lambda: 1
sys.modules[mod.__name__] = mod
entry_point_name = "distributed.comm.backends"
backends_entry_map = pkg_resources.get_entry_map("distributed")
if entry_point_name not in backends_entry_map:
backends_entry_map[entry_point_name] = dict()
backends_entry_map[entry_point_name]["udp"] = pkg_resources.EntryPoint(
"udp", mod.__name__, attrs=["UDPBackend"], dist=dist)
result = get_backend("udp")
assert result == 1
def test_registered():
assert "ucx" in backends
backend = get_backend("ucx")
assert isinstance(backend, ucx.UCXBackend)
def test_registered():
assert "ws" in backends
backend = get_backend("ws")
assert isinstance(backend, ws.WSBackend)
def _get_backend_on_path(path):
sys.path.append(os.fsdecode(path))
return get_backend("udp")
async def connect(
addr, timeout=None, deserialize=True, handshake_overrides=None, **connection_args
):
"""
Connect to the given address (a URI such as ``tcp://127.0.0.1:1234``)
and yield a ``Comm`` object. If the connection attempt fails, it is
retried until the *timeout* is expired.
"""
if timeout is None:
timeout = dask.config.get("distributed.comm.timeouts.connect")
timeout = parse_timedelta(timeout, default="seconds")
scheme, loc = parse_address(addr)
backend = registry.get_backend(scheme)
connector = backend.get_connector()
comm = None
start = time()
def time_left():
deadline = start + timeout
return max(0, deadline - time())
backoff_base = 0.01
attempt = 0
# Prefer multiple small attempts than one long attempt. This should protect
# primarily from DNS race conditions
# gh3104, gh4176, gh4167
intermediate_cap = timeout / 5
active_exception = None
while time_left() > 0:
try:
comm = await asyncio.wait_for(
connector.connect(loc, deserialize=deserialize, **connection_args),
timeout=min(intermediate_cap, time_left()),
)
break
except FatalCommClosedError:
raise
# Note: CommClosed inherits from OSError
except (asyncio.TimeoutError, OSError) as exc:
active_exception = exc
# As descibed above, the intermediate timeout is used to distributed
# initial, bulk connect attempts homogeneously. In particular with
# the jitter upon retries we should not be worred about overloading
# any more DNS servers
intermediate_cap = timeout
# FullJitter see https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
upper_cap = min(time_left(), backoff_base * (2**attempt))
backoff = random.uniform(0, upper_cap)
attempt += 1
logger.debug(
"Could not connect to %s, waiting for %s before retrying", loc, backoff
)
await asyncio.sleep(backoff)
else:
raise OSError(
f"Timed out trying to connect to {addr} after {timeout} s"
) from active_exception
local_info = {
**comm.handshake_info(),
**(handshake_overrides or {}),
}
try:
# This would be better, but connections leak if worker is closed quickly
# write, handshake = await asyncio.gather(comm.write(local_info), comm.read())
handshake = await asyncio.wait_for(comm.read(), time_left())
await asyncio.wait_for(comm.write(local_info), time_left())
except Exception as exc:
with suppress(Exception):
await comm.close()
raise OSError(
f"Timed out during handshake while connecting to {addr} after {timeout} s"
) from exc
comm.remote_info = handshake
comm.remote_info["address"] = comm._peer_addr
comm.local_info = local_info
comm.local_info["address"] = comm._local_addr
comm.handshake_options = comm.handshake_configuration(
comm.local_info, comm.remote_info
)
return comm
