Deploying Multiple Worker Types with Helm

ljstrnadiii · January 17, 2022, 4:58pm

I am getting more and more familiar with dask and have successfully deployed with helm with GKE. It is such a game-changer and we are so excited to be using this.

As a machine learning engineer, I would like to run specific operations on specific node types. For example, it would be nice to pull down zarr/xarray datastores, run transforms on cpu machines, and transfer to GPU nodes for inference.

I have found this page Worker Resources — Dask.distributed 2022.8.1+6.gc15a10e8 documentation that seems to describe how to select with resources to pick for certain tasks. Something like

from distributed import Client
client = Client('scheduler:8786')

data = [client.submit(load, fn) for fn in filenames]
processed = [client.submit(process, d, resources={'GPU': 1}) for d in data]
final = client.submit(aggregate, processed, resources={'MEMORY': 70e9})

However, I am not sure how to deploy multiple worker types with helm. It feels natural to me that worker types could be defined in the values.yaml file like:

scheduler:
  name: scheduler  # Dask scheduler name.
  enabled: true  # Enable/disable scheduler.
  image:
    repository: "daskdev/dask"   # Container image repository.
    tag: 2022.1.0 # Container image tag.
    pullPolicy: IfNotPresent  # Container image pull policy.
    pullSecrets:  # Container image [pull secrets](https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/).
  ...
  
worker:
  name: cpu-worker-group  # Dask worker name. # can we make multiple with different node pools?
  image:
    repository: "daskdev/dask"  # Container image repository.
    tag: 2022.1.0  # Container image tag.
    pullPolicy: IfNotPresent  # Container image pull policy.
    dask_worker: "dask-worker --name cpu-worker"  # Dask worker command. E.g `dask-cuda-worker` for GPU worker.
    pullSecrets:  # Container image [pull secrets](https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/).
    #  - name: regcred
  ...
  nodeSelector:
    cloud.google.com/gke-nodepool: name-of-cpu-node-pool

  
worker:
  name: gpu-worker-group  # Dask worker name. # can we make multiple with different node pools?
  image:
    repository: "daskdev/dask"  # Container image repository.
    tag: 2022.1.0  # Container image tag.
    pullPolicy: IfNotPresent  # Container image pull policy.
    dask_worker: "dask-cuda-worker --name gpu-worker"  # Dask worker command. E.g `dask-cuda-worker` for GPU worker.
    pullSecrets:  # Container image [pull secrets](https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/).
    #  - name: regcred
  ...
  nodeSelector:
    cloud.google.com/gke-nodepool: name-of-gpu-node-pool

Ideally, multiple worker types would be available, but I am having a hard time finding anything about this with a helm deployment. When I try to deploy with the above values.yaml, I only see the first worker and the second seems to be ignored.

Also, it would be great to be able to choose a worker set for a specific task. Above, the second worker has name: gpu-worker-group and an option for a name in the worker command: dask-cuda-worker --name gpu-worker. I don’t think it makes sense to use --name here as it is used to name an individual instance of those workers from what I understand. However, anything like this would be great:

# either pass the worker set name with a pattern
future = client.submit(func, *args, workers=['gpu-worker*'])
# or specify a group name somehow
future = client.submit(func, *args, workers=['gpu-worker-group'])

I have seen that we can specify the host:port, but I don’t want to mess with dask’s ability to chose which worker might already have data, etc. I would like to just specify a worker type for things to run on to take advantage of hardware and resources differences. Ideally, I would have worker sets: small cpu resources, large cpu resources, gpu resources.

A few questions here:

Is there a way values.yaml for a helm deployment can define multiple workers? This would be so nice and seems possible with Ray’s helm deployment.
Is there a way we can select worker types for client.submit, client.map, etc? This seems like an analogy to the nodeSelector.
If there is a way to do this, how would we scale? cluster.scale(5, worker_type='gpu-worker-group)?

Thanks for any thoughts!

ljstrnadiii · January 18, 2022, 12:14am

Some relevant issues:

github.com/dask/distributed

Cluster should support many worker types

opened 12:40PM - 15 Jul 18 UTC

mrocklin

discussion

The various Cluster objects often allow the user to provide specifications of a …worker (cores, memory, software environment, ...) and then provides mechanisms around increasing and decreasing the number of workers. However sometimes a dask deployment has a few different kinds of workers, for example machines with GPUs or high memory, or machines from a queue that is more or less expensive or reliable in some way. This suggests that maybe the Cluster object should accept a list of worker pools, and provide common functionality around them. Things like the widget are easy to scale to multiple pools. Adaptivity is a bit weirder. cc @lesteve @jhamman (dask-jobqueue) @jcrist (dask-yarn) @jacobtomlinson (dask-kubernetes) Credit for this thought goes to @lesteve

github.com/dask/dask

Support Task Annotations within the graph

opened 12:34PM - 19 Jul 18 UTC

sjperkins

scheduler core

I think this issue has some relation to https://github.com/dask/distributed/issu…es/2127, where @kkraus14 wants to run certain tasks on CPU/GPU workers. I've also wanted to run tasks on specific workers, or require resources to be exclusive for certain tasks. Currently, these task dependencies must be specified as additional arguments to `compute`/`persist` etc. rather than at the point of actual construction -- embedding resource/worker dependencies in the graph is not currently possible. To support this, how about adding a `TaskAnnotation` type? This can be a namedtuple, itself containing nested tuples representing key-value pairs. e.g. ```python annot = TaskAnnotation(an=(('resource', ('GPU': '1'), ('worker', 'alice'))) ``` dask array graphs tend to have the following structure: ```python dsk = { (tsk_name, 0) : (fn, arg1, arg2, ..., argn), (tsk_name, 1) : (fn, arg1, arg2, ..., argn), } ``` How about embedding annotations within value tuples? ```python dsk = { (tsk_name, 0) : (fn, arg1, arg2, ..., argn, annotation1), (tsk_name, 1) : (fn, arg1, arg2, ..., argn, annotation2), } ``` If the scheduler discovers an annotation in the tuple, it could remove it from the argument list and attempt to satisfy the requested constraints. In the above example, annotations are placed at the end of the tuple, but the location could be arbitrary and multiple annotations are possible. Alternatively, it might be better to put them at the start. I realise the above example is somewhat specific to dask arrays (I'm not too familiar with the dataframe and bag collections) so there may be issues I'm not seeing. One problem I can immediately identify would be modifying existing graph construction functions to support the above annotations (atop/top support is probably the first place to look).

jacobtomlinson · January 18, 2022, 11:59am

Thanks for opening this @ljstrnadiii.

This isn’t possible today, but is a great idea. I’ve opened an issue on GitHub to track this request.

github.com/dask/helm-chart

Add additional worker groups

opened 10:40AM - 18 Jan 22 UTC

closed 03:40PM - 20 Jan 22 UTC

jacobtomlinson

enhancement help wanted chart/dask

Inspired by dask/dask-kubernetes#384 it would be awesome to support multiple wor…ker groups in the Helm Chart. Currently the chart creates one deployment for workers, but there is nothing to stop us from creating multiple deployments with different resources. For example we could add an additional GPU or high memory deployment which Dask can select via task annotations. Today the worker is configured like this ```yaml worker: name: worker image: repository: "daskdev/dask" tag: 2022.1.0 replicas: 3 resources: requests: cpu: 1 memory: 3G ``` This could be turned into a list to support multiple groups. ```yaml worker_groups: - name: cpu image: repository: "daskdev/dask" tag: 2022.1.0 replicas: 3 resources: limits: cpu: 1 memory: 3G nvidia.com/gpu: 1 requests: cpu: 1 memory: 3G nvidia.com/gpu: 1 - name: gpu image: repository: "daskdev/dask" tag: 2022.1.0 replicas: 2 resources: requests: cpu: 4 memory: 12G nvidia.com/gpu: 1 ``` Or perhaps to maintain backward compatibility we could do something like this. ```yaml worker: name: worker image: repository: "daskdev/dask" tag: 2022.1.0 replicas: 3 resources: requests: cpu: 1 memory: 3G additional_worker_groups: - name: gpu image: repository: "daskdev/dask" tag: 2022.1.0 replicas: 2 resources: requests: cpu: 4 memory: 12G nvidia.com/gpu: 1 ``` Hopefully this couple be implemented by reusing the current worker deployment template and just recycling it over and over.

jacobtomlinson · February 21, 2022, 4:55pm

This is now possible! The Helm Chart can have additional worker groups configured with different resources.

Check out the full writeup here How to run different worker types with the Dask Helm Chart

Topic		Replies	Views
Additional worker deployments using `KubeCluster` Deploying Dask	1	290	May 12, 2022
Installing Dask Workers on Partcular Node Pool Distributed dask-gateway , distributed	3	351	January 31, 2024
Get host:port of additional worker groups Distributed	3	161	August 25, 2022
Define "imagePullSecrets" for scheduler and worker image Deploying Dask dask-gateway , kubernetes , scheduler , worker	1	248	July 14, 2022
Multi-GPU dask gateway pods Deploying Dask	5	133	December 1, 2023

Deploying Multiple Worker Types with Helm

Related topics