Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Disable AWS launch_template from nebari-config schema #2856

Merged
merged 3 commits into from
Nov 18, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 30 additions & 8 deletions src/_nebari/stages/infrastructure/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import tempfile
from typing import Annotated, Any, Dict, List, Literal, Optional, Tuple, Type, Union

from pydantic import ConfigDict, Field, PrivateAttr, field_validator, model_validator
from pydantic import ConfigDict, Field, field_validator, model_validator

from _nebari import constants
from _nebari.provider import terraform
Expand Down Expand Up @@ -136,7 +136,7 @@ class AWSAmiTypes(str, enum.Enum):

class AWSNodeLaunchTemplate(schema.Base):
pre_bootstrap_command: Optional[str] = None
_ami_id: Optional[str] = PrivateAttr(default=None)
ami_id: Optional[str] = None


class AWSNodeGroupInputVars(schema.Base):
Expand All @@ -152,10 +152,23 @@ class AWSNodeGroupInputVars(schema.Base):
launch_template: Optional[AWSNodeLaunchTemplate] = None


def construct_aws_ami_type(gpu_enabled: bool, launch_template: AWSNodeLaunchTemplate):
"""Construct the AWS AMI type based on the provided parameters."""
def construct_aws_ami_type(
gpu_enabled: bool, launch_template: AWSNodeLaunchTemplate
) -> str:
"""
This function selects the Amazon Machine Image (AMI) type for AWS nodes by evaluating
the provided parameters. The selection logic prioritizes the launch template over the
GPU flag.

Returns the AMI type (str) determined by the following rules:
- Returns "CUSTOM" if a `launch_template` is provided and it includes a valid `ami_id`.
- Returns "AL2_x86_64_GPU" if `gpu_enabled` is True and no valid
`launch_template` is provided (None).
- Returns "AL2_x86_64" as the default AMI type if `gpu_enabled` is False and no
valid `launch_template` is provided (None).
"""

if launch_template and launch_template._ami_id:
if launch_template and getattr(launch_template, "ami_id", None):
return "CUSTOM"

if gpu_enabled:
Expand Down Expand Up @@ -480,7 +493,16 @@ class AWSNodeGroup(schema.Base):
gpu: bool = False
single_subnet: bool = False
permissions_boundary: Optional[str] = None
launch_template: Optional[AWSNodeLaunchTemplate] = None
# Disabled as part of 2024.11.1 until #2832 is resolved
# launch_template: Optional[AWSNodeLaunchTemplate] = None

@model_validator(mode="before")
def check_launch_template(cls, values):
if "launch_template" in values:
raise ValueError(
"The 'launch_template' field is currently unavailable and has been removed from the configuration schema.\nPlease omit this field until it is reintroduced in a future update.",
)
return values


DEFAULT_AWS_NODE_GROUPS = {
Expand Down Expand Up @@ -899,10 +921,10 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
max_size=node_group.max_nodes,
single_subnet=node_group.single_subnet,
permissions_boundary=node_group.permissions_boundary,
launch_template=node_group.launch_template,
launch_template=None,
ami_type=construct_aws_ami_type(
gpu_enabled=node_group.gpu,
launch_template=node_group.launch_template,
launch_template=None,
),
)
for name, node_group in self.config.amazon_web_services.node_groups.items()
Expand Down
Loading