Add taint to user and worker nodes

src/_nebari/stages/infrastructure/__init__.py

@@ -41,10 +41,33 @@ class ExistingInputVars(schema.Base):
     kube_context: str
 
 
-class DigitalOceanNodeGroup(schema.Base):
+class NodeGroup(schema.Base):
     instance: str
-    min_nodes: int
-    max_nodes: int
+    min_nodes: Annotated[int, Field(ge=0)] = 0
+    max_nodes: Annotated[int, Field(ge=1)] = 1
+    taints: Optional[List[schema.Taint]] = []
+
+    @field_validator("taints", mode="before")
+    def validate_taint_strings(cls, value: List[str | schema.Taint]):
+        TAINT_STR_REGEX = re.compile(r"(\w+)=(\w+):(\w+)")
+        parsed_taints = []
+        for taint in value:
+            if not isinstance(taint, (str, schema.Taint)):
+                raise ValueError(
+                    f"Unable to parse type: {type(taint)} as taint.  Must be a string or Taint object."
+                )
+
+            if isinstance(taint, schema.Taint):
+                parsed_taint = taint
+            elif isinstance(taint, str):
+                match = TAINT_STR_REGEX.match(taint)
+                if not match:
+                    raise ValueError(f"Invalid taint string: {taint}")
+                key, value, effect = match.groups()
+                parsed_taint = schema.Taint(key=key, value=value, effect=effect)
+            parsed_taints.append(parsed_taint)
+
+        return parsed_taints
 
 
 class DigitalOceanInputVars(schema.Base):
@@ -53,7 +76,7 @@ class DigitalOceanInputVars(schema.Base):
     region: str
     tags: List[str]
     kubernetes_version: str
-    node_groups: Dict[str, DigitalOceanNodeGroup]
+    node_groups: Dict[str, "DigitalOceanNodeGroup"]
     kubeconfig_filename: str = get_kubeconfig_filename()
 
 
@@ -62,10 +85,26 @@ class GCPNodeGroupInputVars(schema.Base):
     instance_type: str
     min_size: int
     max_size: int
+    node_taints: List[dict]
     labels: Dict[str, str]
     preemptible: bool
     guest_accelerators: List["GCPGuestAccelerator"]
 
+    @field_validator("node_taints", mode="before")
+    def convert_taints(cls, value: Optional[List[schema.Taint]]):
+        return [
+            dict(
+                key=taint.key,
+                value=taint.value,
+                effect={
+                    schema.TaintEffectEnum.NoSchedule: "NO_SCHEDULE",
+                    schema.TaintEffectEnum.PreferNoSchedule: "PREFER_NO_SCHEDULE",
+                    schema.TaintEffectEnum.NoExecute: "NO_EXECUTE",
+                }[taint.effect],
+            )
+            for taint in value
+        ]
+
 
 class GCPPrivateClusterConfig(schema.Base):
     enable_private_nodes: bool
@@ -211,16 +250,14 @@ class KeyValueDict(schema.Base):
     value: str
 
 
-class DigitalOceanNodeGroup(schema.Base):
+class DigitalOceanNodeGroup(NodeGroup):
     """Representation of a node group with Digital Ocean
 
     - Kubernetes limits: https://docs.digitalocean.com/products/kubernetes/details/limits/
     - Available instance types: https://slugs.do-api.dev/
     """
 
-    instance: str
     min_nodes: Annotated[int, Field(ge=1)] = 1
-    max_nodes: Annotated[int, Field(ge=1)] = 1
 
 
 DEFAULT_DO_NODE_GROUPS = {
@@ -305,19 +342,26 @@ class GCPGuestAccelerator(schema.Base):
     count: Annotated[int, Field(ge=1)] = 1
 
 
-class GCPNodeGroup(schema.Base):
-    instance: str
-    min_nodes: Annotated[int, Field(ge=0)] = 0
-    max_nodes: Annotated[int, Field(ge=1)] = 1
+class GCPNodeGroup(NodeGroup):
     preemptible: bool = False
     labels: Dict[str, str] = {}
     guest_accelerators: List[GCPGuestAccelerator] = []
 
 
 DEFAULT_GCP_NODE_GROUPS = {
     "general": GCPNodeGroup(instance="e2-highmem-4", min_nodes=1, max_nodes=1),
-    "user": GCPNodeGroup(instance="e2-standard-4", min_nodes=0, max_nodes=5),
-    "worker": GCPNodeGroup(instance="e2-standard-4", min_nodes=0, max_nodes=5),
+    "user": GCPNodeGroup(
+        instance="e2-standard-4",
+        min_nodes=0,
+        max_nodes=5,
+        taints=[schema.Taint(key="dedicated", value="user", effect="NoSchedule")],
+    ),
+    "worker": GCPNodeGroup(
+        instance="e2-standard-4",
+        min_nodes=0,
+        max_nodes=5,
+        taints=[schema.Taint(key="dedicated", value="worker", effect="NoSchedule")],
+    ),
 }
 
 
@@ -355,10 +399,8 @@ def _check_input(cls, data: Any) -> Any:
         return data
 
 
-class AzureNodeGroup(schema.Base):
-    instance: str
-    min_nodes: int
-    max_nodes: int
+class AzureNodeGroup(NodeGroup):
+    pass
 
 
 DEFAULT_AZURE_NODE_GROUPS = {
@@ -426,10 +468,7 @@ def _validate_tags(cls, value: Optional[Dict[str, str]]) -> Dict[str, str]:
         return value if value is None else azure_cloud.validate_tags(value)
 
 
-class AWSNodeGroup(schema.Base):
-    instance: str
-    min_nodes: int = 0
-    max_nodes: int
+class AWSNodeGroup(NodeGroup):
     gpu: bool = False
     single_subnet: bool = False
     permissions_boundary: Optional[str] = None
@@ -738,6 +777,7 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
                         instance_type=node_group.instance,
                         min_size=node_group.min_nodes,
                         max_size=node_group.max_nodes,
+                        node_taints=node_group.taints,
                         preemptible=node_group.preemptible,
                         guest_accelerators=node_group.guest_accelerators,
                     )

src/_nebari/stages/infrastructure/template/gcp/modules/kubernetes/main.tf

@@ -92,6 +92,15 @@ resource "google_container_node_pool" "main" {
 
     oauth_scopes = local.node_group_oauth_scopes
 
+    dynamic "taint" {
+      for_each = local.merged_node_groups[count.index].node_taints
+      content {
+        key    = taint.value.key
+        value  = taint.value.value
+        effect = taint.value.effect
+      }
+    }
+
     metadata = {
       disable-legacy-endpoints = "true"
     }

src/_nebari/stages/infrastructure/template/gcp/modules/kubernetes/variables.tf

@@ -50,20 +50,23 @@ variable "node_groups" {
       min_size      = 1
       max_size      = 1
       labels        = {}
+      node_taints   = []
     },
     {
       name          = "user"
       instance_type = "n1-standard-2"
       min_size      = 0
       max_size      = 2
       labels        = {}
+      node_taints   = [] # TODO: Do this for other cloud providers
     },
     {
       name          = "worker"
       instance_type = "n1-standard-2"
       min_size      = 0
       max_size      = 5
       labels        = {}
+      node_taints   = []
     }
   ]
 }

src/_nebari/stages/kubernetes_services/__init__.py

@@ -382,6 +382,19 @@ class CondaStoreInputVars(schema.Base):
     )
 
 
+class TolerationOperatorEnum(str, enum.Enum):
+    Equal = "Equal"
+    Exists = "Exists"
+
+    @classmethod
+    def to_yaml(cls, representer, node):
+        return representer.represent_str(node.value)
+
+
+class Toleration(schema.Taint):
+    operator: TolerationOperatorEnum = TolerationOperatorEnum.Equal
+
+
 class JupyterhubInputVars(schema.Base):
     jupyterhub_theme: Dict[str, Any] = Field(alias="jupyterhub-theme")
     jupyterlab_image: ImageNameTag = Field(alias="jupyterlab-image")
@@ -405,6 +418,9 @@ class JupyterhubInputVars(schema.Base):
     jhub_apps_enabled: bool = Field(alias="jhub-apps-enabled")
     cloud_provider: str = Field(alias="cloud-provider")
     jupyterlab_preferred_dir: Optional[str] = Field(alias="jupyterlab-preferred-dir")
+    node_taint_tolerations: Optional[List[Toleration]] = Field(
+        alias="node-taint-tolerations"
+    )
 
 
 class DaskGatewayInputVars(schema.Base):
@@ -565,6 +581,12 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
             jupyterlab_default_settings=self.config.jupyterlab.default_settings,
             jupyterlab_gallery_settings=self.config.jupyterlab.gallery_settings,
             jupyterlab_preferred_dir=self.config.jupyterlab.preferred_dir,
+            node_taint_tolerations=[
+                Toleration(**taint.model_dump())
+                for taint in self.config.google_cloud_platform.node_groups[
+                    "user"
+                ].taints
+            ],  # TODO: support other cloud providers
         )
 
         dask_gateway_vars = DaskGatewayInputVars(

src/_nebari/stages/kubernetes_services/template/jupyterhub.tf

@@ -85,6 +85,16 @@ variable "idle-culler-settings" {
   type        = any
 }
 
+variable "node-taint-tolerations" {
+  description = "Node taint toleration"
+  type = list(object({
+    key      = string
+    operator = string
+    value    = string
+    effect   = string
+  }))
+}
+
 module "kubernetes-nfs-server" {
   count = var.jupyterhub-shared-endpoint == null ? 1 : 0
 
@@ -137,6 +147,7 @@ module "jupyterhub" {
   conda-store-service-name                           = module.kubernetes-conda-store-server.service_name
   conda-store-jhub-apps-token                        = module.kubernetes-conda-store-server.service-tokens.jhub-apps
   jhub-apps-enabled                                  = var.jhub-apps-enabled
+  node-taint-tolerations                             = var.node-taint-tolerations
 
   extra-mounts = {
     "/etc/dask" = {

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/files/jupyterhub/03-profiles.py

@@ -241,6 +241,25 @@ def base_profile_extra_mounts():
     }
 
 
+def node_taint_tolerations():
+    tolerations = z2jh.get_config("custom.node-taint-tolerations")
+
+    if not tolerations:
+        return {}
+
+    return {
+        "tolerations": [
+            {
+                "key": taint["key"],
+                "operator": taint["operator"],
+                "value": taint["value"],
+                "effect": taint["effect"],
+            }
+            for taint in tolerations
+        ]
+    }
+
+
 def configure_user_provisioned_repositories(username):
     # Define paths and configurations
     pvc_home_mount_path = f"home/{username}"
@@ -519,6 +538,7 @@ def render_profile(profile, username, groups, keycloak_profilenames):
             configure_user(username, groups),
             configure_user_provisioned_repositories(username),
             profile_kubespawner_override,
+            node_taint_tolerations(),
         ],
         {},
     )

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/main.tf

@@ -78,6 +78,7 @@ resource "helm_release" "jupyterhub" {
         conda-store-jhub-apps-token   = var.conda-store-jhub-apps-token
         jhub-apps-enabled             = var.jhub-apps-enabled
         initial-repositories          = var.initial-repositories
+        node-taint-tolerations        = var.node-taint-tolerations
         skel-mount = {
           name      = kubernetes_config_map.etc-skel.metadata.0.name
           namespace = kubernetes_config_map.etc-skel.metadata.0.namespace

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/variables.tf

@@ -208,3 +208,13 @@ variable "initial-repositories" {
   type        = string
   default     = "[]"
 }
+
+variable "node-taint-tolerations" {
+  description = "Node taint toleration"
+  type = list(object({
+    key      = string
+    operator = string
+    value    = string
+    effect   = string
+  }))
+}

src/nebari/schema.py

@@ -91,3 +91,16 @@ def is_version_accepted(v):
     for deployment with the current Nebari package.
     """
     return Main.is_version_accepted(v)
+
+
+# TODO: Make sure the taint is actually applied to the nodes for each provider
+class TaintEffectEnum(str, enum.Enum):
+    NoSchedule: str = "NoSchedule"
+    PreferNoSchedule: str = "PreferNoSchedule"
+    NoExecute: str = "NoExecute"
+
+
+class Taint(Base):
+    key: str
+    value: str
+    effect: TaintEffectEnum