Repository: vllm-project/aibrix
File: pkg/cache/cache.go (Go)

Code before:
type Cache struct {
	mu                sync.RWMutex
	redisClient       *redis.Client
	prometheusApi     prometheusv1.API
	initialized       bool
	subscribers       []metrics.MetricSubscriber
	metrics           map[string]interface{}
	ModelMetrics      map[string]map[string]interface{}
	Pods              map[string]*v1.Pod
	PodMetrics        map[string]map[string]metrics.MetricValue            // pod_name: map[metric_name]metric_val
	PodModelMetrics   map[string]map[string]map[string]metrics.MetricValue // pod_name: map[model_name]map[metric_name]metric_val
	PodToModelMapping map[string]map[string]struct{}                       // pod_name: map[model_name]struct{}
	ModelToPodMapping map[string]map[string]*v1.Pod                        // model_name: map[pod_name]*v1.Pod
	requestTrace      *sync.Map                                            // model_name: RequestTrace
	numRequestsTraces int32                                                // counter for requestTrace
	pendingRequests   *sync.Map                                            // model_name: *int32
}

type Block struct {
	modelToPods    map[string]map[string]time.Time // model_name: map[pod_name]pod_last_access_time
	lastAccessTime time.Time                       //block_last_access_time
}

const (
	modelIdentifier                       = "model.aibrix.ai/name"
	nodeType                              = "ray.io/node-type"
	podPort                               = 8000
	defaultPodMetricRefreshIntervalInMS   = 50
	expireWriteRequestTraceIntervalInMins = 10
)

var (
	instance                Cache
	counterGaugeMetricNames = []string{
		metrics.NumRequestsRunning,
		metrics.NumRequestsWaiting,
		metrics.NumRequestsSwapped,
		metrics.AvgPromptThroughputToksPerS,
		metrics.AvgGenerationThroughputToksPerS,
		metrics.GPUCacheUsagePerc,
		metrics.CPUCacheUsagePerc,
	}
	// histogram metric example - time_to_first_token_seconds, _sum, _bucket _count.
	histogramMetricNames = []string{
		metrics.IterationTokensTotal,
		metrics.TimeToFirstTokenSeconds,
		metrics.TimePerOutputTokenSeconds,
		metrics.E2ERequestLatencySeconds,
		metrics.RequestQueueTimeSeconds,
		metrics.RequestInferenceTimeSeconds,
		metrics.RequestDecodeTimeSeconds,

Review comment (suggestion):
I am thinking if we want to add a aibrix label like `orchestration.aibrix.ai/server-pod`. autoscaling need that feature as well. We can start with ray specific label and replace with our own later as a separate issue


Code after:
type Cache struct {
	mu                sync.RWMutex
	redisClient       *redis.Client
	prometheusApi     prometheusv1.API
	initialized       bool
	subscribers       []metrics.MetricSubscriber
	metrics           map[string]interface{}
	ModelMetrics      map[string]map[string]interface{}
	Pods              map[string]*v1.Pod
	PodMetrics        map[string]map[string]metrics.MetricValue            // pod_name: map[metric_name]metric_val
	PodModelMetrics   map[string]map[string]map[string]metrics.MetricValue // pod_name: map[model_name]map[metric_name]metric_val
	PodToModelMapping map[string]map[string]struct{}                       // pod_name: map[model_name]struct{}
	ModelToPodMapping map[string]map[string]*v1.Pod                        // model_name: map[pod_name]*v1.Pod
	requestTrace      *sync.Map                                            // model_name: RequestTrace
	numRequestsTraces int32                                                // counter for requestTrace
	pendingRequests   *sync.Map                                            // model_name: *int32
}

type Block struct {
	modelToPods    map[string]map[string]time.Time // model_name: map[pod_name]pod_last_access_time
	lastAccessTime time.Time                       //block_last_access_time
}

const (
	modelIdentifier                       = "model.aibrix.ai/name"
	nodeType                              = "ray.io/node-type"
	nodeWorker                            = "worker"
	podPort                               = 8000
	defaultPodMetricRefreshIntervalInMS   = 50
	expireWriteRequestTraceIntervalInMins = 10
)

var (
	instance                Cache
	counterGaugeMetricNames = []string{
		metrics.NumRequestsRunning,
		metrics.NumRequestsWaiting,
		metrics.NumRequestsSwapped,
		metrics.AvgPromptThroughputToksPerS,
		metrics.AvgGenerationThroughputToksPerS,
		metrics.GPUCacheUsagePerc,
		metrics.CPUCacheUsagePerc,
	}
	// histogram metric example - time_to_first_token_seconds, _sum, _bucket _count.
	histogramMetricNames = []string{
		metrics.IterationTokensTotal,
		metrics.TimeToFirstTokenSeconds,
		metrics.TimePerOutputTokenSeconds,
		metrics.E2ERequestLatencySeconds,
		metrics.RequestQueueTimeSeconds,
		metrics.RequestInferenceTimeSeconds,