Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds Health gRPC Server and Refactor Main() #148

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions pkg/ext-proc/backend/datastore.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At startup, I think we want to ensure that the extension did a sync with the api server and fetched the models, but not declare itself ready only if at least one model is defined.


return ready
}

func RandomWeightedDraw(model *v1alpha1.InferenceModel, seed int64) string {
var weights int32

Expand Down
115 changes: 115 additions & 0 deletions pkg/ext-proc/backend/datastore_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading