Add initial ext proc implementation with LoRA affinity

pkg/ext-proc/Dockerfile

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

pkg/ext-proc/backend/fake.go

@@ -0,0 +1,26 @@
+package backend
+
+import (
+	dto "github.com/prometheus/client_model/go"
+)
+
+type FakePodLister struct {
+	Err  error
+	Pods PodSet
+}
+
+type FakePodMetricsClient struct {
+	Err map[Pod]error
+	Res map[Pod]map[string]*dto.MetricFamily
+}
+
+func (f *FakePodMetricsClient) FetchMetrics(pod Pod) (map[string]*dto.MetricFamily, error) {
+	if err, ok := f.Err[pod]; ok {
+		return nil, err
+	}
+	return f.Res[pod], nil
+}
+
+func (fpl *FakePodLister) List() (PodSet, error) {
+	return fpl.Pods, fpl.Err
+}

pkg/ext-proc/backend/metrics.go

@@ -0,0 +1,135 @@
+package backend
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	dto "github.com/prometheus/client_model/go"
+	"go.uber.org/multierr"
+	klog "k8s.io/klog/v2"
+)
+
+const (
+	ActiveLoRAAdaptersMetricName        = "vllm:info_active_adapters_info"
+	LoRAAdapterPendingRequestMetricName = "vllm:active_lora_adapters"
+	// TODO: Replace these with the num_tokens_running/waiting below once we add those to the fork.
+	RunningQueueSizeMetricName = "vllm:num_requests_running"
+	WaitingQueueSizeMetricName = "vllm:num_requests_waiting"
+	/* TODO: Uncomment this once the following are added to the fork.
+	RunningQueueSizeMetricName        = "vllm:num_tokens_running"
+	WaitingQueueSizeMetricName        = "vllm:num_tokens_waiting"
+	*/
+	KVCacheUsagePercentMetricName     = "vllm:gpu_cache_usage_perc"
+	KvCacheMaxTokenCapacityMetricName = "vllm:gpu_cache_max_token_capacity"
+)
+
+func (p *Provider) refreshMetricsOnce() error {
+	start := time.Now()
+	defer func() {
+		d := time.Now().Sub(start)
+		// TODO: add a metric instead of logging
+		klog.V(3).Infof("Refreshed metrics in %v", d)
+	}()
+	var wg sync.WaitGroup
+	var errs error
+	processOnePod := func(key, value any) bool {
+		pod := key.(Pod)
+		metrics := value.(*PodMetrics)
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			metricFamilies, err := p.pmc.FetchMetrics(pod)
+			if err != nil {
+				multierr.Append(errs, fmt.Errorf("failed to parse metrics from %s: %v", pod, err))
+				return
+			}
+			updated, err := promToPodMetrics(metricFamilies, metrics)
+			klog.V(3).Infof("Updated metrics for pod %s: %v", pod, updated.Metrics)
+			if err != nil {
+				multierr.Append(errs, fmt.Errorf("failed to get all pod metrics updated from prometheus: %v", err))
+			}
+			p.UpdatePodMetrics(pod, updated)
+		}()
+		return true
+	}
+	p.podMetrics.Range(processOnePod)
+	wg.Wait()
+	return errs
+}
+
+// promToPodMetrics updates internal pod metrics with scraped prometheus metrics.
+// A combined error is returned if errors occur in one or more metric processing.
+// it returns a new PodMetrics pointer which can be used to atomically update the pod metrics map.
+func promToPodMetrics(metricFamilies map[string]*dto.MetricFamily, existing *PodMetrics) (*PodMetrics, error) {
+	var errs error
+	updated := existing.Clone()
+	runningQueueSize, _, err := getLatestMetric(metricFamilies, RunningQueueSizeMetricName)
+	multierr.Append(errs, err)
+	if err != nil {
+		updated.RunningQueueSize = int(runningQueueSize.GetCounter().GetValue())
+	}
+	waitingQueueSize, _, err := getLatestMetric(metricFamilies, WaitingQueueSizeMetricName)
+	multierr.Append(errs, err)
+	if err != nil {
+		updated.WaitingQueueSize = int(waitingQueueSize.GetGauge().GetValue())
+	}
+	cachePercent, _, err := getLatestMetric(metricFamilies, KVCacheUsagePercentMetricName)
+	multierr.Append(errs, err)
+	if err != nil {
+		updated.KVCacheUsagePercent = cachePercent.GetGauge().GetValue()
+	}
+	/* TODO: uncomment once this is available in vllm.
+	kvCap, _, err := getGaugeLatestValue(metricFamilies, KvCacheMaxTokenCapacityMetricName)
+	multierr.Append(errs, err)
+	if err != nil {
+		updated.KvCacheMaxTokenCapacity = int(kvCap)
+	}
+	*/
+
+	// Update active loras
+	mf, ok := metricFamilies[ActiveLoRAAdaptersMetricName]
+	if ok {
+		// IMPORTANT: replace the map entries instead of appending to it.
+		updated.CachedModels = make(map[string]int)
+		for _, metric := range mf.GetMetric() {
+			for _, label := range metric.GetLabel() {
+				if label.GetName() == "active_adapters" {
+					if label.GetValue() != "" {
+						adapterList := strings.Split(label.GetValue(), ",")
+						for _, adapter := range adapterList {
+							updated.CachedModels[adapter] = 0
+						}
+					}
+				}
+			}
+		}
+	} else {
+		klog.Warningf("metric family %q not found", ActiveLoRAAdaptersMetricName)
+		multierr.Append(errs, fmt.Errorf("metric family %q not found", ActiveLoRAAdaptersMetricName))
+	}
+
+	return updated, errs
+}
+
+// getLatestMetric gets the latest metric of a family. This should be used to get the latest Gauge metric.
+func getLatestMetric(metricFamilies map[string]*dto.MetricFamily, metricName string) (*dto.Metric, time.Time, error) {
+	mf, ok := metricFamilies[metricName]
+	if !ok {
+		klog.Warningf("metric family %q not found", metricName)
+		return nil, time.Time{}, fmt.Errorf("metric family %q not found", metricName)
+	}
+	if len(mf.GetMetric()) == 0 {
+		return nil, time.Time{}, fmt.Errorf("no metrics available for %q", metricName)
+	}
+	var latestTs int64
+	var latest *dto.Metric
+	for _, m := range mf.GetMetric() {
+		if m.GetTimestampMs() > latestTs {
+			latestTs = m.GetTimestampMs()
+			latest = m
+		}
+	}
+	return latest, time.Unix(0, latestTs*1000), nil
+}

pkg/ext-proc/backend/pod_client.go

@@ -0,0 +1,34 @@
+package backend
+
+import (
+	"fmt"
+	"net/http"
+
+	dto "github.com/prometheus/client_model/go"
+	"github.com/prometheus/common/expfmt"
+	klog "k8s.io/klog/v2"
+)
+
+type PodMetricsClientImpl struct {
+}
+
+// FetchMetrics fetches metrics from a given pod.
+func (p *PodMetricsClientImpl) FetchMetrics(pod Pod) (map[string]*dto.MetricFamily, error) {
+	// Currently the metrics endpoint is hard-coded.
+	// TODO: Consider making this configurable.
+	url := fmt.Sprintf("http://%s/metrics", pod.Address)
+	resp, err := http.Get(url)
+	if err != nil {
+		klog.Errorf("failed to fetch metrics from %s: %v", pod, err)
+		return nil, fmt.Errorf("failed to fetch metrics from %s: %w", pod, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		klog.Errorf("unexpected status code from %s: %v", pod, resp.StatusCode)
+		return nil, fmt.Errorf("unexpected status code from %s: %v", pod, resp.StatusCode)
+	}
+
+	parser := expfmt.TextParser{}
+	return parser.TextToMetricFamilies(resp.Body)
+}

pkg/ext-proc/backend/provider.go

@@ -0,0 +1,122 @@
+package backend
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	dto "github.com/prometheus/client_model/go"
+	klog "k8s.io/klog/v2"
+)
+
+func NewProvider(pmc PodMetricsClient, pl PodLister) *Provider {
+	p := &Provider{
+		podMetrics: sync.Map{},
+		pmc:        pmc,
+		pl:         pl,
+	}
+	return p
+}
+
+// Provider provides backend pods and information such as metrics.
+type Provider struct {
+	// key: Pod, value: *PodMetrics
+	podMetrics sync.Map
+	pmc        PodMetricsClient
+	pl         PodLister
+}
+
+type PodMetricsClient interface {
+	FetchMetrics(pod Pod) (map[string]*dto.MetricFamily, error)
+}
+
+type PodLister interface {
+	List() (PodSet, error)
+}
+
+func (p *Provider) AllPodMetrics() []*PodMetrics {
+	res := []*PodMetrics{}
+	fn := func(k, v any) bool {
+		res = append(res, v.(*PodMetrics))
+		return true
+	}
+	p.podMetrics.Range(fn)
+	return res
+}
+
+func (p *Provider) UpdatePodMetrics(pod Pod, pm *PodMetrics) {
+	p.podMetrics.Store(pod, pm)
+}
+
+func (p *Provider) GetPodMetrics(pod Pod) (*PodMetrics, bool) {
+	val, ok := p.podMetrics.Load(pod)
+	if ok {
+		return val.(*PodMetrics), true
+	}
+	return nil, false
+}
+
+func (p *Provider) Init(refreshPodsInterval, refreshMetricsInterval time.Duration) error {
+	if err := p.refreshPodsOnce(); err != nil {
+		return fmt.Errorf("failed to init pods: %v", err)
+	}
+	if err := p.refreshMetricsOnce(); err != nil {
+		return fmt.Errorf("failed to init metrics: %v", err)
+	}
+
+	klog.V(2).Infof("Initialized pods and metrics: %+v", p.AllPodMetrics())
+
+	// periodically refresh pods
+	go func() {
+		for {
+			time.Sleep(refreshPodsInterval)
+			if err := p.refreshPodsOnce(); err != nil {
+				klog.V(1).Infof("Failed to refresh podslist pods: %v", err)
+			}
+		}
+	}()
+
+	// periodically refresh metrics
+	go func() {
+		for {
+			time.Sleep(refreshMetricsInterval)
+			if err := p.refreshMetricsOnce(); err != nil {
+				klog.V(1).Infof("Failed to refresh metrics: %v", err)
+			}
+		}
+	}()
+
+	return nil
+}
+
+// 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 {
+		if _, ok := p.podMetrics.Load(pod); !ok {
+			new := &PodMetrics{
+				Pod: pod,
+				Metrics: Metrics{
+					CachedModels: make(map[string]int),
+				},
+			}
+			p.podMetrics.Store(pod, new)
+		}
+	}
+	// remove pods that don't exist any more.
+	mergeFn := func(k, v any) bool {
+		pod := k.(Pod)
+		if _, ok := pods[pod]; !ok {
+			p.podMetrics.Delete(pod)
+		}
+		return true
+	}
+	p.podMetrics.Range(mergeFn)
+	return nil
+}

pkg/ext-proc/backend/types.go

@@ -0,0 +1,52 @@
+// Package backend is a library to interact with backend model servers such as probing metrics.
+package backend
+
+import "fmt"
+
+type PodSet map[Pod]bool
+
+type Pod struct {
+	Namespace string
+	Name      string
+	Address   string
+}
+
+func (p Pod) String() string {
+	return p.Namespace + "." + p.Name
+}
+
+type Metrics struct {
+	// CachedModels is a set of models(including LoRA adapters) that are currently cached to GPU.
+	CachedModels            map[string]int
+	RunningQueueSize        int
+	WaitingQueueSize        int
+	KVCacheUsagePercent     float64
+	KvCacheMaxTokenCapacity int
+}
+
+type PodMetrics struct {
+	Pod
+	Metrics
+}
+
+func (pm *PodMetrics) String() string {
+	return fmt.Sprintf("Pod: %+v; Metrics: %+v", pm.Pod, pm.Metrics)
+}
+
+func (pm *PodMetrics) Clone() *PodMetrics {
+	cm := make(map[string]int, len(pm.CachedModels))
+	for k, v := range pm.CachedModels {
+		cm[k] = v
+	}
+	clone := &PodMetrics{
+		Pod: pm.Pod,
+		Metrics: Metrics{
+			CachedModels:            cm,
+			RunningQueueSize:        pm.RunningQueueSize,
+			WaitingQueueSize:        pm.WaitingQueueSize,
+			KVCacheUsagePercent:     pm.KVCacheUsagePercent,
+			KvCacheMaxTokenCapacity: pm.KvCacheMaxTokenCapacity,
+		},
+	}
+	return clone
+}