From 22fc170427156d2a1fa6aeb890f3ecd9bfe9b284 Mon Sep 17 00:00:00 2001 From: viniciusdc Date: Thu, 16 May 2024 17:01:10 -0300 Subject: [PATCH 1/3] expose volume mounts config to schema & present dask-workers multi-image array --- .../stages/kubernetes_services/__init__.py | 8 ++++ .../dask-gateway/files/gateway_config.py | 38 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/src/_nebari/stages/kubernetes_services/__init__.py b/src/_nebari/stages/kubernetes_services/__init__.py index cdc1ae9151..3bd677c782 100644 --- a/src/_nebari/stages/kubernetes_services/__init__.py +++ b/src/_nebari/stages/kubernetes_services/__init__.py @@ -166,6 +166,14 @@ def check_default(cls, value): ) return value +class DaskWorkerMounts(schema.Base): + # Both are required when one is present + volumes: List[Dict[str, str]] = [] + volume_mounts: List[Dict[str, str]] = [] + +class DaskWorker(schema.Base): + worker_images: Optional[Dict[str, str]] = None + extra_mounts: Optional[DaskWorkerMounts] = None class CondaEnvironment(schema.Base): name: str diff --git a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py index 2219d14e56..7bee6a5f4b 100644 --- a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py +++ b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py @@ -226,6 +226,29 @@ def base_username_mount(username, uid=1000, gid=100): }, } +def extra_worker_mount(): + extra_mounts = {} + if config["extra-worker-mounts"]: + volumes = config["extra-worker-mounts"]["volumes"] + volume_mounts = config["extra-worker-mounts"]["volume_mounts"] + if volumes: + extra_mounts.update( + { + "scheduler_extra_pod_config": { + "volumes": volumes, + }, + "worker_extra_pod_config": {"volumes": volumes}, + } + ) + if volume_mounts: + extra_mounts.update( + { + "scheduler_extra_container_config": {"volumeMounts": volume_mounts}, + "worker_extra_container_config": {"volumeMounts": volume_mounts}, + } + ) + return extra_mounts + def worker_profile(options, user): namespace, name = options.conda_environment.split("/") @@ -237,6 +260,12 @@ def worker_profile(options, user): base_username_mount(user.name), config["profiles"][options.profile], {"environment": {**options.environment_vars}}, + { + "image": config["worker-images"][options.image] + if config["worker-images"] + else f"{config['cluster-image']['name']}:{config['cluster-image']['tag']}" + }, + extra_worker_mount() ], {}, ) @@ -272,6 +301,15 @@ def user_options(user): label="Cluster Profile", ) ] + if config["worker-images"]: + args += [ + Select( + "image", + list(config["worker-images"].keys()), + default=list(config["worker-images"].keys())[0], + label="Cluster Docker Image", + ) + ] args += [ Mapping("environment_vars", {}, label="Environment Variables"), From 837bb4d81761b0b72a3841c2ac401b82a79a81e9 Mon Sep 17 00:00:00 2001 From: viniciusdc Date: Thu, 16 May 2024 17:12:31 -0300 Subject: [PATCH 2/3] expose input variables --- .../stages/kubernetes_services/__init__.py | 4 ++++ .../template/dask_gateway.tf | 18 ++++++++++++++++++ .../services/dask-gateway/variables.tf | 15 +++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/src/_nebari/stages/kubernetes_services/__init__.py b/src/_nebari/stages/kubernetes_services/__init__.py index 3bd677c782..11498c2fd4 100644 --- a/src/_nebari/stages/kubernetes_services/__init__.py +++ b/src/_nebari/stages/kubernetes_services/__init__.py @@ -391,6 +391,8 @@ class DaskGatewayInputVars(schema.Base): dask_worker_image: ImageNameTag = Field(alias="dask-worker-image") dask_gateway_profiles: Dict[str, Any] = Field(alias="dask-gateway-profiles") cloud_provider: str = Field(alias="cloud-provider") + extra_worker_mounts: Optional[DaskWorkerMounts] = Field(alias="extra-worker-mounts") + worker_images: Optional[Dict[str, str]] = Field(alias="worker-images") class MonitoringInputVars(schema.Base): @@ -540,6 +542,8 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]): ), dask_gateway_profiles=self.config.profiles.model_dump()["dask_worker"], cloud_provider=cloud_provider, + extra_worker_mounts=self.config.get("dask_worker", {}).get("extra_mounts"), + worker_images=self.config.get("dask_worker", {}).get("worker_images"), ) monitoring_vars = MonitoringInputVars( diff --git a/src/_nebari/stages/kubernetes_services/template/dask_gateway.tf b/src/_nebari/stages/kubernetes_services/template/dask_gateway.tf index b9b0a9c6c3..30de93ac38 100644 --- a/src/_nebari/stages/kubernetes_services/template/dask_gateway.tf +++ b/src/_nebari/stages/kubernetes_services/template/dask_gateway.tf @@ -11,6 +11,21 @@ variable "dask-gateway-profiles" { description = "Dask Gateway profiles to expose to user" } +variable "worker-images" { + description = "Worker images list options to use for dask-gateway" + type = map(string) + default = {} +} + +variable "extra-worker-mounts" { + description = "Extra volumes and volume mounts to add to worker pods" + type = object({ + volumes = list(any) + volume_mounts = list(any) + }) + default = {} +} + # =================== RESOURCES ===================== module "dask-gateway" { source = "./modules/kubernetes/services/dask-gateway" @@ -29,6 +44,9 @@ module "dask-gateway" { # needs to match name in module.jupyterhub.extra-mounts dask-etc-configmap-name = "dask-etc" + worker-images = var.worker-images + extra-worker-mounts = var.extra-worker-mounts + # environments conda-store-pvc = module.conda-store-nfs-mount.persistent_volume_claim.name conda-store-mount = "/home/conda" diff --git a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/variables.tf b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/variables.tf index 7f8a4aa978..48c93c1870 100644 --- a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/variables.tf +++ b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/variables.tf @@ -204,3 +204,18 @@ variable "cloud-provider" { description = "Name of the cloud provider to deploy to." type = string } + +variable "worker-images" { + description = "Worker images list options to use for dask-gateway" + type = map(string) + default = {} +} + +variable "extra-worker-mounts" { + description = "Extra volumes and volume mounts to add to worker pods" + type = object({ + volumes = list(any) + volume_mounts = list(any) + }) + default = {} +} From f35433bce475b1b2a2490efc8baa3cdcc37eb6b8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 May 2024 20:15:43 +0000 Subject: [PATCH 3/3] [pre-commit.ci] Apply automatic pre-commit fixes --- src/_nebari/stages/kubernetes_services/__init__.py | 3 +++ .../services/dask-gateway/files/gateway_config.py | 11 +++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/_nebari/stages/kubernetes_services/__init__.py b/src/_nebari/stages/kubernetes_services/__init__.py index 11498c2fd4..b841c77614 100644 --- a/src/_nebari/stages/kubernetes_services/__init__.py +++ b/src/_nebari/stages/kubernetes_services/__init__.py @@ -166,15 +166,18 @@ def check_default(cls, value): ) return value + class DaskWorkerMounts(schema.Base): # Both are required when one is present volumes: List[Dict[str, str]] = [] volume_mounts: List[Dict[str, str]] = [] + class DaskWorker(schema.Base): worker_images: Optional[Dict[str, str]] = None extra_mounts: Optional[DaskWorkerMounts] = None + class CondaEnvironment(schema.Base): name: str channels: Optional[List[str]] = None diff --git a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py index 7bee6a5f4b..d02b2cc11d 100644 --- a/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py +++ b/src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py @@ -226,6 +226,7 @@ def base_username_mount(username, uid=1000, gid=100): }, } + def extra_worker_mount(): extra_mounts = {} if config["extra-worker-mounts"]: @@ -261,11 +262,13 @@ def worker_profile(options, user): config["profiles"][options.profile], {"environment": {**options.environment_vars}}, { - "image": config["worker-images"][options.image] - if config["worker-images"] - else f"{config['cluster-image']['name']}:{config['cluster-image']['tag']}" + "image": ( + config["worker-images"][options.image] + if config["worker-images"] + else f"{config['cluster-image']['name']}:{config['cluster-image']['tag']}" + ) }, - extra_worker_mount() + extra_worker_mount(), ], {}, )