Auto shutdown idle clusters (#672)

* Add autoCleanupTimeout option

* Add idle_timeout

* Update to use dask/distributed#7642

* Rename spec ketyword

* Updated with changes in distributed

* Nudge CI

* Set idle_timeout longer than resource_timeout

* Rename cluster to avoid collision

* Fix timeout default value

* Reduce logging noise and add warning when autoshudown can't reach scheduler

Author: jacobtomlinson

dask_kubernetes/operator/controller/controller.py

@@ -16,7 +16,8 @@
 from dask_kubernetes.common.networking import get_scheduler_address
 from dask_kubernetes.aiopykube import HTTPClient, KubeConfig
 from dask_kubernetes.aiopykube.dask import DaskCluster
-from distributed.core import rpc
+from distributed.core import rpc, clean_exception
+from distributed.protocol.pickle import dumps
 
 _ANNOTATION_NAMESPACES_TO_IGNORE = (
     "kopf.zalando.org",
@@ -416,7 +417,7 @@ async def retire_workers(
             )
 
     # Otherwise try gracefully retiring via the RPC
-    logger.info(
+    logger.debug(
         f"Scaling {worker_group_name} failed via the HTTP API, falling back to the Dask RPC"
     )
     # Dask version mismatches between the operator and scheduler may cause this to fail in any number of unexpected ways
@@ -435,7 +436,7 @@ async def retire_workers(
             return workers_to_close
 
     # Finally fall back to last-in-first-out scaling
-    logger.info(
+    logger.debug(
         f"Scaling {worker_group_name} failed via the Dask RPC, falling back to LIFO scaling"
     )
     async with kubernetes.client.api_client.ApiClient() as api_client:
@@ -447,6 +448,75 @@ async def retire_workers(
         return [w["metadata"]["name"] for w in workers.items[:-n_workers]]
 
 
+async def check_scheduler_idle(scheduler_service_name, namespace, logger):
+    # Try getting idle time via HTTP API
+    dashboard_address = await get_scheduler_address(
+        scheduler_service_name,
+        namespace,
+        port_name="http-dashboard",
+        allow_external=False,
+    )
+    async with aiohttp.ClientSession() as session:
+        url = f"{dashboard_address}/api/v1/check_idle"
+        async with session.get(url) as resp:
+            if resp.status <= 300:
+                idle_since = (await resp.json())["idle_since"]
+                if idle_since:
+                    logger.debug("Scheduler idle since: %s", idle_since)
+                return idle_since
+            logger.debug(
+                "Received %d response from scheduler API with body %s",
+                resp.status,
+                await resp.text(),
+            )
+
+    # Otherwise try gracefully checking via the RPC
+    logger.debug(
+        f"Checking {scheduler_service_name} idleness failed via the HTTP API, falling back to the Dask RPC"
+    )
+    # Dask version mismatches between the operator and scheduler may cause this to fail in any number of unexpected ways
+    with suppress(Exception):
+        comm_address = await get_scheduler_address(
+            scheduler_service_name,
+            namespace,
+            allow_external=False,
+        )
+        async with rpc(comm_address) as scheduler_comm:
+            idle_since = await scheduler_comm.check_idle()
+            if idle_since:
+                logger.debug("Scheduler idle since: %s", idle_since)
+            return idle_since
+
+    # Finally fall back to code injection via the Dask RPC for distributed<=2023.3.1
+    logger.debug(
+        f"Checking {scheduler_service_name} idleness failed via the Dask RPC, falling back to run_on_scheduler"
+    )
+
+    def idle_since(dask_scheduler=None):
+        if not dask_scheduler.idle_timeout:
+            dask_scheduler.idle_timeout = 300
+        dask_scheduler.check_idle()
+        return dask_scheduler.idle_since
+
+    comm_address = await get_scheduler_address(
+        scheduler_service_name,
+        namespace,
+        allow_external=False,
+    )
+    async with rpc(comm_address) as scheduler_comm:
+        response = await scheduler_comm.run_function(
+            function=dumps(idle_since),
+        )
+        if response["status"] == "error":
+            typ, exc, tb = clean_exception(**response)
+            raise exc.with_traceback(tb)
+        else:
+            idle_since = response["result"]
+            if idle_since:
+                logger.debug("Scheduler idle since: %s", idle_since)
+            return idle_since
+
+
 async def get_desired_workers(scheduler_service_name, namespace, logger):
     # Try gracefully retiring via the HTTP API
     dashboard_address = await get_scheduler_address(
@@ -901,3 +971,23 @@ async def daskautoscaler_adapt(spec, name, namespace, logger, **kwargs):
             logger.debug(
                 "Not autoscaling %s with %d workers", spec["cluster"], current_replicas
             )
+
+
+@kopf.timer("daskcluster.kubernetes.dask.org", interval=5.0)
+async def daskcluster_autoshutdown(spec, name, namespace, logger, **kwargs):
+    if spec["idleTimeout"]:
+        try:
+            idle_since = await check_scheduler_idle(
+                scheduler_service_name=f"{name}-scheduler",
+                namespace=namespace,
+                logger=logger,
+            )
+        except Exception as e:
+            logger.warn("Unable to connect to scheduler, skipping autoshutdown check.")
+            return
+        if idle_since and time.time() > idle_since + spec["idleTimeout"]:
+            api = HTTPClient(KubeConfig.from_env())
+            cluster = await DaskCluster.objects(api, namespace=namespace).get_by_name(
+                name
+            )
+            await cluster.delete()

dask_kubernetes/operator/customresources/templates.yaml

@@ -67,6 +67,10 @@ definitions:
         - scheduler
         - worker
         properties:
+          idleTimeout:
+            type: integer
+            description: Delete cluster if scheduler is idle for longer than this timeout. Set to 0 to never auto cleanup.
+            default: 0
           scheduler:
             $ref: 'python://dask_kubernetes/operator/customresources/templates.yaml#/definitions/dask.k8s.api.v1.DaskScheduler'
           worker:

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

@@ -30,6 +30,9 @@ spec:
         properties:
           spec:
             properties:
+              idleTimeout:
+                default: 0
+                type: integer
               scheduler:
                 properties:
                   metadata:

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

@@ -32,6 +32,9 @@ spec:
                 properties:
                   spec:
                     properties:
+                      idleTimeout:
+                        default: 0
+                        type: integer
                       scheduler:
                         properties:
                           metadata:

dask_kubernetes/operator/kubecluster/kubecluster.py

@@ -102,6 +102,9 @@ class KubeCluster(Cluster):
     shutdown_on_close: bool (optional)
         Whether or not to delete the cluster resource when this object is closed.
         Defaults to ``True`` when creating a cluster and ``False`` when connecting to an existing one.
+    idle_timeout: int (optional)
+        If set Kubernetes will delete the cluster automatically if the scheduler is idle for longer than
+        this timeout in seconds.
     resource_timeout: int (optional)
         Time in seconds to wait for Kubernetes resources to enter their expected state.
         Example: If the ``DaskCluster`` resource that gets created isn't moved into a known ``status.phase``
@@ -170,6 +173,7 @@ def __init__(
         port_forward_cluster_ip=None,
         create_mode=None,
         shutdown_on_close=None,
+        idle_timeout=None,
         resource_timeout=None,
         scheduler_service_type=None,
         custom_cluster_spec=None,
@@ -220,6 +224,9 @@ def __init__(
         self.scheduler_forward_port = dask.config.get(
             "kubernetes.scheduler-forward-port", override_with=scheduler_forward_port
         )
+        self.idle_timeout = dask.config.get(
+            "kubernetes.idle-timeout", override_with=idle_timeout
+        )
 
         if self._custom_cluster_spec:
             if isinstance(self._custom_cluster_spec, str):
@@ -328,6 +335,7 @@ async def _create_cluster(self):
                     n_workers=self.n_workers,
                     image=self.image,
                     scheduler_service_type=self.scheduler_service_type,
+                    idle_timeout=self.idle_timeout,
                 )
             else:
                 data = self._custom_cluster_spec
@@ -722,7 +730,7 @@ def close(self, timeout=3600):
         """Delete the dask cluster"""
         return self.sync(self._close, timeout=timeout)
 
-    async def _close(self, timeout=None):
+    async def _close(self, timeout=3600):
         await super()._close()
         if self.shutdown_on_close:
             async with kubernetes.client.api_client.ApiClient() as api_client:
@@ -881,6 +889,7 @@ def make_cluster_spec(
     env=None,
     worker_command="dask-worker",
     scheduler_service_type="ClusterIP",
+    idle_timeout=0,
 ):
     """Generate a ``DaskCluster`` kubernetes resource.
 
@@ -900,12 +909,15 @@ def make_cluster_spec(
         Environment variables to set on scheduler and workers
     worker_command: str (optional)
         Worker command to use when starting the workers
+    idle_timeout: int (optional)
+        Timeout to cleanup idle cluster
     """
     return {
         "apiVersion": "kubernetes.dask.org/v1",
         "kind": "DaskCluster",
         "metadata": {"name": name},
         "spec": {
+            "idleTimeout": idle_timeout,
             "worker": make_worker_spec(
                 env=env,
                 resources=resources,

dask_kubernetes/operator/kubecluster/tests/test_kubecluster.py

@@ -132,11 +132,11 @@ def test_cluster_without_operator(docker_image):
 def test_cluster_crashloopbackoff(kopf_runner, docker_image):
     with kopf_runner:
         with pytest.raises(SchedulerStartupError, match="Scheduler failed to start"):
-            spec = make_cluster_spec(name="foo", n_workers=1)
+            spec = make_cluster_spec(name="crashloopbackoff", n_workers=1)
             spec["spec"]["scheduler"]["spec"]["containers"][0]["args"][
                 0
             ] = "dask-schmeduler"
-            KubeCluster(custom_cluster_spec=spec, resource_timeout=1)
+            KubeCluster(custom_cluster_spec=spec, resource_timeout=1, idle_timeout=2)
 
 
 def test_adapt(kopf_runner, docker_image):