Dask Working Notes - Posted in 2022

Dask Demo Day November 2022

2022-11-21T00:00:00+00:00

Once a month, the Dask Community team hosts Dask Demo Day: an informal and fun online hangout where folks can showcase new or lesser-known Dask features and the rest of us can learn about all the things we didn’t know Dask could do 😁

November’s Dask Demo Day had five great demos. We learned about:

Visualizing 2-billion lightning flashes with Dask, RAPIDS and Datashader
The new Dask CLI
The Dask-Optuna integration for distributed hyperparameter optimization
Dask-Awkward
Profiling your Dask code with Dask-PySpy

This blog gives you a quick overview of the five demos and demonstrates how they might be useful to you. You can watch the full recording below.

Would it be possible to interactively visualize all the lightning strikes in his dataset, Kevin Tyle (University of Albany) recently asked himself. In this demo, Kevin shows you how he leveraged CUDA, RAPIDS-AI, Dask and Datashader to build a smooth interactive visualization of 8 years’ worth of lightning strikes. That’s over 2 billion rows of data.

Kevin shows you how to finetune performance of such a large-scale data processing workflow by:

Leveraging GPUs
Using a Dask cluster to maximize hardware usage
Making smart choices about file types

Watch the full demo or read more about high-performance visualization strategies with Dask and Datashader.

The New Dask CLI

During the Dask Sprint at SciPy this year, a group of Dask maintainers began work on an upgraded, high-level Dask CLI. Doug Davis (Anaconda) walks us through how the CLI works and all the things you can do with it. After installing dask, you can access the CLI by typing dask into your terminal. The tool is designed to be easily extensible by anyone working on Dask. Doug shows you how to add your own components to the Dask CLI.

Watch the full demo or read the Dask CLI documentation.

XGBoost HPO with Dask and Optuna

Have you ever wanted to speed up your hyperparameter searches by running them in parallel? James Bourbeau (Coiled) shows you how you can use the brand-new dask-optuna integration to run hundreds of hyperparameter searches in parallel on a Dask cluster. Running your Optuna HPO searches on a Dask cluster requires only two changes to your existing optuna code. After making those changes, we’re then able to run 500 HPO iterations in parallel in 25 seconds.

Watch the full demo.

Dask for Awkward Arrays

The PyData ecosystem has historically focused on rectilinear data structures like DataFrames and regular arrays. Awkward Arrays brings NumPy-like operations to non-rectilinear data structures and dask-awkward enables you to work with awkward arrays on a distributed cluster in parallel. Doug Davis (Anaconda) walks you through a quick demo of how to use dask-awkward on a local cluster. This is a helpful tool if you find yourself working with nested data structures at scale.

Watch the full demo.

Profiling Dask on a Cluster with py-spy

py-spy is a Python profiler that lets you dig deeper into your code than just your Python functions. Gabe Joseph (Coiled) shows you how you can use dask-pyspy to profile code on a Dask cluster. By digging down into compiled code, dask-pyspy is able to discover valuable insights about why your Dask code might be running slow and what you might be able to do to resolve this.

Watch the full demo.

Join us for the next Demo Day!

Dask Demo Day is a great opportunity to learn about the latest developments and features in Dask. It’s also a fun hangout where you can ask questions and interact with some of Dask’s core maintainers in an informal, casual online setting. We’d love to see you at the next Demo Day on December 15th!

Curious how you can stay connected and find out about the latest Dask news and events?

You can:

follow us on Twitter @dask_dev
subscribe to the Dask newsletter by sending a blank email to newsletter+subscribe@dask.org
subscribe to the Dask community calendar

Reducing memory usage in Dask workloads by 80%

2022-11-15T00:00:00+00:00

Original version of this post appears on https://www.coiled.io/blog/reducing-dask-memory-usage

There’s a saying in emergency response: “slow is smooth, smooth is fast”.

That saying has always bothered me, because it doesn’t make sense at first, yet it’s entirely correct.

By applying this philosophy to the scheduling algorithm in the latest release of Dask, we’re seeing common workloads use up to 80% less memory than before. This means some workloads that used to be outright un-runnable are now running smoothly—an infinity-X speedup!

The second-most upvoted and commented issue of all time on the dask/distributed repo describes, “tasks early in my graph generate data faster than it can be consumed downstream, causing data to pile up, eventually overwhelming my workers”.

Dask users often struggle with workloads that run out of memory like this. Studying these situations, we realized that the Dask scheduler wasn’t following this “slow is smooth, smooth is fast” adage.

Here’s what the problem was, and how we’ve addressed it:

Historically, Dask tried hard to get each individual task done as fast as possible: if a task could run, it would. So sometimes, tasks would run even if their outputs weren’t going to be used immediately—leaving them sitting around in memory.

If you had thousands of initial tasks loading data—say, fetching Parquet from S3, CSVs from disk, or rows from a database—all those tasks would be scheduled and sent to workers up front.

The workers would churn through them, fetching chunks of data (and accumulating it in memory) as quickly as possible. A worker would tell the scheduler when each chunk was loaded, and the scheduler would respond with what to do with it next—but until that message arrived, there were more data-loading tasks runnable right now, so why not run them?

This slight gap in timing—between the worker immediately starting on a less-useful task, then only later finding out about a more-useful task it should have run instead—allowed this lower-priority data to pile up in memory. We call this “root task overproduction”.

Overall, there could be at least twice as much initial data in memory at once as necessary—and, therefore, twice as many intermediate results. (See this comment for a detailed explanation of why the 2x happens.)

When this put workers under memory pressure, this initial problem would snowball into a bigger one. Workers had to spill data to disk (slow), then read it back from disk to use it or transfer it (slow). Workers might exceed their memory limits and crash, losing progress and requiring tasks to be recomputed on a pool of workers that were already struggling.

In the end, this meant that a whole class of workloads were slow, or even un-runnable, without using hugely oversized clusters.

There were plenty of ways to approach this problem, but we wanted to try the simplest thing first: just don’t tell the workers about more tasks than they can run at once.

We’re calling this mode of scheduling “queuing”, or “root task withholding”. The scheduler puts data-loading tasks in an internal queue, and only drips one out to a worker once it’s finished its current work and there’s nothing more useful to run instead that utilizes the work it just completed.

Queuing adds a cost in latency. Every time workers finish a task, they have to ask the scheduler what to do next and sit under-utilized, or even idle, until they get an answer. (Before, they had a backlog of things they could do before the answer came back.)

For a while, we hadn’t considered this approach, because intuitively, we assumed the latency would be too much of a slow-down.

However, by slowing down the pace of task assignment, and running only the best tasks, scheduling gets much smoother. And with that smoothness, we see that most benchmark workloads use much less memory across the board:

Percent decrease in peak memory used in the latest release. Notice the axes: up to 80% reduction.

This large a reduction in memory use is a big deal!

For many users, this might mean that workloads that weren’t runnable before can now run smoothly.

Julius Busecke reports that a common geoscience task which used to always crash now works out of the box with the new scheduling mode.

Some workloads will also be cheaper to run in the cloud, since they can use instances with less memory. We see some benchmarks that could, in theory, be run for 30-50% less total cost. This is not universal: others would cost more because they get slower. More on that later.

Beyond this, execution is just more predictable. Memory usage is much more consistent and less likely to spike rapidly.

Smooth is fast

In a few cases, it turns out that smooth scheduling can be even faster.

On average, one representative oceanography workload ran 20% faster. A few other workloads showed modest speedups as well. This is mostly because they no longer load too much data into memory and then have to spill it to disk, which creates significant slowdowns.

Additionally, we found that the extra latency we were worried about didn’t actually slow things down in typical cases. There was no measurable change in pure task throughput on a cluster with fast networking and multi-CPU workers, like Coiled clusters or a single-machine LocalCluster. This was a good lesson in trying the simplest thing first.

Sometimes, slow is still slow (but not why you’d expect)

However, we did notice that a few benchmarks run slower with scheduler-side queuing. The typical slowdown is 5-10%, but in the worst case, they are ~50% slower (though they also use about half the memory).

The problem is that implementing queuing meant giving up a scheduling feature introduced last year called co-assignment.

As described in the docs, co-assignment tries to schedule initial tasks on the same worker if their outputs will be combined later. This avoids having to transfer data from one worker to another when the downstream task runs, because all the data is already on one worker.

In a graph like this, we’d want a and b to run on the same worker. Otherwise, one of e or f would have to be transferred between workers before i can run.

  i       j
 / \     / \
e   f   g   h
|   |   |   |
a   b   c   d

Avoiding these transfers speeds things up, because network is relatively slow. It also reduces memory usage by avoiding having to hold replicas of the same data on multiple workers.

Unfortunately, the current implementation of co-assignment isn’t compatible with queuing, and updating it is non-trivial. We plan to do this next, to get the best of both worlds.

But in the short term, we had to decide if queuing was beneficial enough to be worth enabling by default right away, despite the loss of co-assignment.

A new default for scheduling

After running a number of benchmarks, and getting some initial community feedback, we think it is.

Queuing makes things possible that used to not work at all. But it doesn’t break anything that works today: everything will still work, some things just could be slower. We feel that that’s a worthwhile enough tradeoff to enable it by default—especially given how much dask users have struggled with memory issues.

Additionally, to avoid impacting workloads that could become latency-bound, the new algorithm is still doing a little overproduction. It’s pushing a handful of extra root tasks to the worker in advance (as opposed to all of them, like before). This comes at the price of some extra memory use, but prevents painful slowdowns in high-latency clusters.

So in the latest release, queuing is enabled by default. Most memory-intensive Array and DataFrame workloads should see reductions in memory use out of the box, ranging from noticeable to 🤩.

Let us know how it goes

We’ve opened a discussion on GitHub for feedback on this change. Please let us know how it helps (or doesn’t).

Keeping old behavior

For users who are sensitive to runtime and have low memory use, you can deactivate queuing and use the old scheduling mode (including co-assignment) via Dask configuration, by setting the new distributed.scheduler.worker-saturation config value to inf.

You can read more about adjusting this setting in the docs.

On Coiled, you can set it with:

import dask
import coiled

with dask.config.set({"distributed.scheduler.worker-saturation": "inf"}):
    cluster = coiled.Cluster(...)  # coiled sends current dask config automatically

You can see examples of setting this configuration for various deployment systems on the discussion issue (when copy-pasting, be sure to change the 1.0 to inf!). And if you find the need to set worker-saturation back to inf, please let us know on the discussion.

Dask Kubernetes Operator

2022-11-09T00:00:00+00:00

We are excited to announce that the Dask Kubernetes Operator is now generally available 🎉!

Notable new features include:

Dask Clusters are now native custom resources
Clusters can be managed with kubectl or the Python API
Cascaded deletions allow for proper teardown
Multiple worker groups enable heterogenous/tagged deployments
DaskJob: running dask workloads with K8s batched job infrastructure
Clusters can be reused between different Python processes
Autoscaling is handled by a custom Kubernetes controller instead of the user code
Scheduler and worker Pods and Services are fully configurable

$ kubectl get daskcluster
NAME         AGE
my-cluster   4m3s

$ kubectl get all -A -l dask.org/cluster-name=my-cluster
NAMESPACE   NAME                                       READY   STATUS    RESTARTS   AGE
default     pod/my-cluster-default-worker-22bd39e33a   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-5f4f2c989a   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-72418a589f   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-9b00a4e1fd   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-d6fc172526   1/1     Running   0          3m43s
default     pod/my-cluster-scheduler                   1/1     Running   0          4m21s

NAMESPACE   NAME                           TYPE        CLUSTER-IP    EXTERNAL-IP   PORT(S)             AGE
default     service/my-cluster-scheduler   ClusterIP   10.96.33.67   <none>        8786/TCP,8787/TCP   4m21s

At the start of 2022 we began the large undertaking of rewriting the dask-kubernetes package in the operator pattern. This design pattern has become very popular in the Kubernetes community with companies like Red Hat building their whole Kubernetes offering Openshift around it.

If you’ve spent any time in the Kubernetes community you’ll have heard the term operator being thrown around seen projects like Golang’s Operator Framework being used to deploy modern applications.

At it’s core an operator is made up of a data structure for describing the thing you want to deploy (in our case a Dask cluster) and a controller which does the actual deploying. In Kubernetes the templates for these data structures are called Custom Resource Definitions (CRDs) and allow you to extend the Kubernetes API with new resource types of your own design.

For dask-kubernetes we have created a few CRDs to describe things like Dask clusters, groups of Dask workers, adaptive autoscalers and a new Dask powered batch job.

We also built a controller using kopf that handles watching for changes to any of these resources and creates/updates/deletes lower level Kubernetes resources like Pods and Services.

Why did we build this?

The original implementation of dask-kubernetes was started shortly after Kubernetes went 1.0 and before any established design patterns had emerged. Its model was based on spawning Dask workers as subprocesses, except those subprocesses are Pods running in Kubernetes. This is the same way dask-jobqueue launches workers as individual job scheduler allocations or dask-ssh opens many SSH connections to various machines.

Over time this has been refactored, rewritten and extended multiple times. One long-asked-for change was to also place the Dask scheduler inside the Kubernetes cluster to simplify scheduler-worker communication and network connectivity. Naturally this lead to more feature requests around configuring the scheduler service and having more control over the cluster. As we extended more and more the original premise of spawning worker subprocesses on a remote system became less helpful.

The final straw in the original design was folks asking for the ability to leave a cluster running and come back to it later. Either to reuse a cluster between separate jobs, or just different stages in a multi-stage pipeline. The premise of spawning subprocesses leads to an assumption that the parent process will be around for the lifetime of the cluster which makes it a reasonable place to hold state such as the template for launching new workers when scaling up. We attempted to implement this feature but it just wasn’t possible with the current design. Moving to a model where the parent process can die and new processes can pick up means that state needs to be moved elsewhere and things were too entangled to successfully pull this out.

The classic implementation that had served us well for so long was creaking and becoming increasingly difficult to modify and maintain. The time had come to pay down our technical debt by rebuilding from scratch under a new model, the operator pattern.

In this new model a Dask cluster is an abstract object that exists within a Kubernetes cluster. We use custom resources to store the state for each cluster and a custom controller to map that state onto reality by creating the individual components that make up the cluster. Want to scale up your cluster? Instead of having some Python code locally that spawns a new Pod on Kubernetes we just modify the state of the Dask cluster resource to specify the desired number of workers and the controller handles adding/removing Pods to match.

New features

While our primary goal was allowing cluster reuse between Python processes and paying down technical debt switching to the operator pattern has allowed us to add a bunch of nice new features. So let’s explore those.

Python or YAML API

With our new implementation we create Dask clusters by creating a DaskCluster resource on our Kubernetes cluster. The controller sees this appear and spawns child resources for the scheduler, workers, etc.

We modify our cluster by editing the DaskCluster resource and our controller reacts to those changes and updates the child resources accordingly.

We delete our cluster by deleting the DaskCluster resource and Kubernetes handles the rest (see the next section on cascade deletion).

By storing all of our state in the resource and all of our logic in the controller this means the KubeCluster class is now much simpler. It’s actually so simple that it is entirely optional.

The primary purpose of the KubeCluster class now is to provide a nice clean API for creating/scaling/deleting your clusters in Python. It can take a small number of keyword arguments and generate all of the YAML to submit to Kubernetes.

from dask_kubernetes.operator import KubeCluster

cluster = KubeCluster(name="my-cluster", n_workers=3, env={"FOO": "bar"})

The above snippet creates the following resource.

apiVersion: kubernetes.dask.org/v1
kind: DaskCluster
metadata:
  name: my-cluster
spec:
  scheduler:
    service:
      ports:
        - name: tcp-comm
          port: 8786
          protocol: TCP
          targetPort: tcp-comm
        - name: http-dashboard
          port: 8787
          protocol: TCP
          targetPort: http-dashboard
      selector:
        dask.org/cluster-name: my-cluster
        dask.org/component: scheduler
      type: ClusterIP
    spec:
      containers:
        - args:
            - dask-scheduler
            - --host
            - 0.0.0.0
          env:
            - name: FOO
              value: bar
          image: ghcr.io/dask/dask:latest
          livenessProbe:
            httpGet:
              path: /health
              port: http-dashboard
            initialDelaySeconds: 15
            periodSeconds: 20
          name: scheduler
          ports:
            - containerPort: 8786
              name: tcp-comm
              protocol: TCP
            - containerPort: 8787
              name: http-dashboard
              protocol: TCP
          readinessProbe:
            httpGet:
              path: /health
              port: http-dashboard
            initialDelaySeconds: 5
            periodSeconds: 10
          resources: null
  worker:
    cluster: my-cluster
    replicas: 3
    spec:
      containers:
        - args:
            - dask-worker
            - --name
            - $(DASK_WORKER_NAME)
          env:
            - name: FOO
              value: bar
          image: ghcr.io/dask/dask:latest
          name: worker
          resources: null

If I want to scale up my workers to 5 I can do this in Python.

cluster.scale(5)

All this does is apply a patch to the resource and modify the spec.worker.replicas value to be 5 and the controller handles the rest.

Ultimately our Python API is generating YAML and handing it to Kubernetes to action. Everything about our cluster is contained in that YAML. If we prefer we can write and store this YAML ourselves and manage our cluster entirely via kubectl.

If we put the above YAML example into a file called my-cluster.yaml we can create it like this. No Python necessary.

$ kubectl apply -f my-cluster.yaml
daskcluster.kubernetes.dask.org/my-cluster created

We can also scale our cluster with kubectl.

$ kubectl scale --replicas=5 daskworkergroup my-cluster-default
daskworkergroup.kubernetes.dask.org/my-cluster-default

This is extremely powerful for advanced users who want to integrate with existing Kubernetes tooling and really modify everything about their Dask cluster.

You can still construct a KubeCluster object in the future and point it to this existing cluster for convenience.

from dask.distributed import Client
from dask_kubernetes.operator import KubeCluster

cluster = KubeCluster.from_name("my-cluster")
cluster.scale(5)
client = Client(cluster)

Cascade deletion

Having a DaskCluster resource also makes deletion much more pleasant.

In the old implementation your local Python process would spawn a bunch of Pod resources along with supporting ones like Service and PodDisruptionBudget resources. It also had some teardown functionality that was either called directly or via a finalizer that deleted all of these resources when you are done.

One downside of this was that if something went wrong either due to a bug in dask-kubernetes or a more severe failure that caused the Python process to exit without calling finalizers you would be left with a ton of resources that you had to clean up manually. I expect some folks have a label based selector command stored in their snippet manager somewhere but most folks would do this cleanup manually.

With the new model the DaskCluster resource is set as the owner of all of the other resources spawned by the controller. This means we can take advantage of cascade deletion for our cleanup. Regardless of how you create your cluster or whether the initial Python process still exists you can just delete the DaskCluster resource and Kubernetes will know to automatically delete all of its children.

$ kubectl get daskcluster  # Here we see our Dask cluster resource
NAME         AGE
my-cluster   4m3s

$ kubectl get all -A -l dask.org/cluster-name=my-cluster  # and all of its child resources
NAMESPACE   NAME                                       READY   STATUS    RESTARTS   AGE
default     pod/my-cluster-default-worker-22bd39e33a   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-5f4f2c989a   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-72418a589f   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-9b00a4e1fd   1/1     Running   0          3m43s
default     pod/my-cluster-default-worker-d6fc172526   1/1     Running   0          3m43s
default     pod/my-cluster-scheduler                   1/1     Running   0          4m21s

NAMESPACE   NAME                           TYPE        CLUSTER-IP    EXTERNAL-IP   PORT(S)             AGE
default     service/my-cluster-scheduler   ClusterIP   10.96.33.67   <none>        8786/TCP,8787/TCP   4m21s

$ kubectl delete daskcluster my-cluster  # We can delete the daskcluster resource
daskcluster.kubernetes.dask.org "my-cluster" deleted

$ kubectl get all -A -l dask.org/cluster-name=my-cluster  # all of the children are removed
No resources found

Multiple worker groups

We also took this opportunity to add support for multiple worker groups as a first class principle. Some workflows benefit from having a few workers in your cluster with some additional resources. This may be a couple of workers with much higher memory than the rest, or GPUs for accelerated compute. Using resource annotations you can steer certain tasks to those workers, so if you have a single step that creates a large amount of intermediate memory you can ensure that task ends up on a worker with enough memory.

By default when you create a DaskCluster resource it creates a single DaskWorkerGroup which in turn creates the worker Pod resources for our cluster. If you wish you can add more worker group resources yourself with different resource configurations.

Here is an example of creating a cluster with five workers that have 16GB of memory and two additional workers with 64GB of memory.

from dask_kubernetes.operator import KubeCluster

cluster = KubeCluster(name='foo',
                      n_workers=5,
                      resources={
                          "requests": {"memory": "16Gi"},
                          "limits": {"memory": "16Gi"}
                      })

cluster.add_worker_group(name="highmem",
                         n_workers=2,
                         resources={
                             "requests": {"memory": "64Gi"},
                             "limits": {"memory": "64Gi"}
                         })

Autoscaling

One of the much loved features of the classic implementation of KubeCluster was adaptive autoscaling. When enabled the KubeCluster object would regularly communicate with the scheduler and ask if it wanted to change the number of workers and then add/remove pods accordingly.

In the new implementation this logic has moved to the controller so the cluster can autoscale even when there is no KubeCluster object in existence.

The Python API remains the same so you can still use KubeCluster to put your cluster into adaptive mode.

from dask_kubernetes.operator import KubeCluster

cluster = KubeCluster(name="my-cluster", n_workers=5)
cluster.adapt(minimum=1, maximum=100)

This call creates a DaskAutoscaler resource which the controller sees and periodically takes action on by asking the scheduler how many workers it wants and updating the DaskWorkerGroup within the configured bounds.

apiVersion: kubernetes.dask.org/v1
kind: DaskAutoscaler
metadata:
  name: my-cluster
spec:
  cluster: my-cluster
  minimum: 1
  maximum: 100

Calling cluster.scale(5) will also delete this resource and set the number of workers back to 5.

DaskJob

Having composable cluster resources also allows us to put together a new DaskJob resource.

Kubernetes has some built-in batch job style resources which ensure a Pod is run to completion one or more times. You can control how many times is should run and how many concurrent pods there should be. This is useful for fire-and-forget jobs that you want to process a specific workload.

The Dask Operator introduces a DaskJob resource which creates a DaskCluster alongside a single client Pod which it attempts to run to completion. If the Pod exits unhappily it will be restarted until it returns a 0 exit code, at which point the DaskCluster is automatically cleaned up.

The client Pod has all of the configuration for the DaskCluster injected at runtime via environment variables, this means your client code doesn’t need to know anything about how the Dask cluster was constructed it just connects and makes use of it. This allows for excellent separation of concerns between your business logic and your deployment tooling.

from dask.distributed import Client

# We don't need to tell the Client anything about the cluster as
# it will find everything it needs in the environment variables
client = Client()

# Do some work...

This new resource type is useful for some batch workflows, but also demonstrates how you could extend the Dask Operator with your own new resource types and hook them together with a controller plugin.

Extensibility and plugins

By moving to native Kubernetes resources and support for the YAML API power users can treat DaskCluster resources (or any of the new Dask resources) as building blocks in larger applications. One of Kubernetes’s superpowers is managing everything as composable resources that can be combined to create complex and flexible applications.

Does your Kubernetes cluster have an opinionated configuration with additional tools like Istio installed? Have you struggled in the last to integrate dask-kubernetes with your existing tooling because it relied on Python to create clusters?

It’s increasingly common for users to need additional resources to be created alongside their Dask cluster like Istio Gateway resources or cert-manager Certificate resources. Now that everything in dask-kubernetes uses custom resources users can mix and match resources from many different operators to construct their application.

If this isn’t enough you can also extend our custom controller. We built the controller with kopf primarily because the Dask community is strong in Python and less so in Golang (the most common way to build operators). It made sense to play into our strengths rather than using the most popular option.

This also means our users should be able to more easily modify the controller logic and we’ve included a plugin system that allows you to add extra logic rules by installing a custom package into the controller container image and registering them via entry points.

# Source for my_controller_plugin.plugin

import kopf

@kopf.on.create("service", labels={"dask.org/component": "scheduler"})
async def handle_scheduler_service_create(meta, new, namespace, logger, **kwargs):
   # Do something here like create an Istio Gateway
   # See https://kopf.readthedocs.io/en/stable/handlers for documentation on what is possible here

# pyproject.toml for my_controller_plugin

[option.entry_points]
dask_operator_plugin =
   my_controller_plugin = my_controller_plugin.plugin

# Dockerfile

FROM ghcr.io/dask/dask-kubernetes-operator:2022.10.0

RUN pip install my-controller-plugin

That’s it, when the controller starts up it will also import all @kopf methods from modules listed in the dask_operator_plugin entry point alongside the core functionality.

Migrating

One caveat to switching to the operator model is that you need to install the CRDs and controller on your Kubernetes before you can start using it. While a small hurdle this is a break in the user experience compared to the classic implementation.

helm repo add dask https://helm.dask.org && helm repo update
kubectl create ns dask-operator
helm install --namespace dask-operator dask-operator dask/dask-kubernetes-operator

We also took this opportunity to make breaking changes to the constructor of KubeCluster to simplify usage for beginners or folks who are happy with the default options. By adopting the YAML API power users can tinker and tweak to their hearts content without having to modify the Python library, so it made sense to make the Python library simpler and more pleasant to use for the majority of users.

We made an explicit decision not to just replace the old KubeCluster with the new one in place because people’s code will just stop working if we did. Instead we are asking folks to read the migration guide and update your imports and construction code. Users of the classic cluster manager will start seeing a deprecation warning as of 2022.10.0 and at some point the classic implementation will be removed all together. If migrating is challenging to do quickly you can always pin your dask-kubernetes version, and from then on you are clearly not getting bug fixes or enhancements. But in all honesty those have been few and far between for the classic implementation lately anyway.

We are optimistic that the new cleaner implementation, faster cluster startup times and bucket of new features is enough to convince you that it’s worth the migration effort.

If you want some help migrating and the migration guide doesn’t cover your use case then don’t hesitate to reach out on the forum. We’ve also worked hard to ensure the new implementation has feature parity with the classic one, but if anything is missing or broken then please open an issue on GitHub.

Understanding Dask’s meta keyword argument

2022-08-09T00:00:00+00:00

If you have worked with Dask DataFrames or Dask Arrays, you have probably come across the meta keyword argument. Perhaps, while using methods like apply():

import dask
import pandas as pd

ddf = dask.datasets.timeseries()


def my_custom_arithmetic(r):
    if r.id > 1000:
        return r.x * r.x + r.y + r.y
    else:
        return 0


ddf["my_computation"] = ddf.apply(
    my_custom_arithmetic, axis=1, meta=pd.Series(dtype="float64")
)

ddf.head()
# Output:
#
#                        id    name         x         y  my_computation
# timestamp
# 2000-01-01 00:00:00  1055  Victor -0.575374  0.868320        2.067696
# 2000-01-01 00:00:01   994   Zelda  0.963684  0.972240        0.000000
# 2000-01-01 00:00:02   982  George -0.997531 -0.876222        0.000000
# 2000-01-01 00:00:03   981  Ingrid  0.852159 -0.419733        0.000000
# 2000-01-01 00:00:04  1029   Jerry -0.839431 -0.736572       -0.768500

You might have also seen one or more of the following warnings/errors:

UserWarning: You did not provide metadata, so Dask is running your function on a small dataset to guess output types. It is possible that Dask will guess incorrectly.

UserWarning: `meta` is not specified, inferred from partial data. Please provide `meta` if the result is unexpected.

ValueError: Metadata inference failed in …

If the above messages look familiar, this blog post is for you. :)

We will discuss:

what the is meta keyword argument,
why does Dask need meta, and
how to use it effectively.

We will look at meta mainly in the context of Dask DataFrames, however, similar principles also apply to Dask Arrays.

Before answering this, let’s quickly discuss Dask DataFrames.

A Dask DataFrame is a lazy object composed of multiple pandas DataFrames, where each pandas DataFrame is called a “partition”. These are stacked along the index and Dask keeps track of these partitions using “divisions”, which is a tuple representing the start and end index of each partition.

When you create a Dask DataFrame, you usually see something like the following:

>>> import pandas as pd
>>> import dask.dataframe as dd
>>> ddf = dd.DataFrame.from_dict(
...     {
...         "x": range(6),
...         "y": range(10, 16),
...     },
...     npartitions=2,
... )
>>> ddf
Dask DataFrame Structure:
                   x      y
npartitions=2
0              int64  int64
3                ...    ...
5                ...    ...
Dask Name: from_pandas, 2 tasks

Here, Dask has created the structure of the DataFrame using some “metadata” information about the column names and their datatypes. This metadata information is called meta. Dask uses meta for understanding Dask operations and creating accurate task graphs (i.e., the logic of your computation).

The meta keyword argument in various Dask DataFrame functions allows you to explicitly share this metadata information with Dask. Note that the keyword argument is concerned with the metadata of the output of those functions.

Why does Dask need `meta`?

Dask computations are evaluated lazily. This means Dask creates the logic and flow, called task graph, of the computation immediately, but evaluates them only when necessary – usually, on calling .compute().

An example task graph generated to compute the sum of the DataFrame:

>>> s = ddf.sum()
>>> s
Dask Series Structure:
npartitions=1
x    int64
y      ...
dtype: int64
Dask Name: dataframe-sum-agg, 5 tasks

>>> s.visualize()

>>> s.compute()
x    15
y    75
dtype: int64

This is a single operation, but Dask workflows usually have multiple such operation chained together. Therefore, to create the task graph effectively, Dask needs to know the strucutre and datatypes of the DataFame after each operation. Especially because Dask does not know the actual values/structure of the DataFrame yet.

This is where meta is comes in.

In the above example, the Dask DataFrame changed into a Series after sum(). Dask knows this (even before we call compute()) only becasue of meta.

Internally, meta is represented as an empty pandas DataFrame or Series, which has the same structure as the Dask DataFrame. To learn more about how meta is defined internally, check out the DataFrame Internal Design documentation.

To see the actual metadata information for a collection, you can look at the ._meta attribute[1]:

>>> s._meta
Series([], dtype: int64)

How to specify `meta`?

You can specify meta in a few different ways, but the recommended way for Dask DataFrame is:

“An empty pd.DataFrame or pd.Series that matches the dtypes and column names of the output.”

~ DataFrame.apply() docstring

For example:

>>> meta_df = pd.DataFrame(columns=["x", "y"], dtype=int)
>>> meta_df

Empty DataFrame
Columns: [x, y]
Index: []

>>> ddf2 = ddf.apply(lambda x: x, axis=1, meta=meta_df).compute()
>>> ddf2
   x  y
0  0  0
1  1  1
2  2  2
3  0  3
4  1  4
5  2  5

The other ways you can describe meta are:

For a DataFrame, you can specify meta as a:
- Python dictionary: {column_name_1: dtype_1, column_name_2: dtype_2, …}
- Iterable of tuples: [(column_name_1, dtype_1), (columns_name_2, dtype_2, …)]
Note that while describing meta as shown above, using a dictionary or iterable of tuples, the order in which you mention column names is important. Dask will use the same order to create the pandas DataFrame for meta. If the orders don’t match, you will see the following error:
```
ValueError: The columns in the computed data do not match the columns in the provided metadata
```
For a Series output, you can specify meta using a single tuple: (coulmn_name, dtype)

You should not describe meta using just a dtype (like: meta="int64"), even for scalar outputs. If you do, you will see the following warning:

FutureWarning: Meta is not valid, `map_partitions` and `map_overlap` expects output to be a pandas object. Try passing a pandas object as meta or a dict or tuple representing the (name, dtype) of the columns. In the future the meta you passed will not work.

During operations like map_partitions or apply (which uses map_partitions internally), Dask coerces the scalar output of each partition into a pandas object. So, the output of functions that take meta will never be scalar.

For example:

>>> ddf = ddf.repartition(npartitions=1)
>>> result = ddf.map_partitions(lambda x: len(x)).compute()
>>> type(result)
pandas.core.series.Series

Here, the Dask DataFrame ddf has only one partition. Hence, len(x) on that one partition would result in a scalar output of integer dtype. However, when we compute it, we see a pandas Series. This confirms that Dask is coercing the outputs to pandas objects.

Another note, Dask Arrays may not always do this conversion. You can look at the API reference for your particular Array operation for details.

>>> import numpy as np
>>> import dask.array as da
>>> my_arr = da.random.random(10)
>>> my_arr.map_blocks(lambda x: len(x)).compute()
10

`meta` does not force the structure or dtypes

meta can be thought of as a suggestion to Dask. Dask uses this meta to generate the task graph until it can infer the actual metadata from the values. It does not force the output to have the structure or dtype of the specified meta.

Consider the following example, and remember that we defined ddf with x and y column names in the previous sections.

If we provide different column names (a and b) in the meta description, Dask uses these new names to create the task graph:

>>> meta_df = pd.DataFrame(columns=["a", "b"], dtype=int)
>>> result = ddf.apply(lambda x:x*x, axis=1, meta=meta_df)
>>> result
Dask DataFrame Structure:
                   a      b
npartitions=2
0              int64  int64
3                ...    ...
5                ...    ...
Dask Name: apply, 4 tasks

However, if we compute result, we will get the following error:

>>> result.compute()
ValueError: The columns in the computed data do not match the columns in the provided metadata
  Extra:   ['x', 'y']
  Missing: ['a', 'b']

While computing, Dask evaluates the actual metadata with columns x and y. This does not match the meta that we provided, and hence, Dask raises a helpful error message. Notice how Dask does not change the output to have a and b here, rather uses a and b column names only for intermediate task graphs.

Using `._meta` directly

In some rare case, you can also set the ._meta attribute[1] directly for a Dask DataFrame. For example, if the DataFrame was created with incorrect dtypes, like:

>>> ddf = dd.DataFrame.from_dict(
...     {
...         "x": range(6),
...         "y": range(10, 16),
...     },
...     dtype="object", # Note the “object” dtype
...     npartitions=2,
... )
>>> ddf
Dask DataFrame Structure:
                    x       y
npartitions=2
0              object  object
3                 ...     ...
5                 ...     ...
Dask Name: from_pandas, 2 tasks

The values are clearly integers but the dtype says object, so we can’t perform integer operations like addition:

>>> add = ddf + 2
ValueError: Metadata inference failed in `add`.

Original error is below:
------------------------
TypeError('can only concatenate str (not "int") to str')

Here, we can explicitly define ._meta[1]:

>>> ddf._meta = pd.DataFrame(columns=["x", "y"], dtype="int64")

Then, perform the addition:

>>> add = ddf + 2
>>> add
Dask DataFrame Structure:
                   x      y
npartitions=2
0              int64  int64
3                ...    ...
5                ...    ...
Dask Name: add, 4 tasks

Thanks for reading!

Have you run into issues with meta before? Please let us know on Discourse, and we will consider including it here, or updating the Dask documentation. :)

[1] NOTE: ._meta is not a public property, so we recommend using it only when necessary. There is an ongoing discussion around creating public methods to get, set, and view the information in ._meta, and this blog post will be updated to use the public methods when they’re created.

Data Proximate Computation on a Dask Cluster Distributed Between Data Centres

2022-07-19T00:00:00+00:00

This work is a joint venture between the Met Office and the European Weather Cloud, which is a partnership of ECMWF and EUMETSAT.

We have devised a technique for creating a Dask cluster where worker nodes are hosted in different data centres, connected by a mesh VPN that allows the scheduler and workers to communicate and exchange results.

A novel (ab)use of Dask resources allows us to run data processing tasks on the workers in the cluster closest to the source data, so that communication between data centres is minimised. If combined with zarr to give access to huge hyper-cube datasets in object storage, we believe that the technique could realise the potential of data-proximate distributed computing in the Cloud.

Introduction

The UK Met Office has been carrying out a study into data-proximate computing in collaboration with the European Weather Cloud. We identified Dask as a key technology, but whilst existing Dask techniques focus on parallel computation in a single data centre, we looked to extend this to computation across data centres. This way, when the data required is hosted in multiple locations, tasks can be run where the data is rather than copying it.

Dask worker nodes exchange data chunks over a network, coordinated by a scheduler. There is an assumption that all nodes are freely able to communicate, which is not generally true across data centres due to firewalls, NAT, etc, so a truly distributed approach has to solve this problem. In addition, it has to manage data transfer efficiently, because moving data chunks between data centres is much more costly than between workers in the same cloud.

This notebook documents a running proof-of-concept that addresses these problems. It runs a computation in 3 locations:

This computer, where the client and scheduler are running. This was run on AWS during development.
The ECMWF data centre. This has compute resources, and hosts data containing predictions.
The EUMETSAT data centre, with compute resources and data on observations

from IPython.display import Image
Image(filename="images/datacentres.png") # this because GitHub doesn't render markup images in private repos

The idea is that tasks accessing data available in a location should be run there. Meanwhile the computation can be defined, invoked, and the results rendered, elsewhere. All this with minimal hinting to the computation as to how this should be done.

Setup

First some imports and conveniences

import os
from time import sleep
import dask
from dask.distributed import Client
from dask.distributed import performance_report, get_task_stream
from dask_worker_pools import pool, propagate_pools
import pytest
import ipytest
import xarray
import matplotlib.pyplot as plt
from orgs import my_org
from tree import tree

ipytest.autoconfig()

In this case we are using a control plane IPv4 WireGuard network on 10.8.0.0/24 to set up the cluster - this is not a necessity, but simplifies this proof of concept. WireGuard peers are running on ECMWF and EUMETSAT machines already, but we have to start one here:

!./start-wg.sh

4: mo-aws-ec2: <POINTOPOINT,NOARP,UP,LOWER_UP> mtu 8921 qdisc noqueue state UNKNOWN group default qlen 1000
    link/none
    inet 10.8.0.3/24 scope global mo-aws-ec2
       valid_lft forever preferred_lft forever

We have worker machines configured in both ECMWF and EUMETSAT, one in each. They are accessible on the control plane network as

ecmwf_host='10.8.0.4'
%env ECMWF_HOST=$ecmwf_host
eumetsat_host='10.8.0.2'
%env EUMETSAT_HOST=$eumetsat_host

env: ECMWF_HOST=10.8.0.4
env: EUMETSAT_HOST=10.8.0.2

Mount the Data

This machine needs access to the data files over the network in order to read NetCDF metadata. The workers are sharing their data with NFS, so we mount them here. (In this proof of concept, the control plane network is used for NFS, but the data plane network could equally be used, or a more appropriate technology such as zarr accessing object storage.)

%%bash
sudo ./data-reset.sh

mkdir -p /data/ecmwf
mkdir -p /data/eumetsat
sudo mount $ECMWF_HOST:/data/ecmwf /data/ecmwf
sudo mount $EUMETSAT_HOST:/eumetsatdata/ /data/eumetsat

Image(filename="images/datacentres-data.png")

Access to Data

For this demonstration, we have two large data files that we want to process. On ECMWF we have predictions in /data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc. Workers running in ECMWF can see the file

!ssh -i ~/.ssh/id_rsa_rcar_infra  rcar@$ECMWF_HOST 'tree /data/'

/data/
└── ecmwf
    └── 000490262cdd067721a34112963bcaa2b44860ab.nc

1 directory, 1 file

and because that directory is mounted here over NFS, so can this computer

!tree /data/ecmwf

/data/ecmwf
└── 000490262cdd067721a34112963bcaa2b44860ab.nc

0 directories, 1 file

This is a big file

!ls -lh /data/ecmwf

total 2.8G
-rw-rw-r-- 1 ec2-user ec2-user 2.8G Mar 25 13:09 000490262cdd067721a34112963bcaa2b44860ab.nc

On EUMETSAT we have observations.nc

!ssh -i ~/.ssh/id_rsa_rcar_infra  rcar@$EUMETSAT_HOST 'tree /data/eumetsat/ad-hoc'

/data/eumetsat/ad-hoc
└── observations.nc

0 directories, 1 file

similarly visible on this computer

!ls -lh /data/eumetsat/ad-hoc/observations.nc

-rw-rw-r-- 1 613600004 613600004 4.8M May 20 10:57 /data/eumetsat/ad-hoc/observations.nc

Crucially, ECMWF data is not visible in the EUMETSAT data centre, and vice versa.

Our Calculation

We want to compare the predictions against the observations.

We can open the predictions file with xarray

predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')
predictions

<xarray.Dataset>
Dimensions:                  (realization: 18, height: 33, latitude: 960,
longitude: 1280, bnds: 2)
Coordinates:
* realization              (realization) int32 0 18 19 20 21 ... 31 32 33 34
* height                   (height) float32 5.0 10.0 20.0 ... 5.5e+03 6e+03
* latitude                 (latitude) float32 -89.91 -89.72 ... 89.72 89.91
* longitude                (longitude) float32 -179.9 -179.6 ... 179.6 179.9
  forecast_period          timedelta64[ns] 1 days 18:00:00
  forecast_reference_time  datetime64[ns] 2021-11-07T06:00:00
  time                     datetime64[ns] 2021-11-09
  Dimensions without coordinates: bnds
  Data variables:
  air_pressure             (realization, height, latitude, longitude) float32 dask.array<chunksize=(18, 33, 192, 160), meta=np.ndarray>
  latitude_longitude       int32 -2147483647
  latitude_bnds            (latitude, bnds) float32 dask.array<chunksize=(960, 2), meta=np.ndarray>
  longitude_bnds           (longitude, bnds) float32 dask.array<chunksize=(1280, 2), meta=np.ndarray>
  Attributes:
  history:                      2021-11-07T10:27:38Z: StaGE Decoupler
  institution:                  Met Office
  least_significant_digit:      1
  mosg__forecast_run_duration:  PT198H
  mosg__grid_domain:            global
  mosg__grid_type:              standard
  mosg__grid_version:           1.6.0
  mosg__model_configuration:    gl_ens
  source:                       Met Office Unified Model
  title:                        MOGREPS-G Model Forecast on Global 20 km St...
  um_version:                   11.5
  Conventions:                  CF-1.7

Dask code running on this machine has read the metadata for the file via NFS, but has not yet read in the data arrays themselves.

Likewise we can see the observations, so we can perform a calculation locally. Here we average the predictions over the realisations and then compare them with the observations at a particular height. (This is a deliberately inefficient calculation, as we could average at only the required height, but you get the point.)

%%time
def scope():
    client = Client()
    predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')
    observations = xarray.open_dataset('/data/eumetsat/ad-hoc/observations.nc').chunk('auto')

    averages = predictions.mean('realization')
    diff = averages.isel(height=10) - observations
    diff.compute()

#scope()

CPU times: user 10 µs, sys: 2 µs, total: 12 µs
Wall time: 13.8 µs

When we uncomment scope() and actually run this, it takes over 14 minutes to complete! Accessing the data over NFS between data centres (we run this notebook in AWS) is just too slow.

In fact just copying the data files onto the computer running this notebook takes the same sort of time. At least 2.8 GiB + 4.8 MiB of data must pass from the data centres to this machine to perform the calculation.

Instead we should obviously run the Dask tasks where the data is. We can do that on a Dask cluster.

Running Up a Cluster

The cluster is run up with a single command. It takes a while though

import subprocess

scheduler_process = subprocess.Popen([
        '../dask_multicloud/dask-boot.sh',
        f"rcar@{ecmwf_host}",
        f"rcar@{eumetsat_host}"
    ])

[#] ip link add dasklocal type wireguard
[#] wg setconf dasklocal /dev/fd/63
[#] ip -6 address add fda5:c0ff:eeee:0::1/64 dev dasklocal
[#] ip link set mtu 1420 up dev dasklocal
[#] ip -6 route add fda5:c0ff:eeee:2::/64 dev dasklocal
[#] ip -6 route add fda5:c0ff:eeee:1::/64 dev dasklocal
2022-06-29 14:46:57,237 - distributed.scheduler - INFO - -----------------------------------------------
2022-06-29 14:46:58,602 - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
2022-06-29 14:46:58,643 - distributed.scheduler - INFO - -----------------------------------------------
2022-06-29 14:46:58,644 - distributed.scheduler - INFO - Clear task state
2022-06-29 14:46:58,646 - distributed.scheduler - INFO -   Scheduler at:     tcp://172.17.0.2:8786
2022-06-29 14:46:58,646 - distributed.scheduler - INFO -   dashboard at:                     :8787
2022-06-29 14:47:16,104 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:1::11]:37977', name: ecmwf-1-2, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:16,107 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:1::11]:37977
2022-06-29 14:47:16,108 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:16,108 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:1::11]:44575', name: ecmwf-1-3, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:16,109 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:1::11]:44575
2022-06-29 14:47:16,109 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:16,113 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:1::11]:40121', name: ecmwf-1-1, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:16,114 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:1::11]:40121
2022-06-29 14:47:16,114 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:16,119 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:1::11]:40989', name: ecmwf-1-0, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:16,121 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:1::11]:40989
2022-06-29 14:47:16,121 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:23,342 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:2::11]:33423', name: eumetsat-2-0, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:23,343 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:2::11]:33423
2022-06-29 14:47:23,343 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:23,346 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:2::11]:43953', name: eumetsat-2-1, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:23,348 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:2::11]:43953
2022-06-29 14:47:23,348 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:23,350 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:2::11]:46089', name: eumetsat-2-3, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:23,352 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:2::11]:46089
2022-06-29 14:47:23,352 - distributed.core - INFO - Starting established connection
2022-06-29 14:47:23,357 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://[fda5:c0ff:eeee:2::11]:43727', name: eumetsat-2-2, status: undefined, memory: 0, processing: 0>
2022-06-29 14:47:23,358 - distributed.scheduler - INFO - Starting worker compute stream, tcp://[fda5:c0ff:eeee:2::11]:43727
2022-06-29 14:47:23,358 - distributed.core - INFO - Starting established connection

We need to wait for 8 distributed.core - INFO - Starting established connection lines - one from each of 4 worker processes on each of 2 worker machines.

What has happened here is:

start-scheduler.sh runs up a Docker container on this computer.
The container creates a WireGuard IPv6 data plane VPN. This involves generating shared keys for all the nodes and a network interface inside itself. This data plane VPN is transient and unique to this cluster.
The container runs a Dask scheduler, hosted on the data plane network.
It then asks each data centre to provision workers and routing.

Each data centre hosts a control process, accessible over the control plane network. On invocation:

The control process creates a WireGuard network interface on the data plane network. This acts as a router between the workers inside the data centres and the scheduler.
It starts Docker containers on compute instances. These containers have their own WireGuard network interface on the data plane network, routing via the control process instance.
The Docker containers spawn (4) Dask worker processes, each of which connects via the data plane network back to the scheduler created at the beginning.

The result is one container on this computer running the scheduler, talking to a container on each worker machine, over a throw-away data plane WireGuard IPv6 network which allows each of the (in this case 8) Dask worker processes to communicate with each other and the scheduler, even though they are partitioned over 3 data centres.

Something like this

Image(filename="images/datacentres-dask.png")

Key

Data plane network
Dask
NetCDF data

Connecting to the Cluster

The scheduler for the cluster is now running in a Docker container on this machine and is exposed on localhost, so we can create a client talking to it

client = Client("localhost:8786")

2022-06-29 14:47:35,535 - distributed.scheduler - INFO - Receive client connection: Client-69f22f41-f7ba-11ec-a0a2-0acd18a5c05a
2022-06-29 14:47:35,536 - distributed.core - INFO - Starting established connection
/home/ec2-user/miniconda3/envs/jupyter/lib/python3.10/site-packages/distributed/client.py:1287: VersionMismatchWarning: Mismatched versions found

+---------+--------+-----------+---------+
| Package | client | scheduler | workers |
+---------+--------+-----------+---------+
| msgpack | 1.0.4  | 1.0.3     | 1.0.3   |
| numpy   | 1.23.0 | 1.22.3    | 1.22.3  |
| pandas  | 1.4.3  | 1.4.2     | 1.4.2   |
+---------+--------+-----------+---------+
Notes:
-  msgpack: Variation is ok, as long as everything is above 0.6
  warnings.warn(version_module.VersionMismatchWarning(msg[0]["warning"]))

If you click through the client you should see the workers under the Scheduler Info node

# client

You can also click through to the Dashboard on http://localhost:8787/status. There we can show the workers on the task stream

def show_all_workers():
    my_org().compute(workers='ecmwf-1-0')
    my_org().compute(workers='ecmwf-1-1')
    my_org().compute(workers='ecmwf-1-2')
    my_org().compute(workers='ecmwf-1-3')
    my_org().compute(workers='eumetsat-2-0')
    my_org().compute(workers='eumetsat-2-1')
    my_org().compute(workers='eumetsat-2-2')
    my_org().compute(workers='eumetsat-2-3')
    sleep(0.5)

show_all_workers()

Running on the Cluster

Now that there is a Dask client in scope, calculations will be run on the cluster. We can define the tasks to be run

predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')
observations = xarray.open_dataset('/data/eumetsat/ad-hoc/observations.nc').chunk('auto')

averages = predictions.mean('realization')
diff = averages.isel(height=10) - observations

But when we try to perform the calculation it fails

with pytest.raises(FileNotFoundError) as excinfo:
    show_all_workers()
    diff.compute()

str(excinfo.value)

"[Errno 2] No such file or directory: b'/data/eumetsat/ad-hoc/observations.nc'"

It fails because the Dask scheduler has sent some of the tasks to read the data to workers running in EUMETSAT. They cannot see the data in ECMWF, and nor do we want them too, because reading all that data between data centres would be too slow.

Data-Proximate Computation

Dask has the concept of resources. Tasks can be scheduled to run only where a resource (such as a GPU or amount of RAM) is available. We can abuse this mechanism to pin tasks to a data centre, by treating the data centre as a resource.

To do this, when we create the workers we mark them as having a pool-ecmwf or pool-eumetsat resource. Then when we want to create tasks that can only run in one data centre, we annotate them as requiring the appropriate resource

with (dask.annotate(resources={'pool-ecmwf': 1})):
    predictions.mean('realization').isel(height=10).compute()

We can hide that boilerplate inside a Python context manager - pool - and write

with pool('ecmwf'):
    predictions.mean('realization').isel(height=10).compute()

The pool context manager is a collaboration with the Dask developers, and is published on GitHub. You can read more on the evolution of the concept on the Dask Discourse.

We can do better than annotating the computation tasks though. If we load the data inside the context manager block, the data loading tasks will carry the annotation with them

with pool('ecmwf'):
    predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')

In this case we need another context manager in the library, propagate_pools, to ensure that the annotation is not lost when the task graph is processed and executed

with propagate_pools():
    predictions.mean('realization').isel(height=10).compute()

The two context managers allow us to annotate data with its pool, and hence where the loading tasks will run

with pool('ecmwf'):
    predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')

with pool('eumetsat'):
    observations = xarray.open_dataset('/data/eumetsat/ad-hoc/observations.nc').chunk('auto')

define some deferred calculations oblivious to the data’s provenance

averaged_predictions = predictions.mean('realization')
diff = averaged_predictions.isel(height=10) - observations

and then perform the final calculation

%%time
with propagate_pools():
    show_all_workers()
    diff.compute()

CPU times: user 127 ms, sys: 6.34 ms, total: 133 ms
Wall time: 4.88 s

Remember, our aim was to distribute a calculation across data centres, whilst preventing workers reading foreign bulk data.

Here we know that data is only being read by workers in the appropriate location, because neither data centre can read the other’s data. Once data is in memory, Dask prefers to schedule tasks on the workers that have it, so that the local workers will tend to perform follow-on calcuations, and data chunks will tend to stay in the data centre that they were read from.

Ordinarily though, if workers are idle, Dask would use them to perform calculations even if they don’t have the data. If allowed, this work stealing would result in data being moved between data centres unnecessarly, a potentially expensive operation. To prevent this, the propagate_pools context manager installs a scheduler optimisation that disallows work-stealing between workers in different pools.

Once data loaded in one pool needs to be combined with data from another (the substraction in averaged_predictions.isel(height=10) - observations above), this is no longer classified as work stealing, and Dask will move data between data centres as required.

That calculation in one go looks like this

%%time
with pool('ecmwf'):
    predictions = xarray.open_dataset('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc').chunk('auto')

with pool('eumetsat'):
    observations = xarray.open_dataset('/data/eumetsat/ad-hoc/observations.nc').chunk('auto')

averages = predictions.mean('realization')
diff = averages.isel(height=10) - observations

with propagate_pools():
    show_all_workers()
    plt.figure(figsize=(6, 6))
    plt.imshow(diff.to_array()[0,...,0], origin='lower')

CPU times: user 234 ms, sys: 27.6 ms, total: 261 ms
Wall time: 6.04 s

In terms of code, compared with the local version above, this has only added the use of with blocks to label data and manage execution, and executes some 100 times faster.

This is a best case, because in the demonstrator the files are actually hosted on the worker machines, so the speed difference between reading the files locally and reading them over NFS is maximized. Perhaps more persuasive is the measured volume of network traffic.

Method	Time Taken	Measured Network Traffic
Calculation over NFS	> 14 minutes	2.8 GiB
Distributed Calculation	~ 10 seconds	8 MiB

Apart from some control and status messages, only the data required to paint the picture is sent over the network to this computer.

Looking at the task stream we see the ECMWF workers (on the bottom) doing the bulk of the reading and computation, with the red transfer tasks joining this with data on EUMETSAT.

Image(filename="images/task-graph.png")

Catalogs

We can simplify this code even more. Because the data-loading tasks are labelled with their resource pool, this can be opaque to the scientist. So we can write

def load_from_catalog(path):
    with pool(path.split('/')[2]):
        return xarray.open_dataset(path).chunk('auto')

allowing us to ignore where the data came from

predictions = load_from_catalog('/data/ecmwf/000490262cdd067721a34112963bcaa2b44860ab.nc')
observations = load_from_catalog('/data/eumetsat/ad-hoc/observations.nc')

averages = predictions.mean('realization')
diff = averages.isel(height=10) - observations

with propagate_pools():
    show_all_workers()
    diff.compute()

Of course the cluster would have to be provisioned with compute resources in the appropriate data centres, although with some work this could be made dynamic as part of the catalog code.

More Information

This notebook, and the code behind it, are published in a GitHub repository.

For details of the prototype implementation, and ideas for enhancements, see dask-multi-cloud-details.ipynb.

Acknowledgements

Thank you to Armagan Karatosun (EUMETSAT) and Vasileios Baousis (ECMWF) for their help and support with the infrastructure to support this proof of concept. Gabe Joseph (Coiled) wrote the clever pool context managers, and Jacob Tomlinson (NVIDIA) reviewed this document.

Documentation Framework

2022-07-15T00:00:00+00:00

Yesterday at the Dask BOF at SciPy we were talking about the recent docs work and how we can fill holes in our documentation. We want to come up with a strategy to improve things.

For a while, we’ve been exploring moving our documentation to the Diátaxis Framework, and after catching up with other maintainers at SciPy it is clear that many projects are converging on this framework and we are confident about continuing on this journey. This post lays out how we will take the existing docs and apply the framework to make content clearer and easier to find. NOTE: This blog post sketches out where we are going, but the change will happen incrementally.

We want the docs to quickly answer questions like:

I know my workflow is parallizable - can Dask help?
How do I find my logs for a particular worker?
Should I be using .persist()?

Theory
Current Documentation
New Structure
How you can help

Theory

The Diataxis Framework proposes that the documentation be split into 4 entirely separate sections.

Credit: https://diataxis.fr/

Each section serves a unique purpose.

Tutorials provide a narrative that addresses a particular larger objective such as predicting global temperature or analyzing financial data.
How-Tos target people who already know what they want to do and are trying to figure out how to do it. These people might ask questions like:
- How do I apply a rolling mean to a timeseries?
- How do I groupby a column?
- How do I write to a geotiff?
Reference provides the exact arguments and outputs of a particular operation.
Explanation gives context and includes descriptions of how operations work internally.

Current Documentation

There are several different sites that comprise different aspects of dask documentation. Of particular interest are Examples, Tutorials and Docs.

The bulk of the documentation that we currently have on Docs falls under “Explanation” and “Reference” but they are pretty intermingled. There are also some small “How-Tos” sprinkled in, particularly in the API docs.

The material on Tutorials is a mixture of “Tutorial” and “Explanation”. They answer questions like: “What can I do with Dask Dataframes?” and questions like “What is a Dask Dataframe?”. These are styled like lectures in that there is often no motivating example and the assumption is that the audience wants to learn both about how to do specific operations in dask and how those operations work. This type of material can be consumed as standalone content and runs on binder.

Examples pretty much falls under “How-To” but there is a fair amount of setup and each example isn’t split into small enough bits. They answer questions like: “How do I use dask dataframes?” and they have some more longer workflows.

Which pages are most used?

From Google Analytics we can see the most commonly viewed pages.

It is hard to understand whether those pages are the most visible, or if they actually contain the information that people are trying to find. But either way, it conveys the importance of navigation in directing users.

New Structure

Tutorial will be left as is and treated as a long-form overview.

Examples will be presented more as How-Tos with little explanation and more code. This will be similar to gallery or cookbook style documentation that you may see in other projects. Historically one of the current roles of examples is to demonstrate what Dask looks like, that role is now subsumed by “10 Minutes to Dask”.

Docs will be reorganized and the left-nav will be slimmed down dramatically to provide direction. One idea for the left-nav is:

Installation
10 Minutes to Dask (How-To)
Tutorials & Talks (Tutorial)
Examples (How-To)- this will be a landing page that points to individual sections of Examples or API Reference.
Best Practices (Explanation)
API (Reference)
User Guide (Explanation)
- DataFrame - explains what a dataframe is - links out aggressively to reference docs.
- Array
- Bag
- Delayed
- Futures
- Task Graphs
- Scheduling
- Diagnostic Dashboard
Configuration (Reference)
Deploy Dask Clusters (Reference)
Development Guidelines (Reference)
Changelog (Reference)
FAQs (Reference)

Many docstrings (aka Reference) already contain their own short-form How-To docs. I think this is a good place for these and we can thoroughly link from other places to these canonical docs.

How you can help

Please raise issues on the dask issue tracker when you find holes in the docs! The largest gaps we see now are in “how to” which commonly are found via Google. So if you search for how to do something in Dask, and you’re looking for copy-paste examples but can’t find any then let us know.

If you see other gaps please let us know about those too. And if you know how your needs fit into the diataxis framework even better :)

How to run different worker types with the Dask Helm Chart

2022-02-17T00:00:00+00:00

Today, we’ll learn how to deploy Dask on a Kubernetes cluster with the Dask Helm Chart and then run and scale different worker types with annotations.

What is the Dask Helm Chart?

The Dask Helm Chart is a convenient way of deploying Dask using Helm, a package manager for Kubernetes applications. After deploying Dask with the Dask Helm Chart, we can connect to our HelmCluster and begin scaling out workers.

What is Dask Kubernetes?

Dask Kubernetes allows you to deploy and manage your Dask deployment on a Kubernetes cluster. The Dask Kubernetes Python package has a HelmCluster class (among other things) that will enable you to manage your cluster from Python. In this tutorial, we will use the HelmCluster as our cluster manager.

Prerequisites

To have Helm installed and be able to run helm commands
To have a running Kubernetes cluster. It doesn’t matter whether you’re running Kubernetes locally using MiniKube or Kind or you’re using a cloud provider like AWS or GCP. But your cluster will need to have access to GPU nodes to run GPU workers. You’ll also need to install RAPIDS to run the GPU worker example.
To have kubectl installed. Although this is not required.

That’s it, let’s get started!

Install Dask Kubernetes

From the documentation,

pip install dask-kubernetes --upgrade

conda install dask-kubernetes -c conda-forge

Install the Dask Helm Chart

First, deploy Dask on Kubernetes with Helm:

helm repo add dask https://helm.dask.org/
helm repo update
helm install my-dask dask/dask

Now you should have Dask running on your Kubernetes cluster. If you have kubectl installed, you can run kubectl get all -n default

You can see that we’ve created a few resources! The main thing to know is that we start with three dask workers.

Add GPU worker group to our Dask Deployment

The Helm Chart has default values that it uses out of the box to deploy our Dask cluster on Kubernetes. But now, because we want to create some GPU workers, we need to change the default values in the Dask Helm Chart. To do this, we can create a copy of the current values.yaml, update it to add a GPU worker group and then update our helm deployment.

First, you can copy the contents of the values.yaml file in the Dask Helm Chart and create a new file called my-values.yaml
Next, we’re going to update the section in the file called additional_worker_groups. The section looks like this:

additional_worker_groups: [] # Additional groups of workers to create
# - name: high-mem-workers  # Dask worker group name.
#   resources:
#     limits:
#       memory: 32G
#     requests:
#       memory: 32G
# ...
# (Defaults will be taken from the primary worker configuration)

Now we’re going to edit the section to look like this:

additional_worker_groups: # Additional groups of workers to create
  - name: gpu-workers # Dask worker group name.
    replicas: 1
    image:
      repository: rapidsai/rapidsai-core
      tag: 21.12-cuda11.5-runtime-ubuntu20.04-py3.8
      dask_worker: dask-cuda-worker
    extraArgs:
      - --resources
      - "GPU=1"
    resources:
      limits:
        nvidia.com/gpu: 1

Now we can update our deployment with our new values in my-values.yaml

helm upgrade -f my-values.yaml my-dask dask/dask

Again, you can run kubectl get all -n default, and you’ll see our new GPU worker pod running:

Now we can open up a jupyter notebook or any editor to write some code.

Scaling the workers Up/Down

We’ll start by importing the HelmCluster cluster manager from Dask Kubernetes. Next, we connect our cluster manager to our dask cluster by passing the release_name of our Dask cluster as an argument. That’s it, the HelmCluster automatically port-forwards the scheduler to us and can give us quick access to logs. Next, we’re going to scale our Dask cluster.

from dask_kubernetes import HelmCluster
cluster = HelmCluster(release_name="my-dask")
cluster

To scale our cluster, we need to provide our desired number of workers as an argument to the HelmCluster’s scale method. By default, the scale method scales our default worker group. You can see in the first example we scaled the default worker group from three to five workers, giving us six workers in total. In the second example, we use the handy worker_group keyword argument to scale our GPU worker group from one to two workers, giving us seven workers in total.

cluster.scale(5)  # scale the default worker group from 3 to 5 workers
cluster

cluster.scale(2, worker_group = "gpu-workers")  # scale the GPU worker group from 1 to 2 workers
cluster

Example: Finding the average New York City taxi trip distance in April 2020

This example will find the average distance traveled by a yellow taxi in New York City in April 2020 using the NY Taxi Dataset. We’ll compute this distance in two different ways. The first way will employ our default dask workers, and the second way will utilize our GPU worker group. We’ll load the NY Taxi dataset as a data frame in both examples and compute the mean of the trip_distance column. The main difference is that we need to run our GPU-specific computations using our GPU worker group. We can do this by utilizing Dask annotations.

import dask.dataframe as dd
import dask

link = "https://s3.amazonaws.com/nyc-tlc/trip+data/yellow_tripdata_2020-04.csv"
ddf = dd.read_csv(link, assume_missing=True)
avg_trip_distance = ddf['trip_distance'].mean().compute()
print(f"In January 2021, the average trip distance for yellow taxis was {avg_trip_distance} miles.")

with dask.annotate(resources={'GPU': 1}):
    import dask_cudf, cudf
    dask_cdf = ddf.map_partitions(cudf.from_pandas)
    avg_trip_distance = dask_cdf['trip_distance'].mean().compute()
    print(f"In January 2021, the average trip distance for yellow taxis was {avg_trip_distance} miles.")

Closing

That’s it! We’ve deployed Dask with Helm, created an additional GPU worker type, and used our workers to run an example calculation using the NY Taxi dataset. We’ve learned several new things:

The Dask Helm Chart lets you create multiple worker groups with different worker types. We saw this when we made two different groups of Dask Workers: CPU and GPU workers.
You can run specific computations on your workers of choice with annotations. Our example computed the average taxi distance using the RAPIDS libraries cudf and dask_cudf on our GPU worker group.
The HelmCluster cluster manager in Dask Kubernetes lets you scale your worker groups quickly from python. We scaled our GPU worker group by conveniently passing the worker group name as a keyword argument in the HelmCluster scale method.

Future Work

We’re thinking a lot about the concept of worker groups in the Dask community. Until now, most Dask deployments have homogenous workers, but as Dask users push Dask further, there is a growing demand for heterogeneous clusters with special-purpose workers. So we want to add worker groups throughout Dask.

Dask Working Notes - Posted in 2022

Dask Demo Day November 2022

The New Dask CLI

XGBoost HPO with Dask and Optuna

Dask for Awkward Arrays

Profiling Dask on a Cluster with py-spy

Join us for the next Demo Day!

Reducing memory usage in Dask workloads by 80%

Smooth is fast

Sometimes, slow is still slow (but not why you’d expect)

A new default for scheduling

Let us know how it goes

Keeping old behavior

Dask Kubernetes Operator

Why did we build this?

New features

Python or YAML API

Cascade deletion

Multiple worker groups

Autoscaling

DaskJob

Extensibility and plugins

Migrating

Understanding Dask’s meta keyword argument

Why does Dask need meta?

How to specify meta?

meta does not force the structure or dtypes

Using ._meta directly

Data Proximate Computation on a Dask Cluster Distributed Between Data Centres

Introduction

Setup

Mount the Data

Access to Data

Our Calculation

Running Up a Cluster

Connecting to the Cluster

Running on the Cluster

Data-Proximate Computation

Catalogs

More Information

Acknowledgements

Documentation Framework

Contents

Theory

Current Documentation

Which pages are most used?

New Structure

How you can help

How to run different worker types with the Dask Helm Chart

What is the Dask Helm Chart?

What is Dask Kubernetes?

Prerequisites

Install Dask Kubernetes

Install the Dask Helm Chart

Add GPU worker group to our Dask Deployment

Scaling the workers Up/Down

Example: Finding the average New York City taxi trip distance in April 2020

Closing

Future Work

Why does Dask need `meta`?

How to specify `meta`?

`meta` does not force the structure or dtypes

Using `._meta` directly