feat: add compartment-based node assignment

operator/.mockery.yaml

@@ -20,12 +20,12 @@ template: testify
 template-data:
   unroll-variadic: true
 packages:
-  github.com/NVIDIA/skyhook/internal/controller:
+  github.com/NVIDIA/skyhook/operator/internal/controller:
     config:
       all: true
     interfaces:
       SkyhookNodes: {}
-  github.com/NVIDIA/skyhook/internal/dal:
+  github.com/NVIDIA/skyhook/operator/internal/dal:
     config:
       all: true
   k8s.io/client-go/tools/record:

operator/api/v1alpha1/deployment_policy_types.go

@@ -109,6 +109,10 @@ type DeploymentBudget struct {
 	Count *int `json:"count,omitempty"`
 }
 
+const (
+	DefaultCompartmentName = "__default__"
+)
+
 // PolicyDefault defines default budget and strategy for unmatched nodes
 type PolicyDefault struct {
 	// Exactly one of percent or count
@@ -153,6 +157,12 @@ type DeploymentPolicy struct {
 
 // +kubebuilder:object:root=true
 
+type DeploymentPolicyList struct {
+	metav1.TypeMeta `json:",inline"`
+	metav1.ListMeta `json:"metadata,omitempty"`
+	Items           []DeploymentPolicy `json:"items"`
+}
+
 // Default applies default values to DeploymentStrategy
 func (s *DeploymentStrategy) Default() {
 	switch {
@@ -261,5 +271,5 @@ func (b *DeploymentBudget) Validate() error {
 }
 
 func init() {
-	SchemeBuilder.Register(&DeploymentPolicy{})
+	SchemeBuilder.Register(&DeploymentPolicy{}, &DeploymentPolicyList{})
 }

operator/api/v1alpha1/deployment_policy_webhook.go

@@ -59,7 +59,7 @@ func (r *DeploymentPolicyWebhook) Default(ctx context.Context, obj runtime.Objec
 		return fmt.Errorf("object is not a DeploymentPolicy")
 	}
 
-	deploymentPolicylog.Info("default", "name", deploymentPolicy.Name)
+	deploymentPolicylog.Info(DefaultCompartmentName, "name", deploymentPolicy.Name)
 
 	// Apply defaults to the default strategy
 	if deploymentPolicy.Spec.Default.Strategy != nil {
@@ -140,6 +140,11 @@ func (r *DeploymentPolicy) Validate() error {
 	selectors := make(map[string]metav1.LabelSelector)
 
 	for _, compartment := range r.Spec.Compartments {
+		// Validate compartment name is not "__default__" (reserved)
+		if compartment.Name == DefaultCompartmentName {
+			return fmt.Errorf("compartment name %q is reserved and cannot be used", compartment.Name)
+		}
+
 		// Validate unique names
 		if names[compartment.Name] {
 			return fmt.Errorf("compartment name %q is not unique", compartment.Name)

operator/api/v1alpha1/deployment_policy_webhook_test.go

@@ -190,6 +190,36 @@ var _ = Describe("DeploymentPolicy", func() {
 			Expect(err).ToNot(HaveOccurred())
 		})
 
+		It("should reject compartment name '__default__' as reserved", func() {
+			deploymentPolicy := &DeploymentPolicy{
+				ObjectMeta: metav1.ObjectMeta{Name: "foobar"},
+				Spec: DeploymentPolicySpec{
+					Default: PolicyDefault{
+						Budget: DeploymentBudget{Percent: ptr.To(25)},
+						Strategy: &DeploymentStrategy{
+							Fixed: &FixedStrategy{},
+						},
+					},
+					Compartments: []Compartment{
+						{
+							Name:     DefaultCompartmentName, // reserved name
+							Selector: metav1.LabelSelector{MatchLabels: map[string]string{"tier": "web"}},
+							Budget:   DeploymentBudget{Percent: ptr.To(25)},
+						},
+					},
+				},
+			}
+
+			_, err := deploymentPolicyWebhook.ValidateCreate(ctx, deploymentPolicy)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring(`compartment name "__default__" is reserved and cannot be used`))
+
+			// Fixed with different name
+			deploymentPolicy.Spec.Compartments[0].Name = "system"
+			_, err = deploymentPolicyWebhook.ValidateCreate(ctx, deploymentPolicy)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
 		It("should allow different selectors", func() {
 			deploymentPolicy := &DeploymentPolicy{
 				ObjectMeta: metav1.ObjectMeta{Name: "foobar"},

operator/internal/controller/cluster_state_v2.go

@@ -70,7 +70,7 @@ type clusterState struct {
 	skyhooks []SkyhookNodes
 }
 
-func BuildState(skyhooks *v1alpha1.SkyhookList, nodes *corev1.NodeList) (*clusterState, error) {
+func BuildState(skyhooks *v1alpha1.SkyhookList, nodes *corev1.NodeList, deploymentPolicies *v1alpha1.DeploymentPolicyList) (*clusterState, error) {
 
 	ret := &clusterState{
 		tracker:  ObjectTracker{objects: make(map[string]client.Object)},
@@ -82,8 +82,9 @@ func BuildState(skyhooks *v1alpha1.SkyhookList, nodes *corev1.NodeList) (*cluste
 		ret.tracker.Track(skyhook.DeepCopy())
 
 		ret.skyhooks[idx] = &skyhookNodes{
-			skyhook: wrapper.NewSkyhookWrapper(&skyhook),
-			nodes:   make([]wrapper.SkyhookNode, 0),
+			skyhook:      wrapper.NewSkyhookWrapper(&skyhook),
+			nodes:        make([]wrapper.SkyhookNode, 0),
+			compartments: make(map[string]*wrapper.Compartment),
 		}
 		for _, node := range nodes.Items {
 			skyNode, err := wrapper.NewSkyhookNode(&node, &skyhook)
@@ -100,6 +101,26 @@ func BuildState(skyhooks *v1alpha1.SkyhookList, nodes *corev1.NodeList) (*cluste
 				ret.skyhooks[idx].AddNode(skyNode)
 			}
 		}
+
+		// find deployment policy and all compartments + the default one
+		// Skip skyhooks that don't have a deployment policy
+		if skyhook.Spec.DeploymentPolicy == "" {
+			continue
+		}
+
+		for _, deploymentPolicy := range deploymentPolicies.Items {
+			if deploymentPolicy.Name == skyhook.Spec.DeploymentPolicy {
+				for _, compartment := range deploymentPolicy.Spec.Compartments {
+					ret.skyhooks[idx].AddCompartment(compartment.Name, wrapper.NewCompartmentWrapper(&compartment))
+				}
+				// use policy default
+				ret.skyhooks[idx].AddCompartment(v1alpha1.DefaultCompartmentName, wrapper.NewCompartmentWrapper(&v1alpha1.Compartment{
+					Name:     v1alpha1.DefaultCompartmentName,
+					Budget:   deploymentPolicy.Spec.Default.Budget,
+					Strategy: deploymentPolicy.Spec.Default.Strategy,
+				}))
+			}
+		}
 	}
 
 	// Sort by priority (ascending), then by name (ascending) if priorities are equal
@@ -155,15 +176,20 @@ type SkyhookNodes interface {
 	UpdateCondition() bool
 	ReportState()
 	Migrate(logger logr.Logger) error
+
+	GetCompartments() map[string]*wrapper.Compartment
+	AddCompartment(name string, compartment *wrapper.Compartment)
+	AddCompartmentNode(name string, node wrapper.SkyhookNode)
 }
 
 var _ SkyhookNodes = &skyhookNodes{}
 
 // skyhookNodes impl's. SkyhookNodes
 type skyhookNodes struct {
-	skyhook     *wrapper.Skyhook
-	nodes       []wrapper.SkyhookNode
-	priorStatus v1alpha1.Status
+	skyhook      *wrapper.Skyhook
+	nodes        []wrapper.SkyhookNode
+	priorStatus  v1alpha1.Status
+	compartments map[string]*wrapper.Compartment
 }
 
 func (s *skyhookNodes) GetPriorStatus() v1alpha1.Status {
@@ -713,6 +739,18 @@ func (skyhook *skyhookNodes) Migrate(logger logr.Logger) error {
 	return nil
 }
 
+func (skyhook *skyhookNodes) GetCompartments() map[string]*wrapper.Compartment {
+	return skyhook.compartments
+}
+
+func (skyhook *skyhookNodes) AddCompartment(name string, compartment *wrapper.Compartment) {
+	skyhook.compartments[name] = compartment
+}
+
+func (skyhook *skyhookNodes) AddCompartmentNode(name string, node wrapper.SkyhookNode) {
+	skyhook.compartments[name].AddNode(node)
+}
+
 // cleanupNodeMap removes nodes from the given map that no longer exist in currentNodes
 // Returns false if nodeMap is nil, otherwise returns true if any nodes were removed
 func cleanupNodeMap[T any](nodeMap map[string]T, currentNodes map[string]struct{}) bool {

operator/internal/controller/cluster_state_v2_test.go

@@ -161,8 +161,9 @@ var _ = Describe("BuildState ordering", func() {
 				},
 			},
 		}
+		deploymentPolicies := &v1alpha1.DeploymentPolicyList{Items: []v1alpha1.DeploymentPolicy{}}
 		nodes := &corev1.NodeList{Items: []corev1.Node{}}
-		clusterState, err := BuildState(skyhooks, nodes)
+		clusterState, err := BuildState(skyhooks, nodes, deploymentPolicies)
 		Expect(err).ToNot(HaveOccurred())
 		ordered := clusterState.skyhooks
 		// Should be: a (priority 1), b (priority 2, name b), c (priority 2, name c)