Provide dask workers volumes/volume_mounts configuration

src/_nebari/stages/kubernetes_services/__init__.py

@@ -170,6 +170,17 @@ 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
@@ -388,6 +399,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")
     forwardauth_middleware_name: str = _forwardauth_middleware_name
 
 
@@ -543,6 +556,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(

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"

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/files/gateway_config.py

@@ -227,6 +227,30 @@ 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("/")
     return functools.reduce(
@@ -237,6 +261,14 @@ 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 +304,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"),

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/dask-gateway/variables.tf

@@ -205,6 +205,20 @@ variable "cloud-provider" {
   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 = {}
+
 variable "forwardauth_middleware_name" {
   type = string
 }