Create scheduler pod through a Deployment (#711)

Author: Matt711

dask_kubernetes/common/networking.py

@@ -196,8 +196,14 @@ async def get_scheduler_address(
 async def wait_for_scheduler(api, cluster_name, namespace, timeout=None):
     pod_start_time = None
     while True:
+        async with kubernetes.client.api_client.ApiClient() as api_client:
+            k8s_api = kubernetes.client.CoreV1Api(api_client)
+            pods = await k8s_api.list_namespaced_pod(
+                namespace=namespace,
+                label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={cluster_name}",
+            )
         pod = await Pod.objects(api, namespace=namespace).get_by_name(
-            cluster_name + "-scheduler"
+            pods.items[0].metadata.name
         )
         phase = pod.obj["status"]["phase"]
         if phase == "Running":

dask_kubernetes/operator/controller/controller.py

@@ -60,22 +60,34 @@ def _get_labels(meta):
     }
 
 
-def build_scheduler_pod_spec(cluster_name, spec, annotations, labels):
+def build_scheduler_deployment_spec(
+    cluster_name, namespace, pod_spec, annotations, labels
+):
     labels.update(
         **{
             "dask.org/cluster-name": cluster_name,
             "dask.org/component": "scheduler",
             "sidecar.istio.io/inject": "false",
         }
     )
+    metadata = {
+        "name": f"{cluster_name}-scheduler",
+        "labels": labels,
+        "annotations": annotations,
+    }
+    spec = {}
+    spec["replicas"] = 1
+    spec["selector"] = {
+        "matchLabels": labels,
+    }
+    spec["template"] = {
+        "metadata": metadata,
+        "spec": pod_spec,
+    }
     return {
-        "apiVersion": "v1",
-        "kind": "Pod",
-        "metadata": {
-            "name": f"{cluster_name}-scheduler",
-            "labels": labels,
-            "annotations": annotations,
-        },
+        "apiVersion": "apps/v1",
+        "kind": "Deployment",
+        "metadata": metadata,
         "spec": spec,
     }
 
@@ -270,21 +282,21 @@ async def daskcluster_create_components(
                 annotations.update(**scheduler_spec["metadata"]["annotations"])
             if "labels" in scheduler_spec["metadata"]:
                 labels.update(**scheduler_spec["metadata"]["labels"])
-        data = build_scheduler_pod_spec(
-            name, scheduler_spec.get("spec"), annotations, labels
+        data = build_scheduler_deployment_spec(
+            name, namespace, scheduler_spec.get("spec"), annotations, labels
         )
         kopf.adopt(data)
         pod = await api.list_namespaced_pod(
             namespace=namespace,
             label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={name}",
         )
         if not pod.items:
-            await api.create_namespaced_pod(
+            await kubernetes.client.AppsV1Api(api_client).create_namespaced_deployment(
                 namespace=namespace,
                 body=data,
             )
             logger.info(
-                f"Scheduler pod {data['metadata']['name']} created in {namespace}."
+                f"Scheduler deployment {data['metadata']['name']} created in {namespace}."
             )
 
         data = build_scheduler_service_spec(
@@ -295,7 +307,7 @@ async def daskcluster_create_components(
             namespace=namespace,
             label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={name}",
         )
-        if not pod.items:
+        if not service.items:
             await api.create_namespaced_service(
                 namespace=namespace,
                 body=data,
@@ -862,8 +874,12 @@ async def daskautoscaler_adapt(spec, name, namespace, logger, **kwargs):
 
         pod_ready = False
         try:
+            pods = await coreapi.list_namespaced_pod(
+                namespace=namespace,
+                label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={spec['cluster']}",
+            )
             scheduler_pod = await coreapi.read_namespaced_pod(
-                f"{spec['cluster']}-scheduler", namespace
+                pods.items[0].metadata.name, namespace
             )
             if scheduler_pod.status.phase == "Running":
                 pod_ready = True

dask_kubernetes/operator/controller/tests/test_controller.py

@@ -104,27 +104,19 @@ def test_operator_plugins(kopf_runner):
 async def test_simplecluster(k8s_cluster, kopf_runner, gen_cluster):
     with kopf_runner as runner:
         async with gen_cluster() as (cluster_name, ns):
-            scheduler_pod_name = "simple-scheduler"
+            scheduler_deployment_name = "simple-scheduler"
             worker_pod_name = "simple-default-worker"
             service_name = "simple-scheduler"
 
-            while scheduler_pod_name not in k8s_cluster.kubectl(
-                "get", "pods", "-n", ns
+            while scheduler_deployment_name not in k8s_cluster.kubectl(
+                "get", "deployments", "-n", ns
             ):
                 await asyncio.sleep(0.1)
             while service_name not in k8s_cluster.kubectl("get", "svc", "-n", ns):
                 await asyncio.sleep(0.1)
             while worker_pod_name not in k8s_cluster.kubectl("get", "pods", "-n", ns):
                 await asyncio.sleep(0.1)
-            k8s_cluster.kubectl(
-                "wait",
-                "-n",
-                ns,
-                "pods",
-                "--for=condition=Ready",
-                scheduler_pod_name,
-                "--timeout=120s",
-            )
+
             with k8s_cluster.port_forward(
                 f"service/{service_name}", 8786, "-n", ns
             ) as port:
@@ -256,26 +248,18 @@ async def test_simplecluster(k8s_cluster, kopf_runner, gen_cluster):
 async def test_scalesimplecluster(k8s_cluster, kopf_runner, gen_cluster):
     with kopf_runner as runner:
         async with gen_cluster() as (cluster_name, ns):
-            scheduler_pod_name = "simple-scheduler"
+            scheduler_deployment_name = "simple-scheduler"
             worker_pod_name = "simple-default-worker"
             service_name = "simple-scheduler"
-            while scheduler_pod_name not in k8s_cluster.kubectl(
-                "get", "pods", "-n", ns
+            while scheduler_deployment_name not in k8s_cluster.kubectl(
+                "get", "deployments", "-n", ns
             ):
                 await asyncio.sleep(0.1)
             while service_name not in k8s_cluster.kubectl("get", "svc", "-n", ns):
                 await asyncio.sleep(0.1)
             while worker_pod_name not in k8s_cluster.kubectl("get", "pods", "-n", ns):
                 await asyncio.sleep(0.1)
-            k8s_cluster.kubectl(
-                "wait",
-                "pods",
-                "-n",
-                ns,
-                "--for=condition=Ready",
-                scheduler_pod_name,
-                "--timeout=120s",
-            )
+
             with k8s_cluster.port_forward(
                 f"service/{service_name}", 8786, "-n", ns
             ) as port:
@@ -311,26 +295,18 @@ async def test_scalesimplecluster_from_cluster_spec(
 ):
     with kopf_runner as runner:
         async with gen_cluster() as (cluster_name, ns):
-            scheduler_pod_name = "simple-scheduler"
+            scheduler_deployment_name = "simple-scheduler"
             worker_pod_name = "simple-default-worker"
             service_name = "simple-scheduler"
-            while scheduler_pod_name not in k8s_cluster.kubectl(
+            while scheduler_deployment_name not in k8s_cluster.kubectl(
                 "get", "pods", "-n", ns
             ):
                 await asyncio.sleep(0.1)
             while service_name not in k8s_cluster.kubectl("get", "svc", "-n", ns):
                 await asyncio.sleep(0.1)
             while worker_pod_name not in k8s_cluster.kubectl("get", "pods", "-n", ns):
                 await asyncio.sleep(0.1)
-            k8s_cluster.kubectl(
-                "wait",
-                "pods",
-                "-n",
-                ns,
-                "--for=condition=Ready",
-                scheduler_pod_name,
-                "--timeout=120s",
-            )
+
             with k8s_cluster.port_forward(
                 f"service/{service_name}", 8786, "-n", ns
             ) as port:
@@ -360,6 +336,44 @@ async def test_scalesimplecluster_from_cluster_spec(
                     await client.wait_for_workers(3)
 
 
+@pytest.mark.asyncio
+async def test_recreate_scheduler_pod(k8s_cluster, kopf_runner, gen_cluster):
+    with kopf_runner as runner:
+        async with gen_cluster() as (cluster_name, ns):
+            scheduler_deployment_name = "simple-scheduler"
+            worker_pod_name = "simple-default-worker"
+            service_name = "simple-scheduler"
+            while scheduler_deployment_name not in k8s_cluster.kubectl(
+                "get", "pods", "-n", ns
+            ):
+                await asyncio.sleep(0.1)
+            while service_name not in k8s_cluster.kubectl("get", "svc", "-n", ns):
+                await asyncio.sleep(0.1)
+            while worker_pod_name not in k8s_cluster.kubectl("get", "pods", "-n", ns):
+                await asyncio.sleep(0.1)
+            k8s_cluster.kubectl(
+                "delete",
+                "pods",
+                "-l",
+                "dask.org/cluster-name=simple,dask.org/component=scheduler",
+                "-n",
+                ns,
+            )
+            k8s_cluster.kubectl(
+                "wait",
+                "--for=condition=Ready",
+                "-l",
+                "dask.org/cluster-name=simple,dask.org/component=scheduler",
+                "pod",
+                "-n",
+                ns,
+                "--timeout=60s",
+            )
+            assert scheduler_deployment_name in k8s_cluster.kubectl(
+                "get", "pods", "-n", ns
+            )
+
+
 def _get_job_status(k8s_cluster, ns):
     return json.loads(
         k8s_cluster.kubectl(

dask_kubernetes/operator/kubecluster/kubecluster.py

@@ -374,11 +374,15 @@ async def _create_cluster(self):
                 )
             except CrashLoopBackOffError as e:
                 logs = await self._get_logs()
+                pods = await core_api.list_namespaced_pod(
+                    namespace=self.namespace,
+                    label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={self.name}",
+                )
                 await self._close()
                 raise SchedulerStartupError(
                     "Scheduler failed to start.",
                     "Scheduler Pod logs:",
-                    logs[self.name + "-scheduler"],
+                    logs[pods.items[0].metadata.name],
                 ) from e
             self._log("Waiting for scheduler service")
             await wait_for_service(core_api, f"{self.name}-scheduler", self.namespace)
@@ -493,22 +497,30 @@ async def _watch_component_status(self):
 
             # Get Scheduler Pod status
             with suppress(pykube.exceptions.ObjectDoesNotExist):
-                pod = await Pod.objects(
-                    self.k8s_api, namespace=self.namespace
-                ).get_by_name(self.name + "-scheduler")
-                phase = pod.obj["status"]["phase"]
-                if phase == "Running":
-                    conditions = {
-                        c["type"]: c["status"] for c in pod.obj["status"]["conditions"]
-                    }
-                    if "Ready" not in conditions or conditions["Ready"] != "True":
-                        phase = "Health Checking"
-                    if "containerStatuses" in pod.obj["status"]:
-                        for container in pod.obj["status"]["containerStatuses"]:
-                            if "waiting" in container["state"]:
-                                phase = container["state"]["waiting"]["reason"]
-
-                self._startup_component_status["schedulerpod"] = phase
+                async with kubernetes.client.api_client.ApiClient() as api_client:
+                    core_api = kubernetes.client.CoreV1Api(api_client)
+                    pods = await core_api.list_namespaced_pod(
+                        namespace=self.namespace,
+                        label_selector=f"dask.org/component=scheduler,dask.org/cluster-name={self.name}",
+                    )
+                if pods.items:
+                    pod = await Pod.objects(
+                        self.k8s_api, namespace=self.namespace
+                    ).get_by_name(pods.items[0].metadata.name)
+                    phase = pod.obj["status"]["phase"]
+                    if phase == "Running":
+                        conditions = {
+                            c["type"]: c["status"]
+                            for c in pod.obj["status"]["conditions"]
+                        }
+                        if "Ready" not in conditions or conditions["Ready"] != "True":
+                            phase = "Health Checking"
+                        if "containerStatuses" in pod.obj["status"]:
+                            for container in pod.obj["status"]["containerStatuses"]:
+                                if "waiting" in container["state"]:
+                                    phase = container["state"]["waiting"]["reason"]
+
+                    self._startup_component_status["schedulerpod"] = phase
 
             # Get Scheduler Service status
             with suppress(pykube.exceptions.ObjectDoesNotExist):