diff --git a/v1/providers/nebius/image.go b/v1/providers/nebius/image.go index 966ade1..18c1e02 100644 --- a/v1/providers/nebius/image.go +++ b/v1/providers/nebius/image.go @@ -32,19 +32,7 @@ func (c *NebiusClient) GetImages(ctx context.Context, args v1.GetImageArgs) ([]v } } - // Apply architecture filters - default to x86_64 if no architecture specified - architectures := args.Architectures - if len(architectures) == 0 { - architectures = []string{"x86_64"} // Default to x86_64 - } - images = filterImagesByArchitectures(images, architectures) - - // Apply name filter if specified - if len(args.NameFilters) > 0 { - images = filterImagesByNameFilters(images, args.NameFilters) - } - - return images, nil + return applyImageFilters(images, args), nil } // getProjectImages retrieves images specific to the current project @@ -180,7 +168,7 @@ func (c *NebiusClient) getDefaultImages(ctx context.Context) ([]v1.Image, error) ID: image.Metadata.Id, Name: image.Metadata.Name, Description: getImageDescription(image), - Architecture: "x86_64", + Architecture: extractArchitecture(image), } // Set creation time if available @@ -211,28 +199,45 @@ const ( // extractArchitecture extracts architecture information from image metadata func extractArchitecture(image *compute.Image) string { - // Check labels for architecture info - if image.Metadata != nil && image.Metadata.Labels != nil { - if arch, exists := image.Metadata.Labels["architecture"]; exists { - return arch - } - if arch, exists := image.Metadata.Labels["arch"]; exists { - return arch + if image != nil && image.Spec != nil { + switch image.Spec.CpuArchitecture { + case compute.ImageSpec_AMD64: + return string(v1.ArchitectureX86_64) + case compute.ImageSpec_ARM64: + return string(v1.ArchitectureARM64) + case compute.ImageSpec_UNSPECIFIED: + default: + return string(v1.ArchitectureUnknown) } } - // Infer from image name - if image.Metadata != nil { + if image != nil && image.Metadata != nil { + for _, label := range []string{"architecture", "arch"} { + rawArchitecture, ok := image.Metadata.Labels[label] + if !ok { + continue + } + architecture := v1.GetArchitecture(rawArchitecture) + if architecture != v1.ArchitectureUnknown { + return string(architecture) + } + } + name := strings.ToLower(image.Metadata.Name) if strings.Contains(name, ArchitectureArm64) || strings.Contains(name, ArchitectureAArch64) { - return ArchitectureArm64 + return string(v1.ArchitectureARM64) } if strings.Contains(name, ArchitectureX86_64) || strings.Contains(name, ArchitectureAMD64) { - return ArchitectureX86_64 + return string(v1.ArchitectureX86_64) } } - return ArchitectureX86_64 + return string(v1.ArchitectureUnknown) +} + +func applyImageFilters(images []v1.Image, args v1.GetImageArgs) []v1.Image { + images = filterImagesByArchitectures(images, args.Architectures) + return filterImagesByNameFilters(images, args.NameFilters) } // filterImagesByArchitectures filters images by multiple architectures diff --git a/v1/providers/nebius/image_test.go b/v1/providers/nebius/image_test.go new file mode 100644 index 0000000..2391957 --- /dev/null +++ b/v1/providers/nebius/image_test.go @@ -0,0 +1,61 @@ +package v1 + +import ( + "testing" + + common "github.com/nebius/gosdk/proto/nebius/common/v1" + compute "github.com/nebius/gosdk/proto/nebius/compute/v1" + "github.com/stretchr/testify/require" + + cloudv1 "github.com/brevdev/cloud/v1" +) + +func TestExtractArchitecture(t *testing.T) { + t.Parallel() + + image := func( + architecture compute.ImageSpec_CPUArchitecture, + labels map[string]string, + ) *compute.Image { + return &compute.Image{ + Metadata: &common.ResourceMetadata{Labels: labels}, + Spec: &compute.ImageSpec{CpuArchitecture: architecture}, + } + } + + tests := []struct { + name string + image *compute.Image + want cloudv1.Architecture + }{ + {name: "AMD64 enum", image: image(compute.ImageSpec_AMD64, nil), want: cloudv1.ArchitectureX86_64}, + { + name: "ARM64 enum wins over legacy metadata", + image: image( + compute.ImageSpec_ARM64, + map[string]string{"architecture": "amd64"}, + ), + want: cloudv1.ArchitectureARM64, + }, + {name: "legacy label", image: image(compute.ImageSpec_UNSPECIFIED, map[string]string{"arch": "aarch64"}), want: cloudv1.ArchitectureARM64}, + {name: "unknown enum", image: image(compute.ImageSpec_CPUArchitecture(99), map[string]string{"architecture": "amd64"}), want: cloudv1.ArchitectureUnknown}, + {name: "unknown", image: &compute.Image{}, want: cloudv1.ArchitectureUnknown}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, string(tt.want), extractArchitecture(tt.image)) + }) + } +} + +func TestApplyImageFiltersDoesNotDefaultToX86(t *testing.T) { + t.Parallel() + + x86 := cloudv1.Image{ID: "x86", Architecture: string(cloudv1.ArchitectureX86_64)} + arm := cloudv1.Image{ID: "arm", Architecture: string(cloudv1.ArchitectureARM64)} + images := []cloudv1.Image{x86, arm} + + require.Equal(t, images, applyImageFilters(images, cloudv1.GetImageArgs{})) +} diff --git a/v1/providers/nebius/instance_test.go b/v1/providers/nebius/instance_test.go index 7b05dbc..3a5523e 100644 --- a/v1/providers/nebius/instance_test.go +++ b/v1/providers/nebius/instance_test.go @@ -212,6 +212,16 @@ func TestGetGPUMemory(t *testing.T) { gpuType: "B200", expectedGiB: 192, }, + {gpuType: "GB200", expectedGiB: 192}, + { + gpuType: "B300", + expectedGiB: 270, + }, + {gpuType: "GB300", expectedGiB: 270}, + { + gpuType: "RTX6000", + expectedGiB: 96, + }, { gpuType: "UNKNOWN_GPU", expectedGiB: 0, @@ -277,6 +287,18 @@ func TestExtractGPUTypeAndName(t *testing.T) { expectedType: "B200", expectedName: "B200", }, + {platformName: "gpu-gb200", expectedType: "GB200", expectedName: "GB200"}, + { + platformName: "gpu-b300-sxm", + expectedType: "B300", + expectedName: "B300", + }, + {platformName: "gpu-gb300", expectedType: "GB300", expectedName: "GB300"}, + { + platformName: "gpu-rtx6000", + expectedType: "RTX6000", + expectedName: "RTX6000", + }, { platformName: "unknown-platform", expectedType: "GPU", @@ -300,6 +322,22 @@ func TestExtractGPUTypeAndName(t *testing.T) { } } +func TestGetGPUQuotaName(t *testing.T) { + client := &NebiusClient{} + tests := map[string]string{ + "gpu-gb200": "compute.instance.gpu.gb200", + "gpu-gb300": "compute.instance.gpu.gb300", + "gpu-b300-sxm": "compute.instance.gpu.b300", + "gpu-rtx6000": "compute.instance.gpu.rtx6000", + } + + for platformName, expected := range tests { + t.Run(platformName, func(t *testing.T) { + assert.Equal(t, expected, client.getGPUQuotaName(platformName)) + }) + } +} + func TestGetNebiusBootImageFamily(t *testing.T) { tests := []struct { name string @@ -342,6 +380,8 @@ func TestIsPlatformSupported(t *testing.T) { {"gpu-h100-sxm", true, "H100 with gpu prefix"}, {"gpu-h200-sxm", true, "H200 with gpu prefix"}, {"gpu-b200-sxm", true, "B200 with gpu prefix"}, + {"gpu-b300-sxm", true, "B300 with gpu prefix"}, + {"gpu-rtx6000", true, "RTX 6000 with gpu prefix"}, {"gpu-l40s", true, "L40S with gpu prefix"}, {"gpu-a100-sxm4", true, "A100 with gpu prefix"}, {"gpu-v100-sxm2", true, "V100 with gpu prefix"}, diff --git a/v1/providers/nebius/instancetype.go b/v1/providers/nebius/instancetype.go index fc8c6bd..07e9a48 100644 --- a/v1/providers/nebius/instancetype.go +++ b/v1/providers/nebius/instancetype.go @@ -186,17 +186,18 @@ func (c *NebiusClient) getInstanceTypesForLocation(ctx context.Context, platform // Convert Nebius platform preset to our InstanceType format instanceType := v1.InstanceType{ - Location: location.Name, - Type: instanceTypeID, // Same as ID - both use dot-separated format - VCPU: preset.Resources.VcpuCount, - Memory: units.Base2Bytes(preset.Resources.MemoryGibibytes) * units.Gibibyte, - MemoryBytes: v1.NewBytes(v1.BytesValue(preset.Resources.MemoryGibibytes), v1.Gibibyte), // Memory in GiB - NetworkPerformance: "standard", // Default network performance - IsAvailable: isAvailable, - Stoppable: true, // All Nebius instances support stop/start operations - ElasticRootVolume: true, // Nebius supports dynamic disk allocation - SupportedStorage: c.buildSupportedStorage(), - Provider: CloudProviderID, // Nebius is the provider + Location: location.Name, + Type: instanceTypeID, // Same as ID - both use dot-separated format + VCPU: preset.Resources.VcpuCount, + Memory: units.Base2Bytes(preset.Resources.MemoryGibibytes) * units.Gibibyte, + MemoryBytes: v1.NewBytes(v1.BytesValue(preset.Resources.MemoryGibibytes), v1.Gibibyte), // Memory in GiB + NetworkPerformance: "standard", // Default network performance + IsAvailable: isAvailable, + Stoppable: true, // All Nebius instances support stop/start operations + ElasticRootVolume: true, // Nebius supports dynamic disk allocation + SupportedStorage: c.buildSupportedStorage(), + SupportedArchitectures: []v1.Architecture{nebiusPlatformArchitecture(platform.Metadata.Name)}, + Provider: CloudProviderID, // Nebius is the provider } // Add GPU information if available @@ -415,8 +416,11 @@ func (c *NebiusClient) getGPUQuotaName(platformName string) string { // Nebius GPU quota names follow pattern: "compute.instance.gpu.{type}" // Examples: "compute.instance.gpu.h100", "compute.instance.gpu.h200", "compute.instance.gpu.l40s" - platformLower := strings.ToLower(platformName) + if gpuType := parseBlackwellGPUType(platformName); gpuType != "" { + return "compute.instance.gpu." + strings.ToLower(gpuType) + } + platformLower := strings.ToLower(platformName) if strings.Contains(platformLower, "h100") { return "compute.instance.gpu.h100" } @@ -432,20 +436,49 @@ func (c *NebiusClient) getGPUQuotaName(platformName string) string { if strings.Contains(platformLower, "v100") { return "compute.instance.gpu.v100" } - if strings.Contains(platformLower, "b200") { - return "compute.instance.gpu.b200" + if strings.Contains(platformLower, "rtx6000") { + return "compute.instance.gpu.rtx6000" } return "" } +var nebiusPlatformArchitectures = map[string]v1.Architecture{ + "gpu-gb200": v1.ArchitectureARM64, + "gpu-gb300": v1.ArchitectureARM64, + "gpu-b300-sxm": v1.ArchitectureX86_64, + "gpu-b200-sxm": v1.ArchitectureX86_64, + "gpu-b200-sxm-a": v1.ArchitectureX86_64, + "gpu-h200-sxm": v1.ArchitectureX86_64, + "gpu-h100-sxm": v1.ArchitectureX86_64, + "gpu-rtx6000": v1.ArchitectureX86_64, + "gpu-l40s-a": v1.ArchitectureX86_64, + "gpu-l40s-d": v1.ArchitectureX86_64, + "cpu-d3": v1.ArchitectureX86_64, + "cpu-e2": v1.ArchitectureX86_64, +} + +func nebiusPlatformArchitecture(platformName string) v1.Architecture { + architecture, ok := nebiusPlatformArchitectures[strings.ToLower(strings.TrimSpace(platformName))] + if !ok { + return v1.ArchitectureUnknown + } + return architecture +} + // isPlatformSupported checks if a platform should be included in instance types func (c *NebiusClient) isPlatformSupported(platformName string) bool { platformLower := strings.ToLower(platformName) + if parseBlackwellGPUType(platformName) != "" { + return true + } + // For GPU platforms: only accept known GPU types // Check for specific GPU model names (with or without "gpu-" prefix) - knownGPUTypes := []string{"h100", "h200", "l40s", "a100", "v100", "a10", "t4", "l4", "b200"} + knownGPUTypes := []string{ + "h100", "h200", "l40s", "a100", "v100", "a10", "t4", "l4", "rtx6000", + } for _, gpuType := range knownGPUTypes { if strings.Contains(platformLower, gpuType) { return true @@ -493,8 +526,6 @@ func (c *NebiusClient) buildSupportedStorage() []v1.Storage { } // applyInstanceTypeFilters applies various filters to the instance type list -// -//nolint:gocognit // Complex function with multiple filter conditions for instance types func (c *NebiusClient) applyInstanceTypeFilters(instanceTypes []v1.InstanceType, args v1.GetInstanceTypeArgs) []v1.InstanceType { var filtered []v1.InstanceType @@ -513,22 +544,9 @@ func (c *NebiusClient) applyInstanceTypeFilters(instanceTypes []v1.InstanceType, } } - // Apply architecture filter - if args.ArchitectureFilter != nil { - arch := determineInstanceTypeArchitecture(instanceType) - // Check if architecture matches the filter requirements - if len(args.ArchitectureFilter.IncludeArchitectures) > 0 { - found := false - for _, allowedArch := range args.ArchitectureFilter.IncludeArchitectures { - if arch == string(allowedArch) { - found = true - break - } - } - if !found { - continue - } - } + if args.ArchitectureFilter != nil && + !supportsAllowedArchitecture(instanceType, args.ArchitectureFilter) { + continue } filtered = append(filtered, instanceType) @@ -537,10 +555,23 @@ func (c *NebiusClient) applyInstanceTypeFilters(instanceTypes []v1.InstanceType, return filtered } +func supportsAllowedArchitecture(instanceType v1.InstanceType, filter *v1.ArchitectureFilter) bool { + for _, architecture := range instanceType.SupportedArchitectures { + if filter.IsAllowed(architecture) { + return true + } + } + return false +} + // extractGPUTypeAndName extracts GPU type and name from platform name // Note: Returns model name only (e.g., "H100"), not full name with manufacturer // Manufacturer info is stored separately in GPU.Manufacturer field func extractGPUTypeAndName(platformName string) (string, string) { + if gpuType := parseBlackwellGPUType(platformName); gpuType != "" { + return gpuType, gpuType + } + platformLower := strings.ToLower(platformName) if strings.Contains(platformLower, "h100") { @@ -558,26 +589,49 @@ func extractGPUTypeAndName(platformName string) (string, string) { if strings.Contains(platformLower, "v100") { return "V100", "V100" } - if strings.Contains(platformLower, "b200") { - return "B200", "B200" + if strings.Contains(platformLower, "rtx6000") { + return "RTX6000", "RTX6000" } return "GPU", "GPU" // Generic fallback } +func parseBlackwellGPUType(platformName string) string { + platformLower := strings.ToLower(platformName) + + // Keep Grace Blackwell variants before their matching B-series variants: + // "gb300" contains "b300", and "gb200" contains "b200". + switch { + case strings.Contains(platformLower, "gb300"): + return "GB300" + case strings.Contains(platformLower, "gb200"): + return "GB200" + case strings.Contains(platformLower, "b300"): + return "B300" + case strings.Contains(platformLower, "b200"): + return "B200" + default: + return "" + } +} + // getGPUMemory returns the VRAM for a given GPU type in GiB func getGPUMemory(gpuType string) units.Base2Bytes { // Static mapping of GPU types to their VRAM capacities vramMap := map[string]int64{ - "L40S": 48, // 48 GiB VRAM - "H100": 80, // 80 GiB VRAM - "H200": 141, // 141 GiB VRAM - "A100": 80, // 80 GiB VRAM (most common variant) - "V100": 32, // 32 GiB VRAM (most common variant) - "A10": 24, // 24 GiB VRAM - "T4": 16, // 16 GiB VRAM - "L4": 24, // 24 GiB VRAM - "B200": 192, // 192 GiB VRAM + "L40S": 48, // 48 GiB VRAM + "H100": 80, // 80 GiB VRAM + "H200": 141, // 141 GiB VRAM + "A100": 80, // 80 GiB VRAM (most common variant) + "V100": 32, // 32 GiB VRAM (most common variant) + "A10": 24, // 24 GiB VRAM + "T4": 16, // 16 GiB VRAM + "L4": 24, // 24 GiB VRAM + "B200": 192, // 192 GiB VRAM + "GB200": 192, // 192 GiB VRAM + "B300": 270, // 270 GiB VRAM + "GB300": 270, // 270 GiB VRAM + "RTX6000": 96, // 96 GiB VRAM } if vramGiB, exists := vramMap[gpuType]; exists { @@ -588,17 +642,6 @@ func getGPUMemory(gpuType string) units.Base2Bytes { return units.Base2Bytes(0) } -// determineInstanceTypeArchitecture determines architecture from instance type -func determineInstanceTypeArchitecture(instanceType v1.InstanceType) string { - // Check if ARM architecture is indicated in the type or name - typeLower := strings.ToLower(instanceType.Type) - if strings.Contains(typeLower, "arm") || strings.Contains(typeLower, "aarch64") { - return "arm64" - } - - return "x86_64" // Default assumption -} - // getPricingForInstanceType fetches real pricing from Nebius Billing Calculator API // Returns nil if pricing cannot be fetched (non-critical failure) func (c *NebiusClient) getPricingForInstanceType(ctx context.Context, platformName, presetName, _ string) *currency.Amount { diff --git a/v1/providers/nebius/instancetype_test.go b/v1/providers/nebius/instancetype_test.go new file mode 100644 index 0000000..4a7915f --- /dev/null +++ b/v1/providers/nebius/instancetype_test.go @@ -0,0 +1,52 @@ +package v1 + +import ( + "testing" + + "github.com/stretchr/testify/require" + + cloudv1 "github.com/brevdev/cloud/v1" +) + +func TestNebiusPlatformArchitecture(t *testing.T) { + t.Parallel() + + require.Equal(t, cloudv1.ArchitectureX86_64, nebiusPlatformArchitecture(" GPU-H100-SXM ")) + require.Equal(t, cloudv1.ArchitectureX86_64, nebiusPlatformArchitecture("cpu-d3")) + require.Equal(t, cloudv1.ArchitectureARM64, nebiusPlatformArchitecture("gpu-gb200")) + require.Equal(t, cloudv1.ArchitectureARM64, nebiusPlatformArchitecture("gpu-gb300")) + require.Equal(t, cloudv1.ArchitectureUnknown, nebiusPlatformArchitecture("future-platform")) +} + +func TestApplyInstanceTypeFiltersUsesArchitectureMetadata(t *testing.T) { + t.Parallel() + + instanceType := func(name string, architecture cloudv1.Architecture) cloudv1.InstanceType { + return cloudv1.InstanceType{ + Type: name, + SupportedArchitectures: []cloudv1.Architecture{architecture}, + } + } + x86 := instanceType("arm-in-name", cloudv1.ArchitectureX86_64) + arm := instanceType("x86-in-name", cloudv1.ArchitectureARM64) + unknown := instanceType("unknown", cloudv1.ArchitectureUnknown) + instanceTypes := []cloudv1.InstanceType{x86, arm, unknown} + client := &NebiusClient{} + apply := func(filter *cloudv1.ArchitectureFilter) []cloudv1.InstanceType { + return client.applyInstanceTypeFilters( + instanceTypes, + cloudv1.GetInstanceTypeArgs{ArchitectureFilter: filter}, + ) + } + + require.Equal(t, []cloudv1.InstanceType{arm}, apply( + &cloudv1.ArchitectureFilter{ + IncludeArchitectures: []cloudv1.Architecture{cloudv1.ArchitectureARM64}, + }, + )) + require.Equal(t, []cloudv1.InstanceType{x86, unknown}, apply( + &cloudv1.ArchitectureFilter{ + ExcludeArchitectures: []cloudv1.Architecture{cloudv1.ArchitectureARM64}, + }, + )) +} diff --git a/v1/providers/nebius/integration_test.go b/v1/providers/nebius/integration_test.go index 372dbe4..ae2dec5 100644 --- a/v1/providers/nebius/integration_test.go +++ b/v1/providers/nebius/integration_test.go @@ -234,15 +234,29 @@ func TestIntegration_InstanceLifecycle(t *testing.T) { // Step 0: Get available instance types to find one we can use t.Log("Discovering available instance types...") - instanceTypes, err := client.GetInstanceTypes(ctx, v1.GetInstanceTypeArgs{}) + instanceTypes, err := client.GetInstanceTypes(ctx, v1.GetInstanceTypeArgs{ + ArchitectureFilter: &v1.ArchitectureFilter{ + IncludeArchitectures: []v1.Architecture{v1.ArchitectureX86_64}, + }, + }) require.NoError(t, err, "Failed to get instance types") if len(instanceTypes) == 0 { - t.Skip("No instance types available - skipping instance lifecycle test") + t.Skip("No x86_64 instance types available - skipping instance lifecycle test") + } + + selectedInstanceTypeIndex := -1 + for i := range instanceTypes { + if instanceTypes[i].IsAvailable { + selectedInstanceTypeIndex = i + break + } + } + if selectedInstanceTypeIndex == -1 { + t.Skip("No available x86_64 instance types - skipping instance lifecycle test") } - // Use the first available instance type (should have quota) - selectedInstanceType := instanceTypes[0] + selectedInstanceType := instanceTypes[selectedInstanceTypeIndex] t.Logf("Using instance type: %s (Location: %s)", selectedInstanceType.ID, selectedInstanceType.Location) // Step 0.5: Generate SSH key pair for testing (inspired by Shadeform's SSH key handling) @@ -256,11 +270,11 @@ func TestIntegration_InstanceLifecycle(t *testing.T) { createAttrs := v1.CreateInstanceAttrs{ RefID: instanceRefID, Name: instanceName, - InstanceType: string(selectedInstanceType.ID), // Use discovered instance type - ImageID: "ubuntu22.04-cuda12", // Use known-good Nebius image family - DiskSize: 50 * 1024 * 1024 * 1024, // 50 GiB in bytes - Location: selectedInstanceType.Location, // Use the instance type's location - PublicKey: publicKey, // SSH public key for access (like Shadeform) + InstanceType: selectedInstanceType.Type, // Native {platform}.{preset} format + ImageID: "ubuntu22.04-cuda12", // Use known-good Nebius image family + DiskSize: 50 * 1024 * 1024 * 1024, // 50 GiB in bytes + Location: selectedInstanceType.Location, // Use the instance type's location + PublicKey: publicKey, // SSH public key for access (like Shadeform) Tags: map[string]string{ "test": "integration", "created-by": "nebius-integration-test", @@ -565,7 +579,7 @@ func TestIntegration_GetInstanceTypes(t *testing.T) { priceStr := it.BasePrice.Number() var priceFloat float64 if _, err := fmt.Sscanf(priceStr, "%f", &priceFloat); err == nil { - assert.Greater(t, priceFloat, 0.0, "Price should be positive") + assert.GreaterOrEqual(t, priceFloat, 0.0, "Price should not be negative") assert.Less(t, priceFloat, 1000.0, "Price per hour should be reasonable (< $1000/hr)") } } else { diff --git a/v1/providers/testkube/README.md b/v1/providers/testkube/README.md index 15b0c46..3ed1c46 100644 --- a/v1/providers/testkube/README.md +++ b/v1/providers/testkube/README.md @@ -31,8 +31,8 @@ brew install minikube kubectl minikube start --driver=docker --profile testkube kubectl config use-context testkube -docker build -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest ./v1/providers/testkube/images/ubuntu-vm -minikube --profile testkube image load ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest +docker build -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 ./v1/providers/testkube/images/ubuntu-vm +minikube --profile testkube image load ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 kubectl create namespace testkube # In another terminal, keep this running while validation runs. @@ -52,4 +52,4 @@ Clean up: ```bash minikube --profile testkube delete -``` \ No newline at end of file +``` diff --git a/v1/providers/testkube/images/ubuntu-vm/README.md b/v1/providers/testkube/images/ubuntu-vm/README.md index 5a84978..97ef6e1 100644 --- a/v1/providers/testkube/images/ubuntu-vm/README.md +++ b/v1/providers/testkube/images/ubuntu-vm/README.md @@ -1,9 +1,10 @@ # Testkube Ubuntu VM Image -This image backs the `test.ok.cpu` testkube instance type: +This image backs the architecture-specific TestKube instance types: ```text -ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest +test.ok.cpu -> ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 +test.ok.cpu.arm64 -> ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 ``` ## Publish to GHCR @@ -16,22 +17,12 @@ gh auth refresh -h github.com -s write:packages gh auth token | docker login ghcr.io -u "$(gh api user --jq .login)" --password-stdin ``` -Build and push the image from the repository root. For EKS, publish an amd64 image because `test.ok.cpu` advertises `x86_64`: - -```bash -docker buildx build \ - --platform linux/amd64 \ - -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest \ - --push \ - ./v1/providers/testkube/images/ubuntu-vm -``` - -If you need both local Apple Silicon clusters and amd64 EKS nodes to pull the same tag, publish a multi-arch manifest: +Build and push the image from the repository root. For EKS, publish the versioned multi-arch manifest used by both `test.ok.cpu` and `test.ok.cpu.arm64`: ```bash docker buildx build \ --platform linux/amd64,linux/arm64 \ - -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest \ + -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 \ --push \ ./v1/providers/testkube/images/ubuntu-vm ``` @@ -50,6 +41,6 @@ For local minikube or kind validation where the image is loaded directly into th ```bash docker build \ - -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest \ + -t ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2 \ ./v1/providers/testkube/images/ubuntu-vm ``` diff --git a/v1/providers/testkube/instance.go b/v1/providers/testkube/instance.go index 9d981bb..fd4c754 100644 --- a/v1/providers/testkube/instance.go +++ b/v1/providers/testkube/instance.go @@ -133,6 +133,7 @@ func (c *TestKubeClient) createInstanceAsK8sResources(ctx context.Context, attrs Spec: corev1.PodSpec{ RestartPolicy: corev1.RestartPolicyNever, TerminationGracePeriodSeconds: int64Ptr(1), + NodeSelector: maps.Clone(instanceTypeSpec.nodeSelector), Containers: []corev1.Container{ { Name: "vm", diff --git a/v1/providers/testkube/instance_test.go b/v1/providers/testkube/instance_test.go index 10aab7c..906e87c 100644 --- a/v1/providers/testkube/instance_test.go +++ b/v1/providers/testkube/instance_test.go @@ -176,6 +176,66 @@ func TestInstanceUsesBakedImageSpec(t *testing.T) { } } +func TestARM64InstanceUsesVersionedMultiarchImage(t *testing.T) { + ctx := context.Background() + client := newTestClient(t) + + instance, err := client.CreateInstance(ctx, cloudv1.CreateInstanceAttrs{ + RefID: "arm64-image-spec", + Name: "arm64 image spec", + InstanceType: InstanceTypeOKCPUARM64, + }) + require.NoError(t, err) + require.Equal(t, "testkube-ubuntu-vm-arm64", instance.ImageID) + + pod, err := client.k8sClient.CoreV1().Pods(client.namespace).Get(ctx, string(instance.CloudID), metav1.GetOptions{}) + require.NoError(t, err) + require.Equal(t, "ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2", pod.Spec.Containers[0].Image) +} + +func TestInstanceUsesArchitectureNodeSelector(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + instanceType string + nodeArchitecture string + }{ + { + name: "x86_64", + instanceType: InstanceTypeOKCPU, + nodeArchitecture: "amd64", + }, + { + name: "arm64", + instanceType: InstanceTypeOKCPUARM64, + nodeArchitecture: "arm64", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + client := newTestClient(t) + instance, err := client.CreateInstance(ctx, cloudv1.CreateInstanceAttrs{ + RefID: "node-selector-" + tt.name, + Name: "node selector " + tt.name, + InstanceType: tt.instanceType, + }) + require.NoError(t, err) + + pod, err := client.k8sClient.CoreV1().Pods(client.namespace).Get( + ctx, + string(instance.CloudID), + metav1.GetOptions{}, + ) + require.NoError(t, err) + require.Equal(t, map[string]string{ + corev1.LabelArchStable: tt.nodeArchitecture, + }, pod.Spec.NodeSelector) + }) + } +} + func TestPopulateNetworkLoadBalancer(t *testing.T) { instance := &cloudv1.Instance{} populateNetwork(&corev1.Service{ diff --git a/v1/providers/testkube/instancetype.go b/v1/providers/testkube/instancetype.go index 152f35e..73b9c3a 100644 --- a/v1/providers/testkube/instancetype.go +++ b/v1/providers/testkube/instancetype.go @@ -13,11 +13,13 @@ import ( const ( DefaultImageID = "testkube-ubuntu-vm" - DefaultImage = "ghcr.io/brevdev/cloud/testkube-ubuntu-vm:latest" + DefaultImage = "ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2" + ARM64ImageID = "testkube-ubuntu-vm-arm64" DefaultPriceCentsPerHour = 1 InstanceTypeOKCPU = "test.ok.cpu" + InstanceTypeOKCPUARM64 = "test.ok.cpu.arm64" InstanceTypeFailCapacity = "test.fail.capacity" InstanceTypeFailQuota = "test.fail.quota" InstanceTypeFailBuild = "test.fail.build" // TODO: trigger build failure, maybe with a process that monitors build? @@ -29,23 +31,36 @@ type instanceTypeSpec struct { instanceType cloudv1.InstanceType imageID string image string + nodeSelector map[string]string serviceType corev1.ServiceType } var allInstanceTypeSpecs = []instanceTypeSpec{ - makeInstanceTypeSpec(InstanceTypeOKCPU), - makeInstanceTypeSpec(InstanceTypeFailCapacity), - makeInstanceTypeSpec(InstanceTypeFailQuota), - makeInstanceTypeSpec(InstanceTypeFailBuild), + makeInstanceTypeSpec(InstanceTypeOKCPU, cloudv1.ArchitectureX86_64, DefaultImageID), + makeInstanceTypeSpec(InstanceTypeOKCPUARM64, cloudv1.ArchitectureARM64, ARM64ImageID), + makeInstanceTypeSpec(InstanceTypeFailCapacity, cloudv1.ArchitectureX86_64, DefaultImageID), + makeInstanceTypeSpec(InstanceTypeFailQuota, cloudv1.ArchitectureX86_64, DefaultImageID), + makeInstanceTypeSpec(InstanceTypeFailBuild, cloudv1.ArchitectureX86_64, DefaultImageID), } -func makeInstanceTypeSpec(instanceType string) instanceTypeSpec { +func makeInstanceTypeSpec( + instanceType string, + architecture cloudv1.Architecture, + imageID string, +) instanceTypeSpec { estimatedDeployTime := 20 * time.Second + nodeArchitecture := string(architecture) + if architecture == cloudv1.ArchitectureX86_64 { + nodeArchitecture = "amd64" + } return instanceTypeSpec{ - instanceType: makeCPUInstanceType(instanceType, true, &estimatedDeployTime), - imageID: DefaultImageID, + instanceType: makeCPUInstanceType(instanceType, architecture, true, &estimatedDeployTime), + imageID: imageID, image: DefaultImage, - serviceType: corev1.ServiceTypeLoadBalancer, + nodeSelector: map[string]string{ + corev1.LabelArchStable: nodeArchitecture, + }, + serviceType: corev1.ServiceTypeLoadBalancer, } } @@ -118,7 +133,12 @@ func (c *TestKubeClient) instanceTypeSpecToBrevInstanceType(spec instanceTypeSpe return instanceType } -func makeCPUInstanceType(instanceType string, available bool, estimatedDeployTime *time.Duration) cloudv1.InstanceType { +func makeCPUInstanceType( + instanceType string, + architecture cloudv1.Architecture, + available bool, + estimatedDeployTime *time.Duration, +) cloudv1.InstanceType { basePrice, _ := currency.NewAmountFromInt64(DefaultPriceCentsPerHour, "USD") it := cloudv1.InstanceType{ Type: instanceType, @@ -139,7 +159,7 @@ func makeCPUInstanceType(instanceType string, available bool, estimatedDeployTim DefaultCores: 2, VCPU: 2, SupportedArchitectures: []cloudv1.Architecture{ - cloudv1.ArchitectureX86_64, + architecture, }, Stoppable: false, Rebootable: false, diff --git a/v1/providers/testkube/instancetype_test.go b/v1/providers/testkube/instancetype_test.go index 0762a81..5a5a616 100644 --- a/v1/providers/testkube/instancetype_test.go +++ b/v1/providers/testkube/instancetype_test.go @@ -14,7 +14,7 @@ func TestGetInstanceTypes(t *testing.T) { instanceTypes, err := client.GetInstanceTypes(context.Background(), cloudv1.GetInstanceTypeArgs{}) require.NoError(t, err) - require.Len(t, instanceTypes, 4) + require.Len(t, instanceTypes, 5) instanceTypeByName := map[string]cloudv1.InstanceType{} for _, instanceType := range instanceTypes { @@ -26,6 +26,7 @@ func TestGetInstanceTypes(t *testing.T) { InstanceTypeFailCapacity, InstanceTypeFailQuota, InstanceTypeFailBuild, + InstanceTypeOKCPUARM64, } { instanceType, ok := instanceTypeByName[expected] require.True(t, ok, "missing instance type %s", expected) @@ -40,6 +41,49 @@ func TestGetInstanceTypes(t *testing.T) { } } +func TestGetInstanceTypesFiltersByArchitecture(t *testing.T) { + client := newTestClient(t) + + tests := []struct { + name string + architecture cloudv1.Architecture + expected []string + }{ + { + name: "x86_64", + architecture: cloudv1.ArchitectureX86_64, + expected: []string{ + InstanceTypeOKCPU, + InstanceTypeFailCapacity, + InstanceTypeFailQuota, + InstanceTypeFailBuild, + }, + }, + { + name: "arm64", + architecture: cloudv1.ArchitectureARM64, + expected: []string{InstanceTypeOKCPUARM64}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + instanceTypes, err := client.GetInstanceTypes(context.Background(), cloudv1.GetInstanceTypeArgs{ + ArchitectureFilter: &cloudv1.ArchitectureFilter{ + IncludeArchitectures: []cloudv1.Architecture{tt.architecture}, + }, + }) + require.NoError(t, err) + + actual := make([]string, 0, len(instanceTypes)) + for _, instanceType := range instanceTypes { + actual = append(actual, instanceType.Type) + } + require.ElementsMatch(t, tt.expected, actual) + }) + } +} + func TestGetInstanceTypesWithGPUManufacturerFilterIncludesCPU(t *testing.T) { client := newTestClient(t) @@ -49,7 +93,7 @@ func TestGetInstanceTypesWithGPUManufacturerFilterIncludesCPU(t *testing.T) { }, }) require.NoError(t, err) - require.Len(t, instanceTypes, 4) + require.Len(t, instanceTypes, 5) } func TestCapabilitiesDoNotAdvertiseImages(t *testing.T) {