dvt-distributed 2026.3.6__py3-none-any.whl

distributed/__init__.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+# isort: off
+from distributed import config  # load distributed configuration first
+from distributed import widgets  # load distributed widgets second
+
+# isort: on
+
+import atexit
+import weakref
+
+# This finalizer registers an atexit handler that has to happen before
+# distributed registers its handlers, otherwise we observe hangs on
+# cluster shutdown when using the UCX comms backend. See
+# https://github.com/dask/distributed/issues/7726 for more discussion
+# of the problem and the search for long term solutions
+
+weakref.finalize(lambda: None, lambda: None)
+import dask
+from dask.config import config  # type: ignore
+
+from distributed.actor import Actor, ActorFuture, BaseActorFuture
+from distributed.client import (
+    Client,
+    CompatibleExecutor,
+    Future,
+    as_completed,
+    default_client,
+    fire_and_forget,
+    futures_of,
+    get_task_metadata,
+    get_task_stream,
+    performance_report,
+    wait,
+)
+from distributed.core import Status, connect, rpc
+from distributed.deploy import (
+    Adaptive,
+    LocalCluster,
+    SpecCluster,
+    SSHCluster,
+    SubprocessCluster,
+)
+from distributed.diagnostics.plugin import (
+    CondaInstall,
+    Environ,
+    InstallPlugin,
+    NannyPlugin,
+    PipInstall,
+    SchedulerPlugin,
+    UploadDirectory,
+    UploadFile,
+    WorkerPlugin,
+)
+from distributed.diagnostics.progressbar import progress
+from distributed.event import Event
+from distributed.lock import Lock
+from distributed.multi_lock import MultiLock
+from distributed.nanny import Nanny
+from distributed.queues import Queue
+from distributed.scheduler import KilledWorker, Scheduler
+from distributed.security import Security
+from distributed.semaphore import Semaphore
+from distributed.spans import span
+from distributed.threadpoolexecutor import rejoin
+from distributed.utils import CancelledError, TimeoutError, sync
+from distributed.variable import Variable
+from distributed.worker import (
+    Reschedule,
+    Worker,
+    get_client,
+    get_worker,
+    print,
+    secede,
+    warn,
+)
+from distributed.worker_client import local_client, worker_client
+
+try:
+    # Backwards compatibility with versioneer
+    from distributed._version import __commit_id__ as __git_revision__
+    from distributed._version import __version__
+except ImportError:
+    __git_revision__ = "unknown"
+    __version__ = "unknown"
+
+
+__all__ = [
+    "Actor",
+    "ActorFuture",
+    "Adaptive",
+    "BaseActorFuture",
+    "CancelledError",
+    "Client",
+    "CompatibleExecutor",
+    "CondaInstall",
+    "Environ",
+    "Event",
+    "Future",
+    "KilledWorker",
+    "LocalCluster",
+    "Lock",
+    "MultiLock",
+    "Nanny",
+    "NannyPlugin",
+    "InstallPlugin",
+    "PipInstall",
+    "Queue",
+    "Reschedule",
+    "SSHCluster",
+    "Scheduler",
+    "SchedulerPlugin",
+    "Security",
+    "Semaphore",
+    "SpecCluster",
+    "Status",
+    "SubprocessCluster",
+    "TimeoutError",
+    "UploadDirectory",
+    "UploadFile",
+    "Variable",
+    "Worker",
+    "WorkerPlugin",
+    "as_completed",
+    "config",
+    "connect",
+    "dask",
+    "default_client",
+    "fire_and_forget",
+    "futures_of",
+    "get_client",
+    "get_task_metadata",
+    "get_task_stream",
+    "get_worker",
+    "local_client",
+    "performance_report",
+    "print",
+    "progress",
+    "rejoin",
+    "rpc",
+    "secede",
+    "span",
+    "sync",
+    "wait",
+    "warn",
+    "widgets",
+    "worker_client",
+]

distributed/_async_taskgroup.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any, TypeVar
+
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+    R = TypeVar("R")
+    T = TypeVar("T")
+    Coro = Coroutine[Any, Any, T]
+
+
+class _LoopBoundMixin:
+    """Backport of the private asyncio.mixins._LoopBoundMixin from 3.11"""
+
+    _global_lock = threading.Lock()
+
+    _loop = None
+
+    def _get_loop(self):
+        loop = asyncio.get_running_loop()
+
+        if self._loop is None:
+            with self._global_lock:
+                if self._loop is None:
+                    self._loop = loop
+        if loop is not self._loop:
+            raise RuntimeError(f"{self!r} is bound to a different event loop")
+        return loop
+
+
+class AsyncTaskGroupClosedError(RuntimeError):
+    pass
+
+
+def _delayed(corofunc: Callable[P, Coro[T]], delay: float) -> Callable[P, Coro[T]]:
+    """Decorator to delay the evaluation of a coroutine function by the given delay in seconds."""
+
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        await asyncio.sleep(delay)
+        return await corofunc(*args, **kwargs)
+
+    return wrapper
+
+
+class AsyncTaskGroup(_LoopBoundMixin):
+    """Collection tracking all currently running asynchronous tasks within a group"""
+
+    #: If True, the group is closed and does not allow adding new tasks.
+    closed: bool
+
+    def __init__(self) -> None:
+        self.closed = False
+        self._ongoing_tasks: set[asyncio.Task[None]] = set()
+
+    def call_soon(
+        self, afunc: Callable[P, Coro[None]], /, *args: P.args, **kwargs: P.kwargs
+    ) -> None:
+        """Schedule a coroutine function to be executed as an `asyncio.Task`.
+
+        The coroutine function `afunc` is scheduled with `args` arguments and `kwargs` keyword arguments
+        as an `asyncio.Task`.
+
+        Parameters
+        ----------
+        afunc
+            Coroutine function to schedule.
+        *args
+            Arguments to be passed to `afunc`.
+        **kwargs
+            Keyword arguments to be passed to `afunc`
+
+        Returns
+        -------
+            None
+
+        Raises
+        ------
+        AsyncTaskGroupClosedError
+            If the task group is closed.
+        """
+        if self.closed:  # Avoid creating a coroutine
+            raise AsyncTaskGroupClosedError(
+                "Cannot schedule a new coroutine function as the group is already closed."
+            )
+        task = self._get_loop().create_task(afunc(*args, **kwargs))
+        task.add_done_callback(self._ongoing_tasks.remove)
+        self._ongoing_tasks.add(task)
+        return None
+
+    def call_later(
+        self,
+        delay: float,
+        afunc: Callable[P, Coro[None]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> None:
+        """Schedule a coroutine function to be executed after `delay` seconds as an `asyncio.Task`.
+
+        The coroutine function `afunc` is scheduled with `args` arguments and `kwargs` keyword arguments
+        as an `asyncio.Task` that is executed after `delay` seconds.
+
+        Parameters
+        ----------
+        delay
+            Delay in seconds.
+        afunc
+            Coroutine function to schedule.
+        *args
+            Arguments to be passed to `afunc`.
+        **kwargs
+            Keyword arguments to be passed to `afunc`
+
+        Returns
+        -------
+            The None
+
+        Raises
+        ------
+        AsyncTaskGroupClosedError
+            If the task group is closed.
+        """
+        self.call_soon(_delayed(afunc, delay), *args, **kwargs)
+
+    def close(self) -> None:
+        """Closes the task group so that no new tasks can be scheduled.
+
+        Existing tasks continue to run.
+        """
+        self.closed = True
+
+    async def stop(self) -> None:
+        """Close the group and stop all currently running tasks.
+
+        Closes the task group and cancels all tasks. All tasks are cancelled
+        an additional time for each time this task is cancelled.
+        """
+        self.close()
+
+        current_task = asyncio.current_task(self._get_loop())
+        err = None
+        while tasks_to_stop := (self._ongoing_tasks - {current_task}):
+            for task in tasks_to_stop:
+                task.cancel()
+            try:
+                await asyncio.wait(tasks_to_stop)
+            except asyncio.CancelledError as e:
+                err = e
+
+        if err is not None:
+            raise err
+
+    def __len__(self):
+        return len(self._ongoing_tasks)

distributed/_asyncio.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+
+class RLock:
+    """asyncio reentrant lock, which allows the same owner (or two owners that compare
+    as equal) to be inside the critical section at the same time.
+
+    Note that here the owner is an explicit, generic object, whereas in
+    :func:`threading.RLock` and :func:`multiprocessing.RLock` it's hardcoded
+    respectively to the thread ID and process ID.
+
+    **Usage**::
+
+        lock = RLock()
+        async with lock("my-owner"):
+            ...
+
+    **Tip**
+
+    You can mix reentrant and non-reentrant owners; all you need to do is create an
+    owner that doesn't compare as equal to other instances of itself::
+
+        lock = RLock()
+
+        async def non_reentrant():
+            async with lock(object()):
+                ...
+
+        async def reentrant():
+            async with lock("foo"):
+                ...
+
+    In the above example, at any time you may have inside the critical section
+    at most one call to ``non_reentrant`` and no calls to ``reentrant``, or any number
+    of calls to ``reentrant`` but no calls to ``non_reentrant``.
+    """
+
+    _owner: object
+    _count: int
+    _lock: asyncio.Lock
+
+    def __init__(self):
+        self._owner = None
+        self._count = 0
+        self._lock = asyncio.Lock()
+
+    async def acquire(self, owner: object) -> None:
+        if self._count == 0 or self._owner != owner:
+            await self._lock.acquire()
+            self._owner = owner
+        self._count += 1
+
+    def release(self, owner: object) -> None:
+        if self._count == 0 or self._owner != owner:
+            raise RuntimeError("release unlocked lock or mismatched owner")
+        self._count -= 1
+        if self._count == 0:
+            self._owner = None
+            self._lock.release()
+
+    @asynccontextmanager
+    async def __call__(self, owner: object) -> AsyncIterator[None]:
+        await self.acquire(owner)
+        try:
+            yield
+        finally:
+            self.release(owner)

distributed/_concurrent_futures_thread.py

@@ -0,0 +1,166 @@
+# This was copied from CPython 3.6
+
+# Copyright 2009 Brian Quinlan. All Rights Reserved.
+# Licensed to PSF under a Contributor Agreement.
+
+"""Implements ThreadPoolExecutor."""
+
+from __future__ import annotations
+
+__author__ = "Brian Quinlan (brian@sweetapp.com)"
+
+import atexit
+import itertools
+import os
+import queue
+import threading
+import weakref
+from concurrent.futures import _base
+
+# Workers are created as daemon threads. This is done to allow the interpreter
+# to exit when there are still idle threads in a ThreadPoolExecutor's thread
+# pool (i.e. shutdown() was not called). However, allowing workers to die with
+# the interpreter has two undesirable properties:
+#   - The workers would still be running during interpreter shutdown,
+#     meaning that they would fail in unpredictable ways.
+#   - The workers could be killed while evaluating a work item, which could
+#     be bad if the callable being evaluated has external side-effects e.g.
+#     writing to a file.
+#
+# To work around this problem, an exit handler is installed which tells the
+# workers to exit when their work queues are empty and then waits until the
+# threads finish.
+
+_threads_queues: weakref.WeakKeyDictionary[threading.Thread, queue.Queue] = (
+    weakref.WeakKeyDictionary()
+)
+_shutdown = False
+
+
+def _python_exit():
+    global _shutdown
+    _shutdown = True
+    items = list(_threads_queues.items())
+    for _, q in items:
+        q.put(None)
+    for t, _ in items:
+        t.join()
+
+
+atexit.register(_python_exit)
+
+
+class _WorkItem:
+    def __init__(self, future, fn, args, kwargs):
+        self.future = future
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+
+    def run(self):
+        if not self.future.set_running_or_notify_cancel():  # pragma: no cover
+            return
+
+        try:
+            result = self.fn(*self.args, **self.kwargs)
+        except BaseException as e:
+            self.future.set_exception(e)
+        else:
+            self.future.set_result(result)
+
+
+def _worker(executor_reference, work_queue):
+    try:
+        while True:
+            work_item = work_queue.get(block=True)
+            if work_item is not None:
+                work_item.run()
+                # Delete references to object. See issue16284
+                del work_item
+                continue
+            executor = executor_reference()
+            # Exit if:
+            #   - The interpreter is shutting down OR
+            #   - The executor that owns the worker has been collected OR
+            #   - The executor that owns the worker has been shutdown.
+            if _shutdown or executor is None or executor._shutdown:
+                # Notice other workers
+                work_queue.put(None)
+                return
+            del executor
+    except BaseException:
+        _base.LOGGER.critical("Exception in worker", exc_info=True)
+
+
+class ThreadPoolExecutor(_base.Executor):
+    # Used to assign unique thread names when thread_name_prefix is not supplied.
+    _counter = itertools.count()
+
+    def __init__(self, max_workers=None, thread_name_prefix=""):
+        """Initializes a new ThreadPoolExecutor instance.
+
+        Args:
+            max_workers: The maximum number of threads that can be used to
+                execute the given calls.
+            thread_name_prefix: An optional name prefix to give our threads.
+        """
+        if max_workers is None:
+            # Use this number because ThreadPoolExecutor is often
+            # used to overlap I/O instead of CPU work.
+            max_workers = (os.cpu_count() or 1) * 5
+        if max_workers <= 0:
+            raise ValueError("max_workers must be greater than 0")
+
+        self._max_workers = max_workers
+        self._work_queue = queue.Queue()
+        self._threads = set()
+        self._shutdown = False
+        self._shutdown_lock = threading.Lock()
+        self._thread_name_prefix = thread_name_prefix or (
+            "ThreadPoolExecutor-%d" % next(self._counter)
+        )
+
+    def submit(self, fn, *args, **kwargs):
+        with self._shutdown_lock:
+            if self._shutdown:  # pragma: no cover
+                raise RuntimeError("cannot schedule new futures after shutdown")
+
+            f = _base.Future()
+            w = _WorkItem(f, fn, args, kwargs)
+
+            self._work_queue.put(w)
+            self._adjust_thread_count()
+            return f
+
+    submit.__doc__ = _base.Executor.submit.__doc__
+
+    def _adjust_thread_count(self):
+        # When the executor gets lost, the weakref callback will wake up
+        # the worker threads.
+        def weakref_cb(_, q=self._work_queue):
+            q.put(None)
+
+        # TODO(bquinlan): Should avoid creating new threads if there are more
+        # idle threads than items in the work queue.
+        num_threads = len(self._threads)
+        if num_threads < self._max_workers:
+            thread_name = "%s_%d" % (self._thread_name_prefix or self, num_threads)
+            t = threading.Thread(
+                name=thread_name,
+                target=_worker,
+                args=(weakref.ref(self, weakref_cb), self._work_queue),
+            )
+            t.daemon = True
+            t.start()
+            self._threads.add(t)
+            _threads_queues[t] = self._work_queue
+
+    def shutdown(self, wait=True):
+        with self._shutdown_lock:
+            self._shutdown = True
+            self._work_queue.put(None)
+        if wait:
+            for t in self._threads:
+                t.join()
+
+    shutdown.__doc__ = _base.Executor.shutdown.__doc__

distributed/_signals.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import signal
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+async def wait_for_signals() -> int:
+    """Wait for sigint or sigterm by setting global signal handlers"""
+    signals = (signal.SIGINT, signal.SIGTERM)
+    loop = asyncio.get_running_loop()
+    event = asyncio.Event()
+
+    old_handlers: dict[int, Any] = {}
+    caught_signal: int | None = None
+
+    def handle_signal(signum, frame):
+        # *** Do not log or print anything in here
+        # https://stackoverflow.com/questions/45680378/how-to-explain-the-reentrant-runtimeerror-caused-by-printing-in-signal-handlers
+        nonlocal caught_signal
+        caught_signal = signum
+        # Restore old signal handler to allow for quicker exit
+        # if the user sends the signal again.
+        signal.signal(signum, old_handlers[signum])
+        loop.call_soon_threadsafe(event.set)
+
+    for sig in signals:
+        old_handlers[sig] = signal.signal(sig, handle_signal)
+
+    try:
+        await event.wait()
+        assert caught_signal
+        logger.info(
+            "Received signal %s (%d)", signal.Signals(caught_signal).name, caught_signal
+        )
+        return caught_signal
+    finally:
+        for sig in signals:
+            signal.signal(sig, old_handlers[sig])

distributed/_stories.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from collections.abc import Collection, Iterable
+from typing import TYPE_CHECKING
+
+from dask.typing import Key
+
+if TYPE_CHECKING:
+    # Circular import
+    from distributed.scheduler import Transition
+
+
+def scheduler_story(
+    keys_or_stimuli: set[Key | str], transition_log: Iterable[Transition]
+) -> list[Transition]:
+    """Creates a story from the scheduler transition log given a set of keys
+    describing tasks or stimuli.
+
+    Parameters
+    ----------
+    keys_or_stimuli : set[str]
+        Task keys or stimulus_id's
+    log : iterable
+        The scheduler transition log
+
+    Returns
+    -------
+    story : list[tuple]
+    """
+    return [
+        t
+        for t in transition_log
+        if t[0] in keys_or_stimuli or keys_or_stimuli.intersection(t[3])
+    ]
+
+
+def worker_story(keys_or_stimuli: Collection[Key | str], log: Iterable[tuple]) -> list:
+    """Creates a story from the worker log given a set of keys
+    describing tasks or stimuli.
+
+    Parameters
+    ----------
+    keys_or_stimuli : set[str]
+        Task keys or stimulus_id's
+    log : iterable
+        The worker log
+
+    Returns
+    -------
+    story : list[str]
+    """
+    return [
+        msg
+        for msg in log
+        if any(key in msg for key in keys_or_stimuli)
+        or any(
+            key in c
+            for key in keys_or_stimuli
+            for c in msg
+            if isinstance(c, (tuple, list, set))
+        )
+    ]

distributed/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '2026.3.6'
+__version_tuple__ = version_tuple = (2026, 3, 6)
+
+__commit_id__ = commit_id = 'geb11c4c6b'