Add health gRPC server and refactor main() for better lifecycle management

- Introduced a health gRPC server to handle liveness and readiness probes.
- Refactored main() to manage server goroutines using sync.WaitGroup.
- Added graceful shutdown for servers and controller manager.
- Improved logging consistency and ensured datastore readiness checks.

Signed-off-by: Daneyon Hansen <[email protected]>

Author: danehans

pkg/ext-proc/backend/datastore.go

@@ -75,6 +75,29 @@ func (s *K8sDatastore) FetchModelData(modelName string) (returnModel *v1alpha1.I
 	return
 }
 
+func (ds *K8sDatastore) IsReady() bool {
+	ds.poolMu.RLock()
+	defer ds.poolMu.RUnlock()
+
+	// Check if the inference pool is set
+	if ds.inferencePool == nil {
+		return false
+	}
+
+	// Check if there are valid models associated with the pool
+	ready := false
+	ds.InferenceModels.Range(func(key, value any) bool {
+		model, ok := value.(*v1alpha1.InferenceModel)
+		if !ok || model == nil || model.Spec.PoolRef.Name != ds.inferencePool.Name {
+			return true
+		}
+		ready = true
+		return false
+	})
+
+	return ready
+}
+
 func RandomWeightedDraw(model *v1alpha1.InferenceModel, seed int64) string {
 	var weights int32
 

pkg/ext-proc/backend/datastore_test.go

@@ -4,8 +4,123 @@ import (
 	"testing"
 
 	"inference.networking.x-k8s.io/llm-instance-gateway/api/v1alpha1"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
+func TestIsReady(t *testing.T) {
+	tests := []struct {
+		name            string
+		inferencePool   *v1alpha1.InferencePool
+		inferenceModels []*v1alpha1.InferenceModel
+		expectedReady   bool
+	}{
+		{
+			name: "Ready when at least one model matches configured pool",
+			inferencePool: &v1alpha1.InferencePool{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-pool",
+					Namespace: "default",
+				},
+			},
+			inferenceModels: []*v1alpha1.InferenceModel{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Name:      "test-model",
+						Namespace: "default",
+					},
+					Spec: v1alpha1.InferenceModelSpec{
+						PoolRef: v1alpha1.PoolObjectReference{
+							Name: "other-pool",
+						},
+					},
+				},
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Name:      "test-model",
+						Namespace: "default",
+					},
+					Spec: v1alpha1.InferenceModelSpec{
+						PoolRef: v1alpha1.PoolObjectReference{
+							Name: "test-pool",
+						},
+					},
+				},
+			},
+			expectedReady: true,
+		},
+		{
+			name: "Not ready when model references non-matching pool",
+			inferencePool: &v1alpha1.InferencePool{
+				ObjectMeta: v1.ObjectMeta{
+					Name:      "test-pool",
+					Namespace: "default",
+				},
+			},
+			inferenceModels: []*v1alpha1.InferenceModel{
+				{
+					ObjectMeta: v1.ObjectMeta{
+						Name:      "test-model",
+						Namespace: "default",
+					},
+					Spec: v1alpha1.InferenceModelSpec{
+						PoolRef: v1alpha1.PoolObjectReference{
+							Name: "other-pool",
+						},
+					},
+				},
+			},
+			expectedReady: false,
+		},
+		{
+			name:          "Not ready when pool is nil",
+			inferencePool: nil,
+			inferenceModels: []*v1alpha1.InferenceModel{
+				{Spec: v1alpha1.InferenceModelSpec{}},
+				{Spec: v1alpha1.InferenceModelSpec{}},
+			},
+			expectedReady: false,
+		},
+		{
+			name:          "Not ready when models are missing",
+			inferencePool: &v1alpha1.InferencePool{},
+			inferenceModels: []*v1alpha1.InferenceModel{
+				nil,
+			},
+			expectedReady: false,
+		},
+		{
+			name:            "Not ready when models are empty",
+			inferencePool:   &v1alpha1.InferencePool{},
+			inferenceModels: []*v1alpha1.InferenceModel{},
+			expectedReady:   false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			datastore := NewK8sDataStore()
+
+			// Set the inference pool
+			if tt.inferencePool != nil {
+				datastore.setInferencePool(tt.inferencePool)
+			}
+
+			// Set the inference models
+			for _, model := range tt.inferenceModels {
+				if model != nil {
+					datastore.InferenceModels.Store(model.Spec.ModelName, model)
+				}
+			}
+
+			// Check readiness
+			isReady := datastore.IsReady()
+			if isReady != tt.expectedReady {
+				t.Errorf("IsReady() = %v, want %v", isReady, tt.expectedReady)
+			}
+		})
+	}
+}
+
 func TestRandomWeightedDraw(t *testing.T) {
 	tests := []struct {
 		name  string