From c138a22993fbb0dd3fae41af432dfd03047e3dbd Mon Sep 17 00:00:00 2001 From: vinicius douglas cerutti Date: Wed, 16 Oct 2024 19:22:58 -0300 Subject: [PATCH 1/6] refactor ami_type validation logic --- src/_nebari/stages/infrastructure/__init__.py | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index 682e9d50b..da85f9908 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -151,25 +151,21 @@ class AWSNodeGroupInputVars(schema.Base): ami_type: Optional[AWSAmiTypes] = None launch_template: Optional[AWSNodeLaunchTemplate] = None - @field_validator("ami_type", mode="before") - @classmethod - def _infer_and_validate_ami_type(cls, value, values) -> str: - gpu_enabled = values.get("gpu", False) - - # Auto-set ami_type if not provided - if not value: - if values.get("launch_template") and values["launch_template"].ami_id: - return "CUSTOM" - if gpu_enabled: - return "AL2_x86_64_GPU" - return "AL2_x86_64" - - # Explicit validation - if value == "AL2_x86_64" and gpu_enabled: - raise ValueError( - "ami_type 'AL2_x86_64' cannot be used with GPU enabled (gpu=True)." - ) - return value + +def construct_aws_ami_type( + gpu_enabled: bool, launch_template: Dict, ami_type: str = None +): + """Construct the AWS AMI type based on the provided parameters.""" + if ami_type: + return ami_type + + if launch_template and launch_template.get("ami_id"): + return AWSAmiTypes.CUSTOM + + if gpu_enabled: + return AWSAmiTypes.AL2_x86_64_GPU + + return AWSAmiTypes.AL2_x86_64 class AWSInputVars(schema.Base): @@ -864,6 +860,10 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]): single_subnet=node_group.single_subnet, permissions_boundary=node_group.permissions_boundary, launch_template=node_group.launch_template, + ami_type=construct_aws_ami_type( + gpu_enabled=node_group.gpu, + launch_template=node_group.launch_template, + ), ) for name, node_group in self.config.amazon_web_services.node_groups.items() ], From cd6611b8948184d8edd15b63d24931cd6369c061 Mon Sep 17 00:00:00 2001 From: vinicius douglas cerutti Date: Wed, 16 Oct 2024 20:52:34 -0300 Subject: [PATCH 2/6] use serializable ami_type response --- src/_nebari/stages/infrastructure/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index da85f9908..1fcc6aea7 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -160,12 +160,12 @@ def construct_aws_ami_type( return ami_type if launch_template and launch_template.get("ami_id"): - return AWSAmiTypes.CUSTOM + return AWSAmiTypes.CUSTOM.value if gpu_enabled: - return AWSAmiTypes.AL2_x86_64_GPU + return AWSAmiTypes.AL2_x86_64_GPU.value - return AWSAmiTypes.AL2_x86_64 + return AWSAmiTypes.AL2_x86_64.value class AWSInputVars(schema.Base): From eb3ae65f7dfc4185fc7b719f08b5671246c83cae Mon Sep 17 00:00:00 2001 From: vinicius douglas cerutti Date: Wed, 16 Oct 2024 21:12:41 -0300 Subject: [PATCH 3/6] use serializable ami_type response --- src/_nebari/stages/infrastructure/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index 1fcc6aea7..70e68946a 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -160,12 +160,12 @@ def construct_aws_ami_type( return ami_type if launch_template and launch_template.get("ami_id"): - return AWSAmiTypes.CUSTOM.value + return "CUSTOM" if gpu_enabled: - return AWSAmiTypes.AL2_x86_64_GPU.value + return "AL2_x86_64_GPU" - return AWSAmiTypes.AL2_x86_64.value + return "AL2_x86_64" class AWSInputVars(schema.Base): From 3d4017b207175e7c7ee6e315dc05e112268ed40e Mon Sep 17 00:00:00 2001 From: vinicius douglas cerutti Date: Thu, 17 Oct 2024 10:25:44 -0300 Subject: [PATCH 4/6] use serializable ami_type Enum object --- src/_nebari/stages/infrastructure/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index 70e68946a..4d3ca9545 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -128,12 +128,11 @@ class AzureInputVars(schema.Base): workload_identity_enabled: bool = False -class AWSAmiTypes(enum.Enum): +class AWSAmiTypes(str, enum.Enum): AL2_x86_64 = "AL2_x86_64" AL2_x86_64_GPU = "AL2_x86_64_GPU" CUSTOM = "CUSTOM" - class AWSNodeLaunchTemplate(schema.Base): pre_bootstrap_command: Optional[str] = None ami_id: Optional[str] = None From fabd56e5ec9f0f7e1b165ecacca7e1a91a067ab8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 22 Oct 2024 15:04:22 +0000 Subject: [PATCH 5/6] [pre-commit.ci] Apply automatic pre-commit fixes --- src/_nebari/stages/infrastructure/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index 4d3ca9545..b24613e0e 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -133,6 +133,7 @@ class AWSAmiTypes(str, enum.Enum): AL2_x86_64_GPU = "AL2_x86_64_GPU" CUSTOM = "CUSTOM" + class AWSNodeLaunchTemplate(schema.Base): pre_bootstrap_command: Optional[str] = None ami_id: Optional[str] = None From 409ad314fed6a0ff7790ac46707f7b9676c2721a Mon Sep 17 00:00:00 2001 From: "Vinicius D. Cerutti" <51954708+viniciusdc@users.noreply.github.com> Date: Fri, 25 Oct 2024 12:54:36 -0300 Subject: [PATCH 6/6] Address type error with AWSNodeLaunchTemplate --- src/_nebari/stages/infrastructure/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/_nebari/stages/infrastructure/__init__.py b/src/_nebari/stages/infrastructure/__init__.py index b24613e0e..18c3f9981 100644 --- a/src/_nebari/stages/infrastructure/__init__.py +++ b/src/_nebari/stages/infrastructure/__init__.py @@ -153,13 +153,13 @@ class AWSNodeGroupInputVars(schema.Base): def construct_aws_ami_type( - gpu_enabled: bool, launch_template: Dict, ami_type: str = None + gpu_enabled: bool, launch_template: AWSNodeLaunchTemplate, ami_type: str = None ): """Construct the AWS AMI type based on the provided parameters.""" if ami_type: return ami_type - if launch_template and launch_template.get("ami_id"): + if launch_template and launch_template.ami_id: return "CUSTOM" if gpu_enabled: