Dask Distributed

repository·main·Indexed 23 days ago

https://github.com/dask/distributed

A distributed scheduler and orchestration layer for Dask that enables running tasks across a cluster of workers. It provides a Client for task submission and cluster management, various cluster implementations (LocalCluster, SSHCluster, SubprocessCluster), and synchronization primitives like Lock, Semaphore, and Variable. Key features include Actors for stateful distributed objects, an Active Memory Manager (AMM) for optimizing memory usage, and ClientExecutor for integration with the concurrent.futures API.

Tokens
65.1K
Snippets
109
Records
344
Agent score
77%

What's inside distributed

  1. Overview of Dask.distributed

    main

    Dask.distributed is a lightweight library for distributed computing in Python designed for moderate-sized clusters. It extends both the standard library concurrent.futures API and the dask API.

    Key features include:

    • Low Latency: Approximately 1ms of overhead per task.
    • Peer-to-Peer Data Sharing: Workers communicate directly via TCP to share data, reducing central bottlenecks.
    • Complex Scheduling: Supports sophisticated workflows beyond simple map/filter/reduce patterns.
    • Data Locality: Scheduling algorithms prioritize executing computations where the data resides to minimize network traffic.
    • Pure Python: Built in Python, making it pip installable and easy to debug.
  2. Overview of Dask Distributed

    main
    Dask Distributed is a library designed for distributed computation. It provides the infrastructure to scale Python code across multiple cores, machines, or even clusters, enabling parallel execution of tasks and large-scale data processing.
  3. What is the Active Memory Manager (AMM)?

    main

    The Active Memory Manager (AMM) is an experimental daemon that optimizes memory usage across Dask workers in a cluster. It works by running policies that suggest actions to either replicate or delete data (tasks) to balance memory usage and reduce duplication.

    How it works

    • Replication: The AMM can create copies of in-memory tasks on other workers to avoid re-transferring data.
    • Deletion: The AMM can remove excess replicas of tasks to free up memory.
    • No 'Move' operation: Moving data is a two-pass process: a policy first creates a copy, and in a subsequent iteration, a policy deletes the original.
    • Automatic Balancing: Unless constrained by a policy, the AMM automatically replicates data to workers with the lowest memory usage and deletes replicas from workers with the highest memory usage.
  4. How Dask's plugin system works

    main

    Dask's plugin system allows you to execute custom Python code in response to specific lifecycle events within the distributed cluster. Plugins are categorized by the component they target:

    • Scheduler Plugins: Run on the scheduler. They can monitor task transitions or access the scheduler's internal state.
    • Worker Plugins: Run on all workers. They are useful for setting up worker-specific resources (like process pools) when a worker starts.
    • Nanny Plugins: Run on the nanny process, which manages worker processes.

    You can create custom plugins by subclassing the provided base classes or use Dask's built-in plugins for common tasks like installing packages or uploading files.

  5. How Dask Distributed chooses workers for initial tasks

    main

    Dask uses different heuristics for initial task placement depending on whether queuing is enabled. Initial task placement is a 'forward-looking' decision aimed at reducing future data transfers by co-locating related tasks.

    With Queuing Enabled (Default)

    When queuing is enabled, each initial task is scheduled on the least-busy worker available at that moment. If all worker threads are occupied, the task remains in the queue and is not sent to any worker.

    Without Queuing (Worker Saturation set to inf)

    When queuing is disabled (by setting distributed.scheduler.worker-saturation to inf), Dask attempts to co-assign neighboring root tasks to the same worker. This is done by identifying 'root-ish' tasks within a TaskGroup (tasks sharing the same key prefix) and assigning them in batches to workers to ensure that tasks likely to be combined in downstream operations (like dask.array operations) stay on the same worker.

  6. How the Dask Distributed protocol works

    main

    The Dask Distributed protocol enables communication between schedulers, workers, and clients by passing messages that encode commands, status updates, and data.

    Mental Model

    • Messages as Dictionaries: At the application level, messages are represented as Python dictionaries (e.g., {'op': 'compute', 'args': ['x']}).
    • Serialization Layers: To move these dictionaries over the wire, Dask uses a multi-layered approach:
      1. MsgPack: Used for the primary message envelope and small, structured data. It is fast and supports bytestrings.
      2. CloudPickle: Used to serialize Python-specific objects like functions and user-defined classes into bytes before they are passed to MsgPack.
      3. Custom Payload Frames: For large data (like NumPy arrays) or complex types, the message is split. The administrative dictionary is sent via MsgPack, and the heavy data is sent in separate 'frames' (bytestrings) to avoid MsgPack limitations.
    • Frames: The final output sent over the socket is a sequence of frames: a count of frames, the length of each frame, and the actual frame data.
  7. What a Dask Worker does

    main

    A Dask Worker performs two primary functions in a distributed cluster:

    1. Task Computation: Executes tasks as directed by the scheduler.
    2. Data Storage and Serving: Stores computed results locally and serves them to other workers or clients on demand.

    If a worker needs to evaluate a task but lacks the required dependencies, it communicates with peer workers to gather the necessary data. For example, if Worker B needs a result held by Worker A, Worker B will request that data directly from Worker A.

  8. How distributed communication and protocols work

    main

    The dask.distributed system relies on several layers to manage complexity:

    1. Communication Layer: Handles the encoding, sending, and decoding of arbitrary Python objects between distributed endpoints (Client, Scheduler, Worker).
    2. Protocol: While the communication layer is abstract, participants in a cluster follow a specific protocol that defines request-response semantics using a well-defined message format.
    3. Concurrency Model: Nodes (Workers and Schedulers) operate concurrently using Tornado coroutines to handle multiple overlapping requests and computations without blocking.

    Endpoints like Workers, Schedulers, and Nannies are built upon a common Server class. To interact with these servers, developers typically use rpc objects which provide a Pythonic method-call interface.

  9. Understand the criteria for task stealing

    main

    The Dask scheduler uses several heuristics to decide if stealing a task is worth the overhead:

    1. Computation to Communication Ratio: The scheduler prefers tasks where computation_time >> communication_time. It uses sys.getsizeof for dependency sizes and an exponentially weighted moving average for function runtimes.
    2. Saturated Worker Burden: If a worker has a very long backlog and there are many idle workers, the scheduler may steal tasks even if the compute-to-communicate ratio is poor.
    3. Data Replication: Stealing is encouraged if it results in highly-sought-after data being replicated on more workers, improving long-term data locality.
    4. Steal from the Rich: The scheduler targets workers with the largest backlogs rather than those with only a few excess tasks.

    Handling Restrictions

    If a task is restricted to specific workers (e.g., via workers=... with allow_other_workers=False or specific resources), the scheduler will still attempt to steal it to balance load, but it strictly enforces the restrictions. A task will only be stolen by an idle worker that is part of the allowed subset and possesses the required resources.

  10. How Dask Distributed handles data locality and task scheduling

    main

    Dask Distributed uses scheduling policies to minimize data movement, which is critical for performance in analytic computations.

    Task Submission Policy

    When submitting a task f(x) that requires data x, the scheduler attempts to run the task on the worker that already holds x. If the required data is split across multiple workers, the scheduler selects the worker that requires the least amount of data transfer. Data size is determined using sys.getsizeof via the __sizeof__ protocol.

    Data Scatter Policy

    When using client.scatter() to distribute data from a local process to the cluster, data is distributed in a round-robin fashion, grouped by the number of cores available on the workers.

  11. Understand the Worker state machine architecture

    main

    The Dask Worker state machine is organized into three distinct layers of responsibility to separate state logic from side effects (like networking and threading):

    1. WorkerState (Pure Logic):

      • An agnostic data holder that encapsulates the state of the worker and its TaskState objects.
      • It has no knowledge of asyncio, networking, disk I/O, or threads.
      • The only way to mutate state is via handle_stimulus(stimulus), which returns a list of Instruction objects (actions to be taken).
    2. BaseWorker (Async Orchestration):

      • Wraps WorkerState and adds awareness of asyncio.
      • It consumes Instruction objects and manages the lifecycle of asyncio tasks.
      • It defines abstract async methods execute and gather_dep which are intended to be implemented by subclasses.
    3. distributed.Worker (Implementation):

      • Subclasses BaseWorker and provides the actual implementation for networking, threading, and disk I/O.
      • Implements execute and gather_dep to perform real work.

    The Stimulus-Instruction Loop

    1. A StateMachineEvent (stimulus) arrives (from the scheduler or the worker itself).
    2. WorkerState.handle_stimulus processes the event and returns Instruction objects.
    3. BaseWorker executes these instructions (e.g., creating an asyncio task).
    4. When the task finishes, it produces a new StateMachineEvent, which is fed back into the loop.
  12. Handle impure functions with pure=False

    main

    By default, distributed assumes all functions are pure (they always return the same output for the same input and have no side effects). The scheduler uses this assumption to avoid redundant computations by caching results based on a unique key generated from the function and its inputs.

    If your function is impure (e.g., it uses random or modifies global state), you must pass pure=False to client.submit() or client.map(). This tells the scheduler to generate a random key (using uuid4) instead of a deterministic one, ensuring the function is executed every time even if the inputs are identical.