diff --git a/pkg/recipe/criteria.go b/pkg/recipe/criteria.go index 43a0bf304..01b5c5424 100644 --- a/pkg/recipe/criteria.go +++ b/pkg/recipe/criteria.go @@ -78,6 +78,7 @@ const ( CriteriaAcceleratorAny CriteriaAcceleratorType = "any" CriteriaAcceleratorH100 CriteriaAcceleratorType = "h100" CriteriaAcceleratorGB200 CriteriaAcceleratorType = "gb200" + CriteriaAcceleratorB200 CriteriaAcceleratorType = "b200" CriteriaAcceleratorA100 CriteriaAcceleratorType = "a100" CriteriaAcceleratorL40 CriteriaAcceleratorType = "l40" ) @@ -91,6 +92,8 @@ func ParseCriteriaAcceleratorType(s string) (CriteriaAcceleratorType, error) { return CriteriaAcceleratorH100, nil case "gb200": return CriteriaAcceleratorGB200, nil + case "b200": + return CriteriaAcceleratorB200, nil case "a100": return CriteriaAcceleratorA100, nil case "l40": @@ -102,7 +105,7 @@ func ParseCriteriaAcceleratorType(s string) (CriteriaAcceleratorType, error) { // GetCriteriaAcceleratorTypes returns all supported accelerator types sorted alphabetically. func GetCriteriaAcceleratorTypes() []string { - return []string{"a100", "gb200", "h100", "l40"} + return []string{"a100", "b200", "gb200", "h100", "l40"} } // CriteriaIntentType represents the workload intent. @@ -204,7 +207,7 @@ type Criteria struct { // Service is the Kubernetes service type (eks, gke, aks, oke, self-managed). Service CriteriaServiceType `json:"service,omitempty" yaml:"service,omitempty"` - // Accelerator is the GPU/accelerator type (h100, gb200, a100, l40). + // Accelerator is the GPU/accelerator type (h100, gb200, b200, a100, l40). Accelerator CriteriaAcceleratorType `json:"accelerator,omitempty" yaml:"accelerator,omitempty"` // Intent is the workload intent (training, inference). diff --git a/pkg/recipe/criteria_test.go b/pkg/recipe/criteria_test.go index 75825dd80..18f4286e4 100644 --- a/pkg/recipe/criteria_test.go +++ b/pkg/recipe/criteria_test.go @@ -69,6 +69,7 @@ func TestParseCriteriaAcceleratorType(t *testing.T) { {"h100", "h100", CriteriaAcceleratorH100, false}, {"H100 uppercase", "H100", CriteriaAcceleratorH100, false}, {"gb200", "gb200", CriteriaAcceleratorGB200, false}, + {"b200", "b200", CriteriaAcceleratorB200, false}, {"a100", "a100", CriteriaAcceleratorA100, false}, {"l40", "l40", CriteriaAcceleratorL40, false}, {"invalid", "v100", CriteriaAcceleratorAny, true}, @@ -686,7 +687,7 @@ func TestGetCriteriaAcceleratorTypes(t *testing.T) { types := GetCriteriaAcceleratorTypes() // Should return sorted list - expected := []string{"a100", "gb200", "h100", "l40"} + expected := []string{"a100", "b200", "gb200", "h100", "l40"} if len(types) != len(expected) { t.Errorf("GetCriteriaAcceleratorTypes() returned %d types, want %d", len(types), len(expected)) } diff --git a/pkg/recipe/doc.go b/pkg/recipe/doc.go index c1667cf15..ec0e2f79d 100644 --- a/pkg/recipe/doc.go +++ b/pkg/recipe/doc.go @@ -26,7 +26,7 @@ // // type Criteria struct { // Service CriteriaServiceType // eks, gke, aks, any -// Accelerator CriteriaAcceleratorType // h100, gb200, a100, l40, any +// Accelerator CriteriaAcceleratorType // h100, gb200, b200, a100, l40, any // Intent CriteriaIntentType // training, inference, any // OS CriteriaOSType // ubuntu, cos, rhel, any // Nodes int // node count (0 = any) @@ -68,6 +68,7 @@ // Accelerator types for GPU selection: // - CriteriaAcceleratorH100: NVIDIA H100 // - CriteriaAcceleratorGB200: NVIDIA GB200 +// - CriteriaAcceleratorB200: NVIDIA B200 // - CriteriaAcceleratorA100: NVIDIA A100 // - CriteriaAcceleratorL40: NVIDIA L40 // - CriteriaAcceleratorAny: Any accelerator (wildcard) @@ -115,7 +116,7 @@ // // The HTTP handler accepts these query parameters for GET requests: // - service: eks, gke, aks, any (default: any) -// - accelerator: h100, gb200, a100, l40, any (default: any) +// - accelerator: h100, gb200, b200, a100, l40, any (default: any) // - gpu: alias for accelerator (backwards compatibility) // - intent: training, inference, any (default: any) // - os: ubuntu, cos, rhel, any (default: any) diff --git a/pkg/recipe/snapshot.go b/pkg/recipe/snapshot.go index c2069476f..9cac7c7ef 100644 --- a/pkg/recipe/snapshot.go +++ b/pkg/recipe/snapshot.go @@ -101,6 +101,10 @@ func matchAccelerator(model string) CriteriaAcceleratorType { switch { case strings.Contains(lower, "gb200"): return CriteriaAcceleratorGB200 + // b200 must be checked after gb200 to avoid false-matching GB200 models. + // Follow this pattern when adding future Blackwell variants (e.g., check "gb300" before "b300"). + case strings.Contains(lower, "b200"): + return CriteriaAcceleratorB200 case strings.Contains(lower, "h100"): return CriteriaAcceleratorH100 case strings.Contains(lower, "a100"): diff --git a/pkg/recipe/snapshot_test.go b/pkg/recipe/snapshot_test.go index 35e94d22e..3d0d1a087 100644 --- a/pkg/recipe/snapshot_test.go +++ b/pkg/recipe/snapshot_test.go @@ -196,6 +196,29 @@ func TestExtractCriteriaFromSnapshot(t *testing.T) { } }, }, + { + name: "GPU B200", + snapshot: &snapshotter.Snapshot{ + Measurements: []*measurement.Measurement{ + { + Type: measurement.TypeGPU, + Subtypes: []measurement.Subtype{ + { + Name: "device", + Data: map[string]measurement.Reading{ + "model": measurement.Str("NVIDIA-B200"), + }, + }, + }, + }, + }, + }, + validate: func(t *testing.T, c *Criteria) { + if c.Accelerator != CriteriaAcceleratorB200 { + t.Errorf("Accelerator = %v, want %v", c.Accelerator, CriteriaAcceleratorB200) + } + }, + }, { name: "GPU A100", snapshot: &snapshotter.Snapshot{ @@ -384,6 +407,7 @@ func TestMatchAccelerator(t *testing.T) { {"H100 lowercase", "h100-sxm", CriteriaAcceleratorH100}, {"A100", "A100-SXM4-80GB", CriteriaAcceleratorA100}, {"GB200", "NVIDIA GB200", CriteriaAcceleratorGB200}, + {"B200", "NVIDIA-B200", CriteriaAcceleratorB200}, {"L40", "NVIDIA L40S", CriteriaAcceleratorL40}, {"unknown model", "NVIDIA T4", ""}, {"empty string", "", ""},