Add option for AWS node-groups to run in a single subnet/AZ

nebari/initialize.py

@@ -237,8 +237,18 @@ def __getattr__(name):
     "kubernetes_version": "PLACEHOLDER",
     "node_groups": {
         "general": {"instance": "m5.2xlarge", "min_nodes": 1, "max_nodes": 1},
-        "user": {"instance": "m5.xlarge", "min_nodes": 1, "max_nodes": 5},
-        "worker": {"instance": "m5.xlarge", "min_nodes": 1, "max_nodes": 5},
+        "user": {
+            "instance": "m5.xlarge",
+            "min_nodes": 1,
+            "max_nodes": 5,
+            "single_subnet": True,
+        },
+        "worker": {
+            "instance": "m5.xlarge",
+            "min_nodes": 1,
+            "max_nodes": 5,
+            "single_subnet": True,
+        },
     },
 }
 

nebari/schema.py

@@ -291,6 +291,10 @@ def validate_guest_accelerators(cls, v):
                 raise ValueError(assertion_error_message)
 
 
+class AWSNodeGroup(NodeGroup):
+    single_subnet: typing.Optional[bool] = False
+
+
 class DigitalOceanProvider(Base):
     region: str
     kubernetes_version: str
@@ -320,7 +324,7 @@ class AmazonWebServicesProvider(Base):
     region: str
     availability_zones: typing.Optional[typing.List[str]]
     kubernetes_version: str
-    node_groups: typing.Dict[str, NodeGroup]
+    node_groups: typing.Dict[str, AWSNodeGroup]
     terraform_overrides: typing.Any
 
 

nebari/stages/input_vars.py

@@ -110,6 +110,7 @@ def stage_02_infrastructure(stage_outputs, config):
                     "max_size": value["max_nodes"],
                     "gpu": value.get("gpu", False),
                     "instance_type": value["instance"],
+                    "single_subnet": value.get("single_subnet", False),
                 }
                 for key, value in config["amazon_web_services"]["node_groups"].items()
             ],

nebari/template/stages/02-infrastructure/aws/modules/kubernetes/main.tf

@@ -22,7 +22,7 @@ resource "aws_eks_node_group" "main" {
   cluster_name    = aws_eks_cluster.main.name
   node_group_name = var.node_groups[count.index].name
   node_role_arn   = aws_iam_role.node-group.arn
-  subnet_ids      = var.cluster_subnets
+  subnet_ids      = var.node_groups[count.index].single_subnet ? [element(var.cluster_subnets, 0)] : var.cluster_subnets
 
   instance_types = [var.node_groups[count.index].instance_type]
   ami_type       = var.node_groups[count.index].gpu == true ? "AL2_x86_64_GPU" : "AL2_x86_64"

nebari/template/stages/02-infrastructure/aws/modules/kubernetes/variables.tf

@@ -50,6 +50,7 @@ variable "node_groups" {
     min_size      = number
     desired_size  = number
     max_size      = number
+    single_subnet = bool
   }))
 }
 

nebari/template/stages/02-infrastructure/aws/variables.tf

@@ -27,6 +27,7 @@ variable "node_groups" {
     min_size      = number
     desired_size  = number
     max_size      = number
+    single_subnet = bool
   }))
 }