Skip to content

Commit 0c33c45

Browse files
maxtwardowskiMax Twardowski
maxtwardowski
and
Max Twardowski
authored
feat: add AWS capacity reservations support to node templates (#650)
* regen sdk * fix edge location resource after sdk regen * nt resource changes * regen docs * update examples * fix TestNodeTemplateResourceReadContext * remove resource group arn from TestAccEKS_ResourceNodeTemplate_basic * fix cr count expectation --------- Co-authored-by: Max Twardowski <max@cast.ai>
1 parent 976e12a commit 0c33c45

File tree

9 files changed

+668
-45
lines changed

9 files changed

+668
-45
lines changed

‎castai/resource_edge_location.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -765,7 +765,7 @@ func (r *edgeLocationResource) toGCP(ctx context.Context, plan, config *gcpModel
765765
ProjectId: plan.ProjectID.ValueString(),
766766
InstanceServiceAccount: instanceServiceAccount,
767767
Credentials: &omni.GCPParamCredentials{
768-
ClientServiceAccountJsonBase64: config.ClientServiceAccountJSONBase64WO.ValueString(),
768+
ClientServiceAccountJsonBase64: config.ClientServiceAccountJSONBase64WO.ValueStringPointer(),
769769
},
770770
Networking: &omni.GCPParamGCPNetworking{
771771
NetworkName: plan.NetworkName.ValueString(),

‎castai/resource_node_template.go‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,11 @@ const (
9292
FieldNodeTemplatePriceAdjustmentConfiguration = "price_adjustment_configuration"
9393
FieldNodeTemplateInstanceTypeAdjustments = "instance_type_adjustments"
9494
FieldNodeTemplateUserManagedGPUDrivers = "user_managed_gpu_drivers"
95+
FieldNodeTemplateAWSConstraints = "aws"
96+
FieldNodeTemplateCapacityReservations = "capacity_reservations"
97+
FieldNodeTemplateCapacityReservationId = "id"
98+
FieldNodeTemplateCapacityResourceGroupArn = "capacity_resource_group_arn"
99+
FieldNodeTemplateCapacityReservationType = "type"
95100
)
96101

97102
const (
@@ -121,6 +126,11 @@ const (
121126
Unspecified = "unspecified"
122127
)
123128

129+
const (
130+
CapacityReservationTypeOnDemand = "ON_DEMAND_CAPACITY_RESERVATION"
131+
CapacityReservationTypeCapacityBlock = "CAPACITY_BLOCK"
132+
)
133+
124134
type nodeSelectorOperatorsSlice []string
125135

126136
var nodeSelectorOperators = nodeSelectorOperatorsSlice{
@@ -609,6 +619,41 @@ func resourceNodeTemplate() *schema.Resource {
609619
},
610620
},
611621
},
622+
FieldNodeTemplateAWSConstraints: {
623+
Type: schema.TypeList,
624+
MaxItems: 1,
625+
Optional: true,
626+
Elem: &schema.Resource{
627+
Schema: map[string]*schema.Schema{
628+
FieldNodeTemplateCapacityReservations: {
629+
Type: schema.TypeList,
630+
Optional: true,
631+
Elem: &schema.Resource{
632+
Schema: map[string]*schema.Schema{
633+
FieldNodeTemplateCapacityReservationId: {
634+
Type: schema.TypeString,
635+
Optional: true,
636+
Description: "AWS capacity reservation ID.",
637+
},
638+
FieldNodeTemplateCapacityResourceGroupArn: {
639+
Type: schema.TypeString,
640+
Optional: true,
641+
Description: "Capacity resource group ARN for UltraServer capacity blocks.",
642+
},
643+
FieldNodeTemplateCapacityReservationType: {
644+
Type: schema.TypeString,
645+
Optional: true,
646+
ValidateDiagFunc: validation.ToDiagFunc(validation.StringInSlice([]string{CapacityReservationTypeOnDemand, CapacityReservationTypeCapacityBlock}, false)),
647+
Description: fmt.Sprintf("Type of capacity reservation. Allowed values: %s, %s.", CapacityReservationTypeOnDemand, CapacityReservationTypeCapacityBlock),
648+
},
649+
},
650+
},
651+
Description: "Capacity reservations that this template can use for provisioning.",
652+
},
653+
},
654+
},
655+
Description: "AWS-specific constraints for the node template.",
656+
},
612657
},
613658
},
614659
},
@@ -1043,6 +1088,9 @@ func flattenConstraints(c *sdk.NodetemplatesV1TemplateConstraints) ([]map[string
10431088
out[FieldNodeTemplateArchitecturePriority] = lo.FromPtr(c.ArchitecturePriority)
10441089
}
10451090
out[FieldNodeTemplateResourceLimits] = flattenResourceLimits(c.ResourceLimits)
1091+
if c.Aws != nil {
1092+
out[FieldNodeTemplateAWSConstraints] = flattenAWSConstraints(c.Aws)
1093+
}
10461094
if c.BareMetal != nil {
10471095
if lo.FromPtr(c.BareMetal) {
10481096
out[FieldNodeTemplateBareMetal] = True
@@ -1102,6 +1150,33 @@ func flattenResourceLimits(resourceLimits *sdk.NodetemplatesV1TemplateConstraint
11021150
return []map[string]any{out}
11031151
}
11041152

1153+
func flattenAWSConstraints(aws *sdk.NodetemplatesV1TemplateConstraintsAWSConstraints) []map[string]any {
1154+
if aws == nil {
1155+
return nil
1156+
}
1157+
out := map[string]any{}
1158+
if aws.CapacityReservations == nil {
1159+
return []map[string]any{out}
1160+
}
1161+
1162+
reservations := make([]map[string]any, 0, len(*aws.CapacityReservations))
1163+
for _, r := range *aws.CapacityReservations {
1164+
m := map[string]any{}
1165+
if r.Id != nil {
1166+
m[FieldNodeTemplateCapacityReservationId] = *r.Id
1167+
}
1168+
if r.CapacityResourceGroupArn != nil {
1169+
m[FieldNodeTemplateCapacityResourceGroupArn] = *r.CapacityResourceGroupArn
1170+
}
1171+
if r.Type != nil {
1172+
m[FieldNodeTemplateCapacityReservationType] = string(*r.Type)
1173+
}
1174+
reservations = append(reservations, m)
1175+
}
1176+
out[FieldNodeTemplateCapacityReservations] = reservations
1177+
return []map[string]any{out}
1178+
}
1179+
11051180
func flattenGpu(gpu *sdk.NodetemplatesV1TemplateConstraintsGPUConstraints) []map[string]any {
11061181
if gpu == nil {
11071182
return nil
@@ -1711,6 +1786,11 @@ func toTemplateConstraints(obj map[string]any) *sdk.NodetemplatesV1TemplateConst
17111786
out.ResourceLimits = toTemplateConstraintsResourceLimits(val)
17121787
}
17131788
}
1789+
if v, ok := obj[FieldNodeTemplateAWSConstraints].([]any); ok && len(v) > 0 {
1790+
if val, ok := v[0].(map[string]any); ok {
1791+
out.Aws = toTemplateConstraintsAWSConstraints(val)
1792+
}
1793+
}
17141794

17151795
if v, ok := obj[FieldNodeTemplateBareMetal].(string); ok {
17161796
switch v {
@@ -1829,6 +1909,37 @@ func toTemplateConstraintsResourceLimits(o map[string]any) *sdk.NodetemplatesV1T
18291909
return out
18301910
}
18311911

1912+
func toTemplateConstraintsAWSConstraints(o map[string]any) *sdk.NodetemplatesV1TemplateConstraintsAWSConstraints {
1913+
if o == nil {
1914+
return nil
1915+
}
1916+
out := &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{}
1917+
v, ok := o[FieldNodeTemplateCapacityReservations].([]any)
1918+
if !ok || len(v) == 0 {
1919+
return out
1920+
}
1921+
1922+
reservations := lo.FilterMap(v, func(item any, _ int) (sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation, bool) {
1923+
m, ok := item.(map[string]any)
1924+
if !ok {
1925+
return sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{}, false
1926+
}
1927+
r := sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{}
1928+
if id, ok := m[FieldNodeTemplateCapacityReservationId].(string); ok && id != "" {
1929+
r.Id = toPtr(id)
1930+
}
1931+
if arn, ok := m[FieldNodeTemplateCapacityResourceGroupArn].(string); ok && arn != "" {
1932+
r.CapacityResourceGroupArn = toPtr(arn)
1933+
}
1934+
if t, ok := m[FieldNodeTemplateCapacityReservationType].(string); ok && t != "" {
1935+
r.Type = toPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationType(t))
1936+
}
1937+
return r, true
1938+
})
1939+
out.CapacityReservations = &reservations
1940+
return out
1941+
}
1942+
18321943
func toTemplateConstraintsGpuConstraints(o map[string]any) *sdk.NodetemplatesV1TemplateConstraintsGPUConstraints {
18331944
if o == nil {
18341945
return nil

‎castai/resource_node_template_test.go‎

Lines changed: 163 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ constraints.0.spot_interruption_predictions_type = aws-rebalance-recommendations
247247
constraints.0.storage_optimized = false
248248
constraints.0.storage_optimized_state = enabled
249249
constraints.0.use_spot_fallbacks = false
250+
constraints.0.aws.# = 0
250251
constraints.0.bare_metal = unspecified
251252
custom_instances_enabled = true
252253
custom_instances_with_extended_memory_enabled = true
@@ -794,6 +795,10 @@ func TestAccEKS_ResourceNodeTemplate_basic(t *testing.T) {
794795
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.#", "1"),
795796
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.0.cpu_limit_enabled", "false"),
796797
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.0.cpu_limit_max_cores", "0"),
798+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.#", "1"),
799+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.#", "1"),
800+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.0.id", "cr-12345678901234567"),
801+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.0.type", "ON_DEMAND_CAPACITY_RESERVATION"),
797802
resource.TestCheckResourceAttr(resourceName, "edge_location_ids.#", "2"),
798803
resource.TestCheckResourceAttrSet(resourceName, "edge_location_ids.0"),
799804
resource.TestCheckResourceAttrSet(resourceName, "edge_location_ids.1"),
@@ -880,6 +885,10 @@ func TestAccEKS_ResourceNodeTemplate_basic(t *testing.T) {
880885
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.#", "1"),
881886
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.0.cpu_limit_enabled", "true"),
882887
resource.TestCheckResourceAttr(resourceName, "constraints.0.resource_limits.0.cpu_limit_max_cores", "50"),
888+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.#", "1"),
889+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.#", "1"),
890+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.0.id", "cr-98765432109876543"),
891+
resource.TestCheckResourceAttr(resourceName, "constraints.0.aws.0.capacity_reservations.0.type", "ON_DEMAND_CAPACITY_RESERVATION"),
883892
resource.TestCheckResourceAttr(resourceName, "gpu.0.default_shared_clients_per_gpu", "1"),
884893
resource.TestCheckResourceAttr(resourceName, "gpu.0.enable_time_sharing", "false"),
885894
resource.TestCheckResourceAttr(resourceName, "edge_location_ids.#", "1"),
@@ -982,6 +991,13 @@ func testAccNodeTemplateConfig(rName, clusterName string) string {
982991
983992
cpu_manufacturers = ["INTEL", "AMD"]
984993
architecture_priority = ["amd64"]
994+
995+
aws {
996+
capacity_reservations {
997+
id = "cr-12345678901234567"
998+
type = "ON_DEMAND_CAPACITY_RESERVATION"
999+
}
1000+
}
9851001
}
9861002
}
9871003
`, rName))
@@ -1137,6 +1153,13 @@ func testNodeTemplateUpdated(rName, clusterName string) string {
11371153
cpu_limit_enabled = true
11381154
cpu_limit_max_cores = 50
11391155
}
1156+
1157+
aws {
1158+
capacity_reservations {
1159+
id = "cr-98765432109876543"
1160+
type = "ON_DEMAND_CAPACITY_RESERVATION"
1161+
}
1162+
}
11401163
}
11411164
}
11421165
`, rName))
@@ -1368,9 +1391,9 @@ func Test_toTemplateGpu(t *testing.T) {
13681391
FieldNodeTemplateUserManagedGPUDrivers: true,
13691392
},
13701393
want: &sdk.NodetemplatesV1GPU{
1371-
EnableTimeSharing: lo.ToPtr(false),
1372-
SharingConfiguration: &map[string]sdk.NodetemplatesV1SharedGPU{},
1373-
SharingStrategy: lo.ToPtr(sdk.GPUSHARINGSTRATEGYMPS),
1394+
EnableTimeSharing: lo.ToPtr(false),
1395+
SharingConfiguration: &map[string]sdk.NodetemplatesV1SharedGPU{},
1396+
SharingStrategy: lo.ToPtr(sdk.GPUSHARINGSTRATEGYMPS),
13741397
UserManagedGpuDrivers: lo.ToPtr(true),
13751398
},
13761399
},
@@ -1572,3 +1595,140 @@ func Test_flattenGpuSettings(t *testing.T) {
15721595
})
15731596
}
15741597
}
1598+
1599+
func Test_flattenAWSConstraints(t *testing.T) {
1600+
tests := []struct {
1601+
name string
1602+
input *sdk.NodetemplatesV1TemplateConstraintsAWSConstraints
1603+
want []map[string]any
1604+
}{
1605+
{
1606+
name: "nil input returns nil",
1607+
input: nil,
1608+
want: nil,
1609+
},
1610+
{
1611+
name: "empty reservations",
1612+
input: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{},
1613+
want: []map[string]any{{}},
1614+
},
1615+
{
1616+
name: "reservation with id",
1617+
input: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{
1618+
CapacityReservations: &[]sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{
1619+
{
1620+
Id: lo.ToPtr("cr-123"),
1621+
Type: lo.ToPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationTypeONDEMANDCAPACITYRESERVATION),
1622+
},
1623+
},
1624+
},
1625+
want: []map[string]any{{
1626+
FieldNodeTemplateCapacityReservations: []map[string]any{
1627+
{
1628+
FieldNodeTemplateCapacityReservationId: "cr-123",
1629+
FieldNodeTemplateCapacityReservationType: CapacityReservationTypeOnDemand,
1630+
},
1631+
},
1632+
}},
1633+
},
1634+
{
1635+
name: "reservation with ARN",
1636+
input: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{
1637+
CapacityReservations: &[]sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{
1638+
{
1639+
CapacityResourceGroupArn: lo.ToPtr("arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group"),
1640+
Type: lo.ToPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationTypeCAPACITYBLOCK),
1641+
},
1642+
},
1643+
},
1644+
want: []map[string]any{{
1645+
FieldNodeTemplateCapacityReservations: []map[string]any{
1646+
{
1647+
FieldNodeTemplateCapacityResourceGroupArn: "arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group",
1648+
FieldNodeTemplateCapacityReservationType: CapacityReservationTypeCapacityBlock,
1649+
},
1650+
},
1651+
}},
1652+
},
1653+
}
1654+
for _, tt := range tests {
1655+
t.Run(tt.name, func(t *testing.T) {
1656+
got := flattenAWSConstraints(tt.input)
1657+
require.Equal(t, tt.want, got)
1658+
})
1659+
}
1660+
}
1661+
1662+
func Test_toTemplateConstraintsAWSConstraints(t *testing.T) {
1663+
tests := map[string]struct {
1664+
input map[string]any
1665+
want *sdk.NodetemplatesV1TemplateConstraintsAWSConstraints
1666+
}{
1667+
"nil input": {
1668+
input: nil,
1669+
want: nil,
1670+
},
1671+
"reservation with id": {
1672+
input: map[string]any{
1673+
FieldNodeTemplateCapacityReservations: []any{
1674+
map[string]any{
1675+
FieldNodeTemplateCapacityReservationId: "cr-123",
1676+
FieldNodeTemplateCapacityReservationType: CapacityReservationTypeOnDemand,
1677+
},
1678+
},
1679+
},
1680+
want: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{
1681+
CapacityReservations: &[]sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{
1682+
{
1683+
Id: lo.ToPtr("cr-123"),
1684+
Type: lo.ToPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationTypeONDEMANDCAPACITYRESERVATION),
1685+
},
1686+
},
1687+
},
1688+
},
1689+
"reservation with ARN": {
1690+
input: map[string]any{
1691+
FieldNodeTemplateCapacityReservations: []any{
1692+
map[string]any{
1693+
FieldNodeTemplateCapacityResourceGroupArn: "arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group",
1694+
FieldNodeTemplateCapacityReservationType: CapacityReservationTypeCapacityBlock,
1695+
},
1696+
},
1697+
},
1698+
want: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{
1699+
CapacityReservations: &[]sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{
1700+
{
1701+
CapacityResourceGroupArn: lo.ToPtr("arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group"),
1702+
Type: lo.ToPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationTypeCAPACITYBLOCK),
1703+
},
1704+
},
1705+
},
1706+
},
1707+
"reservation with both id and ARN": {
1708+
input: map[string]any{
1709+
FieldNodeTemplateCapacityReservations: []any{
1710+
map[string]any{
1711+
FieldNodeTemplateCapacityReservationId: "cr-456",
1712+
FieldNodeTemplateCapacityResourceGroupArn: "arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group",
1713+
FieldNodeTemplateCapacityReservationType: CapacityReservationTypeCapacityBlock,
1714+
},
1715+
},
1716+
},
1717+
want: &sdk.NodetemplatesV1TemplateConstraintsAWSConstraints{
1718+
CapacityReservations: &[]sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservation{
1719+
{
1720+
Id: lo.ToPtr("cr-456"),
1721+
CapacityResourceGroupArn: lo.ToPtr("arn:aws:ec2:us-east-1:123456789012:capacity-reservation-group/my-group"),
1722+
Type: lo.ToPtr(sdk.NodetemplatesV1TemplateConstraintsAWSConstraintsCapacityReservationTypeCAPACITYBLOCK),
1723+
},
1724+
},
1725+
},
1726+
},
1727+
}
1728+
for name, tt := range tests {
1729+
t.Run(name, func(t *testing.T) {
1730+
got := toTemplateConstraintsAWSConstraints(tt.input)
1731+
require.Equal(t, tt.want, got)
1732+
})
1733+
}
1734+
}

0 commit comments

Comments
 (0)