Change MachineClass capabilities to support arbitrary resources

broker/machinebroker/server/status.go

@@ -80,15 +80,18 @@ func (s *Server) filterIronCoreMachineClasses(
 }
 
 func (s *Server) convertIronCoreMachineClassStatus(machineClass *computev1alpha1.MachineClass, quantity *resource.Quantity) *iri.MachineClassStatus {
-	cpu := machineClass.Capabilities.CPU()
-	memory := machineClass.Capabilities.Memory()
+	resources := map[string]int64{}
+	resourceList := machineClass.Capabilities
+
+	for resource, quantity := range resourceList {
+		resources[string(resource)] = quantity.Value()
+	}
 
 	return &iri.MachineClassStatus{
 		MachineClass: &iri.MachineClass{
 			Name: machineClass.Name,
 			Capabilities: &iri.MachineClassCapabilities{
-				CpuMillis:   cpu.MilliValue(),
-				MemoryBytes: memory.Value(),
+				Resources: resources,
 			},
 		},
 		Quantity: quantity.Value(),

internal/controllers/compute/machine_scheduler_test.go

@@ -434,4 +434,75 @@ var _ = Describe("MachineScheduler", func() {
 			HaveField("Spec.MachinePoolRef", Equal(&corev1.LocalObjectReference{Name: machinePool.Name})),
 		))
 	})
+
+	It("should schedule machine on pool with correctly allocatable resources", func(ctx SpecContext) {
+		By("creating a machine pool")
+		machinePool := &computev1alpha1.MachinePool{
+			ObjectMeta: metav1.ObjectMeta{
+				GenerateName: "test-pool-",
+			},
+		}
+		Expect(k8sClient.Create(ctx, machinePool)).To(Succeed(), "failed to create machine pool")
+
+		By("patching the machine pool status to contain a machine class")
+		Eventually(UpdateStatus(machinePool, func() {
+			machinePool.Status.AvailableMachineClasses = []corev1.LocalObjectReference{{Name: machineClass.Name}}
+			machinePool.Status.Allocatable = corev1alpha1.ResourceList{
+				corev1alpha1.ClassCountFor(corev1alpha1.ClassTypeMachineClass, machineClass.Name): resource.MustParse("5"),
+			}
+		})).Should(Succeed())
+
+		By("creating a second machine pool")
+		secondMachinePool := &computev1alpha1.MachinePool{
+			ObjectMeta: metav1.ObjectMeta{
+				GenerateName: "second-test-pool-",
+			},
+		}
+		Expect(k8sClient.Create(ctx, secondMachinePool)).To(Succeed(), "failed to create the second machine pool")
+
+		By("creating a second machine class")
+		secondMachineClass := &computev1alpha1.MachineClass{
+			ObjectMeta: metav1.ObjectMeta{
+				GenerateName: "second-machine-class-",
+			},
+			Capabilities: corev1alpha1.ResourceList{
+				corev1alpha1.ResourceCPU:    resource.MustParse("1"),
+				corev1alpha1.ResourceMemory: resource.MustParse("1Gi"),
+				"type-a.vendor.com/gpu":     resource.MustParse("1"),
+			},
+		}
+		Expect(k8sClient.Create(ctx, secondMachineClass)).To(Succeed(), "failed to create second machine class")
+
+		By("patching the second machine pool status to contain a both machine classes")
+		Eventually(UpdateStatus(secondMachinePool, func() {
+			secondMachinePool.Status.AvailableMachineClasses = []corev1.LocalObjectReference{
+				{Name: machineClass.Name},
+				{Name: secondMachineClass.Name},
+			}
+			secondMachinePool.Status.Allocatable = corev1alpha1.ResourceList{
+				corev1alpha1.ClassCountFor(corev1alpha1.ClassTypeMachineClass, machineClass.Name):       resource.MustParse("5"),
+				corev1alpha1.ClassCountFor(corev1alpha1.ClassTypeMachineClass, secondMachineClass.Name): resource.MustParse("5"),
+			}
+		})).Should(Succeed())
+
+		By("creating a machine")
+		machine := &computev1alpha1.Machine{
+			ObjectMeta: metav1.ObjectMeta{
+				Namespace:    ns.Name,
+				GenerateName: "test-machine-",
+			},
+			Spec: computev1alpha1.MachineSpec{
+				Image: "my-image",
+				MachineClassRef: corev1.LocalObjectReference{
+					Name: secondMachineClass.Name,
+				},
+			},
+		}
+		Expect(k8sClient.Create(ctx, machine)).To(Succeed(), "failed to create the machine")
+
+		By("checking that the machine is scheduled onto the machine pool")
+		Eventually(Object(machine)).Should(SatisfyAll(
+			HaveField("Spec.MachinePoolRef.Name", Equal(secondMachinePool.Name)),
+		))
+	})
 })

iri/apis/machine/v1alpha1/api.proto

@@ -46,8 +46,7 @@ message EventFilter {
 }
 
 message MachineClassCapabilities {
-  int64 cpu_millis = 1;
-  int64 memory_bytes = 2;
+  map<string, int64> resources = 1;
 }
 
 message Machine {

irictl-machine/tableconverters/status.go

@@ -4,6 +4,7 @@
 package tableconverters
 
 import (
+	corev1alpha1 "github.com/ironcore-dev/ironcore/api/core/v1alpha1"
 	iri "github.com/ironcore-dev/ironcore/iri/apis/machine/v1alpha1"
 	"github.com/ironcore-dev/ironcore/irictl/api"
 	"github.com/ironcore-dev/ironcore/irictl/tableconverter"
@@ -23,8 +24,8 @@ var (
 		Rows: tableconverter.SingleRowFrom(func(status *iri.MachineClassStatus) (api.Row, error) {
 			return api.Row{
 				status.MachineClass.Name,
-				resource.NewMilliQuantity(status.MachineClass.Capabilities.CpuMillis, resource.DecimalSI).String(),
-				resource.NewQuantity(status.MachineClass.Capabilities.MemoryBytes, resource.DecimalSI).String(),
+				resource.NewMilliQuantity(status.MachineClass.Capabilities.Resources[string(corev1alpha1.ResourceCPU)], resource.DecimalSI).String(),
+				resource.NewQuantity(status.MachineClass.Capabilities.Resources[string(corev1alpha1.ResourceMemory)], resource.DecimalSI).String(),
 				resource.NewQuantity(status.Quantity, resource.DecimalSI).String(),
 			}, nil
 		}),

poollet/machinepoollet/controllers/controllers_suite_test.go

@@ -189,8 +189,10 @@ func SetupTest() (*corev1.Namespace, *computev1alpha1.MachinePool, *computev1alp
 					MachineClass: &iri.MachineClass{
 						Name: mc.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   mc.Capabilities.CPU().MilliValue(),
-							MemoryBytes: mc.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    mc.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): mc.Capabilities.Memory().Value(),
+							},
 						},
 					},
 				},

poollet/machinepoollet/controllers/machine_controller.go

@@ -717,12 +717,15 @@ func (r *MachineReconciler) prepareIRIMachineClass(ctx context.Context, machine
 }
 
 func getIRIMachineClassCapabilities(machineClass *computev1alpha1.MachineClass) *iri.MachineClassCapabilities {
-	cpu := machineClass.Capabilities.CPU()
-	memory := machineClass.Capabilities.Memory()
+	resources := map[string]int64{}
+	resourceList := machineClass.Capabilities
+
+	for resource, quantity := range resourceList {
+		resources[string(resource)] = quantity.Value()
+	}
 
 	return &iri.MachineClassCapabilities{
-		CpuMillis:   cpu.MilliValue(),
-		MemoryBytes: memory.Value(),
+		Resources: resources,
 	}
 }
 

poollet/machinepoollet/controllers/machinepool_controller_test.go

@@ -43,8 +43,10 @@ var _ = Describe("MachinePoolController", func() {
 					MachineClass: &iri.MachineClass{
 						Name: machineClass.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   machineClass.Capabilities.CPU().MilliValue(),
-							MemoryBytes: machineClass.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    machineClass.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): machineClass.Capabilities.Memory().Value(),
+							},
 						},
 					},
 					Quantity: machineClassCapacity,
@@ -55,8 +57,10 @@ var _ = Describe("MachinePoolController", func() {
 					MachineClass: &iri.MachineClass{
 						Name: machineClass2.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   machineClass2.Capabilities.CPU().MilliValue(),
-							MemoryBytes: machineClass2.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    machineClass2.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): machineClass2.Capabilities.Memory().Value(),
+							},
 						},
 					},
 					Quantity: machineClass2Capacity,
@@ -140,6 +144,7 @@ var _ = Describe("MachinePoolController", func() {
 			Capabilities: corev1alpha1.ResourceList{
 				corev1alpha1.ResourceCPU:    resource.MustParse("2"),
 				corev1alpha1.ResourceMemory: resource.MustParse("2Gi"),
+				"type-a.vendor.com/gpu":     resource.MustParse("1"),
 			},
 		}
 		Expect(k8sClient.Create(ctx, machineClass)).To(Succeed(), "failed to create test machine class")
@@ -150,8 +155,11 @@ var _ = Describe("MachinePoolController", func() {
 					MachineClass: &iri.MachineClass{
 						Name: machineClass.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   machineClass.Capabilities.CPU().MilliValue(),
-							MemoryBytes: machineClass.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    machineClass.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): machineClass.Capabilities.Memory().Value(),
+								"type-a.vendor.com/gpu":             1,
+							},
 						},
 					},
 				},
@@ -189,8 +197,11 @@ var _ = Describe("MachinePoolController", func() {
 					MachineClass: &iri.MachineClass{
 						Name: machineClass.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   machineClass.Capabilities.CPU().MilliValue(),
-							MemoryBytes: machineClass.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    machineClass.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): machineClass.Capabilities.Memory().Value(),
+								"type-a.vendor.com/gpu":             1,
+							},
 						},
 					},
 				},
@@ -200,8 +211,10 @@ var _ = Describe("MachinePoolController", func() {
 					MachineClass: &iri.MachineClass{
 						Name: machineClass2.Name,
 						Capabilities: &iri.MachineClassCapabilities{
-							CpuMillis:   machineClass2.Capabilities.CPU().MilliValue(),
-							MemoryBytes: machineClass2.Capabilities.Memory().Value(),
+							Resources: map[string]int64{
+								string(corev1alpha1.ResourceCPU):    machineClass2.Capabilities.CPU().Value(),
+								string(corev1alpha1.ResourceMemory): machineClass2.Capabilities.Memory().Value(),
+							},
 						},
 					},
 				},

poollet/machinepoollet/mcm/generic.go

@@ -6,6 +6,7 @@ package mcm
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"sync"
 	"time"
 
@@ -20,16 +21,21 @@ import (
 	ctrl "sigs.k8s.io/controller-runtime"
 )
 
-type capabilities struct {
-	cpuMillis   int64
-	memoryBytes int64
+type capabilities map[string]int64
+
+type machineClassQuantity struct {
+	*iri.MachineClass
+	quantity int64
 }
 
-func getCapabilities(iriCaps *iri.MachineClassCapabilities) capabilities {
-	return capabilities{
-		cpuMillis:   iriCaps.CpuMillis,
-		memoryBytes: iriCaps.MemoryBytes,
+func getMachineClassByCapabilities(machineClassByCapabilities map[machineClassQuantity]capabilities, capabilities capabilities) []machineClassQuantity {
+	matchingMachineClassStatus := []machineClassQuantity{}
+	for machineClassStatus, machineClassCapabilities := range machineClassByCapabilities {
+		if reflect.DeepEqual(machineClassCapabilities, capabilities) {
+			matchingMachineClassStatus = append(matchingMachineClassStatus, machineClassStatus)
+		}
 	}
+	return matchingMachineClassStatus
 }
 
 type Generic struct {
@@ -40,8 +46,8 @@ type Generic struct {
 
 	listener sets.Set[*listener]
 
-	machineClassByName         map[string]*iri.MachineClassStatus
-	machineClassByCapabilities map[capabilities][]*iri.MachineClassStatus
+	machineClassByName       map[string]*iri.MachineClassStatus
+	machineClassCapabilities map[machineClassQuantity]capabilities
 
 	machineRuntime machine.RuntimeService
 
@@ -97,19 +103,19 @@ func (g *Generic) relist(ctx context.Context, log logr.Logger) error {
 	oldMachineClassByName := maps.Clone(g.machineClassByName)
 
 	maps.Clear(g.machineClassByName)
-	maps.Clear(g.machineClassByCapabilities)
+	maps.Clear(g.machineClassCapabilities)
 
 	var notify bool
 	for _, machineClassStatus := range res.MachineClassStatus {
 		machineClass := machineClassStatus.GetMachineClass()
 		notify = notify || shouldNotify(oldMachineClassByName, machineClassStatus)
 
-		caps := capabilities{
-			cpuMillis:   machineClass.Capabilities.CpuMillis,
-			memoryBytes: machineClass.Capabilities.MemoryBytes,
-		}
 		g.machineClassByName[machineClass.Name] = machineClassStatus
-		g.machineClassByCapabilities[caps] = append(g.machineClassByCapabilities[caps], machineClassStatus)
+		machineClassWithQuantity := machineClassQuantity{
+			MachineClass: machineClass,
+			quantity:     machineClassStatus.Quantity,
+		}
+		g.machineClassCapabilities[machineClassWithQuantity] = machineClass.Capabilities.Resources
 	}
 
 	if notify {
@@ -121,12 +127,12 @@ func (g *Generic) relist(ctx context.Context, log logr.Logger) error {
 
 	for _, machineClassStatus := range res.MachineClassStatus {
 		machineClass := machineClassStatus.GetMachineClass()
-		caps := capabilities{
-			cpuMillis:   machineClass.Capabilities.CpuMillis,
-			memoryBytes: machineClass.Capabilities.MemoryBytes,
-		}
 		g.machineClassByName[machineClass.Name] = machineClassStatus
-		g.machineClassByCapabilities[caps] = append(g.machineClassByCapabilities[caps], machineClassStatus)
+		machineClassWithQuantity := machineClassQuantity{
+			MachineClass: machineClass,
+			quantity:     machineClassStatus.Quantity,
+		}
+		g.machineClassCapabilities[machineClassWithQuantity] = machineClass.Capabilities.Resources
 	}
 
 	if !g.sync {
@@ -151,23 +157,21 @@ func (g *Generic) GetMachineClassFor(ctx context.Context, name string, caps *iri
 	g.mu.RLock()
 	defer g.mu.RUnlock()
 
-	expected := getCapabilities(caps)
-	if byName, ok := g.machineClassByName[name]; ok && getCapabilities(byName.MachineClass.Capabilities) == expected {
+	expected := caps.Resources
+	if byName, ok := g.machineClassByName[name]; ok && reflect.DeepEqual(byName.MachineClass.Capabilities.Resources, expected) {
 		return byName.MachineClass, byName.Quantity, nil
 	}
 
-	if byCaps, ok := g.machineClassByCapabilities[expected]; ok {
-		switch len(byCaps) {
-		case 0:
-			return nil, 0, ErrNoMatchingMachineClass
-		case 1:
-			return byCaps[0].MachineClass, byCaps[0].Quantity, nil
-		default:
-			return nil, 0, ErrAmbiguousMatchingMachineClass
-		}
+	byCaps := getMachineClassByCapabilities(g.machineClassCapabilities, expected)
+	switch len(byCaps) {
+	case 0:
+		return nil, 0, ErrNoMatchingMachineClass
+	case 1:
+		return byCaps[0].MachineClass, byCaps[0].quantity, nil
+	default:
+		return nil, 0, ErrAmbiguousMatchingMachineClass
 	}
 
-	return nil, 0, ErrNoMatchingMachineClass
 }
 
 func (g *Generic) WaitForSync(ctx context.Context) error {
@@ -192,12 +196,12 @@ func setGenericOptionsDefaults(o *GenericOptions) {
 func NewGeneric(runtime machine.RuntimeService, opts GenericOptions) MachineClassMapper {
 	setGenericOptionsDefaults(&opts)
 	return &Generic{
-		synced:                     make(chan struct{}),
-		machineClassByName:         map[string]*iri.MachineClassStatus{},
-		machineClassByCapabilities: map[capabilities][]*iri.MachineClassStatus{},
-		listener:                   sets.New[*listener](),
-		machineRuntime:             runtime,
-		relistPeriod:               opts.RelistPeriod,
+		synced:                   make(chan struct{}),
+		machineClassByName:       map[string]*iri.MachineClassStatus{},
+		machineClassCapabilities: map[machineClassQuantity]capabilities{},
+		listener:                 sets.New[*listener](),
+		machineRuntime:           runtime,
+		relistPeriod:             opts.RelistPeriod,
 	}
 }