Rearchitecture: Kubernetes Backend

[jcrist]

In #195 server architecture was redone to improve resiliency/performance and make the kubernetes backend more kubernetes-native. In that PR we dropped support entirely for kubernetes, support will be added back in a subsequent PR. Here I'll outline how I think the new Kubernetes backend should be architected.

---

When running on Kubernetes, dask-gateway will be composed of the following components:

• A dask-gateway CRD object to describe specification/state about a single dask cluster
• One or more instances of the dask-gateway api server, running behind a ClusterIP Service. These provide the internal api routing, but users won't (usually) connect directly to the api server.
• One or more instances of traefik configured to use the Kubernetes CRD provider (https://docs.traefik.io/providers/kubernetes-crd/). These will handle proxying both the scheduler and dashboard routes.
• One instance of a controller, managing dask-gateway CRD objects.

When initially installed, only a single IngressRoute object will be created to expose the dask-gateway api server through the proxy. As dask-clusters come and go, additional routes will be added/removed. Users will only communicate with the traefik proxy service directly, remaining network traffic doesn't need to be exposed directly.

Walking through the components individually.

Dask-Gateway CRD

An instance of a Dask-Gateway CRD object needs to contain information on how to run a single dask cluster. Here's a sketch of a proposed schema (in pseudocode):

spec:
    schedulerTemplate: ...  # A pod template for what the scheduler pod should look like
    workerTemplate: ...  # A pod template for what a single worker pod should look like
    ingressRoutes: ...  # A list of IngressRoute objects to create after the scheduler pod is running
    replicas: 0  # The number of workers to create, initially 0
status:
    # Status tracking things go here, not sure what we'd need to track

Each dask cluster created by the gateway will have one of these objects created.

Backend Class

To plugin to the dask-gateway api server, we'd need to write a kubernetes backend. We may need to tweak the interface exposed by the dask_gateway_server.backends.base.Backend class, but I think the existing interface is pretty close to what we'll want.

Walking through the methods:

• start_cluster(self, user, cluster_options)

This is called when a new cluster is requested for a user. This should validate the request, then compose and create the dask-gateway CRD object and corresponding Secret (storing the TLS credentials and api token). It should return the object name (or some proxy for it) so that it can be looked up again. Note that when this method returns the cluster doesn't need to be running, only submitted.

• stop_cluster(self, cluster_name, failed=False)

This is called when a cluster is requested to stop. This should do whatever is needed to stop the CRD object, and should be a no-op if it's already stopped. As above, this doesn't need to wait for the cluster to stop, only for it to be stopping.

If we want to keep some job history then stop_cluster can't delete the CRD object, since that's the only record that the job happened. If we don't care about this then deleting the CRD is fine. Other backends keep around a short job history, but that may not matter here.

• get_cluster(self, cluster_name, wait=False)

This is called when a user wants to lookup or connect to a cluster. It provides the cluster name (as returned by start_cluster. It should return a dask_gateway_server.models.Cluster object.

I think what we'll want is for each instance of the api server to watch for objects required to compose the Cluster model (dask-gateway CRD objects, secrets with a label selector, etc...). If the needed object is already reflected locally, we can use that. If it's not, we can try querying the k8s server in case the object was created and hasn't propagated to us yet (and then store it in the reflector locally). It should be fine to return slightly outdated information from this call.

• list_clusters(self, username=None, statuses=None)

This is used to query clusters, optionally filtering on username and statuses. Like get_cluster above, this can be slightly out of date, so I think returning only the reflected view is fine. If not, this should map pretty easily onto a k8s query with a label selector (a cluster's username will be stored in a label).

• on_cluster_heartbeat(self, cluster_name, msg)

This responds to messages from the cluster to the gateway. For the k8s backend, we only really need to handle messages that contain changes to the cluster worker count. These will translate to patch updates to the CRD objects, updating the replicas field.

Controller

The controller handles the bulk of the logic around when/how to run pods. It should watch for updates to the CRD objects, and update the child resources (pods, ...) accordingly. Walking through the logic:

• When a new CRD is created, the controller will create a scheduler pod from the scheduler template, along with any IngressRoute objects specified in the CRD. The scheduler pod should have the CRD as an owner reference so it'll be automatically cleaned up on CRD deletion. Likewise the IngressRoute objects should have the scheduler pod as an owner reference.

• When the scheduler pod is running, it will transition the CRD status to RUNNING. This signals to the API server that the cluster is now available.

• When the replicas value is higher than the current number of child workers (scale up event), additional worker pods will be created. The worker pods should be created with the scheduler pod as an owner reference so that they'll be automatically cleaned up.

• When the replicas value is lower than the current number of child workers the controller should do nothing. All worker pods should be configured with restartPolicy: onFailure. When a cluster is scaling down, the workers will exit successfully and the pods will stop themselves.

• When the CRD is signaled to shutdown, the scheduler pod should be deleted, which will delete the worker pods as well. We'll also need to delete the corresponding Secret. Perhaps start_cluster sets the CRD as the owner reference of the secret, and the controller is robust to CRDs being created before their corresponding Secret? Not sure.

We want to support running clusters in multiple namespaces in the same gateway. I'm not sure what the performance implications of watching all namespaces are vs a single namespace - if it's significant we may want to make this configurable, otherwise we may watch all namespaces always even if we only deploy in a single namespace.

We'll also want to be resilient to pod creation being denied due to a Resource Quota.

I've written a controller with the Python k8s api and it was pretty simple, but I'd be happy writing this using operator-sdk or kubebuilder as well. No strong opinions.

---

Since this splits up the kubernetes implementation into a couple of smaller services, any meaningful testing will likely require all of them running. I'm not familiar with the tooling people use for this (see #196), advice here would be useful. We can definitely write unit tests for the smaller components, but an integration test will require everything running.

[jcrist]

cc @yuvipanda, @droctothorpe, @gforsyth

[mmccarty]

cc @dankerrigan

[droctothorpe]

When a new CRD CR is created, the controller will create a scheduler pod from the scheduler template, along with any IngressRoute objects specified in the CRD CR. The scheduler pod should have the CRD CR as an owner reference so it'll be automatically cleaned up on CRD CR deletion. Likewise the IngressRoute objects should have the scheduler pod as an owner reference.

Small, admittedly pedantic, point of clarification because it can lead to some confusion: it's important to disambiguate between the custom resource definition (CRD) and individual custom resources (CRs). The former are like a class, the latter, instances of the class. Another way to think about it is that the CRD is like a table and CRs are like rows in that table.

The above is currently implemented in the operator.

When the scheduler pod is running, it will transition the CRD CR status to RUNNING. This signals to the API server that the cluster is now available.

When the replicas value is higher than the current number of child workers (scale up event), additional worker pods will be created. The worker pods should be created with the scheduler pod as an owner reference so that they'll be automatically cleaned up.

The above is urrently implemented in the operator.

When the replicas value is lower than the current number of child workers the controller should do nothing. All worker pods should be configured with restartPolicy: onFailure. When a cluster is scaling down, the workers will exit successfully and the pods will stop themselves.

The above is currently implemented in the operator.

When the CRD CR is signaled to shutdown, the scheduler pod should be deleted, which will delete the worker pods as well. We'll also need to delete the corresponding Secret. Perhaps start_cluster sets the CRD CR as the owner reference of the secret, and the controller is robust to CRDs CRs being created before their corresponding Secret? Not sure.

Easy to implement owner reference for secret. It's already being done for the deployments, services, and ingresses.

We want to support running clusters in multiple namespaces in the same gateway. I'm not sure what the performance implications of watching all namespaces are vs a single namespace - if it's significant we may want to make this configurable, otherwise we may watch all namespaces always even if we only deploy in a single namespace.

OSDK supports both workflows. Here's some documentation that touches on the pros and cons. If we're not relying on Kubernetes-native horizontal pod autoscaling, watching a single namespace may be sufficient, since scaling directives will always flow from the CRs down. All of the CRs could reside in a single namespace. That's all the controller would need to watch. The child deployments, on the other hand, could be deployed to any namespace, and the native controller could handle recovery in the event of pod failure or rescheduling.

Since this splits up the kubernetes implementation into a couple of smaller services, any meaningful testing will likely require all of them running. I'm not familiar with the tooling people use for this (see #196), advice here would be useful. We can definitely write unit tests for the smaller components, but an integration test will require everything running.

OSDK provides a testing framework. We've written some pretty comprehensive unit and acceptance tests. Those will need to be refactored alongside the forthcoming changes of course.

[jcrist]

It's already being done for the deployments, services, and ingresses.

Oh, your controller implementation uses Deployment objects? I was thinking we'd avoid using a Deployment per cluster. These objects have behavior on scaling and restarting that we don't necessarily want (restartPolicy must be Always, scaling down won't honor the cluster choosing which pod to close, ...). Rather, I think we should implement some deployment-like behaviors ourselves in the controller. This would mean that we'd need to watch child pods for each cluster CR as well.

[droctothorpe]

We're using deployments currently, but we can definitely strip that and use replica sets or even just pods instead.

The main drawback of managing pods directly would be that we would need to recreate some fantastic, universally adopted, and thoroughly tested benefits provided by the ReplicationController.

If there was a way to leverage the Kubernetes native functionality that doesn't infringe on the level of control that Dask Gateway requires, that would be ideal in my opinion. Definitely open to trying out / implementing any approach though. I'm definitely not a go native or go home zealot.

One other advantage of using deployments / replica sets is that it would permit the operator to be modular and function as a standalone solution that would provide an alternative to dask-kubernetes. I could see frameworks like Kubeflow or Flyte using it in a pluggable way (engineers from both products expressed interest in the Dask operator being open-sourced at Kubecon), which would ultimately promote Dask adoption. That being said, we could always maintain two versions of the operator, one that plugs into Dask Gateway and another that functions as a standalone solution.

[jcrist]

The main drawback of managing pods directly would be that we would need to recreate some fantastic, universally adopted, and thoroughly tested benefits provided by the ReplicationController.

I also would vastly prefer to reuse things, but am not sure the existing abstractions would work for our use case. We don't need all the fancy features that deployments provide either, just pod management and rescheduling, so implementing the subset we care about may not be too terrible.

What we want is a replicaset like thing, with the following behavior:

• Pods are run with restartPolicy: OnFailure
• If replicas > len(managed_pods), start new pods
• If replicas < len(managed_pods), do nothing, wait for pods to exit successfully

One other advantage of using deployments / replica sets is that it would permit the operator to be modular and function as a standalone solution that would provide an alternative to dask-kubernetes.

Right now I think we should focus on solving the dask-gateway use case. If after we get things working something seems generalizable, then we may extract things out.

[mrocklin]

This comes up most times that we interact with Kubernetes. When Dask scales things down, it wants to carefully select which dask-workers to kill, move data off of those workers, and then kill them. This sort of stateful handling isn't something that most Kubernetes Deployments were comfortable thinking about, so historically we've managed this ourself. Fortunately, our implementation is also decently well tested and used (although obviously less well than Kubernetes') and so ends up handling this process ok.

My guess is that Operators give us more flexibility to handle this in a more native Kubernetes fashion. I'm not familiar enough with them to say one way or the other though.

[droctothorpe]

What criteria does Dask use when selecting workers to kill?

I haven't used it before, but we may be able to use the preStop life cycle hook to ensure data transfer before termination.

Happy to work towards either implementation! Thanks for taking my input into consideration.

[mrocklin]

We want to close workers that aren't working on anything right now, and aren't storing much. In principle this could get more complex if there are some data that we can't move easily (maybe it requires a GPU for example).

The full implementation is here, you might be interested in the docstring

https://github.com/dask/distributed/blob/f561e9646a121f74d1aecc6d1ee31baeabffad49/distributed/scheduler.py#L3213-L3344

The implementation of actually retiring those workers gracefully is just below.

In my ideal world a Kubernetes Operator would query the Scheduler pod and ask it to gracefully close the right things down, and then would clean up the underlying pods. This is the kind of logic that we may want to change in the future, and I think that it would be good to keep that logic in one place.

[jcrist]

In my ideal world a Kubernetes Operator would query the Scheduler pod and ask it to gracefully close the right things down, and then would clean up the underlying pods. This is the kind of logic that we may want to change in the future, and I think that it would be good to keep that logic in one place.

This isn't exactly ideal. What if the scheduler pod is unresponsive, or slow to respond? When scaling down many clusters you'd have to limit the number of open requests. Scaling down a worker nicely can also take time.

Instead, in dask-gateway we have the scheduler manage its own n_workers setpoint. If it needs to scale up it asks the gateway for more workers. If it needs to scale down it tells the gateway what workers it's going to close, then closes them. This way the gateway is always reacting to inputs from the scheduler, not the other way around. The design outlined above in the original issue comment removes any communication between the scheduler and the operator, while still benefiting from the scale-down logic implemented in dask.

I haven't used it before, but we may be able to use the preStop life cycle hook to ensure data transfer before termination.

preStop is a hook executed before terminating a pod, not a way to choose which pods to terminate unfortunately.

[mrocklin]

Yeah, I guess I was thinking of this in the Kubernetes Operator perspective, where a user might want to scale a Dask deployment thing to n replicas. In that case it would be nice for the operator to ask the scheduler what to do. You're right though that I wasn't thinking of this in the full context of Dask-Gateway. Happy to retract my comment.

[mmccarty]

preStop is a hook executed before terminating a pod, not a way to choose which pods to terminate unfortunately.

Maybe not for choosing which pod to terminate, but what about ensuring data is transferred to other workers before termination?

[jcrist]

It could be used for that, but it's not ideal. We'd like to select a pod that has the least data/work to minimize the cost of shutting it down. Kubernetes could select a bad pod (say one with most of the data) and the cluster could rebalance itself before scaling down, but it'd be less efficient than if we let dask make the decision itself.

[mmccarty]

That makes sense. I see where the data transfer is initiated in dask now https://github.com/dask/distributed/blob/f561e9646a121f74d1aecc6d1ee31baeabffad49/distributed/scheduler.py#L3423-L3429