From 381b3a5b3b9aba3a462d83ef474c8992927a16fe Mon Sep 17 00:00:00 2001 From: Jan Sykora Date: Fri, 8 Mar 2024 12:37:02 +0100 Subject: [PATCH] feat: Add node affinity to Node Template (#283) --- castai/resource_node_template.go | 79 +++++++++++++++++++++++++++ castai/resource_node_template_test.go | 14 +++++ castai/sdk/api.gen.go | 31 +++++++---- docs/resources/node_template.md | 11 ++++ 4 files changed, 124 insertions(+), 11 deletions(-) diff --git a/castai/resource_node_template.go b/castai/resource_node_template.go index 9442572e..0bc76427 100644 --- a/castai/resource_node_template.go +++ b/castai/resource_node_template.go @@ -56,6 +56,9 @@ const ( FieldNodeTemplateStorageOptimized = "storage_optimized" FieldNodeTemplateUseSpotFallbacks = "use_spot_fallbacks" FieldNodeTemplateCustomPriority = "custom_priority" + FieldNodeTemplateNodeAffinity = "node_affinity" + FieldNodeTemplateAzName = "az_name" + FieldNodeTemplateInstanceTypes = "instance_types" ) const ( @@ -345,6 +348,32 @@ func resourceNodeTemplate() *schema.Resource { }, }, }, + FieldNodeTemplateNodeAffinity: { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + FieldNodeTemplateInstanceTypes: { + Type: schema.TypeList, + Required: true, + Elem: &schema.Schema{ + Type: schema.TypeString, + }, + Description: "Instance types in this node group.", + }, + FieldNodeTemplateAzName: { + Required: true, + Type: schema.TypeString, + Description: "Availability zone name.", + }, + FieldNodeTemplateName: { + Required: true, + Type: schema.TypeString, + Description: "Name of node group.", + }, + }, + }, + }, }, }, }, @@ -489,6 +518,9 @@ func flattenConstraints(c *sdk.NodetemplatesV1TemplateConstraints) ([]map[string if c.CustomPriority != nil && len(*c.CustomPriority) > 0 { out[FieldNodeTemplateCustomPriority] = flattenCustomPriority(*c.CustomPriority) } + if c.NodeAffinity != nil && len(*c.NodeAffinity) > 0 { + out[FieldNodeTemplateNodeAffinity] = flattenNodeAffinity(*c.NodeAffinity) + } if c.InstanceFamilies != nil { out[FieldNodeTemplateInstanceFamilies] = flattenInstanceFamilies(c.InstanceFamilies) } @@ -596,6 +628,19 @@ func flattenCustomPriority(priorities []sdk.NodetemplatesV1TemplateConstraintsCu }) } +func flattenNodeAffinity(affinities []sdk.NodetemplatesV1TemplateConstraintsNodeAffinity) any { + return lo.Map(affinities, func(item sdk.NodetemplatesV1TemplateConstraintsNodeAffinity, index int) map[string]any { + result := map[string]any{} + if item.InstanceTypes != nil { + result[FieldNodeTemplateInstanceTypes] = *item.InstanceTypes + } + + result[FieldNodeTemplateName] = lo.FromPtr(item.Name) + result[FieldNodeTemplateAzName] = lo.FromPtr(item.AzName) + return result + }) +} + func resourceNodeTemplateDelete(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics { client := meta.(*ProviderConfig).api clusterID := d.Get(FieldClusterID).(string) @@ -1019,6 +1064,21 @@ func toTemplateConstraints(obj map[string]any) *sdk.NodetemplatesV1TemplateConst })) } } + if v, ok := obj[FieldNodeTemplateNodeAffinity].([]any); ok && len(v) > 0 { + if ok { + out.NodeAffinity = lo.ToPtr(lo.FilterMap(v, func(item any, _ int) (sdk.NodetemplatesV1TemplateConstraintsNodeAffinity, bool) { + val, ok := item.(map[string]any) + if !ok { + return sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{}, false + } + res := toTemplateConstraintsNodeAffinity(val) + if res == nil { + return sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{}, false + } + return *res, true + })) + } + } return out } @@ -1083,3 +1143,22 @@ func toTemplateConstraintsCustomPriority(o map[string]any) *sdk.NodetemplatesV1T return &out } + +func toTemplateConstraintsNodeAffinity(o map[string]any) *sdk.NodetemplatesV1TemplateConstraintsNodeAffinity { + if o == nil { + return nil + } + + out := sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{} + if v, ok := o[FieldNodeTemplateName].(string); ok { + out.Name = toPtr(v) + } + if v, ok := o[FieldNodeTemplateAzName].(string); ok { + out.AzName = toPtr(v) + } + if v, ok := o[FieldNodeTemplateInstanceTypes].([]any); ok { + out.InstanceTypes = toPtr(toStringList(v)) + } + + return &out +} diff --git a/castai/resource_node_template_test.go b/castai/resource_node_template_test.go index 236d7d88..af1c439c 100644 --- a/castai/resource_node_template_test.go +++ b/castai/resource_node_template_test.go @@ -87,6 +87,13 @@ func TestNodeTemplateResourceReadContext(t *testing.T) { "spot": true, "onDemand": true } + ], + "nodeAffinity": [ + { + "name": "foo", + "azName": "eu-central-1a", + "instanceTypes": ["m5.24xlarge"] + } ] }, "version": "3", @@ -171,6 +178,11 @@ constraints.0.max_cpu = 10000 constraints.0.max_memory = 0 constraints.0.min_cpu = 10 constraints.0.min_memory = 0 +constraints.0.node_affinity.# = 1 +constraints.0.node_affinity.0.az_name = eu-central-1a +constraints.0.node_affinity.0.instance_types.# = 1 +constraints.0.node_affinity.0.instance_types.0 = m5.24xlarge +constraints.0.node_affinity.0.name = foo constraints.0.on_demand = true constraints.0.os.# = 1 constraints.0.os.0 = linux @@ -434,6 +446,7 @@ func TestAccResourceNodeTemplate_basic(t *testing.T) { resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.instance_families.1", "d"), resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.spot", "true"), resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.on_demand", "true"), + resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.#", "0"), ), }, { @@ -491,6 +504,7 @@ func TestAccResourceNodeTemplate_basic(t *testing.T) { resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.instance_families.1", "d"), resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.spot", "true"), resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.on_demand", "true"), + resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.#", "0"), ), }, }, diff --git a/castai/sdk/api.gen.go b/castai/sdk/api.gen.go index ef11170f..a2b994e2 100644 --- a/castai/sdk/api.gen.go +++ b/castai/sdk/api.gen.go @@ -1495,12 +1495,13 @@ type ExternalclusterV1RegisterClusterRequest struct { // ExternalclusterV1Resources defines model for externalcluster.v1.Resources. type ExternalclusterV1Resources struct { - CpuAllocatableMilli *int32 `json:"cpuAllocatableMilli,omitempty"` - CpuCapacityMilli *int32 `json:"cpuCapacityMilli,omitempty"` - CpuRequestsMilli *int32 `json:"cpuRequestsMilli,omitempty"` - MemAllocatableMib *int32 `json:"memAllocatableMib,omitempty"` - MemCapacityMib *int32 `json:"memCapacityMib,omitempty"` - MemRequestsMib *int32 `json:"memRequestsMib,omitempty"` + BandwidthCapacityMbps *int32 `json:"bandwidthCapacityMbps,omitempty"` + CpuAllocatableMilli *int32 `json:"cpuAllocatableMilli,omitempty"` + CpuCapacityMilli *int32 `json:"cpuCapacityMilli,omitempty"` + CpuRequestsMilli *int32 `json:"cpuRequestsMilli,omitempty"` + MemAllocatableMib *int32 `json:"memAllocatableMib,omitempty"` + MemCapacityMib *int32 `json:"memCapacityMib,omitempty"` + MemRequestsMib *int32 `json:"memRequestsMib,omitempty"` } // Subnet represents cluster subnet. @@ -1991,11 +1992,12 @@ type NodetemplatesV1TemplateConstraints struct { // This template is gpu only. Setting this to true, will result in only instances with GPUs being considered. // In addition, this ensures that all of the added instances for this template won't have any nvidia taints. - IsGpuOnly *bool `json:"isGpuOnly"` - MaxCpu *int32 `json:"maxCpu"` - MaxMemory *int32 `json:"maxMemory"` - MinCpu *int32 `json:"minCpu"` - MinMemory *int32 `json:"minMemory"` + IsGpuOnly *bool `json:"isGpuOnly"` + MaxCpu *int32 `json:"maxCpu"` + MaxMemory *int32 `json:"maxMemory"` + MinCpu *int32 `json:"minCpu"` + MinMemory *int32 `json:"minMemory"` + NodeAffinity *[]NodetemplatesV1TemplateConstraintsNodeAffinity `json:"nodeAffinity,omitempty"` // Should include on-demand instances in the considered pool. OnDemand *bool `json:"onDemand"` @@ -2050,6 +2052,13 @@ type NodetemplatesV1TemplateConstraintsInstanceFamilyConstraints struct { Include *[]string `json:"include,omitempty"` } +// NodetemplatesV1TemplateConstraintsNodeAffinity defines model for nodetemplates.v1.TemplateConstraints.NodeAffinity. +type NodetemplatesV1TemplateConstraintsNodeAffinity struct { + AzName *string `json:"azName,omitempty"` + InstanceTypes *[]string `json:"instanceTypes,omitempty"` + Name *string `json:"name,omitempty"` +} + // NodetemplatesV1UpdateNodeTemplate defines model for nodetemplates.v1.UpdateNodeTemplate. type NodetemplatesV1UpdateNodeTemplate struct { ConfigurationId *string `json:"configurationId,omitempty"` diff --git a/docs/resources/node_template.md b/docs/resources/node_template.md index ac86b1d7..750e6eb3 100644 --- a/docs/resources/node_template.md +++ b/docs/resources/node_template.md @@ -55,6 +55,7 @@ Optional: - `max_memory` (Number) Max Memory (Mib) per node. - `min_cpu` (Number) Min CPU cores per node. - `min_memory` (Number) Min Memory (Mib) per node. +- `node_affinity` (Block List) (see [below for nested schema](#nestedblock--constraints--node_affinity)) - `on_demand` (Boolean) Should include on-demand instances in the considered pool. - `os` (List of String) List of acceptable instance Operating Systems, the default is linux. Allowed values: linux, windows. - `spot` (Boolean) Should include spot instances in the considered pool. @@ -95,6 +96,16 @@ Optional: - `include` (List of String) Instance families to exclude when filtering (includes all other families). + +### Nested Schema for `constraints.node_affinity` + +Required: + +- `az_name` (String) Availability zone name. +- `instance_types` (List of String) Instance types in this node group. +- `name` (String) Name of node group. + + ### Nested Schema for `custom_taints`