Add Autoscaling to Operator (#451)

.github/workflows/helmcluster.yaml

@@ -22,6 +22,9 @@ jobs:
       matrix:
         python-version: ["3.8", "3.9", "3.10"]
 
+    env:
+      KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
+
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
@@ -30,6 +33,7 @@ jobs:
       - name: Install deps
         run: ./ci/install-deps.sh
       - name: Run tests
-        env:
-          KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
         run: pytest dask_kubernetes/common/tests dask_kubernetes/helm/tests
+      - name: Debug k8s resources
+        if: success() || failure()
+        run: kubectl get all -A

.github/workflows/kubecluster.yaml

@@ -22,6 +22,9 @@ jobs:
       matrix:
         python-version: ["3.8", "3.9", "3.10"]
 
+    env:
+      KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
+
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
@@ -30,6 +33,7 @@ jobs:
       - name: Install deps
         run: ./ci/install-deps.sh
       - name: Run tests
-        env:
-          KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
         run: pytest dask_kubernetes/common/tests dask_kubernetes/classic/tests
+      - name: Debug k8s resources
+        if: success() || failure()
+        run: kubectl get all -A

.github/workflows/operator.yaml

@@ -24,6 +24,9 @@ jobs:
       matrix:
         python-version: ["3.8", "3.9", "3.10"]
 
+    env:
+      KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
+
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
@@ -33,6 +36,8 @@ jobs:
         run: ./ci/install-deps.sh
       - name: Run tests
         env:
-          KUBECONFIG: .pytest-kind/pytest-kind/kubeconfig
           TEST_ISTIO: "true"
         run: pytest dask_kubernetes/common/tests dask_kubernetes/operator/tests dask_kubernetes/experimental/tests
+      - name: Debug k8s resources
+        if: success() || failure()
+        run: kubectl get all -A

ci/pre-commit-crd.py

@@ -83,6 +83,7 @@ def main(version, *args):
             run_action("daskcluster.yaml", temp_dir, crd_path, output_paths)
             run_action("daskworkergroup.yaml", temp_dir, crd_path, output_paths)
             run_action("daskjob.yaml", temp_dir, crd_path, output_paths)
+            run_action("daskautoscaler.yaml", temp_dir, crd_path, output_paths)
 
         else:
             run_action(changed_file, temp_dir, crd_path, output_paths)

dask_kubernetes/conftest.py

@@ -67,7 +67,7 @@ def customresources(k8s_cluster):
     temp_dir = tempfile.TemporaryDirectory()
     crd_path = os.path.join(DIR, "operator", "customresources")
 
-    for crd in ["daskcluster", "daskworkergroup", "daskjob"]:
+    for crd in ["daskcluster", "daskworkergroup", "daskjob", "daskautoscaler"]:
         run_generate(
             os.path.join(crd_path, f"{crd}.yaml"),
             os.path.join(crd_path, f"{crd}.patch.yaml"),

dask_kubernetes/experimental/kubecluster.py

@@ -491,6 +491,15 @@ async def _scale(self, n, worker_group="default"):
             custom_objects_api.api_client.set_default_header(
                 "content-type", "application/merge-patch+json"
             )
+            # Disable adaptivity if enabled
+            with suppress(kubernetes.client.ApiException):
+                await custom_objects_api.delete_namespaced_custom_object(
+                    group="kubernetes.dask.org",
+                    version="v1",
+                    plural="daskautoscalers",
+                    namespace=self.namespace,
+                    name=self.name,
+                )
             await custom_objects_api.patch_namespaced_custom_object_scale(
                 group="kubernetes.dask.org",
                 version="v1",
@@ -500,11 +509,59 @@ async def _scale(self, n, worker_group="default"):
                 body={"spec": {"replicas": n}},
             )
 
-    def adapt(self, *args, **kwargs):
-        """Turn on adaptivity"""
-        raise NotImplementedError(
-            "Adaptive mode is not supported yet for this KubeCluster."
-        )
+    def adapt(self, minimum=None, maximum=None):
+        """Turn on adaptivity
+
+        Parameters
+        ----------
+        minimum : int
+            Minimum number of workers
+        minimum : int
+            Maximum number of workers
+
+        Examples
+        --------
+        >>> cluster.adapt()  # Allow scheduler to add/remove workers within k8s cluster resource limits
+        >>> cluster.adapt(minimum=1, maximum=10) # Allow scheduler to add/remove workers within 1-10 range
+        """
+        return self.sync(self._adapt, minimum, maximum)
+
+    async def _adapt(self, minimum=None, maximum=None):
+        async with kubernetes.client.api_client.ApiClient() as api_client:
+            custom_objects_api = kubernetes.client.CustomObjectsApi(api_client)
+            custom_objects_api.api_client.set_default_header(
+                "content-type", "application/merge-patch+json"
+            )
+            try:
+                await custom_objects_api.patch_namespaced_custom_object_scale(
+                    group="kubernetes.dask.org",
+                    version="v1",
+                    plural="daskautoscalers",
+                    namespace=self.namespace,
+                    name=self.name,
+                    body={"spec": {"minimum": minimum, "maximum": maximum}},
+                )
+            except kubernetes.client.ApiException:
+                await custom_objects_api.create_namespaced_custom_object(
+                    group="kubernetes.dask.org",
+                    version="v1",
+                    plural="daskautoscalers",
+                    namespace=self.namespace,
+                    body={
+                        "apiVersion": "kubernetes.dask.org/v1",
+                        "kind": "DaskAutoscaler",
+                        "metadata": {
+                            "name": self.name,
+                            "dask.org/cluster-name": self.cluster_name,
+                            "dask.org/component": "autoscaler",
+                        },
+                        "spec": {
+                            "cluster": self.cluster_name,
+                            "minimum": minimum,
+                            "maximum": maximum,
+                        },
+                    },
+                )
 
     def _build_scheduler_spec(self, cluster_name):
         # TODO: Take the values provided in the current class constructor

dask_kubernetes/experimental/tests/test_kubecluster.py

@@ -84,3 +84,19 @@ def test_additional_worker_groups(kopf_runner, docker_image):
 def test_cluster_without_operator(docker_image):
     with pytest.raises(TimeoutError, match="is the Dask Operator running"):
         KubeCluster(name="noop", n_workers=1, image=docker_image, resource_timeout=1)
+
+
+def test_adapt(kopf_runner, docker_image):
+    with kopf_runner:
+        with KubeCluster(
+            name="adaptive",
+            image=docker_image,
+            n_workers=0,
+        ) as cluster:
+            cluster.adapt(minimum=0, maximum=1)
+            with Client(cluster) as client:
+                assert client.submit(lambda x: x + 1, 10).result() == 11
+
+            # Need to clean up the DaskAutoscaler object
+            # See https://github.com/dask/dask-kubernetes/issues/546
+            cluster.scale(0)

dask_kubernetes/operator/customresources/daskautoscaler.patch.yaml

@@ -0,0 +1 @@
+[]

dask_kubernetes/operator/customresources/daskautoscaler.yaml

@@ -0,0 +1,37 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: daskautoscalers.kubernetes.dask.org
+spec:
+  scope: Namespaced
+  group: kubernetes.dask.org
+  names:
+    kind: DaskAutoscaler
+    plural: daskautoscalers
+    singular: daskautoscaler
+  versions:
+    - name: v1
+      served: true
+      storage: true
+      schema:
+        openAPIV3Schema:
+          type: object
+          properties:
+            spec:
+              type: object
+              required:
+              - cluster
+              - minimum
+              - maximum
+              properties:
+                cluster:
+                  type: string
+                  description: Name of the cluster to associate this worker group with
+                minimum:
+                  type: integer
+                  description: Minimum number of workers
+                maximum:
+                  type: integer
+                  description: Maximum number of workers
+            status:
+              type: object

dask_kubernetes/operator/deployment/helm/dask-kubernetes-operator/crds/daskautoscaler.yaml

@@ -0,0 +1,37 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: daskautoscalers.kubernetes.dask.org
+spec:
+  group: kubernetes.dask.org
+  names:
+    kind: DaskAutoscaler
+    plural: daskautoscalers
+    singular: daskautoscaler
+  scope: Namespaced
+  versions:
+  - name: v1
+    schema:
+      openAPIV3Schema:
+        properties:
+          spec:
+            properties:
+              cluster:
+                description: Name of the cluster to associate this worker group with
+                type: string
+              maximum:
+                description: Maximum number of workers
+                type: integer
+              minimum:
+                description: Minimum number of workers
+                type: integer
+            required:
+            - cluster
+            - minimum
+            - maximum
+            type: object
+          status:
+            type: object
+        type: object
+    served: true
+    storage: true

dask_kubernetes/operator/deployment/manifests/daskautoscaler.yaml

@@ -0,0 +1,37 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: daskautoscalers.kubernetes.dask.org
+spec:
+  group: kubernetes.dask.org
+  names:
+    kind: DaskAutoscaler
+    plural: daskautoscalers
+    singular: daskautoscaler
+  scope: Namespaced
+  versions:
+  - name: v1
+    schema:
+      openAPIV3Schema:
+        properties:
+          spec:
+            properties:
+              cluster:
+                description: Name of the cluster to associate this worker group with
+                type: string
+              maximum:
+                description: Maximum number of workers
+                type: integer
+              minimum:
+                description: Minimum number of workers
+                type: integer
+            required:
+            - cluster
+            - minimum
+            - maximum
+            type: object
+          status:
+            type: object
+        type: object
+    served: true
+    storage: true