Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Migrate KubeCluster scale and adapt to kr8s #787

Merged
merged 3 commits into from
Jul 31, 2023
Merged

Migrate KubeCluster scale and adapt to kr8s #787

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6f47793
Migrate KubeCluster scale and adapt to kr8s
jacobtomlinson Jul 28, 2023
f8d5a9d
Tweak names so aiopykube and kr8s can coexist for now
jacobtomlinson Jul 28, 2023
f2112d9
Import typo
jacobtomlinson Jul 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 40 additions & 62 deletions dask_kubernetes/operator/kubecluster/kubecluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import pykube.exceptions
import kubernetes_asyncio as kubernetes
import yaml
import kr8s

import dask.config
from distributed.core import Status, rpc
Expand All @@ -43,9 +44,13 @@
)
from dask_kubernetes.common.utils import get_current_namespace
from dask_kubernetes.aiopykube import HTTPClient, KubeConfig
from dask_kubernetes.aiopykube.dask import DaskCluster, DaskWorkerGroup
from dask_kubernetes.aiopykube.dask import (
DaskCluster,
DaskWorkerGroup as AIODaskWorkerGroup,
)
from dask_kubernetes.aiopykube.objects import Pod, Service
from dask_kubernetes.exceptions import CrashLoopBackOffError, SchedulerStartupError
from dask_kubernetes.operator._objects import DaskWorkerGroup, DaskAutoscaler

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -541,7 +546,7 @@ async def _watch_component_status(self):

# Get DaskWorkerGroup status
with suppress(pykube.exceptions.ObjectDoesNotExist):
await DaskWorkerGroup.objects(
await AIODaskWorkerGroup.objects(
self.k8s_api, namespace=self.namespace
).get_by_name(self.name + "-default")
self._startup_component_status["workergroup"] = "Created"
Expand Down Expand Up @@ -799,31 +804,18 @@ def scale(self, n, worker_group="default"):
return self.sync(self._scale, n, worker_group)

async def _scale(self, n, worker_group="default"):
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"
)
# 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",
plural="daskworkergroups",
namespace=self.namespace,
name=f"{self.name}-{worker_group}",
body={"spec": {"replicas": n}},
)
for instance in self._instances:
if instance.name == self.name:
instance.scheduler_info = self.scheduler_info
# Disable adaptivity if enabled
with suppress(kr8s.NotFoundError):
autoscaler = await DaskAutoscaler(self.name, self.namespace)
await autoscaler.delete()

wg = await DaskWorkerGroup(
f"{self.name}-{worker_group}", namespace=self.namespace
)
await wg.scale(n)
for instance in self._instances:
if instance.name == self.name:
instance.scheduler_info = self.scheduler_info

def adapt(self, minimum=None, maximum=None):
"""Turn on adaptivity
Expand All @@ -843,41 +835,27 @@ def adapt(self, minimum=None, maximum=None):
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.name,
"dask.org/component": "autoscaler",
},
"spec": {
"cluster": self.name,
"minimum": minimum,
"maximum": maximum,
},
},
)
autoscaler = await DaskAutoscaler(
{
"apiVersion": "kubernetes.dask.org/v1",
"kind": "DaskAutoscaler",
"metadata": {
"name": self.name,
"dask.org/cluster-name": self.name,
"dask.org/component": "autoscaler",
},
"spec": {
"cluster": self.name,
"minimum": minimum,
"maximum": maximum,
},
},
self.namespace,
)
try:
await autoscaler.patch({"spec": {"minimum": minimum, "maximum": maximum}})
except kr8s.NotFoundError:
await autoscaler.create()

def __enter__(self):
return self
Expand Down