Dm/armsupport

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,14 @@ func TestGetGPUMemory(t *testing.T) {
 			gpuType:     "B200",
 			expectedGiB: 192,
 		},
+		{
+			gpuType:     "B300",
+			expectedGiB: 270,
+		},
+		{
+			gpuType:     "RTX6000",
+			expectedGiB: 96,
+		},
 		{
 			gpuType:      "UNKNOWN_GPU",
 			expectedGiB:  0,
@@ -277,6 +285,16 @@ func TestExtractGPUTypeAndName(t *testing.T) {
 			expectedType: "B200",
 			expectedName: "B200",
 		},
+		{
+			platformName: "gpu-b300-sxm",
+			expectedType: "B300",
+			expectedName: "B300",
+		},
+		{
+			platformName: "gpu-rtx6000",
+			expectedType: "RTX6000",
+			expectedName: "RTX6000",
+		},
 		{
 			platformName: "unknown-platform",
 			expectedType: "GPU",
@@ -342,6 +360,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"},

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
@@ -435,17 +436,50 @@ func (c *NebiusClient) getGPUQuotaName(platformName string) string {
 	if strings.Contains(platformLower, "b200") {
 		return "compute.instance.gpu.b200"
 	}
+	if strings.Contains(platformLower, "gb300") {
+		return "compute.instance.gpu.gb300"
+	}
+	if strings.Contains(platformLower, "b300") {
+		return "compute.instance.gpu.b300"
+	}
+	if strings.Contains(platformLower, "rtx6000") {
+		return "compute.instance.gpu.rtx6000"
+	}
 
 	return ""
 }
 
+var nebiusPlatformArchitectures = map[string]v1.Architecture{
+	"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)
 
 	// 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", "b200", "b300", "rtx6000",
+	}
 	for _, gpuType := range knownGPUTypes {
 		if strings.Contains(platformLower, gpuType) {
 			return true
@@ -493,8 +527,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 +545,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,6 +556,15 @@ 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
@@ -561,6 +589,12 @@ func extractGPUTypeAndName(platformName string) (string, string) {
 	if strings.Contains(platformLower, "b200") {
 		return "B200", "B200"
 	}
+	if strings.Contains(platformLower, "b300") {
+		return "B300", "B300"
+	}
+	if strings.Contains(platformLower, "rtx6000") {
+		return "RTX6000", "RTX6000"
+	}
 
 	return "GPU", "GPU" // Generic fallback
 }
@@ -569,15 +603,17 @@ func extractGPUTypeAndName(platformName string) (string, string) {
 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
+		"B300":    270, // 270 GiB VRAM
+		"RTX6000": 96,  // 96 GiB VRAM
 	}
 
 	if vramGiB, exists := vramMap[gpuType]; exists {
@@ -588,17 +624,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 {