diff --git a/distributed/_submission_permit_client.py b/distributed/_submission_permit_client.py new file mode 100644 index 0000000000..20dcdccaf6 --- /dev/null +++ b/distributed/_submission_permit_client.py @@ -0,0 +1,551 @@ +"""Explicit client-side submission permits for a single Dask collection. + +This is deliberately a private, opt-in API. It holds a graph while it is +prepared, then holds the operation's Futures until scheduler admission. + +The caller supplies finite lease/network limits and a clock-rate bound plus a +safety margin. Timing validity assumes continuously advancing client/server +clocks whose relative rate respects that bound; it is not a guarantee for an +arbitrarily suspended or misconfigured clock. No renewal or fallback is used. + +Optimization and serialization keep the ordinary Client calling thread. Async +cancellation is cooperative and cannot interrupt that synchronous preparation. +After a dispatch attempt, cancellation or an indeterminate error can leave work +running: the graph is never retransmitted and an abort RPC is never issued. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from contextvars import ContextVar +from math import isfinite +from time import monotonic +from typing import TYPE_CHECKING, Any + +import dask + +if TYPE_CHECKING: + from distributed.client import Client, Future + + +class SubmissionPermitClientError(RuntimeError): + """Base class for client-side submission permit failures.""" + + +class SubmissionPermitUnsupportedError(SubmissionPermitClientError): + """The connected scheduler did not advertise a compatible permit protocol.""" + + +class SubmissionPermitExpiredError(SubmissionPermitClientError): + """The conservative local permit interval ended before admission.""" + + +class SubmissionPermitRejectedError(SubmissionPermitClientError): + """The scheduler rejected the tagged graph.""" + + +class SubmissionPermitIndeterminateError(SubmissionPermitClientError): + """The graph may have been sent, but its admission outcome was not observed.""" + + +# Short names make these useful in small caller-side error handlers too. +PermitUnsupportedError = SubmissionPermitUnsupportedError +PermitExpiredError = SubmissionPermitExpiredError +PermitRejectedError = SubmissionPermitRejectedError +PermitIndeterminateError = SubmissionPermitIndeterminateError + + +_current_submission: ContextVar[SubmissionPermitOperation | None] = ContextVar( + "current_submission", default=None +) + + +def _finite( + name: str, value: object, *, minimum: float, inclusive: bool = False +) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not isfinite(value) + ): + raise ValueError(f"{name} must be a finite number") + value = float(value) + if value < minimum or (not inclusive and value == minimum): + comparison = "at least" if inclusive else "greater than" + raise ValueError(f"{name} must be {comparison} {minimum}") + return value + + +class SubmissionPermitOperation: + """The one graph and Futures owned by one protected submission.""" + + def __init__( + self, + *, + duration: float, + timeout: float, + max_clock_rate: float, + clock_margin: float, + clock: Callable[[], float], + ) -> None: + self.duration = duration + self.timeout = timeout + self.max_clock_rate = max_clock_rate + self.clock_margin = clock_margin + self.clock = clock + self.started: float | None = None + self._last_clock: float | None = None + self.client: Client | None = None + self.generation: int | None = None + self.carrier: Any = None + self.capabilities: dict[str, Any] | None = None + self.message: dict[str, Any] | None = None + self.futures: list[Future] = [] + self.graph_started = False + self.sequence: int | None = None + self.epoch: str | None = None + self.granted_duration: float | None = None + self.dispatch_started = False + + def begin_graph(self, client: Client) -> None: + if self.client is not client or self.graph_started: + raise RuntimeError("a protected submission may capture exactly one graph") + self.graph_started = True + + def own(self, future: Future) -> None: + self.futures.append(future) + + def capture(self, message: dict[str, Any]) -> None: + if self.message is not None: + raise RuntimeError("a protected submission may capture exactly one graph") + self.message = message + + def remaining(self) -> float: + if self.granted_duration is None or self.started is None: + raise RuntimeError("submission permit was not acquired") + return ( + self.granted_duration + - self.max_clock_rate * (self._clock_now() - self.started) + - self.clock_margin + ) + + def _clock_now(self) -> float: + now = self.clock() + if ( + isinstance(now, bool) + or not isinstance(now, (int, float)) + or not isfinite(now) + ): + raise SubmissionPermitRejectedError("clock returned a non-finite value") + now = float(now) + if self._last_clock is not None and now < self._last_clock: + raise SubmissionPermitRejectedError("clock moved backwards") + self._last_clock = now + return now + + def ensure_valid(self) -> None: + if self.remaining() <= 0: + raise SubmissionPermitExpiredError("submission permit expired locally") + + def ensure_origin(self, client: Client) -> None: + capabilities = client._submission_permit_capabilities + if ( + self.client is not client + or self.generation != client.generation + or self.carrier is not client.scheduler_comm + or self.carrier is None + or self.carrier.closed() + or client.status != "running" + or capabilities is not self.capabilities + or capabilities is None + or capabilities.get("epoch") != self.epoch + ): + raise SubmissionPermitRejectedError( + "submission permit connection changed before graph dispatch" + ) + + async def acquire(self, client: Client) -> None: + if self.sequence is not None: + raise RuntimeError("a submission operation cannot acquire another permit") + capabilities = client._submission_permit_capabilities + if ( + not isinstance(capabilities, dict) + or capabilities.get("version") != 1 + or not isinstance(capabilities.get("epoch"), str) + ): + raise SubmissionPermitUnsupportedError( + "scheduler does not support submission permits" + ) + epoch = capabilities["epoch"] + maximum = capabilities.get("max_duration") + if ( + isinstance(maximum, bool) + or not isinstance(maximum, (int, float)) + or not isfinite(maximum) + or maximum <= 0 + or self.duration > maximum + ): + raise SubmissionPermitUnsupportedError( + "requested permit duration is not supported" + ) + + async with client._submission_permit_acquire_lock: + # Capabilities can change while another acquisition owns the lock. + if client._submission_permit_capabilities is not capabilities: + raise SubmissionPermitRejectedError( + "submission permit connection changed during acquire" + ) + client._submission_permit_sequence += 1 + sequence = client._submission_permit_sequence + self.client = client + self.generation = client.generation + self.carrier = client.scheduler_comm + self.capabilities = capabilities + self.epoch = epoch + self.sequence = sequence + self.started = self._clock_now() + self.ensure_origin(client) + connection_changed = client._submission_permit_changed + request = asyncio.create_task( + client.scheduler.submission_permit_acquire( + client=client.id, + epoch=epoch, + sequence=sequence, + duration=self.duration, + ) + ) + changed = asyncio.create_task(connection_changed.wait()) + try: + done, _ = await asyncio.wait( + (request, changed), + timeout=self.timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if connection_changed.is_set(): + raise SubmissionPermitRejectedError( + "connection changed before the grant was received" + ) + if request not in done: + raise asyncio.TimeoutError( + "timed out waiting for a submission permit" + ) + reply = request.result() + finally: + request.cancel() + changed.cancel() + await asyncio.gather(request, changed, return_exceptions=True) + + if ( + not isinstance(reply, dict) + or reply.get("sequence") != sequence + or reply.get("state") != "pending" + or isinstance(reply.get("duration"), bool) + or not isinstance(reply.get("duration"), (int, float)) + or not isfinite(reply["duration"]) + or reply["duration"] <= 0 + or reply["duration"] > self.duration + ): + await self.abort(client) + raise SubmissionPermitRejectedError( + "scheduler did not grant a pending submission permit" + ) + self.granted_duration = float(reply["duration"]) + try: + self.ensure_origin(client) + self.ensure_valid() + except BaseException: + await self.abort(client) + raise + + async def abort(self, client: Client) -> None: + if self.dispatch_started or self.epoch is None or self.sequence is None: + return + try: + await asyncio.wait_for( + client.scheduler.submission_permit_abort( + client=client.id, epoch=self.epoch, sequence=self.sequence + ), + timeout=self.timeout, + ) + except BaseException: + # An acquire response can have been lost; cleanup is best effort only. + pass + + async def commit(self, client: Client) -> None: + if self.message is None or self.epoch is None or self.sequence is None: + raise SubmissionPermitRejectedError("protected graph was not captured") + self.ensure_origin(client) + self.ensure_valid() + message = dict( + self.message, submission_epoch=self.epoch, submission_sequence=self.sequence + ) + key = (self.epoch, self.sequence) + waiter = asyncio.get_running_loop().create_future() + if key in client._submission_permit_pending: + raise RuntimeError("duplicate submission permit admission waiter") + client._submission_permit_pending[key] = waiter + try: + # send may enqueue before raising or being interrupted. From this + # point onward, never race an abort RPC against a possibly queued graph. + self.dispatch_started = True + self.carrier.send(message) + try: + reply = await asyncio.wait_for( + asyncio.shield(waiter), timeout=self.timeout + ) + except asyncio.TimeoutError as exc: + raise SubmissionPermitIndeterminateError( + "timed out waiting for graph admission" + ) from exc + except asyncio.CancelledError: + raise + except SubmissionPermitClientError: + raise + except Exception as exc: + raise SubmissionPermitIndeterminateError( + "graph dispatch or its admission response was interrupted" + ) from exc + finally: + if client._submission_permit_pending.get(key) is waiter: + del client._submission_permit_pending[key] + if not waiter.done(): + waiter.cancel() + elif not waiter.cancelled(): + waiter.exception() + if ( + reply.get("status") != "accepted" + or reply.get("epoch") != self.epoch + or reply.get("sequence") != self.sequence + ): + raise SubmissionPermitRejectedError( + reply.get("reason") or "scheduler rejected protected graph" + ) + if ( + self.client is not client + or self.generation != client.generation + or self.carrier is not client.scheduler_comm + or self.carrier is None + or self.carrier.closed() + or client.status != "running" + or client._submission_permit_capabilities is not self.capabilities + or client._submission_permit_capabilities is None + or client._submission_permit_capabilities.get("epoch") != self.epoch + ): + raise SubmissionPermitIndeterminateError( + "connection changed after graph admission" + ) + + def release_owned(self) -> None: + for future in self.futures: + try: + future.release() + except BaseException: + pass + self.futures.clear() + self.message = None + + def publish(self) -> None: + self.futures.clear() + self.message = None + + async def cleanup(self, client: Client) -> None: + # Future.release queues refcount work onto this same event loop. + await asyncio.sleep(0) + + +def _operation( + *, + duration: object, + timeout: object, + max_clock_rate: object, + clock_margin: object, + clock: Callable[[], float], +) -> SubmissionPermitOperation: + if not callable(clock): + raise TypeError("clock must be callable") + return SubmissionPermitOperation( + duration=_finite("duration", duration, minimum=0), + timeout=_finite("timeout", timeout, minimum=0), + max_clock_rate=_finite( + "max_clock_rate", max_clock_rate, minimum=1, inclusive=True + ), + clock_margin=_finite("clock_margin", clock_margin, minimum=0, inclusive=True), + clock=clock, + ) + + +def _prepare( + client: Client, + operation: SubmissionPermitOperation, + collection: Any, + persist: bool, + kwargs: dict[str, Any], +) -> Any: + token = _current_submission.set(operation) + try: + result = ( + client.persist(collection, **kwargs) + if persist + else client.compute(collection, sync=False, **kwargs) + ) + finally: + _current_submission.reset(token) + if operation.message is None: + raise RuntimeError("protected submission did not capture a graph") + return result + + +async def _run( + client: Client, + operation: SubmissionPermitOperation, + collection: Any, + persist: bool, + kwargs: dict[str, Any], +) -> Any: + try: + await operation.acquire(client) + operation.ensure_origin(client) + operation.ensure_valid() + result = _prepare(client, operation, collection, persist, kwargs) + # Preparation is synchronous; allow a queued cancellation to win before + # the carrier is allowed to dispatch the graph. + await asyncio.sleep(0) + operation.ensure_valid() + await operation.commit(client) + operation.publish() + return result + except BaseException: + try: + operation.release_owned() + except BaseException: + pass + if not operation.dispatch_started: + try: + await operation.abort(client) + except BaseException: + pass + try: + await operation.cleanup(client) + except BaseException: + pass + raise + + +def _protected( + client: Client, + collection: Any, + *, + persist: bool, + duration: object, + timeout: object, + max_clock_rate: object, + clock_margin: object, + clock: Callable[[], float], + kwargs: dict[str, Any], +) -> Any: + if not dask.is_dask_collection(collection): + raise TypeError("protected submission requires one Dask collection") + if _current_submission.get() is not None: + raise RuntimeError("nested protected submissions are not supported") + operation = _operation( + duration=duration, + timeout=timeout, + max_clock_rate=max_clock_rate, + clock_margin=clock_margin, + clock=clock, + ) + if client.asynchronous: + return _run(client, operation, collection, persist, kwargs) + # Keep graph optimization and serialization in the calling thread. Only + # scheduler RPCs and the direct carrier dispatch run on Client's loop. + try: + client.sync(operation.acquire, client) + operation.ensure_origin(client) + operation.ensure_valid() + result = _prepare(client, operation, collection, persist, kwargs) + operation.ensure_valid() + client.sync(operation.commit, client) + operation.publish() + return result + except BaseException: + try: + operation.release_owned() + except BaseException: + pass + if not operation.dispatch_started: + try: + client.sync(operation.abort, client) + except BaseException: + pass + try: + client.sync(operation.cleanup, client) + except BaseException: + pass + raise + + +def protected_compute( + client: Client, + collection: Any, + *, + duration: float, + timeout: float, + max_clock_rate: float, + clock_margin: float, + clock: Callable[[], float] = monotonic, + **kwargs: Any, +) -> Any: + """Return one collection's Future after acknowledged protected submission. + + For an asynchronous Client, await this function to obtain the Future, then + await that Future for the computation result. For a synchronous Client, + this function blocks through admission and returns the ordinary Future. + ``sync=True`` and nested protected submissions are unsupported. + + ``timeout`` bounds each acquisition/admission/cleanup network phase; it + does not interrupt graph preparation. ``duration`` is the requested server + lease in seconds. The local usable interval subtracts ``clock_margin`` and + ``max_clock_rate`` times elapsed client-clock time from the granted lease. + All limits are explicit experimental choices, without production defaults. + """ + return _protected( + client, + collection, + persist=False, + duration=duration, + timeout=timeout, + max_clock_rate=max_clock_rate, + clock_margin=clock_margin, + clock=clock, + kwargs=kwargs, + ) + + +def protected_persist( + client: Client, + collection: Any, + *, + duration: float, + timeout: float, + max_clock_rate: float, + clock_margin: float, + clock: Callable[[], float] = monotonic, + **kwargs: Any, +) -> Any: + """Return one persisted collection after acknowledged protected submission. + + The mode, timing and failure semantics match :func:`protected_compute`. + Other Client calls remain unprotected before graph dispatch. + """ + return _protected( + client, + collection, + persist=True, + duration=duration, + timeout=timeout, + max_clock_rate=max_clock_rate, + clock_margin=clock_margin, + clock=clock, + kwargs=kwargs, + ) diff --git a/distributed/_submission_permit_extension.py b/distributed/_submission_permit_extension.py new file mode 100644 index 0000000000..5b7bb87483 --- /dev/null +++ b/distributed/_submission_permit_extension.py @@ -0,0 +1,101 @@ +"""Explicitly enabled prototype of scheduler-side submission protection. + +This extension is not part of DEFAULT_EXTENSIONS and does not change Client APIs. +Its limits must be supplied explicitly while the public contract is evaluated. +""" + +from __future__ import annotations + +from collections.abc import Callable +from time import monotonic +from typing import TYPE_CHECKING + +from distributed._submission_permits import ( + ClosedPermitError, + SubmissionPermitRegistry, +) +from distributed.core import Status + +if TYPE_CHECKING: + from distributed.scheduler import Scheduler + + +class SubmissionPermitExtension: + def __init__( + self, + scheduler: Scheduler, + *, + max_duration: float, + max_pending_per_client: int, + max_pending: int, + max_outcomes_per_client: int, + clock: Callable[[], float] = monotonic, + ) -> None: + self.scheduler = scheduler + self.registry = SubmissionPermitRegistry( + max_duration=max_duration, + max_pending_per_client=max_pending_per_client, + max_pending=max_pending, + max_outcomes_per_client=max_outcomes_per_client, + clock=clock, + ) + scheduler.handlers.update( + { + "submission_permit_acquire": self.acquire, + "submission_permit_status": self.status, + "submission_permit_abort": self.abort, + } + ) + + def register_client(self, client: str) -> str: + return self.registry.register(client) + + def unregister_client(self, client: str, epoch: str) -> None: + self.registry.unregister(client, epoch) + + def capabilities(self, epoch: str) -> dict[str, int | float | str]: + return { + "version": 1, + "epoch": epoch, + "max_duration": self.registry.max_duration, + } + + def acquire( + self, client: str, epoch: str, sequence: int, duration: float + ) -> dict[str, int | float | str]: + comm = self.scheduler.client_comms.get(client) + if self.scheduler.status != Status.running or comm is None or comm.closed(): + raise ClosedPermitError("client connection or scheduler is not running") + permit = self.registry.acquire(client, epoch, sequence, duration) + if permit.state == "pending": + self.scheduler.idle_since = None + return permit.to_dict() + + def status( + self, client: str, epoch: str, sequence: int + ) -> dict[str, int | float | str]: + return self.registry.status(client, epoch, sequence).to_dict() + + def abort( + self, client: str, epoch: str, sequence: int + ) -> dict[str, int | float | str]: + return self.registry.abort(client, epoch, sequence).to_dict() + + def transfer(self, client: str, epoch: str, sequence: int) -> bool: + # The caller has entered _active_graph_updates, with no await between + # acquiring that guard and consuming this permit. + comm = self.scheduler.client_comms.get(client) + if self.scheduler.status != Status.running or comm is None or comm.closed(): + raise ClosedPermitError("client connection or scheduler is not running") + return self.registry.transfer(client, epoch, sequence) + + def has_pending(self) -> bool: + return self.registry.has_pending() + + def commit_idle_shutdown(self) -> None: + # check_idle queues close before Scheduler.status changes. Fence grants + # synchronously at the idle-close decision, not when the coroutine runs. + self.registry.close() + + def teardown(self) -> None: + self.registry.close() diff --git a/distributed/_submission_permits.py b/distributed/_submission_permits.py new file mode 100644 index 0000000000..32fd4fc18b --- /dev/null +++ b/distributed/_submission_permits.py @@ -0,0 +1,327 @@ +"""Bounded, generation-scoped submission permits. + +This module deliberately has no Scheduler dependency. The extension that uses it owns +transport and idle-shutdown policy; this registry owns only the one-shot admission state. +""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass, field +from heapq import heapify, heappop, heappush +from math import isfinite +from time import monotonic +from uuid import uuid4 + +_HEAP_SLACK = 64 + + +class SubmissionPermitError(ValueError): + """Base class for a permit that cannot be used.""" + + +class UnknownPermitError(SubmissionPermitError): + pass + + +class ExpiredPermitError(SubmissionPermitError): + pass + + +class AbortedPermitError(SubmissionPermitError): + pass + + +class RetiredPermitError(SubmissionPermitError): + pass + + +class ClosedPermitError(SubmissionPermitError): + pass + + +class PermitCapacityError(SubmissionPermitError): + pass + + +@dataclass(frozen=True, slots=True) +class PermitSnapshot: + sequence: int + state: str + duration: float + remaining: float + + def to_dict(self) -> dict[str, float | int | str]: + return { + "sequence": self.sequence, + "state": self.state, + "duration": self.duration, + "remaining": self.remaining, + } + + +@dataclass(slots=True) +class _Permit: + sequence: int + duration: float + deadline: float + state: str = "pending" + + +@dataclass(slots=True) +class _Generation: + epoch: str + high_watermark: int = 0 + active: dict[int, _Permit] = field(default_factory=dict) + outcomes: OrderedDict[int, _Permit] = field(default_factory=OrderedDict) + + +class SubmissionPermitRegistry: + """A finite permit registry for a single scheduler. + + Sequences are monotonic *per registered epoch*. Tombstones are deliberately + finite: a sequence older than the high-water mark that is no longer retained is + ``retired`` and never becomes live again. + """ + + def __init__( + self, + max_duration: float, + max_pending_per_client: int, + max_pending: int, + max_outcomes_per_client: int, + clock: Callable[[], float] = monotonic, + ) -> None: + if ( + not isinstance(max_duration, (int, float)) + or isinstance(max_duration, bool) + or not isfinite(max_duration) + or max_duration <= 0 + ): + raise ValueError("max_duration must be a positive finite number") + for name, value in ( + ("max_pending_per_client", max_pending_per_client), + ("max_pending", max_pending), + ("max_outcomes_per_client", max_outcomes_per_client), + ): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + self.max_duration = float(max_duration) + self.max_pending_per_client = max_pending_per_client + self.max_pending = max_pending + self.max_outcomes_per_client = max_outcomes_per_client + self._clock = clock + self._generations: dict[str, _Generation] = {} + # Entries deliberately contain only scalar identity data. Finished, + # unregistered, and replaced generations leave stale entries behind; + # _expire and _compact_heap validate them against active state. + self._deadlines: list[tuple[float, str, str, int]] = [] + self._pending = 0 + self._closed = False + + def register(self, client: str) -> str: + self._validate_client(client) + if self._closed: + raise ClosedPermitError("submission permit registry is closed") + # Replacing a generation invalidates, and releases, all old pending permits. + old = self._generations.get(client) + if old is not None: + self._pending -= len(old.active) + epoch = uuid4().hex + self._generations[client] = _Generation(epoch) + self._compact_heap() + return epoch + + def is_current(self, client: str, epoch: str) -> bool: + self._validate_client(client) + self._validate_epoch(epoch) + generation = self._generations.get(client) + return generation is not None and generation.epoch == epoch + + def unregister(self, client: str, epoch: str) -> bool: + self._validate_client(client) + self._validate_epoch(epoch) + generation = self._generations.get(client) + if generation is None or generation.epoch != epoch: + return False + self._pending -= len(generation.active) + del self._generations[client] + self._compact_heap() + return True + + def acquire( + self, client: str, epoch: str, sequence: int, duration: float + ) -> PermitSnapshot: + generation = self._current_generation(client, epoch) + self._validate_sequence(sequence) + self._validate_duration(duration) + if self._closed: + raise ClosedPermitError("submission permit registry is closed") + now = self._clock() + self._expire(now) + existing = self._lookup(generation, sequence) + if existing is not None: + return self._snapshot(existing, now) + self._raise_if_retired(generation, sequence) + if len(generation.active) >= self.max_pending_per_client: + raise PermitCapacityError("client pending permit capacity reached") + if self._pending >= self.max_pending: + raise PermitCapacityError("registry pending permit capacity reached") + permit = _Permit(sequence, float(duration), now + float(duration)) + generation.active[sequence] = permit + generation.high_watermark = sequence + self._pending += 1 + heappush(self._deadlines, (permit.deadline, client, epoch, sequence)) + self._compact_heap() + return self._snapshot(permit, now) + + def status(self, client: str, epoch: str, sequence: int) -> PermitSnapshot: + generation = self._current_generation(client, epoch) + self._validate_sequence(sequence) + now = self._clock() + self._expire(now) + permit = self._lookup(generation, sequence) + if permit is not None: + return self._snapshot(permit, now) + return PermitSnapshot( + sequence, + "retired" if sequence <= generation.high_watermark else "unknown", + 0.0, + 0.0, + ) + + def abort(self, client: str, epoch: str, sequence: int) -> PermitSnapshot: + generation = self._current_generation(client, epoch) + self._validate_sequence(sequence) + now = self._clock() + self._expire(now) + permit = self._lookup(generation, sequence) + if permit is None: + return self.status(client, epoch, sequence) + if permit.state == "pending": + self._finish(generation, permit, "aborted") + self._compact_heap() + return self._snapshot(permit, now) + + def transfer(self, client: str, epoch: str, sequence: int) -> bool: + generation = self._current_generation(client, epoch) + self._validate_sequence(sequence) + self._expire(self._clock()) + permit = self._lookup(generation, sequence) + if permit is None: + if sequence <= generation.high_watermark: + raise RetiredPermitError(f"submission permit {sequence} is retired") + raise UnknownPermitError(f"unknown submission permit {sequence}") + if permit.state == "accepted": + return False + if permit.state == "expired": + raise ExpiredPermitError(f"submission permit {sequence} expired") + if permit.state == "aborted": + raise AbortedPermitError(f"submission permit {sequence} was aborted") + self._finish(generation, permit, "accepted") + self._compact_heap() + return True + + def has_pending(self) -> bool: + self._expire(self._clock()) + return self._pending > 0 + + def close(self) -> None: + if self._closed: + return + self._closed = True + for generation in self._generations.values(): + for permit in list(generation.active.values()): + self._finish(generation, permit, "aborted") + self._deadlines.clear() + + def _current_generation(self, client: str, epoch: str) -> _Generation: + self._validate_client(client) + self._validate_epoch(epoch) + generation = self._generations.get(client) + if generation is None or generation.epoch != epoch: + raise UnknownPermitError("unknown or stale submission permit epoch") + return generation + + def _expire(self, now: float) -> None: + while self._deadlines and self._deadlines[0][0] <= now: + deadline, client, epoch, sequence = heappop(self._deadlines) + generation = self._generations.get(client) + if generation is None or generation.epoch != epoch: + continue + permit = generation.active.get(sequence) + if permit is not None and permit.deadline == deadline: + self._finish(generation, permit, "expired") + self._compact_heap() + + def _compact_heap(self) -> None: + if len(self._deadlines) <= 2 * self._pending + _HEAP_SLACK: + return + self._deadlines = [ + entry for entry in self._deadlines if self._is_live_entry(entry) + ] + heapify(self._deadlines) + + def _is_live_entry(self, entry: tuple[float, str, str, int]) -> bool: + deadline, client, epoch, sequence = entry + generation = self._generations.get(client) + if generation is None or generation.epoch != epoch: + return False + permit = generation.active.get(sequence) + return permit is not None and permit.deadline == deadline + + def _finish(self, generation: _Generation, permit: _Permit, state: str) -> None: + del generation.active[permit.sequence] + self._pending -= 1 + permit.state = state + generation.outcomes[permit.sequence] = permit + generation.outcomes.move_to_end(permit.sequence) + while len(generation.outcomes) > self.max_outcomes_per_client: + generation.outcomes.popitem(last=False) + + @staticmethod + def _snapshot(permit: _Permit, now: float) -> PermitSnapshot: + remaining = ( + max(0.0, permit.deadline - now) if permit.state == "pending" else 0.0 + ) + return PermitSnapshot(permit.sequence, permit.state, permit.duration, remaining) + + @staticmethod + def _lookup(generation: _Generation, sequence: int) -> _Permit | None: + return generation.active.get(sequence) or generation.outcomes.get(sequence) + + @staticmethod + def _raise_if_retired(generation: _Generation, sequence: int) -> None: + if sequence <= generation.high_watermark: + raise RetiredPermitError(f"submission permit {sequence} is retired") + + @staticmethod + def _validate_client(client: str) -> None: + if not isinstance(client, str) or not client: + raise ValueError("client must be a non-empty string") + + @staticmethod + def _validate_epoch(epoch: str) -> None: + if not isinstance(epoch, str) or not epoch: + raise ValueError("epoch must be a non-empty string") + + @staticmethod + def _validate_sequence(sequence: int) -> None: + if ( + not isinstance(sequence, int) + or isinstance(sequence, bool) + or not 0 < sequence <= 2**63 - 1 + ): + raise ValueError("sequence must be an integer between 1 and 2**63 - 1") + + def _validate_duration(self, duration: float) -> None: + if ( + not isinstance(duration, (int, float)) + or isinstance(duration, bool) + or not isfinite(duration) + or not 0 < duration <= self.max_duration + ): + raise ValueError( + "duration must be positive, finite, and no greater than max_duration" + ) diff --git a/distributed/client.py b/distributed/client.py index 5f14cb4911..97b0ed1dc7 100644 --- a/distributed/client.py +++ b/distributed/client.py @@ -1132,6 +1132,13 @@ def __init__( # Communication self.scheduler_comm = None + self._submission_permit_capabilities: dict[str, Any] | None = None + self._submission_permit_pending: dict[ + tuple[str, int], asyncio.Future[dict[str, Any]] + ] = {} + self._submission_permit_sequence = 0 + self._submission_permit_acquire_lock = asyncio.Lock() + self._submission_permit_changed = asyncio.Event() if address is None: address = dask.config.get("scheduler-address", None) @@ -1221,6 +1228,7 @@ def __init__( "error": self._handle_error, "event": self._handle_event, "adjust-heartbeat-interval": self._adjust_heartbeat_intervals, + "submission-permit-admission": self._handle_submission_permit_admission, } self._state_handlers = { @@ -1513,8 +1521,10 @@ async def _start(self, timeout=no_default, **kwargs): async def _reconnect(self): assert self.scheduler_comm.comm.closed() + self._fail_submission_permits("scheduler connection lost") self.status = "connecting" self.scheduler_comm = None + self._submission_permit_capabilities = None for st in self.futures.values(): st.cancel( @@ -1588,6 +1598,7 @@ async def _ensure_connected(self, timeout=None): msg = await comm.read() assert len(msg) == 1 assert msg[0]["op"] == "stream-start" + self._submission_permit_capabilities = msg[0].get("submission-permits") if msg[0].get("error"): raise ImportError(msg[0]["error"]) @@ -1713,6 +1724,35 @@ def _inc_ref(self, key): with self._refcount_lock: self.refcount[key] += 1 + def _handle_submission_permit_admission( + self, epoch=None, sequence=None, status=None, reason=None, detail=None + ): + if ( + not isinstance(epoch, str) + or not isinstance(sequence, int) + or isinstance(sequence, bool) + ): + return + pending = self._submission_permit_pending.get((epoch, sequence)) + if pending is not None and not pending.done(): + pending.set_result( + { + "epoch": epoch, + "sequence": sequence, + "status": status, + "reason": reason, + "detail": detail, + } + ) + + def _fail_submission_permits(self, reason): + self._submission_permit_changed.set() + self._submission_permit_changed = asyncio.Event() + for pending in self._submission_permit_pending.values(): + if not pending.done(): + pending.set_exception(CommClosedError(reason)) + self._submission_permit_pending.clear() + def _dec_ref(self, key): with self._refcount_lock: self.refcount[key] -= 1 @@ -1824,6 +1864,7 @@ def _handle_task_erred(self, key=None, exception=None, traceback=None): state.set_error(exception, traceback) def _handle_restart(self): + self._fail_submission_permits("scheduler restarted") logger.info("Receive restart signal from scheduler") for state in self.futures.values(): state.cancel( @@ -1873,6 +1914,8 @@ async def _close(self, fast: bool = False) -> None: self.status = "closing" + self._fail_submission_permits("client closing") + await self.preloads.teardown() with suppress(AttributeError): @@ -3384,6 +3427,11 @@ def _graph_to_futures( fifo_timeout=0, actors=None, ): + from distributed._submission_permit_client import _current_submission + + submission = _current_submission.get() + if submission is not None: + submission.begin_graph(self) with self._refcount_lock: if actors is not None and actors is not True and actors is not False: actors = list(self._expand_key(actors)) @@ -3415,7 +3463,14 @@ def _graph_to_futures( validate_key(key) # Create futures before sending graph (helps avoid contention) - futures = {key: Future(key, self) for key in keyset} + if submission is None: + futures = {key: Future(key, self) for key in keyset} + else: + futures = {} + for key in keyset: + future = Future(key, self) + submission.own(future) + futures[key] = future # This is done manually here to get better exception messages on # scheduler side and be able to produce the below warning about @@ -3437,20 +3492,22 @@ def _graph_to_futures( computations = self._get_computation_code( nframes=dask.config.get("distributed.diagnostics.computations.nframes") ) - self._send_to_scheduler( - { - "op": "update-graph", - "expr_ser": expr_ser, - "keys": set(keys), - "internal_priority": internal_priority, - "submitting_task": getattr(thread_state, "key", None), - "fifo_timeout": fifo_timeout, - "actors": actors, - "code": ToPickle(computations), - "annotations": ToPickle(annotations), - "span_metadata": ToPickle(span_metadata), - } - ) + message = { + "op": "update-graph", + "expr_ser": expr_ser, + "keys": set(keys), + "internal_priority": internal_priority, + "submitting_task": getattr(thread_state, "key", None), + "fifo_timeout": fifo_timeout, + "actors": actors, + "code": ToPickle(computations), + "annotations": ToPickle(annotations), + "span_metadata": ToPickle(span_metadata), + } + if submission is None: + self._send_to_scheduler(message) + else: + submission.capture(message) return futures def get( diff --git a/distributed/scheduler.py b/distributed/scheduler.py index 92f22b807a..fa9206f5b9 100644 --- a/distributed/scheduler.py +++ b/distributed/scheduler.py @@ -65,7 +65,7 @@ from dask._expr import LLGExpr from dask._task_spec import GraphNode, convert_legacy_graph from dask.core import istask, validate_key -from dask.typing import Key, no_default +from dask.typing import Key, NoDefault, no_default from dask.utils import ( format_bytes, format_time, @@ -151,6 +151,7 @@ from dask._expr import Expr + from distributed._submission_permit_extension import SubmissionPermitExtension from distributed.diagnostics.task_stream import TaskStreamPlugin FuncT = TypeVar("FuncT", bound=Callable[..., Any]) @@ -5008,13 +5009,58 @@ async def update_graph( code: tuple[SourceCode, ...] = (), annotations: dict | None = None, stimulus_id: str | None = None, + submission_epoch: str | None | NoDefault = no_default, + submission_sequence: int | None | NoDefault = no_default, ) -> None: start = time() stimulus_id = stimulus_id or f"update-graph-{start}" self._active_graph_updates += 1 evt_msg: dict[str, Any] + tagged = ( + submission_epoch is not no_default or submission_sequence is not no_default + ) + permit_ext: SubmissionPermitExtension | None = self.extensions.get( + "submission-permits" + ) + submission_comm = self.client_comms.get(client) try: + if tagged: + admission: dict[str, Any] = { + "op": "submission-permit-admission", + "epoch": None + if submission_epoch is no_default + else submission_epoch, + "sequence": ( + None + if submission_sequence is no_default + else submission_sequence + ), + "status": "rejected", + } + try: + if ( + permit_ext is None + or not isinstance(submission_epoch, str) + or not isinstance(submission_sequence, int) + ): + raise ValueError("invalid or unsupported submission permit tag") + if permit_ext.transfer( + client, submission_epoch, submission_sequence + ): + admission["status"] = "accepted" + else: + admission["reason"] = "sequence-already-consumed" + except ValueError as e: + admission["reason"] = type(e).__name__ + admission["detail"] = str(e) + self.client_send(client, admission) + if admission["status"] != "accepted": + # Admission failure is per submission, never per task key. + # A key-wide error could corrupt another existing Future. + return + # Acceptance means entry to this guarded handler, not successful + # graph preparation, TaskState creation, or computation. logger.debug("Received new graph. Deserializing...") try: expr = deserialize(expr_ser.header, expr_ser.frames) @@ -5058,6 +5104,17 @@ async def update_graph( # has to happen in the same event loop. # ************************************* + if tagged and ( + self.status != Status.running + or submission_comm is None + or submission_comm.closed() + or self.client_comms.get(client) is not submission_comm + or permit_ext is None + or not isinstance(submission_epoch, str) + or not permit_ext.registry.is_current(client, submission_epoch) + ): + return + if self._find_lost_dependencies(dsk, keys): self.report( { @@ -5116,6 +5173,16 @@ async def update_graph( self.log_event(["scheduler", client], evt_msg) logger.debug("Task state created. %i new tasks", len(self.tasks) - before) except Exception as e: + if tagged and ( + self.status != Status.running + or submission_comm is None + or submission_comm.closed() + or self.client_comms.get(client) is not submission_comm + or permit_ext is None + or not isinstance(submission_epoch, str) + or not permit_ext.registry.is_current(client, submission_epoch) + ): + return evt_msg = { "action": "update-graph", "stimulus_id": stimulus_id, @@ -5909,6 +5976,24 @@ async def add_client( We listen to all future messages from this Comm. """ assert client is not None + permit_ext: SubmissionPermitExtension | None = self.extensions.get( + "submission-permits" + ) + submission_epoch = None + if permit_ext is not None: + # The old stream may still be awaiting graph preparation or cleanup. + # Preserve its ClientState and ownership until that lifetime ends. + # Closing just the new comm lets Client._reconnect retry normally. + if client in self.client_comms: + await comm.close() + return + from distributed._submission_permits import ClosedPermitError + + try: + submission_epoch = permit_ext.register_client(client) + except ClosedPermitError: + await comm.close() + return comm.name = "Scheduler->Client" logger.info("Receive client connection: %s", client) self.log_event(["all", client], {"action": "add-client", "client": client}) @@ -5925,7 +6010,10 @@ async def add_client( bcomm = BatchedSend(interval="2ms", loop=self.loop) bcomm.start(comm) self.client_comms[client] = bcomm - msg = {"op": "stream-start"} + msg: dict[str, Any] = {"op": "stream-start"} + if permit_ext is not None: + assert submission_epoch is not None + msg["submission-permits"] = permit_ext.capabilities(submission_epoch) version_warning = version_module.error_message( version_module.get_versions(), {w: ws.versions for w, ws in self.workers.items()}, @@ -5937,19 +6025,38 @@ async def add_client( try: await self.handle_stream(comm=comm, extra={"client": client}) finally: + if permit_ext is not None: + assert submission_epoch is not None + permit_ext.unregister_client(client, submission_epoch) self.remove_client(client=client, stimulus_id=f"remove-client-{time()}") logger.debug("Finished handling client %s", client) finally: - if not comm.closed(): - self.client_comms[client].send({"op": "stream-closed"}) - try: - if not self._is_finalizing(): - await self.client_comms[client].close() - del self.client_comms[client] - if self.status == Status.running: - logger.info("Close client connection: %s", client) - except TypeError: # comm becomes None during GC - pass + if permit_ext is None: + # Preserve legacy same-ID overlap behavior. Only the optional + # permit path serializes stream ownership before registration. + if not comm.closed(): + self.client_comms[client].send({"op": "stream-closed"}) + try: + if not self._is_finalizing(): + await self.client_comms[client].close() + del self.client_comms[client] + if self.status == Status.running: + logger.info("Close client connection: %s", client) + except TypeError: # comm becomes None during GC + pass + else: + if not comm.closed(): + bcomm.send({"op": "stream-closed"}) + try: + if not self._is_finalizing(): + await bcomm.close() + if self.status == Status.running: + logger.info("Close client connection: %s", client) + except TypeError: # comm becomes None during GC + pass + finally: + if self.client_comms.get(client) is bcomm: + del self.client_comms[client] def remove_client(self, client: str, stimulus_id: str | None = None) -> None: """Remove client from network""" @@ -8673,6 +8780,13 @@ def check_idle(self) -> float | None: self.idle_since = None return None + permit_ext: SubmissionPermitExtension | None = self.extensions.get( + "submission-permits" + ) + if permit_ext is not None and permit_ext.has_pending(): + self.idle_since = None + return None + if ( self.queued or self.unrunnable @@ -8696,6 +8810,8 @@ def check_idle(self) -> float | None: if self.idle_timeout: if time() > self.idle_since + self.idle_timeout: assert self.idle_since + if permit_ext is not None: + permit_ext.commit_idle_shutdown() logger.info( "Scheduler closing after being idle for %s", format_time(self.idle_timeout), diff --git a/distributed/tests/test_submission_permit_client.py b/distributed/tests/test_submission_permit_client.py new file mode 100644 index 0000000000..47455e8993 --- /dev/null +++ b/distributed/tests/test_submission_permit_client.py @@ -0,0 +1,640 @@ +from __future__ import annotations + +import asyncio +import threading +from functools import partial +from time import sleep + +import pytest + +from dask import delayed + +import distributed.client as client_module +from distributed import Client +from distributed._submission_permit_client import ( + PermitExpiredError, + PermitIndeterminateError, + PermitRejectedError, + PermitUnsupportedError, + _current_submission, + protected_compute, + protected_persist, +) +from distributed._submission_permit_extension import SubmissionPermitExtension +from distributed.core import CommClosedError +from distributed.metrics import time +from distributed.scheduler import DEFAULT_EXTENSIONS +from distributed.utils_test import async_poll_for, cluster, gen_cluster, inc + + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class ClockedExtension(SubmissionPermitExtension): + def __init__(self, scheduler): + self.clock = Clock() + super().__init__( + scheduler, + max_duration=10, + max_pending_per_client=5, + max_pending=10, + max_outcomes_per_client=10, + clock=self.clock, + ) + + +class IdleTimeoutAfterPermit(SubmissionPermitExtension): + def acquire( + self, client: str, epoch: str, sequence: int, duration: float + ) -> dict[str, int | float | str]: + result = super().acquire(client, epoch, sequence, duration) + if result["state"] == "pending": + # Arm the timeout atomically with the grant. Client startup and + # acquisition latency are outside the interval protected by a permit. + self.scheduler.idle_timeout = 0.05 + self.scheduler._idle_transition_counter = self.scheduler.transition_counter + self.scheduler.idle_since = time() - 1 + return result + + +SCHEDULER_KWARGS = { + "extensions": {**DEFAULT_EXTENSIONS, "submission-permits": ClockedExtension} +} +OPTIONS = {"duration": 5, "timeout": 1, "max_clock_rate": 1, "clock_margin": 0.1} + + +def watch_graph_sends(c, monkeypatch): + messages = [] + send = c.scheduler_comm.send + + def capture(*msgs): + messages.extend(msg.copy() for msg in msgs if msg.get("op") == "update-graph") + send(*msgs) + + monkeypatch.setattr(c.scheduler_comm, "send", capture) + return messages + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_async_protected_compute_and_persist(c, s, a, monkeypatch): + messages = watch_graph_sends(c, monkeypatch) + value = delayed(inc)(1) + future = await protected_compute(c, value, **OPTIONS) + assert await future == 2 + persisted = await protected_persist(c, delayed(inc)(4), **OPTIONS) + assert await c.compute(persisted) == 5 + assert [m.get("submission_sequence") for m in messages] == [1, 2, None] + assert all( + m["submission_epoch"] == c._submission_permit_capabilities["epoch"] + for m in messages[:2] + ) + assert not c._submission_permit_pending + assert _current_submission.get() is None + assert not s.extensions["submission-permits"].has_pending() + s.validate_state() + + +@gen_cluster(client=True, nthreads=[]) +async def test_unsupported_peer_never_prepares(c, s, monkeypatch): + def unexpected(*args, **kwargs): + pytest.fail("unsupported peer started graph preparation") + + monkeypatch.setattr(c, "compute", unexpected) + with pytest.raises(PermitUnsupportedError): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert not c.futures + assert not c._submission_permit_pending + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_late_acquire_ack_never_prepares(c, s, monkeypatch): + clock = Clock() + acquire = s.handlers["submission_permit_acquire"] + + def late_acquire(client, epoch, sequence, duration): + result = acquire(client, epoch, sequence, duration) + clock.now = 6 + return result + + def unexpected(*args, **kwargs): + pytest.fail("unusable grant started graph preparation") + + monkeypatch.setitem(s.handlers, "submission_permit_acquire", late_acquire) + monkeypatch.setattr(c, "compute", unexpected) + with pytest.raises(PermitExpiredError): + await protected_compute(c, delayed(inc)(1), clock=clock, **OPTIONS) + assert not c.futures + assert not s.extensions["submission-permits"].has_pending() + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_serialization_error_releases_owned_futures(c, s, monkeypatch): + class SerializationFailure(RuntimeError): + pass + + def fail(*args, **kwargs): + assert c.refcount # Failure happens after private Future construction. + raise SerializationFailure("original serialization error") + + messages = watch_graph_sends(c, monkeypatch) + monkeypatch.setattr(client_module, "serialize", fail) + with pytest.raises(SerializationFailure, match="original serialization error"): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert not messages + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + assert not s.extensions["submission-permits"].has_pending() + assert _current_submission.get() is None + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_server_rejection_preserves_existing_same_key_future( + c, s, a, monkeypatch +): + value = delayed(inc)(1) + existing = c.compute(value) + assert await existing == 2 + refs_before = dict(c.refcount) + state = existing._state + task_state = s.tasks[existing.key] + compute = c.compute + ext = s.extensions["submission-permits"] + + def expire_after_preparation(*args, **kwargs): + result = compute(*args, **kwargs) + assert c.refcount[existing.key] == refs_before[existing.key] + 1 + ext.clock.now = 5 + return result + + monkeypatch.setattr(c, "compute", expire_after_preparation) + with pytest.raises(PermitRejectedError): + await protected_compute(c, value, **OPTIONS) + assert dict(c.refcount) == refs_before + assert existing._state is state + assert s.tasks[existing.key] is task_state + assert task_state.state == "memory" + assert await existing == 2 + assert not c._submission_permit_pending + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_local_expiry_during_preparation_never_sends(c, s, monkeypatch): + clock = Clock() + compute = c.compute + messages = watch_graph_sends(c, monkeypatch) + + def expire(*args, **kwargs): + result = compute(*args, **kwargs) + clock.now = 6 + return result + + monkeypatch.setattr(c, "compute", expire) + with pytest.raises(PermitExpiredError): + await protected_compute(c, delayed(inc)(1), clock=clock, **OPTIONS) + assert not messages + assert not c.refcount + assert not c.futures + assert not s.tasks + assert not s.extensions["submission-permits"].has_pending() + + +@pytest.mark.parametrize("cancel", [False, True]) +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_lost_admission_or_cancel_after_send_never_retries_or_aborts( + c, s, a, cancel, monkeypatch +): + messages = watch_graph_sends(c, monkeypatch) + admissions = [] + arrived = asyncio.Event() + client_send = s.client_send + abort = s.handlers["submission_permit_abort"] + aborts = [] + + def hold_admission(client, msg): + if msg["op"] == "submission-permit-admission": + admissions.append(msg.copy()) + arrived.set() + else: + client_send(client, msg) + + def watch_abort(*args, **kwargs): + aborts.append((args, kwargs)) + return abort(*args, **kwargs) + + monkeypatch.setattr(s, "client_send", hold_admission) + monkeypatch.setitem(s.handlers, "submission_permit_abort", watch_abort) + options = dict(OPTIONS, timeout=0.1) + task = asyncio.create_task(protected_compute(c, delayed(inc)(1), **options)) + await arrived.wait() + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(PermitIndeterminateError): + await task + assert len(messages) == 1 + assert len(admissions) == 1 + assert not aborts + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + # Delayed and duplicate results cannot recreate a removed operation. + response = {k: v for k, v in admissions[0].items() if k != "op"} + c._handle_submission_permit_admission(**response) + c._handle_submission_permit_admission(**response) + assert not c._submission_permit_pending + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_cancel_queued_during_preparation_prevents_send(c, s, monkeypatch): + compute = c.compute + messages = watch_graph_sends(c, monkeypatch) + + def cancel_after_preparation(*args, **kwargs): + result = compute(*args, **kwargs) + asyncio.current_task().cancel() + return result + + monkeypatch.setattr(c, "compute", cancel_after_preparation) + task = asyncio.create_task(protected_compute(c, delayed(inc)(1), **OPTIONS)) + with pytest.raises(asyncio.CancelledError): + await task + assert not messages + assert not c.refcount + assert not c.futures + assert not s.extensions["submission-permits"].has_pending() + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_reconnect_during_preparation_cannot_buffer_old_graph(c, s, monkeypatch): + epoch = c._submission_permit_capabilities["epoch"] + compute = c.compute + messages = watch_graph_sends(c, monkeypatch) + + def lose_connection(*args, **kwargs): + result = compute(*args, **kwargs) + c.scheduler_comm.comm.abort() + return result + + monkeypatch.setattr(c, "compute", lose_connection) + with pytest.raises(PermitRejectedError): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + await async_poll_for( + lambda: ( + c.status == "running" + and c._submission_permit_capabilities is not None + and c._submission_permit_capabilities["epoch"] != epoch + ) + ) + assert not messages + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + assert not [msg for msg in c._pending_msg_buffer if msg["op"] == "update-graph"] + assert not s.tasks + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_overlapping_operations_are_independent(c, s, a, monkeypatch): + messages = watch_graph_sends(c, monkeypatch) + first, second = await asyncio.gather( + protected_compute(c, delayed(inc)(1), **OPTIONS), + protected_compute(c, delayed(inc)(8), **OPTIONS), + ) + assert await first == 2 + assert await second == 9 + assert {msg["submission_sequence"] for msg in messages} == {1, 2} + assert not c._submission_permit_pending + assert not s.extensions["submission-permits"].has_pending() + + +def test_sync_client_preserves_preparation_thread_and_survives_idle_timeout( + monkeypatch, +): + extensions = { + **DEFAULT_EXTENSIONS, + "submission-permits": partial( + IdleTimeoutAfterPermit, + max_duration=10, + max_pending_per_client=5, + max_pending=10, + max_outcomes_per_client=10, + ), + } + with cluster(nworkers=1, scheduler_kwargs={"extensions": extensions}) as ( + s, + workers, + ): + with Client(s["address"]) as c: + original = c.compute + caller_thread = threading.get_ident() + + def slow_preparation(*args, **kwargs): + assert threading.get_ident() == caller_thread + sleep(0.4) + return original(*args, **kwargs) + + monkeypatch.setattr(c, "compute", slow_preparation) + + future = protected_compute(c, delayed(inc)(1), **OPTIONS) + assert future.result() == 2 + assert not c._submission_permit_pending + assert _current_submission.get() is None + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_cancel_while_grant_response_is_pending(c, s, monkeypatch): + entered = asyncio.Event() + release = asyncio.Event() + acquire = s.handlers["submission_permit_acquire"] + + async def hold_grant(client, epoch, sequence, duration): + grant = acquire(client, epoch, sequence, duration) + entered.set() + await release.wait() + return grant + + monkeypatch.setitem(s.handlers, "submission_permit_acquire", hold_grant) + task = asyncio.create_task(protected_compute(c, delayed(inc)(1), **OPTIONS)) + try: + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not s.extensions["submission-permits"].has_pending() + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + assert not s.tasks + finally: + release.set() + + +@pytest.mark.parametrize("failure", ["reconnect", "restart", "close", "timeout"]) +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_pending_grant_is_bounded_and_wakes_on_connection_change( + c, s, failure, monkeypatch +): + entered = asyncio.Event() + release = asyncio.Event() + acquire = s.handlers["submission_permit_acquire"] + old_epoch = c._submission_permit_capabilities["epoch"] + + async def hold_grant(client, epoch, sequence, duration): + grant = acquire(client, epoch, sequence, duration) + entered.set() + await release.wait() + return grant + + def unexpected(*args, **kwargs): + pytest.fail("an unconfirmed grant started graph preparation") + + monkeypatch.setitem(s.handlers, "submission_permit_acquire", hold_grant) + monkeypatch.setattr(c, "compute", unexpected) + options = dict(OPTIONS, timeout=0.05 if failure == "timeout" else 5) + task = asyncio.create_task(protected_compute(c, delayed(inc)(1), **options)) + try: + await entered.wait() + if failure == "reconnect": + c.scheduler_comm.comm.abort() + elif failure == "restart": + c._handle_restart() + elif failure == "close": + await c.close() + if failure == "timeout": + with pytest.raises(asyncio.TimeoutError, match="submission permit"): + await asyncio.wait_for(task, 1) + else: + with pytest.raises(PermitRejectedError): + await asyncio.wait_for(task, 1) + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + if failure == "reconnect": + await async_poll_for( + lambda: ( + c.status == "running" + and c._submission_permit_capabilities is not None + and c._submission_permit_capabilities["epoch"] != old_epoch + ) + ) + assert not s.extensions["submission-permits"].has_pending() + finally: + release.set() + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_restart_before_grant_reply_prevents_preparation(c, s, monkeypatch): + acquire = s.handlers["submission_permit_acquire"] + + def restart_during_grant(client, epoch, sequence, duration): + grant = acquire(client, epoch, sequence, duration) + c._handle_restart() + return grant + + def unexpected(*args, **kwargs): + pytest.fail("a pre-restart grant started preparation") + + monkeypatch.setitem(s.handlers, "submission_permit_acquire", restart_during_grant) + monkeypatch.setattr(c, "compute", unexpected) + with pytest.raises(PermitRejectedError): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert not c.refcount + assert not s.tasks + assert not s.extensions["submission-permits"].has_pending() + + +@pytest.mark.parametrize("failure", ["reconnect", "restart", "close"]) +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_connection_lifecycle_wakes_admission_waiter( + c, s, a, failure, monkeypatch +): + arrived = asyncio.Event() + original_send = s.client_send + old_epoch = c._submission_permit_capabilities["epoch"] + messages = watch_graph_sends(c, monkeypatch) + aborts = [] + abort = s.handlers["submission_permit_abort"] + + def hold(client, msg): + if msg["op"] == "submission-permit-admission": + arrived.set() + else: + original_send(client, msg) + + def watch_abort(*args, **kwargs): + aborts.append((args, kwargs)) + return abort(*args, **kwargs) + + monkeypatch.setattr(s, "client_send", hold) + monkeypatch.setitem(s.handlers, "submission_permit_abort", watch_abort) + task = asyncio.create_task(protected_compute(c, delayed(inc)(1), **OPTIONS)) + await arrived.wait() + if failure == "reconnect": + c.scheduler_comm.comm.abort() + elif failure == "restart": + c._handle_restart() + else: + await c.close() + with pytest.raises(PermitIndeterminateError): + await task + if failure == "reconnect": + await async_poll_for( + lambda: ( + c.status == "running" + and c._submission_permit_capabilities is not None + and c._submission_permit_capabilities["epoch"] != old_epoch + ) + ) + assert len(messages) == 1 + assert not aborts + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_epoch_change_after_ack_does_not_publish_stale_futures( + c, s, a, monkeypatch +): + capabilities = c._submission_permit_capabilities + handler = c._stream_handlers["submission-permit-admission"] + + def change_epoch(**msg): + handler(**msg) + c._submission_permit_capabilities = dict(capabilities, epoch="replacement") + + monkeypatch.setitem(c._stream_handlers, "submission-permit-admission", change_epoch) + try: + with pytest.raises(PermitIndeterminateError): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + finally: + c._submission_permit_capabilities = capabilities + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_accepted_reply_may_arrive_after_original_permit_duration( + c, s, a, monkeypatch +): + original_send = s.client_send + + def delay_ack(client, msg): + if msg["op"] == "submission-permit-admission": + assert msg["status"] == "accepted" + s.loop.call_later(0.3, original_send, client, msg) + else: + original_send(client, msg) + + monkeypatch.setattr(s, "client_send", delay_ack) + future = await protected_compute( + c, delayed(inc)(1), **dict(OPTIONS, duration=0.2, clock_margin=0.01) + ) + assert await future == 2 + assert not c._submission_permit_pending + + +@pytest.mark.parametrize("captures", [0, 2]) +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_zero_or_multiple_graphs_fail_without_send(c, s, captures, monkeypatch): + original = c.compute + messages = watch_graph_sends(c, monkeypatch) + + def invalid_preparation(collection, **kwargs): + if captures == 0: + return None + original(collection, **kwargs) + return original(collection, **kwargs) + + monkeypatch.setattr(c, "compute", invalid_preparation) + with pytest.raises(RuntimeError, match="graph"): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert not messages + assert not c.refcount + assert not c.futures + assert not s.tasks + assert not s.extensions["submission-permits"].has_pending() + assert _current_submission.get() is None + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_enqueue_then_send_failure_is_indeterminate_and_never_aborts( + c, s, monkeypatch +): + send = c.scheduler_comm.send + messages = [] + aborts = [] + abort = s.handlers["submission_permit_abort"] + + def interrupted_send(*msgs): + send(*msgs) + for msg in msgs: + if msg["op"] == "update-graph": + messages.append(msg) + raise CommClosedError("interrupted after enqueue") + + def watch_abort(*args, **kwargs): + aborts.append((args, kwargs)) + return abort(*args, **kwargs) + + monkeypatch.setattr(c.scheduler_comm, "send", interrupted_send) + monkeypatch.setitem(s.handlers, "submission_permit_abort", watch_abort) + with pytest.raises(PermitIndeterminateError): + await protected_compute(c, delayed(inc)(1), **OPTIONS) + assert len(messages) == 1 + assert not aborts + assert not c.refcount + assert not c.futures + assert not c._submission_permit_pending + + +@gen_cluster( + client=True, nthreads=[("127.0.0.1", 1)], scheduler_kwargs=SCHEDULER_KWARGS +) +async def test_rejected_persist_releases_only_its_multiple_owned_futures( + c, s, a, monkeypatch +): + import dask.bag as db + + collection = db.from_sequence(range(6), npartitions=3).map(inc) + persisted = await protected_persist(c, collection, **OPTIONS) + assert await c.compute(persisted.sum()) == 21 + # The temporary reduction Future may release asynchronously. + await asyncio.sleep(0) + refs = dict(c.refcount) + assert len(refs) == 3 + persist = c.persist + + def expire(*args, **kwargs): + result = persist(*args, **kwargs) + assert dict(c.refcount) == {key: count + 1 for key, count in refs.items()} + s.extensions["submission-permits"].clock.now = 5 + return result + + monkeypatch.setattr(c, "persist", expire) + with pytest.raises(PermitRejectedError): + await protected_persist(c, collection, **OPTIONS) + assert dict(c.refcount) == refs + assert await c.compute(persisted.sum()) == 21 diff --git a/distributed/tests/test_submission_permit_client_unit.py b/distributed/tests/test_submission_permit_client_unit.py new file mode 100644 index 0000000000..69c52b276b --- /dev/null +++ b/distributed/tests/test_submission_permit_client_unit.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import asyncio +from typing import Any, cast + +import pytest + +from distributed._submission_permit_client import ( + SubmissionPermitExpiredError, + SubmissionPermitOperation, + SubmissionPermitRejectedError, + SubmissionPermitUnsupportedError, + _operation, +) +from distributed.client import Client as DaskClient + + +class Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +class RPC: + async def submission_permit_acquire(self, **kwargs: Any) -> dict[str, Any]: + return { + "sequence": kwargs["sequence"], + "state": "pending", + "duration": kwargs["duration"], + } + + async def submission_permit_abort(self, **kwargs: Any) -> dict[str, Any]: + return {"state": "aborted", **kwargs} + + +class Carrier: + def __init__(self, client: FakeClient) -> None: + self.client = client + self.messages: list[dict[str, Any]] = [] + self.is_closed = False + + def closed(self) -> bool: + return self.is_closed + + def send(self, message: dict[str, Any]) -> None: + self.messages.append(message) + + def admit() -> None: + key = (message["submission_epoch"], message["submission_sequence"]) + self.client._submission_permit_pending[key].set_result( + {"epoch": key[0], "sequence": key[1], "status": "accepted"} + ) + + asyncio.get_running_loop().call_soon(admit) + + +class FakeClient: + asynchronous: bool = True + generation: int = 4 + status: str = "running" + id: str = "unit-client" + + def __init__(self) -> None: + self._submission_permit_capabilities: dict[str, Any] | None = { + "version": 1, + "epoch": "epoch", + "max_duration": 10, + } + self._submission_permit_sequence = 0 + self._submission_permit_acquire_lock = asyncio.Lock() + self._submission_permit_changed = asyncio.Event() + self._submission_permit_pending: dict[ + tuple[str, int], asyncio.Future[dict[str, Any]] + ] = {} + self.scheduler: RPC = RPC() + self.scheduler_comm: Carrier = Carrier(self) + + +def as_dask_client(client: FakeClient) -> DaskClient: + """Limit the real Client type boundary to these lightweight unit doubles.""" + return cast(DaskClient, client) + + +def test_operation_tags_only_its_captured_carrier_message(): + async def run() -> None: + clock = Clock() + client = FakeClient() + dask_client = as_dask_client(client) + operation = SubmissionPermitOperation( + duration=5, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=clock, + ) + await operation.acquire(dask_client) + operation.begin_graph(dask_client) + operation.capture({"op": "update-graph", "keys": {"x"}}) + await operation.commit(dask_client) + + assert client.scheduler_comm.messages == [ + { + "op": "update-graph", + "keys": {"x"}, + "submission_epoch": "epoch", + "submission_sequence": 1, + } + ] + assert not client._submission_permit_pending + + asyncio.run(run()) + + +def test_expiry_before_commit_does_not_send_graph(): + async def run() -> None: + clock = Clock() + client = FakeClient() + dask_client = as_dask_client(client) + operation = SubmissionPermitOperation( + duration=2, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=clock, + ) + await operation.acquire(dask_client) + operation.begin_graph(dask_client) + operation.capture({"op": "update-graph"}) + clock.now = 2 + + with pytest.raises(SubmissionPermitExpiredError): + await operation.commit(dask_client) + assert not client.scheduler_comm.messages + + asyncio.run(run()) + + +def test_mismatched_acquire_reply_is_rejected_and_aborted(): + class BadRPC(RPC): + async def submission_permit_acquire(self, **kwargs: Any) -> dict[str, Any]: + return {"sequence": kwargs["sequence"], "state": "accepted", "duration": 1} + + async def run() -> None: + clock = Clock() + client = FakeClient() + dask_client = as_dask_client(client) + client.scheduler = BadRPC() + operation = SubmissionPermitOperation( + duration=1, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=clock, + ) + + with pytest.raises(SubmissionPermitRejectedError): + await operation.acquire(dask_client) + + asyncio.run(run()) + + +def test_argument_validation_and_unsupported_capability(): + with pytest.raises(ValueError, match="duration"): + _operation( + duration=0, timeout=1, max_clock_rate=1, clock_margin=0, clock=Clock() + ) + with pytest.raises(ValueError, match="max_clock_rate"): + _operation( + duration=1, timeout=1, max_clock_rate=0.5, clock_margin=0, clock=Clock() + ) + + operation = SubmissionPermitOperation( + duration=1, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=Clock(), + ) + client = FakeClient() + client._submission_permit_capabilities = None + + async def acquire() -> None: + with pytest.raises(SubmissionPermitUnsupportedError): + await operation.acquire(as_dask_client(client)) + + asyncio.run(acquire()) + + +def test_acquire_snapshot_rejects_a_reconnected_client_before_graph_work(): + async def run() -> None: + clock = Clock() + client = FakeClient() + dask_client = as_dask_client(client) + operation = SubmissionPermitOperation( + duration=1, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=clock, + ) + await operation.acquire(dask_client) + client.generation += 1 + client.scheduler_comm = Carrier(client) + client._submission_permit_capabilities = { + "version": 1, + "epoch": "new-epoch", + "max_duration": 10, + } + + with pytest.raises(SubmissionPermitRejectedError, match="connection changed"): + operation.ensure_origin(dask_client) + with pytest.raises(RuntimeError, match="exactly one graph"): + operation.begin_graph(as_dask_client(FakeClient())) + + asyncio.run(run()) + + +def test_clock_must_not_move_backwards_or_exceed_granted_interval(): + async def run() -> None: + clock = Clock() + client = FakeClient() + dask_client = as_dask_client(client) + operation = SubmissionPermitOperation( + duration=2, + timeout=1, + max_clock_rate=1, + clock_margin=0, + clock=clock, + ) + await operation.acquire(dask_client) + clock.now = -1 + with pytest.raises(SubmissionPermitRejectedError, match="backwards"): + operation.ensure_valid() + + asyncio.run(run()) diff --git a/distributed/tests/test_submission_permit_extension.py b/distributed/tests/test_submission_permit_extension.py new file mode 100644 index 0000000000..6aae8a7aaa --- /dev/null +++ b/distributed/tests/test_submission_permit_extension.py @@ -0,0 +1,396 @@ +from __future__ import annotations + +import asyncio + +import pytest + +import distributed.scheduler as scheduler_module +from distributed._submission_permit_extension import SubmissionPermitExtension +from distributed._submission_permits import ClosedPermitError, UnknownPermitError +from distributed.core import Status +from distributed.metrics import time +from distributed.scheduler import DEFAULT_EXTENSIONS +from distributed.utils_test import async_poll_for, gen_cluster, inc + + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class ClockedExtension(SubmissionPermitExtension): + def __init__(self, scheduler): + self.clock = Clock() + super().__init__( + scheduler, + max_duration=10, + max_pending_per_client=2, + max_pending=4, + max_outcomes_per_client=2, + clock=self.clock, + ) + + +SCHEDULER_KWARGS = { + "extensions": {**DEFAULT_EXTENSIONS, "submission-permits": ClockedExtension} +} + + +async def acquire(c, sequence=1, duration=5): + return await c.scheduler.submission_permit_acquire( + client=c.id, + epoch=c._submission_permit_capabilities["epoch"], + sequence=sequence, + duration=duration, + ) + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_rpc_acknowledges_live_permit_and_retry_does_not_extend(c, s): + ext = s.extensions["submission-permits"] + s.periodic_callbacks["idle-timeout"].stop() + assert c._submission_permit_capabilities == { + "version": 1, + "epoch": c._submission_permit_capabilities["epoch"], + "max_duration": 10, + } + first = await acquire(c) + assert first == {"sequence": 1, "state": "pending", "duration": 5, "remaining": 5} + assert ext.has_pending() + s.idle_timeout = 0.01 + s.idle_since = time() - 1 + assert s.check_idle() is None + assert s.status == Status.running + ext.clock.now = 2 + retry = await acquire(c) + assert retry["remaining"] == 3 + assert retry["state"] == "pending" + epoch = c._submission_permit_capabilities["epoch"] + result = await c.scheduler.submission_permit_abort( + client=c.id, epoch=epoch, sequence=1 + ) + assert result["state"] == "aborted" + assert not ext.has_pending() + assert ( + await c.scheduler.submission_permit_status(client=c.id, epoch=epoch, sequence=1) + == result + ) + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_expiry_allows_idle_shutdown_and_fences_grants_before_close(c, s): + ext = s.extensions["submission-permits"] + s.periodic_callbacks["idle-timeout"].stop() + await acquire(c) + assert s.check_idle() is None + ext.clock.now = 5 + assert s.check_idle() is not None + assert not ext.has_pending() + s.idle_timeout = 0.01 + s.idle_since = time() - 1 + s.check_idle() + # close has been queued but cannot have run until this coroutine yields. + assert s.status == Status.running + with pytest.raises(ClosedPermitError): + ext.acquire(c.id, c._submission_permit_capabilities["epoch"], 2, 1) + await s.finished() + assert s.status == Status.closed + + +@gen_cluster(client=True, nthreads=[], scheduler_kwargs=SCHEDULER_KWARGS) +async def test_rpc_rejects_unknown_or_stale_connection(c, s): + with pytest.raises(UnknownPermitError, match="stale"): + await c.scheduler.submission_permit_acquire( + client=c.id, epoch="old-epoch", sequence=1, duration=1 + ) + with pytest.raises(ClosedPermitError, match="not running"): + await c.scheduler.submission_permit_acquire( + client="not-connected", epoch="missing", sequence=1, duration=1 + ) + assert not s.extensions["submission-permits"].has_pending() + + +@gen_cluster( + client=True, + nthreads=[("127.0.0.1", 1)], + scheduler_kwargs=SCHEDULER_KWARGS, +) +async def test_tagged_graph_transfers_into_active_guard_and_computes(c, s, a): + ext = s.extensions["submission-permits"] + s.periodic_callbacks["idle-timeout"].stop() + epoch = c._submission_permit_capabilities["epoch"] + await acquire(c) + accepted = asyncio.Queue() + c._stream_handlers["submission-permit-admission"] = lambda **msg: ( + accepted.put_nowait(msg) + ) + transfer = ext.transfer + observations = [] + + def observe_transfer(*args): + fresh = transfer(*args) + observations.append( + (s._active_graph_updates, ext.has_pending(), s.check_idle()) + ) + return fresh + + ext.transfer = observe_transfer + send = c._send_to_scheduler + + def tagged_send(msg): + if msg["op"] == "update-graph": + msg = dict(msg, submission_epoch=epoch, submission_sequence=1) + send(msg) + + c._send_to_scheduler = tagged_send + future = c.submit(inc, 1, key="protected") + assert await accepted.get() == {"epoch": epoch, "sequence": 1, "status": "accepted"} + assert await future == 2 + assert observations == [(1, False, None)] + assert ext.registry.status(c.id, epoch, 1).state == "accepted" + assert s._active_graph_updates == 0 + s.validate_state() + + +@gen_cluster(client=True, nthreads=[("127.0.0.1", 1)]) +async def test_legacy_scheduler_has_no_capability_and_computes(c, s, a): + assert c._submission_permit_capabilities is None + assert "submission-permits" not in s.extensions + assert "submission_permit_acquire" not in s.handlers + assert await c.submit(inc, 1) == 2 + + +@gen_cluster( + client=True, + nthreads=[("127.0.0.1", 1)], + scheduler_kwargs=SCHEDULER_KWARGS, +) +async def test_untagged_graph_works_with_optional_extension(c, s, a): + assert await c.submit(inc, 1) == 2 + assert not s.extensions["submission-permits"].has_pending() + + +@pytest.mark.parametrize( + "case,reason", + [ + ("expired", "ExpiredPermitError"), + ("aborted", "AbortedPermitError"), + ("retired", "RetiredPermitError"), + ("unknown", "UnknownPermitError"), + ("stale", "UnknownPermitError"), + ("missing-epoch", "ValueError"), + ("missing-sequence", "ValueError"), + ("invalid-sequence", "ValueError"), + ("null-tags", "ValueError"), + ("invalid-epoch", "ValueError"), + ("closed", "AbortedPermitError"), + ], +) +@gen_cluster( + client=True, + nthreads=[("127.0.0.1", 1)], + scheduler_kwargs=SCHEDULER_KWARGS, +) +async def test_rejected_admission_does_not_process_graph_or_poison_key( + c, s, a, case, reason, monkeypatch +): + future = c.submit(inc, 1, key="existing") + assert await future == 2 + state = s.tasks[future.key] + ext = s.extensions["submission-permits"] + epoch = c._submission_permit_capabilities["epoch"] + await acquire(c) + tag = {"submission_epoch": epoch, "submission_sequence": 1} + if case == "expired": + ext.clock.now = 5 + elif case == "aborted": + ext.abort(c.id, epoch, 1) + elif case == "retired": + for sequence in (1, 2, 3): + ext.acquire(c.id, epoch, sequence, 1) + ext.abort(c.id, epoch, sequence) + elif case == "unknown": + tag["submission_sequence"] = 99 + elif case == "stale": + tag["submission_epoch"] = "previous-connection" + elif case == "missing-epoch": + tag.pop("submission_epoch") + elif case == "missing-sequence": + tag.pop("submission_sequence") + elif case == "invalid-sequence": + tag["submission_sequence"] = True + elif case == "null-tags": + tag = {"submission_epoch": None, "submission_sequence": None} + elif case == "invalid-epoch": + tag["submission_epoch"] = [epoch] + elif case == "closed": + ext.commit_idle_shutdown() + results = asyncio.Queue() + c._stream_handlers["submission-permit-admission"] = lambda **msg: ( + results.put_nowait(msg) + ) + graph_calls = [] + key_messages = [] + + def unexpected_deserialize(*args): + graph_calls.append(args) + raise AssertionError("rejected request entered graph deserialization") + + report = s.report + + def observe_report(msg, *args, **kwargs): + if msg["op"] in ("task-erred", "cancelled-keys"): + key_messages.append(msg) + report(msg, *args, **kwargs) + + monkeypatch.setattr(scheduler_module, "deserialize", unexpected_deserialize) + monkeypatch.setattr(s, "report", observe_report) + c._send_to_scheduler( + { + "op": "update-graph", + "expr_ser": None, + "keys": {"existing", "unsubmitted"}, + "span_metadata": {}, + "internal_priority": None, + "submitting_task": None, + **tag, + } + ) + result = await results.get() + assert result["status"] == "rejected" + assert result["reason"] == reason + assert results.empty() + assert not graph_calls + assert not key_messages + assert "unsubmitted" not in s.tasks + assert s.tasks[future.key] is state + assert state.state == "memory" + assert await future == 2 + assert s._active_graph_updates == 0 + + +@pytest.mark.parametrize("same_graph", [False, True]) +@gen_cluster( + client=True, + nthreads=[("127.0.0.1", 1)], + scheduler_kwargs=SCHEDULER_KWARGS, +) +async def test_consumed_sequence_rejects_every_replay(c, s, a, same_graph): + await acquire(c) + epoch = c._submission_permit_capabilities["epoch"] + results = asyncio.Queue() + c._stream_handlers["submission-permit-admission"] = lambda **msg: ( + results.put_nowait(msg) + ) + send = c._send_to_scheduler + messages = [] + + def tagged_send(msg): + if msg["op"] == "update-graph": + msg = dict(msg, submission_epoch=epoch, submission_sequence=1) + messages.append(msg.copy()) + send(msg) + + c._send_to_scheduler = tagged_send + future = c.submit(inc, 1, key="original") + assert (await results.get())["status"] == "accepted" + assert await future == 2 + replay = messages[0] + if not same_graph: + replay = dict(replay, expr_ser=None, keys={"other-graph"}) + send(replay) + assert await results.get() == { + "epoch": epoch, + "sequence": 1, + "status": "rejected", + "reason": "sequence-already-consumed", + } + assert results.empty() + assert "other-graph" not in s.tasks + assert await future == 2 + assert s._active_graph_updates == 0 + + +@gen_cluster(client=True, nthreads=[("127.0.0.1", 1)]) +async def test_unsupported_tag_rejects_without_graph_work(c, s, a): + results = asyncio.Queue() + c._stream_handlers["submission-permit-admission"] = lambda **msg: ( + results.put_nowait(msg) + ) + c._send_to_scheduler( + { + "op": "update-graph", + "expr_ser": None, + "keys": {"unsupported"}, + "span_metadata": {}, + "internal_priority": None, + "submitting_task": None, + "submission_epoch": "unknown", + "submission_sequence": 1, + } + ) + result = await results.get() + assert result["status"] == "rejected" + assert result["reason"] == "ValueError" + assert not s.tasks + assert s._active_graph_updates == 0 + + +@pytest.mark.parametrize("disconnect", [False, True]) +@gen_cluster( + client=True, + nthreads=[("127.0.0.1", 1)], + scheduler_kwargs=SCHEDULER_KWARGS, +) +async def test_admission_precedes_materialization_and_checks_connection_afterward( + c, s, a, disconnect, monkeypatch +): + ext = s.extensions["submission-permits"] + epoch = c._submission_permit_capabilities["epoch"] + await acquire(c) + results = asyncio.Queue() + c._stream_handlers["submission-permit-admission"] = lambda **msg: ( + results.put_nowait(msg) + ) + entered = asyncio.Event() + release = asyncio.Event() + offload = scheduler_module.offload + + async def paused_offload(func, *args, **kwargs): + if func is scheduler_module._materialize_graph: + entered.set() + await release.wait() + return await offload(func, *args, **kwargs) + + monkeypatch.setattr(scheduler_module, "offload", paused_offload) + send = c._send_to_scheduler + + def tagged_send(msg): + if msg["op"] == "update-graph": + msg = dict(msg, submission_epoch=epoch, submission_sequence=1) + send(msg) + + c._send_to_scheduler = tagged_send + try: + future = c.submit(inc, 1, key="paused-graph") + assert (await results.get())["status"] == "accepted" + await entered.wait() + assert s._active_graph_updates == 1 + assert ext.registry.status(c.id, epoch, 1).state == "accepted" + assert not s.tasks + assert s.check_idle() is None + if disconnect: + # Close the captured scheduler-side transport while the handler is + # awaiting graph preparation. Reconnect must wait for this handler. + s.client_comms[c.id].comm.abort() + release.set() + if disconnect: + await async_poll_for(lambda: s._active_graph_updates == 0) + assert "paused-graph" not in s.tasks + else: + assert await future == 2 + assert results.empty() + finally: + release.set() diff --git a/distributed/tests/test_submission_permit_reconnect.py b/distributed/tests/test_submission_permit_reconnect.py new file mode 100644 index 0000000000..db7a53d18b --- /dev/null +++ b/distributed/tests/test_submission_permit_reconnect.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +import asyncio +from functools import partial +from typing import Any + +import pytest + +from distributed._submission_permit_extension import SubmissionPermitExtension +from distributed.comm.core import Comm, CommClosedError, connect +from distributed.scheduler import DEFAULT_EXTENSIONS +from distributed.utils import wait_for +from distributed.utils_test import async_poll_for, gen_cluster + +pytestmark = pytest.mark.ci1 + + +PERMIT_EXTENSIONS = { + **DEFAULT_EXTENSIONS, + "submission-permits": partial( + SubmissionPermitExtension, + max_duration=10, + max_pending_per_client=2, + max_pending=4, + max_outcomes_per_client=2, + ), +} + + +@gen_cluster(nthreads=[]) +async def test_legacy_overlapping_client_cleanup_still_closes_replacement( + s, monkeypatch +): + client = "legacy-same-client" + first = await connect(s.address) + second = await connect(s.address) + replacement = None + finished = { + first.local_address: asyncio.Event(), + second.local_address: asyncio.Event(), + } + add_client = s.handlers["register-client"] + + async def track_client(comm: Comm, client: str, versions: dict[str, Any]) -> None: + peer = comm.peer_address + try: + await add_client(comm=comm, client=client, versions=versions) + finally: + if peer in finished: + finished[peer].set() + + async def register(comm: Comm) -> None: + await comm.write( + {"op": "register-client", "client": client, "reply": False, "versions": {}} + ) + messages = await wait_for(comm.read(), 1) + assert messages[0]["op"] == "stream-start" + assert "submission-permits" not in messages[0] + + monkeypatch.setitem(s.handlers, "register-client", track_client) + try: + await register(first) + await register(second) + await first.close() + await wait_for(finished[first.local_address].wait(), 1) + assert client not in s.client_comms + with pytest.raises(CommClosedError): + await wait_for(second.read(), 1) + # Both finalizers must finish before a third registration can race them. + await wait_for( + asyncio.gather(*(event.wait() for event in finished.values())), 1 + ) + assert client not in s.clients + assert client not in s.client_comms + + replacement = await connect(s.address) + await register(replacement) + assert client in s.clients + assert client in s.client_comms + assert not s.client_comms[client].closed() + finally: + for comm in (replacement, second, first): + if comm is not None and not comm.closed(): + await comm.close() + + +async def _register(comm: Comm, client: str) -> dict: + await comm.write( + { + "op": "register-client", + "client": client, + "reply": False, + "versions": {}, + } + ) + msg = await wait_for(comm.read(), 1) + assert msg[0]["op"] == "stream-start" + return msg[0]["submission-permits"] + + +async def _assert_registration_is_refused(comm: Comm, client: str) -> None: + await comm.write( + { + "op": "register-client", + "client": client, + "reply": False, + "versions": {}, + } + ) + with pytest.raises(CommClosedError): + await wait_for(comm.read(), 1) + + +@gen_cluster(nthreads=[], scheduler_kwargs={"extensions": PERMIT_EXTENSIONS}) +async def test_submission_permits_reject_overlapping_client_registration(s): + client = "same-client" + first_comm = await connect(s.address) + second_comm = None + third_comm = None + replacement_comm = None + allow_close = None + try: + first_capabilities = await _register(first_comm, client) + old_epoch = first_capabilities["epoch"] + old_bcomm = s.client_comms[client] + added_total = s._client_connections_added_total + old_client_state = s.clients[client] + + second_comm = await connect(s.address) + await _assert_registration_is_refused(second_comm, client) + assert s.clients[client] is old_client_state + assert s.client_comms[client] is old_bcomm + assert s.extensions["submission-permits"].registry.is_current(client, old_epoch) + assert s._client_connections_added_total == added_total + + close_started = asyncio.Event() + allow_close = asyncio.Event() + original_close = old_bcomm.close + + async def delayed_close() -> None: + close_started.set() + await allow_close.wait() + await original_close() + + old_bcomm.close = delayed_close + await first_comm.close() + await wait_for(close_started.wait(), 1) + + # The old BatchedSend remains the connection fence until its close completes. + assert s.client_comms[client] is old_bcomm + third_comm = await connect(s.address) + await _assert_registration_is_refused(third_comm, client) + assert s.client_comms[client] is old_bcomm + + allow_close.set() + await async_poll_for(lambda: client not in s.client_comms) + + replacement_comm = await connect(s.address) + new_capabilities = await _register(replacement_comm, client) + new_epoch = new_capabilities["epoch"] + assert new_epoch != old_epoch + extension = s.extensions["submission-permits"] + assert extension.registry.is_current(client, new_epoch) + + # A delayed old disconnect cleanup must be fenced by the old epoch. + extension.unregister_client(client, old_epoch) + assert extension.registry.is_current(client, new_epoch) + assert extension.acquire(client, new_epoch, 1, 1)["state"] == "pending" + finally: + if allow_close is not None: + allow_close.set() + for comm in (second_comm, third_comm, replacement_comm, first_comm): + if comm is not None and not comm.closed(): + await comm.close() + + +@gen_cluster( + client=True, nthreads=[], scheduler_kwargs={"extensions": PERMIT_EXTENSIONS} +) +async def test_submission_permits_client_reconnect_waits_for_old_comm_cleanup(c, s): + client = c.id + old_bcomm = s.client_comms[client] + old_epoch = c._submission_permit_capabilities["epoch"] + added_total = s._client_connections_added_total + close_started = asyncio.Event() + allow_close = asyncio.Event() + original_close = old_bcomm.close + + async def delayed_close() -> None: + close_started.set() + await allow_close.wait() + await original_close() + + old_bcomm.close = delayed_close + try: + await c.scheduler_comm.comm.close() + await wait_for(close_started.wait(), 1) + + # The client's automatic retries use the same ID but cannot replace the + # still-closing connection or receive a fresh capability epoch yet. + await async_poll_for(lambda: c.status == "connecting") + assert c._submission_permit_capabilities is None + assert s.client_comms[client] is old_bcomm + assert s._client_connections_added_total == added_total + + allow_close.set() + await async_poll_for(lambda: c.status == "running") + new_epoch = c._submission_permit_capabilities["epoch"] + assert new_epoch != old_epoch + assert s.extensions["submission-permits"].registry.is_current(client, new_epoch) + + reply = await c.scheduler.submission_permit_acquire( + client=client, epoch=new_epoch, sequence=1, duration=1 + ) + assert reply["state"] == "pending" + finally: + allow_close.set() + + +@gen_cluster(nthreads=[], scheduler_kwargs={"extensions": PERMIT_EXTENSIONS}) +async def test_submission_permits_reject_client_registration_after_idle_commit(s): + extension = s.extensions["submission-permits"] + extension.commit_idle_shutdown() + + client = "after-idle-commit" + comm = await connect(s.address) + try: + await _assert_registration_is_refused(comm, client) + assert client not in s.clients + assert client not in s.client_comms + assert not extension.registry.is_current(client, "not-an-epoch") + assert s._client_connections_added_total == 0 + finally: + if not comm.closed(): + await comm.close() diff --git a/distributed/tests/test_submission_permits.py b/distributed/tests/test_submission_permits.py new file mode 100644 index 0000000000..44fdd68993 --- /dev/null +++ b/distributed/tests/test_submission_permits.py @@ -0,0 +1,301 @@ +from __future__ import annotations + +from dataclasses import asdict + +import pytest + +from distributed._submission_permits import ( + _HEAP_SLACK, + AbortedPermitError, + ClosedPermitError, + ExpiredPermitError, + PermitCapacityError, + RetiredPermitError, + SubmissionPermitRegistry, + UnknownPermitError, +) + + +class Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +@pytest.fixture +def registry() -> tuple[SubmissionPermitRegistry, Clock, str]: + clock = Clock() + permits = SubmissionPermitRegistry(10, 2, 3, 2, clock) + return permits, clock, permits.register("client") + + +def test_acquire_retry_and_transfer_are_one_shot( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, clock, epoch = registry + first = permits.acquire("client", epoch, 1, 5) + clock.now = 2 + retry = permits.acquire("client", epoch, 1, 5) + + assert first == first.__class__(1, "pending", 5.0, 5.0) + assert retry == retry.__class__(1, "pending", 5.0, 3.0) + assert asdict(retry) == retry.to_dict() + assert permits.transfer("client", epoch, 1) is True + assert permits.transfer("client", epoch, 1) is False + assert permits.status("client", epoch, 1).state == "accepted" + + +def test_overlap_abort_and_expiry( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, clock, epoch = registry + permits.acquire("client", epoch, 1, 2) + permits.acquire("client", epoch, 2, 4) + assert permits.has_pending() + assert permits.abort("client", epoch, 1).state == "aborted" + assert permits.abort("client", epoch, 1).state == "aborted" + with pytest.raises(AbortedPermitError): + permits.transfer("client", epoch, 1) + + clock.now = 4 + assert not permits.has_pending() + assert permits.status("client", epoch, 2).state == "expired" + with pytest.raises(ExpiredPermitError): + permits.transfer("client", epoch, 2) + + +def test_terminal_outcome_eviction_retires_old_sequences( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, _, epoch = registry + for sequence in (1, 2, 3): + permits.acquire("client", epoch, sequence, 1) + assert permits.transfer("client", epoch, sequence) + + assert permits.status("client", epoch, 1).state == "retired" + with pytest.raises(RetiredPermitError): + permits.acquire("client", epoch, 1, 1) + with pytest.raises(RetiredPermitError): + permits.transfer("client", epoch, 1) + assert permits.status("client", epoch, 4).state == "unknown" + + +def test_out_of_order_sequence_is_rejected_after_a_higher_admission( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, _, epoch = registry + permits.acquire("client", epoch, 2, 1) + with pytest.raises(RetiredPermitError): + permits.acquire("client", epoch, 1, 1) + + +def test_pending_capacities_do_not_consume_admission_sequence() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(5, 1, 1, 1, clock) + first_epoch = permits.register("first") + second_epoch = permits.register("second") + permits.acquire("first", first_epoch, 1, 1) + with pytest.raises(PermitCapacityError): + permits.acquire("first", first_epoch, 2, 1) + with pytest.raises(PermitCapacityError): + permits.acquire("second", second_epoch, 1, 1) + assert permits.abort("first", first_epoch, 1).state == "aborted" + assert permits.acquire("second", second_epoch, 1, 1).state == "pending" + + +def test_reconnect_invalidates_old_pending_and_stale_cleanup( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, _, old_epoch = registry + permits.acquire("client", old_epoch, 1, 1) + new_epoch = permits.register("client") + + assert not permits.is_current("client", old_epoch) + assert permits.is_current("client", new_epoch) + assert not permits.unregister("client", old_epoch) + assert permits.acquire("client", new_epoch, 1, 1).state == "pending" + with pytest.raises(UnknownPermitError): + permits.status("client", old_epoch, 1) + assert permits.unregister("client", new_epoch) + assert not permits.has_pending() + + +def test_close_aborts_live_permits_and_rejects_new_registration_and_grants( + registry: tuple[SubmissionPermitRegistry, Clock, str], +) -> None: + permits, _, epoch = registry + permits.acquire("client", epoch, 1, 1) + permits.close() + + assert permits.status("client", epoch, 1).state == "aborted" + assert not permits.has_pending() + with pytest.raises(ClosedPermitError): + permits.acquire("client", epoch, 2, 1) + with pytest.raises(ClosedPermitError): + permits.register("new-client") + + +@pytest.mark.parametrize( + "kwargs", + [ + {"max_duration": 0}, + {"max_duration": float("inf")}, + {"max_duration": True}, + {"max_pending_per_client": 0}, + {"max_pending": False}, + {"max_outcomes_per_client": -1}, + ], +) +def test_constructor_rejects_invalid_bounds(kwargs: dict[str, object]) -> None: + options: dict[str, object] = { + "max_duration": 1, + "max_pending_per_client": 1, + "max_pending": 1, + "max_outcomes_per_client": 1, + } + options.update(kwargs) + with pytest.raises(ValueError): + SubmissionPermitRegistry(**options) # type: ignore[arg-type] + + +@pytest.mark.parametrize("sequence", [0, -1, True, 2**63]) +def test_rejects_invalid_sequences( + registry: tuple[SubmissionPermitRegistry, Clock, str], sequence: object +) -> None: + permits, _, epoch = registry + with pytest.raises(ValueError): + permits.acquire("client", epoch, sequence, 1) # type: ignore[arg-type] + + +@pytest.mark.parametrize("duration", [0, -1, True, float("nan"), float("inf"), 11]) +def test_rejects_invalid_durations( + registry: tuple[SubmissionPermitRegistry, Clock, str], duration: object +) -> None: + permits, _, epoch = registry + with pytest.raises(ValueError): + permits.acquire("client", epoch, 1, duration) # type: ignore[arg-type] + + +def assert_heap_bound(permits: SubmissionPermitRegistry) -> None: + assert len(permits._deadlines) <= 2 * permits._pending + _HEAP_SLACK + assert all(len(entry) == 4 for entry in permits._deadlines) + + +def test_deadline_heap_stays_bounded_through_long_terminal_churn() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(100, 1, 2, 8, clock) + epoch = permits.register("client") + for sequence in range(1, 20_001): + permits.acquire("client", epoch, sequence, 50) + if sequence % 2: + assert permits.transfer("client", epoch, sequence) + else: + assert permits.abort("client", epoch, sequence).state == "aborted" + assert_heap_bound(permits) + + generation = permits._generations["client"] + assert permits._pending == 0 + assert not generation.active + assert len(generation.outcomes) == 8 + assert generation.high_watermark == 20_000 + + +def test_deadline_heap_ignores_replaced_and_unregistered_epochs() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(100, 1, 10, 4, clock) + old_epoch = permits.register("client") + permits.acquire("client", old_epoch, 1, 1) + replacement = permits.register("client") + permits.acquire("client", replacement, 1, 50) + transient = permits.register("transient") + permits.acquire("transient", transient, 1, 1) + assert permits.unregister("transient", transient) + + clock.now = 1 + assert permits.status("client", replacement, 1).state == "pending" + assert permits._pending == 1 + assert_heap_bound(permits) + assert all(entry[2] == replacement for entry in permits._deadlines) + + +def test_deadline_heap_mixed_live_and_churn_has_one_live_entry_per_permit() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(100, 1, 32, 4, clock) + live = [] + for index in range(20): + client = f"live-{index}" + epoch = permits.register(client) + permits.acquire(client, epoch, 1, 50) + live.append((client, epoch)) + + churn_epoch = permits.register("churn") + for sequence in range(1, 20_001): + permits.acquire("churn", churn_epoch, sequence, 50) + permits.abort("churn", churn_epoch, sequence) + assert permits._pending == len(live) + assert_heap_bound(permits) + live_entries = [ + entry for entry in permits._deadlines if permits._is_live_entry(entry) + ] + assert len(live_entries) == len(live) + assert {(client, epoch, 1) for _, client, epoch, _ in live_entries} == { + (client, epoch, 1) for client, epoch in live + } + + +def test_deadline_heap_long_reregister_and_unregister_churn_ignores_old_epochs() -> ( + None +): + clock = Clock() + permits = SubmissionPermitRegistry(100, 1, 3, 4, clock) + anchor_epoch = permits.register("anchor") + permits.acquire("anchor", anchor_epoch, 1, 50) + for _ in range(20_000): + old_epoch = permits.register("churn") + permits.acquire("churn", old_epoch, 1, 1) + replacement_epoch = permits.register("churn") + permits.acquire("churn", replacement_epoch, 1, 1) + assert permits.unregister("churn", replacement_epoch) + current_epoch = permits.register("churn") + permits.acquire("churn", current_epoch, 1, 1) + assert_heap_bound(permits) + + due_epoch = permits.register("due") + permits.acquire("due", due_epoch, 1, 1) + clock.now = 1 + assert permits.status("anchor", anchor_epoch, 1).state == "pending" + assert permits.status("due", due_epoch, 1).state == "expired" + assert permits._pending == 1 + assert_heap_bound(permits) + + +def test_pending_retry_does_not_add_heap_entry_and_close_clears_heap() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(10, 1, 1, 2, clock) + epoch = permits.register("client") + permits.acquire("client", epoch, 1, 5) + before = list(permits._deadlines) + permits.acquire("client", epoch, 1, 5) + assert permits._deadlines == before + permits.close() + assert not permits._deadlines + assert permits._pending == 0 + + +def test_equal_deadlines_expire_together_after_capacity_rejection() -> None: + clock = Clock() + permits = SubmissionPermitRegistry(10, 1, 2, 2, clock) + first = permits.register("first") + second = permits.register("second") + third = permits.register("third") + permits.acquire("first", first, 1, 1) + permits.acquire("second", second, 1, 1) + with pytest.raises(PermitCapacityError): + permits.acquire("third", third, 1, 1) + clock.now = 1 + assert not permits.has_pending() + assert permits.status("first", first, 1).state == "expired" + assert permits.status("second", second, 1).state == "expired" + assert not permits._deadlines