feat(BRE2-1008): architecture / arm support

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

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{}))
+}

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"},