Add utility methods to kr8s objects (#741)

Author: jacobtomlinson

dask_kubernetes/operator/controller/tests/test_controller.py

@@ -9,10 +9,12 @@
 import yaml
 from dask.distributed import Client
 
+from kr8s.asyncio.objects import Pod, Deployment, Service
 from dask_kubernetes.operator.controller import (
     KUBERNETES_DATETIME_FORMAT,
     get_job_runner_pod_name,
 )
+from dask_kubernetes.operator.objects import DaskCluster, DaskWorkerGroup, DaskJob
 
 DIR = pathlib.Path(__file__).parent.absolute()
 
@@ -590,3 +592,70 @@ async def test_failed_job(k8s_cluster, kopf_runner, gen_job):
 
     assert "A DaskJob has been created" in runner.stdout
     assert "Job failed, deleting Dask cluster." in runner.stdout
+
+
+@pytest.mark.asyncio
+async def test_object_dask_cluster(k8s_cluster, kopf_runner, gen_cluster):
+    with kopf_runner as runner:
+        async with gen_cluster() as (cluster_name, ns):
+            cluster = await DaskCluster.get(cluster_name, namespace=ns)
+
+            worker_groups = []
+            while not worker_groups:
+                worker_groups = await cluster.worker_groups()
+                await asyncio.sleep(0.1)
+            assert len(worker_groups) == 1  # Just the default worker group
+            wg = worker_groups[0]
+            assert isinstance(wg, DaskWorkerGroup)
+
+            scheduler_pod = await cluster.scheduler_pod()
+            assert isinstance(scheduler_pod, Pod)
+
+            scheduler_deployment = await cluster.scheduler_deployment()
+            assert isinstance(scheduler_deployment, Deployment)
+
+            scheduler_service = await cluster.scheduler_service()
+            assert isinstance(scheduler_service, Service)
+
+
+@pytest.mark.asyncio
+async def test_object_dask_worker_group(k8s_cluster, kopf_runner, gen_cluster):
+    with kopf_runner as runner:
+        async with gen_cluster() as (cluster_name, ns):
+            cluster = await DaskCluster.get(cluster_name, namespace=ns)
+
+            worker_groups = []
+            while not worker_groups:
+                worker_groups = await cluster.worker_groups()
+                await asyncio.sleep(0.1)
+            assert len(worker_groups) == 1  # Just the default worker group
+            wg = worker_groups[0]
+            assert isinstance(wg, DaskWorkerGroup)
+
+            pods = []
+            while not pods:
+                pods = await wg.pods()
+                await asyncio.sleep(0.1)
+            assert all([isinstance(p, Pod) for p in pods])
+
+            deployments = []
+            while not deployments:
+                deployments = await wg.deployments()
+                await asyncio.sleep(0.1)
+            assert all([isinstance(d, Deployment) for d in deployments])
+
+            assert (await wg.cluster()).name == cluster.name
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="Flaky in CI")
+async def test_object_dask_job(k8s_cluster, kopf_runner, gen_job):
+    with kopf_runner as runner:
+        async with gen_job("simplejob.yaml") as (job_name, ns):
+            job = await DaskJob.get(job_name, namespace=ns)
+
+            job_pod = await job.pod()
+            assert isinstance(job_pod, Pod)
+
+            cluster = await job.cluster()
+            assert isinstance(cluster, DaskCluster)

dask_kubernetes/operator/objects.py

@@ -1,4 +1,7 @@
-from kr8s.asyncio.objects import APIObject
+from __future__ import annotations
+from typing import List
+
+from kr8s.asyncio.objects import APIObject, Pod, Deployment, Service
 
 
 class DaskCluster(APIObject):
@@ -11,6 +14,61 @@ class DaskCluster(APIObject):
     scalable = True
     scalable_spec = "worker.replicas"
 
+    async def worker_groups(self) -> List[DaskWorkerGroup]:
+        return await self.api.get(
+            DaskWorkerGroup.endpoint,
+            label_selector=f"dask.org/cluster-name={self.name}",
+            namespace=self.namespace,
+        )
+
+    async def scheduler_pod(self) -> Pod:
+        pods = []
+        while not pods:
+            pods = await self.api.get(
+                Pod.endpoint,
+                label_selector=",".join(
+                    [
+                        f"dask.org/cluster-name={self.name}",
+                        "dask.org/component=scheduler",
+                    ]
+                ),
+                namespace=self.namespace,
+            )
+        assert len(pods) == 1
+        return pods[0]
+
+    async def scheduler_deployment(self) -> Deployment:
+        deployments = []
+        while not deployments:
+            deployments = await self.api.get(
+                Deployment.endpoint,
+                label_selector=",".join(
+                    [
+                        f"dask.org/cluster-name={self.name}",
+                        "dask.org/component=scheduler",
+                    ]
+                ),
+                namespace=self.namespace,
+            )
+        assert len(deployments) == 1
+        return deployments[0]
+
+    async def scheduler_service(self) -> Service:
+        services = []
+        while not services:
+            services = await self.api.get(
+                Service.endpoint,
+                label_selector=",".join(
+                    [
+                        f"dask.org/cluster-name={self.name}",
+                        "dask.org/component=scheduler",
+                    ]
+                ),
+                namespace=self.namespace,
+            )
+        assert len(services) == 1
+        return services[0]
+
 
 class DaskWorkerGroup(APIObject):
     version = "kubernetes.dask.org/v1"
@@ -21,6 +79,35 @@ class DaskWorkerGroup(APIObject):
     namespaced = True
     scalable = True
 
+    async def pods(self) -> List[Pod]:
+        return await self.api.get(
+            Pod.endpoint,
+            label_selector=",".join(
+                [
+                    f"dask.org/cluster-name={self.spec['cluster']}",
+                    "dask.org/component=worker",
+                    f"dask.org/workergroup-name={self.name}",
+                ]
+            ),
+            namespace=self.namespace,
+        )
+
+    async def deployments(self) -> List[Deployment]:
+        return await self.api.get(
+            Deployment.endpoint,
+            label_selector=",".join(
+                [
+                    f"dask.org/cluster-name={self.spec['cluster']}",
+                    "dask.org/component=worker",
+                    f"dask.org/workergroup-name={self.name}",
+                ]
+            ),
+            namespace=self.namespace,
+        )
+
+    async def cluster(self) -> DaskCluster:
+        return await DaskCluster.get(self.spec["cluster"], namespace=self.namespace)
+
 
 class DaskAutoscaler(APIObject):
     version = "kubernetes.dask.org/v1"
@@ -30,6 +117,9 @@ class DaskAutoscaler(APIObject):
     singular = "daskautoscaler"
     namespaced = True
 
+    async def cluster(self) -> DaskCluster:
+        return await DaskCluster.get(self.spec["cluster"], namespace=self.namespace)
+
 
 class DaskJob(APIObject):
     version = "kubernetes.dask.org/v1"
@@ -38,3 +128,22 @@ class DaskJob(APIObject):
     plural = "daskjobs"
     singular = "daskjob"
     namespaced = True
+
+    async def cluster(self) -> DaskCluster:
+        return await DaskCluster.get(self.name, namespace=self.namespace)
+
+    async def pod(self) -> Pod:
+        pods = []
+        while not pods:
+            pods = await self.api.get(
+                Pod.endpoint,
+                label_selector=",".join(
+                    [
+                        f"dask.org/cluster-name={self.name}",
+                        "dask.org/component=job-runner",
+                    ]
+                ),
+                namespace=self.namespace,
+            )
+        assert len(pods) == 1
+        return pods[0]