Fix a race condition where LLMServerPool has not been initialized yet

Author: liu-cong

pkg/ext-proc/backend/datastore.go

@@ -1,6 +1,7 @@
 package backend
 
 import (
+	"fmt"
 	"math/rand"
 	"sync"
 
@@ -9,24 +10,65 @@ import (
 	"k8s.io/klog/v2"
 )
 
+func NewK8sDataStore(options ...K8sDatastoreOption) *K8sDatastore {
+	store := &K8sDatastore{
+		poolMu:      sync.RWMutex{},
+		llmServices: &sync.Map{},
+		pods:        &sync.Map{},
+	}
+	for _, opt := range options {
+		opt(store)
+	}
+	return store
+}
+
 // The datastore is a local cache of relevant data for the given LLMServerPool (currently all pulled from k8s-api)
 type K8sDatastore struct {
-	LLMServerPool *v1alpha1.LLMServerPool
-	LLMServices   *sync.Map
-	Pods          *sync.Map
+	// poolMu is used to synchronize access to the llmServerPool.
+	poolMu        sync.RWMutex
+	llmServerPool *v1alpha1.LLMServerPool
+	llmServices   *sync.Map
+	pods          *sync.Map
+}
+
+type K8sDatastoreOption func(*K8sDatastore)
+
+// WithPods can be used in tests to override the pods.
+func WithPods(pods []*PodMetrics) K8sDatastoreOption {
+	return func(store *K8sDatastore) {
+		store.pods = &sync.Map{}
+		for _, pod := range pods {
+			store.pods.Store(pod.Pod, true)
+		}
+	}
+}
+
+func (ds *K8sDatastore) SetLLMServerPool(pool *v1alpha1.LLMServerPool) {
+	ds.poolMu.Lock()
+	defer ds.poolMu.Unlock()
+	ds.llmServerPool = pool
+}
+
+func (ds *K8sDatastore) GetLLMServerPool() (*v1alpha1.LLMServerPool, error) {
+	ds.poolMu.RLock()
+	defer ds.poolMu.RUnlock()
+	if ds.llmServerPool == nil {
+		return nil, fmt.Errorf("LLMServerPool hasn't been initialized yet")
+	}
+	return ds.llmServerPool, nil
 }
 
 func (ds *K8sDatastore) GetPodIPs() []string {
 	var ips []string
-	ds.Pods.Range(func(name, pod any) bool {
+	ds.pods.Range(func(name, pod any) bool {
 		ips = append(ips, pod.(*corev1.Pod).Status.PodIP)
 		return true
 	})
 	return ips
 }
 
 func (s *K8sDatastore) FetchModelData(modelName string) (returnModel *v1alpha1.Model) {
-	s.LLMServices.Range(func(k, v any) bool {
+	s.llmServices.Range(func(k, v any) bool {
 		service := v.(*v1alpha1.LLMService)
 		klog.V(3).Infof("Service name: %v", service.Name)
 		for _, model := range service.Spec.Models {

pkg/ext-proc/backend/datastore_test.go

@@ -10,10 +10,9 @@ var ()
 
 func TestRandomWeightedDraw(t *testing.T) {
 	tests := []struct {
-		name      string
-		datastore K8sDatastore
-		model     *v1alpha1.Model
-		want      string
+		name  string
+		model *v1alpha1.Model
+		want  string
 	}{
 		{
 			name: "'random' distribution",

pkg/ext-proc/backend/endpointslice_reconciler.go

@@ -48,9 +48,9 @@ func (c *EndpointSliceReconciler) updateDatastore(slice *discoveryv1.EndpointSli
 	for _, endpoint := range slice.Endpoints {
 		klog.V(4).Infof("Zone: %v \n endpoint: %+v \n", c.Zone, endpoint)
 		if c.validPod(endpoint) {
-			pod := Pod{Name: *&endpoint.TargetRef.Name, Address: endpoint.Addresses[0] + ":" + fmt.Sprint(c.Datastore.LLMServerPool.Spec.TargetPort)}
+			pod := Pod{Name: *&endpoint.TargetRef.Name, Address: endpoint.Addresses[0] + ":" + fmt.Sprint(c.Datastore.llmServerPool.Spec.TargetPort)}
 			podMap[pod] = true
-			c.Datastore.Pods.Store(pod, true)
+			c.Datastore.pods.Store(pod, true)
 		}
 	}
 
@@ -61,14 +61,22 @@ func (c *EndpointSliceReconciler) updateDatastore(slice *discoveryv1.EndpointSli
 			return false
 		}
 		if _, ok := podMap[pod]; !ok {
-			c.Datastore.Pods.Delete(pod)
+			c.Datastore.pods.Delete(pod)
 		}
 		return true
 	}
-	c.Datastore.Pods.Range(removeOldPods)
+	c.Datastore.pods.Range(removeOldPods)
 }
 
 func (c *EndpointSliceReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	llmServerPoolAvailable := func(object client.Object) bool {
+		_, err := c.Datastore.GetLLMServerPool()
+		if err != nil {
+			klog.Warningf("Skipping reconciling EndpointSlice because LLMServerPool is not available yet: %v", err)
+		}
+		return err == nil
+	}
+
 	ownsEndPointSlice := func(object client.Object) bool {
 		// Check if the object is an EndpointSlice
 		endpointSlice, ok := object.(*discoveryv1.EndpointSlice)
@@ -80,7 +88,7 @@ func (c *EndpointSliceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	}
 
 	return ctrl.NewControllerManagedBy(mgr).
-		For(&discoveryv1.EndpointSlice{}, builder.WithPredicates(predicate.NewPredicateFuncs(ownsEndPointSlice))).
+		For(&discoveryv1.EndpointSlice{}, builder.WithPredicates(predicate.NewPredicateFuncs(llmServerPoolAvailable), predicate.NewPredicateFuncs(ownsEndPointSlice))).
 		Complete(c)
 }
 

pkg/ext-proc/backend/endpointslice_reconcilier_test.go

@@ -18,15 +18,15 @@ var (
 func TestUpdateDatastore_EndpointSliceReconciler(t *testing.T) {
 	tests := []struct {
 		name          string
-		datastore     K8sDatastore
+		datastore     *K8sDatastore
 		incomingSlice *discoveryv1.EndpointSlice
-		want          K8sDatastore
+		wantPods      *sync.Map
 	}{
 		{
 			name: "Add new pod",
-			datastore: K8sDatastore{
-				Pods: populateMap(basePod1, basePod2),
-				LLMServerPool: &v1alpha1.LLMServerPool{
+			datastore: &K8sDatastore{
+				pods: populateMap(basePod1, basePod2),
+				llmServerPool: &v1alpha1.LLMServerPool{
 					Spec: v1alpha1.LLMServerPoolSpec{
 						TargetPort: int32(8000),
 					},
@@ -66,15 +66,13 @@ func TestUpdateDatastore_EndpointSliceReconciler(t *testing.T) {
 					},
 				},
 			},
-			want: K8sDatastore{
-				Pods: populateMap(basePod1, basePod2, basePod3),
-			},
+			wantPods: populateMap(basePod1, basePod2, basePod3),
 		},
 		{
 			name: "New pod, but its not ready yet. Do not add.",
-			datastore: K8sDatastore{
-				Pods: populateMap(basePod1, basePod2),
-				LLMServerPool: &v1alpha1.LLMServerPool{
+			datastore: &K8sDatastore{
+				pods: populateMap(basePod1, basePod2),
+				llmServerPool: &v1alpha1.LLMServerPool{
 					Spec: v1alpha1.LLMServerPoolSpec{
 						TargetPort: int32(8000),
 					},
@@ -114,15 +112,13 @@ func TestUpdateDatastore_EndpointSliceReconciler(t *testing.T) {
 					},
 				},
 			},
-			want: K8sDatastore{
-				Pods: populateMap(basePod1, basePod2),
-			},
+			wantPods: populateMap(basePod1, basePod2),
 		},
 		{
 			name: "Existing pod not ready, new pod added, and is ready",
-			datastore: K8sDatastore{
-				Pods: populateMap(basePod1, basePod2),
-				LLMServerPool: &v1alpha1.LLMServerPool{
+			datastore: &K8sDatastore{
+				pods: populateMap(basePod1, basePod2),
+				llmServerPool: &v1alpha1.LLMServerPool{
 					Spec: v1alpha1.LLMServerPoolSpec{
 						TargetPort: int32(8000),
 					},
@@ -162,18 +158,16 @@ func TestUpdateDatastore_EndpointSliceReconciler(t *testing.T) {
 					},
 				},
 			},
-			want: K8sDatastore{
-				Pods: populateMap(basePod3, basePod2),
-			},
+			wantPods: populateMap(basePod3, basePod2),
 		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			endpointSliceReconciler := &EndpointSliceReconciler{Datastore: &test.datastore, Zone: ""}
+			endpointSliceReconciler := &EndpointSliceReconciler{Datastore: test.datastore, Zone: ""}
 			endpointSliceReconciler.updateDatastore(test.incomingSlice)
 
-			if mapsEqual(endpointSliceReconciler.Datastore.Pods, test.want.Pods) {
-				t.Errorf("Unexpected output pod mismatch. \n Got %v \n Want: %v \n", endpointSliceReconciler.Datastore.Pods, test.want.Pods)
+			if mapsEqual(endpointSliceReconciler.Datastore.pods, test.wantPods) {
+				t.Errorf("Unexpected output pod mismatch. \n Got %v \n Want: %v \n", endpointSliceReconciler.Datastore.pods, test.wantPods)
 			}
 		})
 	}

pkg/ext-proc/backend/llmlserverpool_reconciler_test.go

@@ -41,18 +41,12 @@ func (c *LLMServerPoolReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 		return ctrl.Result{}, err
 	}
 
-	c.updateDatastore(serverPool)
+	klog.V(2).Infof("Updated LLMServerPool: %+v", serverPool)
+	c.Datastore.SetLLMServerPool(serverPool)
 
 	return ctrl.Result{}, nil
 }
 
-func (c *LLMServerPoolReconciler) updateDatastore(serverPool *v1alpha1.LLMServerPool) {
-	if c.Datastore.LLMServerPool == nil || serverPool.ObjectMeta.ResourceVersion != c.Datastore.LLMServerPool.ObjectMeta.ResourceVersion {
-		klog.V(2).Infof("Updated LLMServerPool: %+v", serverPool)
-		c.Datastore.LLMServerPool = serverPool
-	}
-}
-
 func (c *LLMServerPoolReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha1.LLMServerPool{}).

pkg/ext-proc/backend/llmserverpool_reconciler.go

@@ -44,12 +44,12 @@ func (c *LLMServiceReconciler) updateDatastore(service *v1alpha1.LLMService) {
 	for _, ref := range service.Spec.PoolRef {
 		if strings.Contains(strings.ToLower(ref.Kind), strings.ToLower("LLMServerPool")) && ref.Name == c.ServerPoolName {
 			klog.V(2).Infof("Adding/Updating service: %+v", service)
-			c.Datastore.LLMServices.Store(service.Name, service)
+			c.Datastore.llmServices.Store(service.Name, service)
 			return
 		}
 	}
 	klog.V(2).Infof("Removing/Not adding service: %+v", service)
 	// The LLMService may have changed to a different pool. Remove such services.
 	// Otherwise this is a noop.
-	c.Datastore.LLMServices.Delete(service.Name)
+	c.Datastore.llmServices.Delete(service.Name)
 }