Dask Working Notes - Posts tagged deployment

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.

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.