Merge pull request #85 from MikeSpreitzer/remember-readiness-relay

Milestone 1 stragglers

Author: MikeSpreitzer

pkg/controller/dual-pods/controller.go

@@ -71,16 +71,17 @@ func NewController(
 	numWorkers int,
 ) (*controller, error) {
 	ctl := &controller{
-		enqueueLogger: logger.WithName(ControllerName),
-		coreclient:    coreClient,
-		namespace:     namespace,
-		podInformer:   corev1PreInformers.Pods().Informer(),
-		podLister:     corev1PreInformers.Pods().Lister(),
-		cmInformer:    corev1PreInformers.ConfigMaps().Informer(),
-		cmLister:      corev1PreInformers.ConfigMaps().Lister(),
-		nodeInformer:  corev1PreInformers.Nodes().Informer(),
-		nodeLister:    corev1PreInformers.Nodes().Lister(),
-		requesters:    make(map[string]*requesterData),
+		enqueueLogger:    logger.WithName(ControllerName),
+		coreclient:       coreClient,
+		namespace:        namespace,
+		podInformer:      corev1PreInformers.Pods().Informer(),
+		podLister:        corev1PreInformers.Pods().Lister(),
+		cmInformer:       corev1PreInformers.ConfigMaps().Informer(),
+		cmLister:         corev1PreInformers.ConfigMaps().Lister(),
+		nodeInformer:     corev1PreInformers.Nodes().Informer(),
+		nodeLister:       corev1PreInformers.Nodes().Lister(),
+		requesters:       make(map[string]*requesterData),
+		inferenceServers: make(map[apitypes.UID]*serverData),
 	}
 	ctl.gpuMap.Store(&map[string]GpuLocation{})
 	ctl.QueueAndWorkers = genctlr.NewQueueAndWorkers(string(ControllerName), numWorkers, ctl.process)
@@ -114,8 +115,13 @@ type controller struct {
 
 	// requesters maps sever-requesting Pod name to data
 	requesters map[string]*requesterData
+
+	// inferenceServers maps UID of serve-requesting Pod to data
+	inferenceServers map[apitypes.UID]*serverData
 }
 
+var _ Controller = &controller{}
+
 type GpuLocation struct {
 	Node  string
 	Index uint
@@ -126,7 +132,11 @@ type requesterData struct {
 	GPUIndices *string
 }
 
-var _ Controller = &controller{}
+// Internal state about an inference server
+type serverData struct {
+	requestingPodName string
+	ReadinessRelayed  *bool
+}
 
 type typedRef struct {
 	Kind string
@@ -294,25 +304,45 @@ func (ctl *controller) processConfigMap(ctx context.Context, cmRef cache.ObjectN
 		}
 		return err, true
 	}
+	oldMap := ctl.gpuMap.Load()
 	newMap := map[string]GpuLocation{}
 	nodeCount := 0
+	additions := 0
 	for nodeName, mapStr := range cm.Data {
-		var nodesMap map[string]uint
-		err = json.Unmarshal([]byte(mapStr), &nodesMap)
+		var newNodesMap map[string]uint
+		err = json.Unmarshal([]byte(mapStr), &newNodesMap)
 		if err != nil {
 			logger.Error(err, "A GPU map entry failed to parse as JSON", "nodeName", nodeName)
 			continue
 		}
-		for uuid, index := range nodesMap {
-			newMap[uuid] = GpuLocation{Node: nodeName, Index: index}
+		for uuid, index := range newNodesMap {
+			newLoc := GpuLocation{Node: nodeName, Index: index}
+			if oldMap == nil || (*oldMap)[uuid] != newLoc {
+				additions++
+			}
+			newMap[uuid] = newLoc
 		}
 		nodeCount += 1
 	}
-	logger.V(1).Info("Parsed GPU map", "numNodes", nodeCount, "numGPUs", len(newMap))
+	logger.V(1).Info("Parsed GPU map", "numNodes", nodeCount, "numGPUs", len(newMap), "additions", additions)
 	ctl.gpuMap.Store(&newMap)
+	if additions > 0 {
+		ctl.enqueueRequesters(ctx)
+	}
 	return nil, false
 }
 
+func (ctl *controller) enqueueRequesters(ctx context.Context) {
+	ctl.mutex.Lock()
+	defer ctl.mutex.Unlock()
+	logger := klog.FromContext(ctx)
+	for reqPodName := range ctl.requesters {
+		ref := typedRef{Kind: podKind, ObjectName: cache.ObjectName{Namespace: ctl.namespace, Name: reqPodName}}
+		logger.V(5).Info("Enqueuing server-requesting Pod because of change to GPU map", "ref", ref)
+		ctl.Queue.Add(ref)
+	}
+}
+
 func (ctl *controller) getRequesterData(name string, podUID apitypes.UID, insist bool) *requesterData {
 	ctl.mutex.Lock()
 	defer ctl.mutex.Unlock()
@@ -329,3 +359,24 @@ func (ctl *controller) clearRequesterData(name string) {
 	defer ctl.mutex.Unlock()
 	delete(ctl.requesters, name)
 }
+
+func (ctl *controller) getServerData(reqName string, reqUID apitypes.UID, insist bool) *serverData {
+	ctl.mutex.Lock()
+	defer ctl.mutex.Unlock()
+	ans := ctl.inferenceServers[reqUID]
+	if ans == nil && insist {
+		ans = &serverData{requestingPodName: reqName}
+		ctl.inferenceServers[reqUID] = ans
+	}
+	return ans
+}
+
+func (ctl *controller) clearServerData(reqName string) {
+	ctl.mutex.Lock()
+	defer ctl.mutex.Unlock()
+	for uid, serveDat := range ctl.inferenceServers {
+		if serveDat.requestingPodName == reqName {
+			delete(ctl.inferenceServers, uid)
+		}
+	}
+}

pkg/controller/dual-pods/server-requesting.go

@@ -105,20 +105,20 @@ func (ctl *controller) processServerRequestingPod(ctx context.Context, requestin
 
 	// use the server patch to build the server-running pod
 	logger.V(5).Info("Building server-running pod from patch", "name", requestingPod.Name, "patch", serverPatch)
-	serverRunningPod, err := composeServerRunningPod(ctx, requestingPod, serverPatch, *reqDat.GPUIndices, api.RunnerData{
+	desiredRunningPod, err := composeServerRunningPod(ctx, requestingPod, serverPatch, *reqDat.GPUIndices, api.RunnerData{
 		NodeName: requestingPod.Spec.NodeName,
 	})
 	if err != nil {
 		return ctl.ensureReqStatus(ctx, requestingPod, fmt.Sprintf("failed to construct the nominal server-running Pod: %s", err.Error()))
 	}
 
-	got, err := ctl.podLister.Pods(serverRunningPod.Namespace).Get(serverRunningPod.Name)
+	actualRunningPod, err := ctl.podLister.Pods(desiredRunningPod.Namespace).Get(desiredRunningPod.Name)
 	if err != nil && !apierrors.IsNotFound(err) {
-		logger.Error(err, "Failed to get existing server-running pod", "name", serverRunningPod.Name)
+		logger.Error(err, "Failed to get existing server-running pod", "name", desiredRunningPod.Name)
 		return err, true
 	}
-	if got != nil {
-		logger.V(5).Info("Server-running pod exists", "name", serverRunningPod.Name)
+	if actualRunningPod != nil {
+		logger.V(5).Info("Server-running pod exists", "name", desiredRunningPod.Name)
 
 		// TODO: we should reconcile the existing server-running pod with the one we just built
 		return nil, false
@@ -131,8 +131,8 @@ func (ctl *controller) processServerRequestingPod(ctx context.Context, requestin
 		return err, false
 	}
 
-	logger.V(2).Info("Creating server-running pod", "name", serverRunningPod.Name, "namespace", serverRunningPod.Namespace, "annotations", serverRunningPod.Annotations, "labels", serverRunningPod.Labels)
-	echo, err := ctl.coreclient.Pods(serverRunningPod.Namespace).Create(ctx, serverRunningPod, metav1.CreateOptions{})
+	logger.V(2).Info("Creating server-running pod", "name", desiredRunningPod.Name, "namespace", desiredRunningPod.Namespace, "annotations", desiredRunningPod.Annotations, "labels", desiredRunningPod.Labels)
+	echo, err := ctl.coreclient.Pods(desiredRunningPod.Namespace).Create(ctx, desiredRunningPod, metav1.CreateOptions{})
 	if err != nil {
 		errMsg := err.Error()
 		if invalidPodRE.MatchString(errMsg) {
@@ -144,7 +144,7 @@ func (ctl *controller) processServerRequestingPod(ctx context.Context, requestin
 		}
 		return err, true
 	}
-	logger.V(5).Info("Created server-running pod", "name", serverRunningPod.Name, "annotations", echo.Annotations, "labels", echo.Labels, "resourceVersion", echo.ResourceVersion)
+	logger.V(5).Info("Created server-running pod", "name", desiredRunningPod.Name, "annotations", echo.Annotations, "labels", echo.Labels, "resourceVersion", echo.ResourceVersion)
 
 	return ctl.ensureReqStatus(ctx, requestingPod)
 }

pkg/controller/dual-pods/server-running.go

@@ -68,6 +68,10 @@ func (ctl *controller) processServerRunningPod(ctx context.Context, runningPod *
 		}
 	}
 
+	if requestingPod == nil {
+		ctl.clearServerData(requestingPodName)
+	}
+
 	if requestingPod == nil || requestingPod.DeletionTimestamp != nil || runningPod.DeletionTimestamp != nil {
 		podOps := ctl.coreclient.Pods(runningPod.Namespace)
 		// Deletion requested, so remove finalizer and delete server-requesting Pod
@@ -90,27 +94,34 @@ func (ctl *controller) processServerRunningPod(ctx context.Context, runningPod *
 		return err, false
 	}
 
+	serverDat := ctl.getServerData(requestingPodName, requestingPod.UID, true)
+
 	// relay the readiness
 	port := requestingPod.Annotations[api.AdminPortAnnotationName]
 	if port == "" {
 		port = api.AdminPortDefaultValue
 	}
-	url, readiness := fmt.Sprintf("http://%s:%s", requestingPod.Status.PodIP, port), ""
-	if isPodReady(runningPod) {
-		logger.V(5).Info("Server-running pod is ready", "name", runningPod.Name)
-		url += stubapi.BecomeReadyPath
-		readiness = "ready"
-	} else {
-		logger.V(5).Info("Server-running pod is not ready", "name", runningPod.Name)
-		url += stubapi.BecomeUnreadyPath
-		readiness = "unready"
-	}
-	err = postToReadiness(url)
-	if err != nil {
-		logger.Error(err, "Failed to relay the readiness", "name", runningPod.Name, "readiness", readiness)
-		return err, true
+	ready := isPodReady(runningPod)
+	if serverDat.ReadinessRelayed == nil || ready != *serverDat.ReadinessRelayed {
+		// TODO: use cache
+		url, readiness := fmt.Sprintf("http://%s:%s", requestingPod.Status.PodIP, port), ""
+		if ready {
+			logger.V(5).Info("Server-running pod is ready", "name", runningPod.Name)
+			url += stubapi.BecomeReadyPath
+			readiness = "ready"
+		} else {
+			logger.V(5).Info("Server-running pod is not ready", "name", runningPod.Name)
+			url += stubapi.BecomeUnreadyPath
+			readiness = "unready"
+		}
+		err = postToReadiness(url)
+		if err != nil {
+			logger.Error(err, "Failed to relay the readiness", "name", runningPod.Name, "readiness", readiness)
+			return err, true
+		}
+		serverDat.ReadinessRelayed = &ready
+		logger.V(5).Info("Successfully relayed the readiness", "name", runningPod.Name, "readiness", readiness)
 	}
-	logger.V(5).Info("Successfully relayed the readiness", "name", runningPod.Name, "readiness", readiness)
 
 	logger.V(5).Info("Processed server-running pod", "name", runningPod.Name)
 	return nil, false