add default node groups to config (#2398)

nebari-dev · Apr 12, 2024 · a26b8df · a26b8df
1 parent f3028b7
commit a26b8df
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 30 deletions.
diff --git a/src/_nebari/initialize.py b/src/_nebari/initialize.py
@@ -18,6 +18,12 @@
 )
 from _nebari.provider.oauth.auth0 import create_client
 from _nebari.stages.bootstrap import CiEnum
+from _nebari.stages.infrastructure import (
+    DEFAULT_AWS_NODE_GROUPS,
+    DEFAULT_AZURE_NODE_GROUPS,
+    DEFAULT_DO_NODE_GROUPS,
+    DEFAULT_GCP_NODE_GROUPS,
+)
 from _nebari.stages.kubernetes_ingress import CertificateEnum
 from _nebari.stages.kubernetes_keycloak import AuthenticationEnum
 from _nebari.stages.terraform_state import TerraformStateEnum
@@ -30,6 +36,10 @@
 WELCOME_HEADER_TEXT = "Your open source data science platform, hosted"
 
 
+def _node_groups_to_dict(node_groups):
+    return {ng_name: ng.dict() for ng_name, ng in node_groups.items()}
+
+
 def render_config(
     project_name: str,
     nebari_domain: str = None,
@@ -117,6 +127,7 @@ def render_config(
         config["digital_ocean"] = {
             "kubernetes_version": do_kubernetes_versions,
             "region": do_region,
+            "node_groups": _node_groups_to_dict(DEFAULT_DO_NODE_GROUPS),
         }
 
         config["theme"]["jupyterhub"][
@@ -131,6 +142,7 @@ def render_config(
         config["google_cloud_platform"] = {
             "kubernetes_version": gcp_kubernetes_version,
             "region": gcp_region,
+            "node_groups": _node_groups_to_dict(DEFAULT_GCP_NODE_GROUPS),
         }
 
         config["theme"]["jupyterhub"][
@@ -152,6 +164,7 @@ def render_config(
             "kubernetes_version": azure_kubernetes_version,
             "region": azure_region,
             "storage_account_postfix": random_secure_string(length=4),
+            "node_groups": _node_groups_to_dict(DEFAULT_AZURE_NODE_GROUPS),
         }
 
         config["theme"]["jupyterhub"][
@@ -170,6 +183,7 @@ def render_config(
         config["amazon_web_services"] = {
             "kubernetes_version": aws_kubernetes_version,
             "region": aws_region,
+            "node_groups": _node_groups_to_dict(DEFAULT_AWS_NODE_GROUPS),
         }
         config["theme"]["jupyterhub"][
             "hub_subtitle"

diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py
@@ -223,21 +223,18 @@ class DigitalOceanNodeGroup(schema.Base):
     max_nodes: pydantic.conint(ge=1) = 1
 
 
+DEFAULT_DO_NODE_GROUPS = {
+    "general": DigitalOceanNodeGroup(instance="g-8vcpu-32gb", min_nodes=1, max_nodes=1),
+    "user": DigitalOceanNodeGroup(instance="g-4vcpu-16gb", min_nodes=1, max_nodes=5),
+    "worker": DigitalOceanNodeGroup(instance="g-4vcpu-16gb", min_nodes=1, max_nodes=5),
+}
+
+
 class DigitalOceanProvider(schema.Base):
     region: str
     kubernetes_version: str
     # Digital Ocean image slugs are listed here https://slugs.do-api.dev/
-    node_groups: Dict[str, DigitalOceanNodeGroup] = {
-        "general": DigitalOceanNodeGroup(
-            instance="g-8vcpu-32gb", min_nodes=1, max_nodes=1
-        ),
-        "user": DigitalOceanNodeGroup(
-            instance="g-4vcpu-16gb", min_nodes=1, max_nodes=5
-        ),
-        "worker": DigitalOceanNodeGroup(
-            instance="g-4vcpu-16gb", min_nodes=1, max_nodes=5
-        ),
-    }
+    node_groups: Dict[str, DigitalOceanNodeGroup] = DEFAULT_DO_NODE_GROUPS
     tags: Optional[List[str]] = []
 
     @pydantic.validator("region")
@@ -329,17 +326,20 @@ class GCPNodeGroup(schema.Base):
     guest_accelerators: List[GCPGuestAccelerator] = []
 
 
+DEFAULT_GCP_NODE_GROUPS = {
+    "general": GCPNodeGroup(instance="n1-standard-8", min_nodes=1, max_nodes=1),
+    "user": GCPNodeGroup(instance="n1-standard-4", min_nodes=0, max_nodes=5),
+    "worker": GCPNodeGroup(instance="n1-standard-4", min_nodes=0, max_nodes=5),
+}
+
+
 class GoogleCloudPlatformProvider(schema.Base):
     region: str
     project: str
     kubernetes_version: str
     availability_zones: Optional[List[str]] = []
     release_channel: str = constants.DEFAULT_GKE_RELEASE_CHANNEL
-    node_groups: Dict[str, GCPNodeGroup] = {
-        "general": GCPNodeGroup(instance="n1-standard-8", min_nodes=1, max_nodes=1),
-        "user": GCPNodeGroup(instance="n1-standard-4", min_nodes=0, max_nodes=5),
-        "worker": GCPNodeGroup(instance="n1-standard-4", min_nodes=0, max_nodes=5),
-    }
+    node_groups: Dict[str, GCPNodeGroup] = DEFAULT_GCP_NODE_GROUPS
     tags: Optional[List[str]] = []
     networking_mode: str = "ROUTE"
     network: str = "default"
@@ -381,16 +381,19 @@ class AzureNodeGroup(schema.Base):
     max_nodes: int
 
 
+DEFAULT_AZURE_NODE_GROUPS = {
+    "general": AzureNodeGroup(instance="Standard_D8_v3", min_nodes=1, max_nodes=1),
+    "user": AzureNodeGroup(instance="Standard_D4_v3", min_nodes=0, max_nodes=5),
+    "worker": AzureNodeGroup(instance="Standard_D4_v3", min_nodes=0, max_nodes=5),
+}
+
+
 class AzureProvider(schema.Base):
     region: str
     kubernetes_version: str
     storage_account_postfix: str
     resource_group_name: str = None
-    node_groups: Dict[str, AzureNodeGroup] = {
-        "general": AzureNodeGroup(instance="Standard_D8_v3", min_nodes=1, max_nodes=1),
-        "user": AzureNodeGroup(instance="Standard_D4_v3", min_nodes=0, max_nodes=5),
-        "worker": AzureNodeGroup(instance="Standard_D4_v3", min_nodes=0, max_nodes=5),
-    }
+    node_groups: Dict[str, AzureNodeGroup] = DEFAULT_AZURE_NODE_GROUPS
     storage_account_postfix: str
     vnet_subnet_id: Optional[Union[str, None]] = None
     private_cluster_enabled: bool = False
@@ -442,19 +445,22 @@ class AWSNodeGroup(schema.Base):
     permissions_boundary: Optional[str] = None
 
 
+DEFAULT_AWS_NODE_GROUPS = {
+    "general": AWSNodeGroup(instance="m5.2xlarge", min_nodes=1, max_nodes=1),
+    "user": AWSNodeGroup(
+        instance="m5.xlarge", min_nodes=0, max_nodes=5, single_subnet=False
+    ),
+    "worker": AWSNodeGroup(
+        instance="m5.xlarge", min_nodes=0, max_nodes=5, single_subnet=False
+    ),
+}
+
+
 class AmazonWebServicesProvider(schema.Base):
     region: str
     kubernetes_version: str
     availability_zones: Optional[List[str]]
-    node_groups: Dict[str, AWSNodeGroup] = {
-        "general": AWSNodeGroup(instance="m5.2xlarge", min_nodes=1, max_nodes=1),
-        "user": AWSNodeGroup(
-            instance="m5.xlarge", min_nodes=0, max_nodes=5, single_subnet=False
-        ),
-        "worker": AWSNodeGroup(
-            instance="m5.xlarge", min_nodes=0, max_nodes=5, single_subnet=False
-        ),
-    }
+    node_groups: Dict[str, AWSNodeGroup] = DEFAULT_AWS_NODE_GROUPS
     existing_subnet_ids: List[str] = None
     existing_security_group_id: str = None
     vpc_cidr_block: str = "10.10.0.0/16"