Add ComputeDomainStatus and set it as its deployment pods come online.

Signed-off-by: Kevin Klues <[email protected]>

Author: klueska

api/nvidia.com/resource/v1beta1/computedomain.go

@@ -24,13 +24,15 @@ import (
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
 // +k8s:openapi-gen=true
 // +kubebuilder:resource:scope=Namespaced
+// +kubebuilder:subresource:status
 
 // ComputeDomain prepares a set of nodes to run a multi-node workload in.
 type ComputeDomain struct {
 	metav1.TypeMeta   `json:",inline"`
 	metav1.ObjectMeta `json:"metadata,omitempty"`
 
-	Spec ComputeDomainSpec `json:"spec,omitempty"`
+	Spec   ComputeDomainSpec   `json:"spec,omitempty"`
+	Status ComputeDomainStatus `json:"status,omitempty"`
 }
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
@@ -51,3 +53,26 @@ type ComputeDomainSpec struct {
 	ResourceClaimName string `json:"resourceClaimName,omitempty"`
 	DeviceClassName   string `json:"deviceClassName,omitempty"`
 }
+
+// ComputeDomainStatus provides the status for a ComputeDomain.
+type ComputeDomainStatus struct {
+	// +listType=map
+	// +listMapKey=name
+	Nodes []*ComputeDomainNode `json:"nodes,omitempty"`
+}
+
+// ComputeDomainNode provides information about each node added to a ComputeDomain.
+type ComputeDomainNode struct {
+	Name      string `json:"name"`
+	IPAddress string `json:"ipAddress"`
+	CliqueID  string `json:"cliqueID"`
+}
+
+// GetNodesAsMap returns the list of nodes in a ComputeDomainStatus as a map.
+func (s *ComputeDomainStatus) GetNodesAsMap() map[string]*ComputeDomainNode {
+	m := make(map[string]*ComputeDomainNode)
+	for _, node := range s.Nodes {
+		m[node.Name] = node
+	}
+	return m
+}

api/nvidia.com/resource/v1beta1/zz_generated.deepcopy.go

@@ -308,7 +308,7 @@ func (m *DeploymentManager) addPodManager(ctx context.Context, labelSelector *me
 		return nil
 	}
 
-	podManager := NewDeploymentPodManager(m.config, m.imexChannelManager, labelSelector, numPods)
+	podManager := NewDeploymentPodManager(m.config, m.imexChannelManager, labelSelector, numPods, m.getComputeDomain)
 
 	if err := podManager.Start(ctx); err != nil {
 		return fmt.Errorf("error creating Pod manager: %w", err)

cmd/nvidia-dra-imex-controller/deployment.go

@@ -29,6 +29,12 @@ import (
 	corev1listers "k8s.io/client-go/listers/core/v1"
 	"k8s.io/client-go/tools/cache"
 	"k8s.io/klog/v2"
+
+	nvapi "github.com/NVIDIA/k8s-dra-driver/api/nvidia.com/resource/v1beta1"
+)
+
+const (
+	CliqueIDLabelKey = "nvidia.com/gpu.clique"
 )
 
 type DeploymentPodManager struct {
@@ -40,14 +46,16 @@ type DeploymentPodManager struct {
 	informer cache.SharedInformer
 	lister   corev1listers.PodLister
 
-	nodeSelector       corev1.NodeSelector
+	getComputeDomain   GetComputeDomainFunc
+	computeDomainNodes []*nvapi.ComputeDomainNode
 	computeDomainLabel string
 	numPods            int
+	nodeSelector       corev1.NodeSelector
 
 	imexChannelManager *ImexChannelManager
 }
 
-func NewDeploymentPodManager(config *ManagerConfig, imexChannelManager *ImexChannelManager, labelSelector *metav1.LabelSelector, numPods int) *DeploymentPodManager {
+func NewDeploymentPodManager(config *ManagerConfig, imexChannelManager *ImexChannelManager, labelSelector *metav1.LabelSelector, numPods int, getComputeDomain GetComputeDomainFunc) *DeploymentPodManager {
 	factory := informers.NewSharedInformerFactoryWithOptions(
 		config.clientsets.Core,
 		informerResyncPeriod,
@@ -79,9 +87,10 @@ func NewDeploymentPodManager(config *ManagerConfig, imexChannelManager *ImexChan
 		factory:            factory,
 		informer:           informer,
 		lister:             lister,
-		nodeSelector:       nodeSelector,
+		getComputeDomain:   getComputeDomain,
 		computeDomainLabel: labelSelector.MatchLabels[computeDomainLabelKey],
 		numPods:            numPods,
+		nodeSelector:       nodeSelector,
 		imexChannelManager: imexChannelManager,
 	}
 
@@ -150,23 +159,68 @@ func (m *DeploymentPodManager) onPodAddOrUpdate(ctx context.Context, obj any) er
 
 	klog.Infof("Processing added or updated Pod: %s/%s", p.Namespace, p.Name)
 
+	cd, err := m.getComputeDomain(p.Labels[computeDomainLabelKey])
+	if err != nil {
+		return fmt.Errorf("error getting ComputeDomain: %w", err)
+	}
+	if cd == nil {
+		return nil
+	}
+
 	if p.Spec.NodeName == "" {
 		return fmt.Errorf("pod not yet scheduled: %s/%s", p.Namespace, p.Name)
 	}
 
-	hostnameLabels := m.nodeSelector.NodeSelectorTerms[0].MatchExpressions[0].Values
-	if !slices.Contains(hostnameLabels, p.Spec.NodeName) {
-		hostnameLabels = append(hostnameLabels, p.Spec.NodeName)
+	var nodeNames []string
+	for _, node := range m.computeDomainNodes {
+		nodeNames = append(nodeNames, node.Name)
 	}
-	m.nodeSelector.NodeSelectorTerms[0].MatchExpressions[0].Values = hostnameLabels
 
-	if len(hostnameLabels) != m.numPods {
-		return fmt.Errorf("node selector not yet complete")
+	if !slices.Contains(nodeNames, p.Spec.NodeName) {
+		node, err := m.GetComputeDomainNode(ctx, p.Spec.NodeName)
+		if err != nil {
+			return fmt.Errorf("error getting ComputeDomainNode: %w", err)
+		}
+		nodeNames = append(nodeNames, node.Name)
+		m.computeDomainNodes = append(m.computeDomainNodes, node)
 	}
 
+	if len(nodeNames) != m.numPods {
+		return fmt.Errorf("not all pods scheduled yet")
+	}
+
+	cd.Status.Nodes = m.computeDomainNodes
+	if _, err = m.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(cd.Namespace).UpdateStatus(ctx, cd, metav1.UpdateOptions{}); err != nil {
+		return fmt.Errorf("error updating nodes in ComputeDomain status: %w", err)
+	}
+
+	m.nodeSelector.NodeSelectorTerms[0].MatchExpressions[0].Values = nodeNames
 	if err := m.imexChannelManager.CreateOrUpdatePool(m.computeDomainLabel, &m.nodeSelector); err != nil {
 		return fmt.Errorf("failed to create or update IMEX channel pool: %w", err)
 	}
 
 	return nil
 }
+
+func (m *DeploymentPodManager) GetComputeDomainNode(ctx context.Context, nodeName string) (*nvapi.ComputeDomainNode, error) {
+	node, err := m.config.clientsets.Core.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	if err != nil {
+		return nil, fmt.Errorf("error getting Node '%s': %w", nodeName, err)
+	}
+
+	var ipAddress string
+	for _, addr := range node.Status.Addresses {
+		if addr.Type == corev1.NodeInternalIP {
+			ipAddress = addr.Address
+			break
+		}
+	}
+
+	n := &nvapi.ComputeDomainNode{
+		Name:      nodeName,
+		IPAddress: ipAddress,
+		CliqueID:  node.Labels[CliqueIDLabelKey],
+	}
+
+	return n, nil
+}

cmd/nvidia-dra-imex-controller/deploymentpods.go

@@ -53,6 +53,32 @@ spec:
             - message: Exactly one of 'resourceClaimName' or 'deviceClassName' must
                 be set.
               rule: '(has(self.resourceClaimName) ? !has(self.deviceClassName) : has(self.deviceClassName))'
+          status:
+            description: ComputeDomainStatus provides the status for a ComputeDomain.
+            properties:
+              nodes:
+                items:
+                  description: ComputeDomainNode provides information about each node
+                    added to a ComputeDomain.
+                  properties:
+                    cliqueID:
+                      type: string
+                    ipAddress:
+                      type: string
+                    name:
+                      type: string
+                  required:
+                  - cliqueID
+                  - ipAddress
+                  - name
+                  type: object
+                type: array
+                x-kubernetes-list-map-keys:
+                - name
+                x-kubernetes-list-type: map
+            type: object
         type: object
     served: true
     storage: true
+    subresources:
+      status: {}

deployments/helm/k8s-dra-gpu-driver/crds/resource.nvidia.com_computedomains.yaml

@@ -53,6 +53,32 @@ spec:
             - message: Exactly one of 'resourceClaimName' or 'deviceClassName' must
                 be set.
               rule: '(has(self.resourceClaimName) ? !has(self.deviceClassName) : has(self.deviceClassName))'
+          status:
+            description: ComputeDomainStatus provides the status for a ComputeDomain.
+            properties:
+              nodes:
+                items:
+                  description: ComputeDomainNode provides information about each node
+                    added to a ComputeDomain.
+                  properties:
+                    cliqueID:
+                      type: string
+                    ipAddress:
+                      type: string
+                    name:
+                      type: string
+                  required:
+                  - cliqueID
+                  - ipAddress
+                  - name
+                  type: object
+                type: array
+                x-kubernetes-list-map-keys:
+                - name
+                x-kubernetes-list-type: map
+            type: object
         type: object
     served: true
     storage: true
+    subresources:
+      status: {}

deployments/helm/k8s-dra-imex-driver/crds/resource.nvidia.com_computedomains.yaml

@@ -8,6 +8,9 @@ rules:
 - apiGroups: ["resource.nvidia.com"]
   resources: ["computedomains"]
   verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: ["resource.nvidia.com"]
+  resources: ["computedomains/status"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
 - apiGroups: ["resource.k8s.io"]
   resources: ["resourceclaims"]
   verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]