Implementing LLMServerPool controller and Pod controller

Author: kfswain

Diff for: .gitignore

@@ -15,6 +15,7 @@ Dockerfile.cross
 
 # Go workspace file
 go.work
+go.work.sum
 
 # Kubernetes Generated files - skip generated files, except for vendored files
 !vendor/**/zz_generated.*

Diff for: pkg/ext-proc/Dockerfile renamed to Dockerfile

@@ -1,11 +1,12 @@
 ## Multistage build
-FROM golang:1.22.5-alpine as build
+FROM golang:1.23-alpine as build
 ENV CGO_ENABLED=0
 ENV GOOS=linux
 ENV GOARCH=amd64
 
 WORKDIR /src
 COPY . .
+WORKDIR /src/pkg/ext-proc
 RUN go mod download
 RUN go build -o /ext-proc
 FROM alpine:latest
@@ -16,4 +17,4 @@ FROM gcr.io/distroless/base-debian10
 WORKDIR /
 COPY --from=build /ext-proc /ext-proc
 
-ENTRYPOINT ["/ext-proc"]
+ENTRYPOINT ["/ext-proc"]

Diff for: examples/poc/ext-proc/go.mod

@@ -1,4 +1,4 @@
-module ext-proc
+module ext-proc-poc
 
 go 1.21
 

Diff for: pkg/ext-proc/backend/datastore.go

@@ -0,0 +1,38 @@
+package backend
+
+import (
+	"sync"
+
+	"inference.networking.x-k8s.io/llm-instance-gateway/api/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/klog/v2"
+)
+
+type Datastore struct {
+	LLMServerPool *v1alpha1.LLMServerPool
+	Pods          sync.Map
+	Port          string
+}
+
+func (ds *Datastore) GetPodIPs() []string {
+	var ips []string
+	ds.Pods.Range(func(name, pod any) bool {
+
+		ips = append(ips, pod.(*corev1.Pod).Status.PodIP)
+		return true
+	})
+	return ips
+}
+
+func (ds *Datastore) LabelsMatch(podLabels map[string]string) bool {
+	selector, err := metav1.LabelSelectorAsSelector(&ds.LLMServerPool.Spec.ModelServerSelector)
+	if err != nil {
+		klog.Error(err.Error())
+		return false
+	}
+	set := labels.Set(podLabels)
+	return selector.Matches(set)
+
+}

Diff for: pkg/ext-proc/backend/llmserverpool_controller.go

@@ -0,0 +1,53 @@
+package backend
+
+import (
+	"context"
+
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"inference.networking.x-k8s.io/llm-instance-gateway/api/v1alpha1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/tools/record"
+	"k8s.io/klog/v2"
+	ctrl "sigs.k8s.io/controller-runtime"
+)
+
+const (
+	controllerNamePrefix = "instance-gateway-"
+)
+
+// LLMServerPoolController is the controller implementation for Instance Gateway resources
+type LLMServerPoolController struct {
+	client.Client
+	Scheme         *runtime.Scheme
+	Record         record.EventRecorder
+	ServerPoolName string
+	Datastore      *Datastore
+}
+
+func (c *LLMServerPoolController) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	logger := klog.FromContext(ctx)
+	logger.V(1).Info("reconciling LLMServerPool")
+
+	var serverPool *v1alpha1.LLMServerPool
+	if err := c.Get(ctx, req.NamespacedName, serverPool); err != nil {
+		logger.Error(err, "unable to get LLMServerPool")
+		return ctrl.Result{}, err
+	}
+
+	c.updateDatastore(serverPool)
+
+	return ctrl.Result{}, nil
+}
+
+func (c *LLMServerPoolController) updateDatastore(serverPool *v1alpha1.LLMServerPool) {
+	if c.Datastore.LLMServerPool == nil || serverPool.ObjectMeta.ResourceVersion != c.Datastore.LLMServerPool.ObjectMeta.ResourceVersion {
+		c.Datastore.LLMServerPool = serverPool
+	}
+}
+
+func (c *LLMServerPoolController) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&v1alpha1.LLMServerPool{}).
+		Complete(c)
+}

Diff for: pkg/ext-proc/backend/pod_controller.go

@@ -0,0 +1,54 @@
+package backend
+
+import (
+	"context"
+
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"k8s.io/client-go/tools/record"
+	"k8s.io/klog/v2"
+)
+
+type PodController struct {
+	client.Client
+	Scheme    *runtime.Scheme
+	Record    record.EventRecorder
+	Datastore *Datastore
+}
+
+func (c *PodController) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	logger := klog.FromContext(ctx)
+	logger.V(1).Info("reconciling Pod")
+
+	var pod *v1.Pod
+	if err := c.Get(ctx, req.NamespacedName, pod); err != nil {
+		logger.Error(err, "unable to get Pod")
+		return ctrl.Result{}, err
+	}
+
+	c.updateDatastore(pod)
+
+	return ctrl.Result{}, nil
+}
+
+func (c *PodController) updateDatastore(pod *v1.Pod) {
+	if pod.ObjectMeta.DeletionTimestamp != nil {
+		klog.Info("Kellen; Pod deletion timestamp is not nil %v", pod.ObjectMeta.DeletionTimestamp)
+	}
+
+	// if labels don't match or pod is scheduled to be deleted, remove in case the pod was part of this pool, otherwise, add.
+	if !c.Datastore.LabelsMatch(pod.ObjectMeta.Labels) || pod.ObjectMeta.DeletionTimestamp != nil {
+		c.Datastore.Pods.Delete(pod.Name)
+	} else {
+		c.Datastore.Pods.Store(pod.Name, pod)
+	}
+}
+
+func (c *PodController) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&v1.Pod{}).
+		Complete(c)
+}

Diff for: pkg/ext-proc/backend/provider.go

@@ -7,18 +7,19 @@ import (
 	"time"
 
 	"go.uber.org/multierr"
+	corev1 "k8s.io/api/core/v1"
 	klog "k8s.io/klog/v2"
 )
 
 const (
 	fetchMetricsTimeout = 5 * time.Second
 )
 
-func NewProvider(pmc PodMetricsClient, pl PodLister) *Provider {
+func NewProvider(pmc PodMetricsClient, datastore *Datastore) *Provider {
 	p := &Provider{
 		podMetrics: sync.Map{},
 		pmc:        pmc,
-		pl:         pl,
+		datastore:  datastore,
 	}
 	return p
 }
@@ -28,17 +29,13 @@ type Provider struct {
 	// key: Pod, value: *PodMetrics
 	podMetrics sync.Map
 	pmc        PodMetricsClient
-	pl         PodLister
+	datastore  *Datastore
 }
 
 type PodMetricsClient interface {
 	FetchMetrics(ctx context.Context, pod Pod, existing *PodMetrics) (*PodMetrics, error)
 }
 
-type PodLister interface {
-	List() (PodSet, error)
-}
-
 func (p *Provider) AllPodMetrics() []*PodMetrics {
 	res := []*PodMetrics{}
 	fn := func(k, v any) bool {
@@ -108,13 +105,11 @@ func (p *Provider) Init(refreshPodsInterval, refreshMetricsInterval time.Duratio
 // refreshPodsOnce lists pods and updates keys in the podMetrics map.
 // Note this function doesn't update the PodMetrics value, it's done separately.
 func (p *Provider) refreshPodsOnce() error {
-	pods, err := p.pl.List()
-	if err != nil {
-		return err
-	}
 	// merge new pods with cached ones.
 	// add new pod to the map
-	for pod := range pods {
+	addNewPods := func(k, v any) bool {
+		k8sPod := v.(*corev1.Pod)
+		pod := Pod{Name: k8sPod.Name, Namespace: k8sPod.Namespace, Address: k8sPod.Status.PodIP + ":" + p.datastore.Port}
 		if _, ok := p.podMetrics.Load(pod); !ok {
 			new := &PodMetrics{
 				Pod: pod,
@@ -124,16 +119,18 @@ func (p *Provider) refreshPodsOnce() error {
 			}
 			p.podMetrics.Store(pod, new)
 		}
+		return true
 	}
 	// remove pods that don't exist any more.
 	mergeFn := func(k, v any) bool {
 		pod := k.(Pod)
-		if _, ok := pods[pod]; !ok {
+		if _, ok := p.datastore.Pods.Load(pod.Name); !ok {
 			p.podMetrics.Delete(pod)
 		}
 		return true
 	}
 	p.podMetrics.Range(mergeFn)
+	p.datastore.Pods.Range(addNewPods)
 	return nil
 }
 

Diff for: pkg/ext-proc/go.mod

@@ -1,6 +1,6 @@
 module ext-proc
 
-go 1.21
+go 1.22.0
 
 require (
 	github.com/bojand/ghz v0.120.0
@@ -9,12 +9,14 @@ require (
 	github.com/jhump/protoreflect v1.15.1
 	github.com/prometheus/client_model v0.6.1
 	github.com/prometheus/common v0.55.0
-	go.uber.org/multierr v1.9.0
+	go.uber.org/multierr v1.11.0
 	google.golang.org/grpc v1.67.0
 	google.golang.org/protobuf v1.34.2
 	k8s.io/klog/v2 v2.130.1
 )
 
+require sigs.k8s.io/controller-runtime v0.19.0 // indirect
+
 require (
 	cel.dev/expr v0.16.0 // indirect
 	cloud.google.com/go/compute/metadata v0.5.0 // indirect
@@ -27,31 +29,60 @@ require (
 	github.com/census-instrumentation/opencensus-proto v0.4.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cncf/xds/go v0.0.0-20240723142845-024c85f92f20 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
 	github.com/envoyproxy/protoc-gen-validate v1.1.0 // indirect
-	github.com/go-logr/logr v1.4.1 // indirect
+	github.com/fxamacker/cbor/v2 v2.7.0 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-openapi/jsonpointer v0.19.6 // indirect
+	github.com/go-openapi/jsonreference v0.20.2 // indirect
+	github.com/go-openapi/swag v0.22.4 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
+	github.com/google/gnostic-models v0.6.8 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
+	github.com/google/gofuzz v1.2.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/huandu/xstrings v1.3.3 // indirect
 	github.com/imdario/mergo v0.3.11 // indirect
 	github.com/jinzhu/configor v1.2.1 // indirect
-	github.com/kr/text v0.2.0 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mitchellh/copystructure v1.0.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.1 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
 	github.com/shopspring/decimal v1.2.0 // indirect
 	github.com/spf13/cast v1.4.1 // indirect
-	go.uber.org/atomic v1.7.0 // indirect
+	github.com/x448/float16 v0.8.4 // indirect
 	golang.org/x/crypto v0.26.0 // indirect
 	golang.org/x/net v0.28.0 // indirect
 	golang.org/x/oauth2 v0.22.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.24.0 // indirect
+	golang.org/x/term v0.23.0 // indirect
 	golang.org/x/text v0.17.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+	inference.networking.x-k8s.io/llm-instance-gateway v0.0.0
+	k8s.io/api v0.31.1 // indirect
+	k8s.io/apimachinery v0.31.1
+	k8s.io/client-go v0.31.1
+	k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340 // indirect
+	k8s.io/utils v0.0.0-20240711033017-18e509b52bc8 // indirect
+	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
+	sigs.k8s.io/yaml v1.4.0 // indirect
 )
+
+//replace inference.networking.x-k8s.io/llm-instance-gateway v0.0.0 => /usr/local/google/home/kfswain/repos/kubernetes/llm-instance-gateway
+replace inference.networking.x-k8s.io/llm-instance-gateway v0.0.0 => /src