dask / dask/dask-kubernetes

Improve advanced KubeCluster customization API

Open
#899 2 comments 0 reactions 0 assignees View on GitHub
enhancement idea
Dominant language
Python
Stars
324
Forks
157
PR merge metrics
No merged PRs in 30d

Description

## Background

When creating a new `KubeCluster` the `__init__()`/`_start()` process happens in a [couple of steps](
https://github.com/dask/dask-kubernetes/blob/734d001836fad228cca809093feec3d07cd634a6/dask_kubernetes/operator/kubecluster/kubecluster.py#L343-L361).

1. Generate a dictionary of the Kubernetes spec from provided kwargs using `make_cluster_spec()`
2. Create a `DaskCluster` class (which was created by `kr8s.asyncio.objects.make_class()`) from that dictionary
3. Call `DaskCluster.create()`

If users want to customize their cluster spec before creating it they can use `make_cluster_spec()` directly which takes the same kwargs as `KubeCluster` and returns the dictionary that is generated in step 1, then they can modify the dict before passing it directly to step 2 via the kwarg `KubeCluster(custom_cluster_spec=...)` .

## Problem statement

When it comes to adding new kwargs to `make_cluster_spec()` we are very conservative because it is easy for that method to become extremely overwhelming. However as highlighted in #898 there may be customizations we want to do regularly like setting the memory limit on the `dask worker` CLI args which is not very ergonomic to do by modifying the dict directly. It would be nice of the object returned from `make_cluster_spec()` was easier to customize than a simple dictionary.

For example

```python
from dask_kubernetes.operator import KubeCluster, make_cluster_spec

cluster_spec = make_cluster_spec("foo", ...)

# Currently we do this
cluster_spec["spec"]["worker"]["spec"]["containers"][0]["args"] += ["--memory-limit", "4GB"]
# It would be nicer if we could do something like this instead
# cluster_spec.set_worker_memory_limit("4GB")

cluster = KubeCluster(custom_cluster_spec=cluster_spec)
```

## Option 1

Instead of returning a dict from `make_cluster_spec()` we could return the `DaskCluster` object from step 2. This object is dict-like so any existing code that modifies the dict should work as expected (although this will need to be verified).

### Pros

The existing `DaskCluster` class is a convenient place to add methods to simplify some of the common customizations. It will be created anyway, we would just be changing the step at which the user can intervene from 1 to 2.

### Cons

The `DaskCluster` object that would be returned by `make_cluster_spec()` would be an async `kr8s` object and potentially `KubeCluster` is being used in a sync context, and we would then be extending it with sync methods to modify things.

Returning the `DaskCluster` object leaks the implementation detail that we are using `kr8s` to interact with Kubernetes.

It also means in theory a user could call `await cluster_spec.create()` which would create the Dask cluster, but it wouldn't handle any of the port forwarding, log retrieval, etc and couldn't be passed to `dask.distributed.Client`.

Having `KubeCluster` and `DaskCluster`, which serve different purposes but have similar names, may cause confusion with users.

## Option 2

We could introduce some new intermediate object with a name like `KubeClusterSpec` which would effectively be a subclass of `dict` but with whatever convenience methods we want. We would return this from `make_cluster_spec()`.

We would need to ensure that `DaskCluster` can be passed a `KubeClusterSpec` instead of a `dict`. There is [already support in `kr8s` for passing a few different things](https://github.com/kr8s-org/kr8s/blob/8a1dbdc03bcba100f5e90ff9b039dea8f3209867/kr8s/_objects.py#L60-L69), so we should be able to just conform to one of those.

### Pros

Adding a new class would avoid any user confusion as users would continue to never interact with `kr8s` resources directly.

The purpose of the new class is well scoped, it is a dictionary of the spec with some utility methods to simplify modifying the spec.

### Cons

Introducing a new class into the chain increases complexity.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.