AIBrix Cache 缓存包深入解析Pod 元数据、模型路由与 KV 事件同步【免费下载链接】aibrixCost-efficient and pluggable Infrastructure components for GenAI inference项目地址: https://gitcode.com/GitHub_Trending/ai/aibrix导读AIBrix 的pkg/cache包是整个 LLM 推理系统的中枢神经它以内存缓存为核心向上为 Envoy Gateway 插件与各类路由算法提供 Pod 可用性、负载、前缀缓存命中率等实时决策数据向下通过 Kubernetes Informer 与 ZMQ 事件流打通 vLLM Pod 的 KV Cache 变更。本文以 pkg/cache/README.md 为骨架结合源码逐层剖析其缓存存储、请求追踪、输出预测、GPU Profile 与 KV 事件同步的实现原理并给出完整的配置项与可运行示例帮助你掌握 AIBrix 路由决策的数据基础。一、Cache 包定位路由决策的数据中枢AIBrix 的网关插件在每次请求路由前都需要回答几个关键问题某个模型当前有哪些可用 Pod哪个 Pod 负载最低哪些 Pod 的前缀缓存命中率最高目标 Pod 的 GPU 利用率如何pkg/cache正是为回答这些问题而设计的集中式缓存与路由支撑组件。从 cache_api.go 可以看到Cache是一个聚合接口它嵌入了 8 个子接口type Cache interface { PodCache ModelCache MetricCache RequestTracker RequestTrackerRegistry ProfileCache types.OutputPredictorProvider types.RouterProvider }这从结构上决定了缓存包的五大核心职责能力接口说明Pod 与模型元数据缓存PodCache/ModelCache按 Pod 查模型、按模型查 Pod以及 LoRA 适配器与基础模型的映射请求路由与负载追踪MetricCache/RequestTrackerPod 指标查询、跨网关运行中请求数统计、请求生命周期追踪KV 缓存事件同步事件管道 前缀索引器vLLM Pod 的 KV 块事件经 ZMQ 同步进前缀缓存索引器GPU Profile 管理ProfileCache按 Pod / 按 Deployment 查询模型 GPU 画像输出预测OutputPredictorProvider基于历史滑窗预测请求的输出 token 数二、核心组件Store 与它的存储结构2.1 Store全局单例缓存仓库Store是缓存包的中心数据结构定义于 cache_init.go。它采用全局单例 sync.Once初始化的模式cache_init.govar ( store Store{} // Global cache store instance once sync.Once // Singleton pattern control lock )Store内部的核心字段包括Pod 存储metaPodsSyncMap[string, *Pod]以namespace/name为键存放所有被发现的 PodrecentlyDeletedPods保存刚删除 Pod 的实时计数器快照支持 Pod 快速重建后恢复请求追踪计数而不是从零开始。模型存储metaModels模型名 →*ModelmodelClaims单独存放不可路由端口为 0的 ModelClaim 广播状态避免污染正常路由。Deployment 画像deploymentProfiles键为aibrix:profile_[model_name]_[deployment_name]。KV 同步syncPrefixIndexer仅在启用 KV 同步时创建与kvEventManager。运行时统计enableTracing是否开启 GPU Optimizer 追踪、podMetricsWorkerCount与podMetricsJobs组成的 Pod 指标更新 worker 池、requestTrackers注册的请求追踪器列表。2.2 初始化流程与三类服务身份InitWithOptions 是生产环境的统一入口它通过InitOptions判断当前调用方身份从而决定启用哪些能力type InitOptions struct { IsGateway bool // 标记调用方为 gateway-plugins 服务 EnableKVSync bool // 是否启动 ZMQ KV 事件同步 RedisClient *redis.Client // KV 同步等特性必需 ModelRouterProvider ModelRouterProviderFunc // 仅网关需要 DiscoveryProvider discovery.Provider // 服务发现 Provider默认为 Kubernetes Informer }根据serviceIdentitycache_init.go的判断逻辑gatewayIsGatewaytrue启用追踪与 GPU Profile 缓存支撑路由metadata有 RedisClient 但非网关关闭 GPU Optimizer 追踪与 Profile 缓存controllers其余场景同样关闭追踪与 Profile 缓存。初始化时若EnableKVSynctrue但RedisClient为 nil会直接klog.Fatalf退出——这是硬性依赖校验。此外初始化还会视 Redis 可用性启动网关快照同步initGatewaySnapshotSync与运行中请求活性心跳initRunningRequestsLiveness后者用于跨网关精确统计每个 Pod 当前在途请求数。2.3 服务发现Kubernetes Informer 与可插拔 Provider缓存数据从何而来答案是服务发现层。所有 Pod / ModelAdapter 的增删改事件都通过统一的discovery.Provider接口进入缓存cache_init.goprovider : opts.DiscoveryProvider if provider nil { // Default: Kubernetes informer-based discovery provider discovery.NewKubernetesProvider(config) }discovery子包pkg/cache/discovery提供了两种开箱即用的 ProviderKubernetes Provider基于 Informer 监听 Pod 生命周期与 ModelAdapter 资源对应 kubernetes.goStatic Provider用于独立/开发模式无需 Kubernetes对应 static.go另有 static_test.go 验证其行为。事件进入handleDiscoveryObject后被分派到addPod/updatePod/deletePod/addModelAdapter等内部方法完成内存索引的维护。informers.go中注册的 Informer 则负责监听 Pod 生命周期事件、ModelAdapter 资源与配置更新。三、KV 缓存事件同步从 vLLM Pod 到前缀索引器KV 事件同步是pkg/cache最具特色的能力它让网关能实时感知每个 vLLM Pod 中 KV Cache 块的存储与移除从而在路由时优先选择前缀缓存命中率高的 Pod显著降低 TTFT首 token 延迟。3.1 事件链路全景README 中给出了完整的五步事件流Pod Discovery缓存监听带有 KV 事件标签的 PodSubscription事件管理器为符合条件的 Pod 创建 ZMQ 客户端Event Processing进入的事件更新前缀缓存索引器Query Routing路由器向索引器查询前缀匹配Request Dispatch请求被路由到具有匹配前缀的 Pod。3.2 KVEventManager带构建标签的双实现KVEventManager是同步的入口它通过Go build tags区分两种实现非 ZMQ 构建默认kv_event_manager.go 中是一个 stubvalidateConfiguration()直接返回需要-tagszmq构建的错误所有OnPodAdd/Update/Delete均为空操作ZMQ 构建-tagszmqkv_event_manager_zmq.go 中它内嵌了*kvevent.Manager通过NewStoreProviderAdapter(store)将 Store 适配为 Pod Provider 与同步 Provider复用 pkg/kvevent 的完整事件管理器实现。其配置校验validateKVEventConfiguration见 kv_event_manager_zmq.go揭示了启用 KV 同步的四项硬性前提环境变量要求说明AIBRIX_PREFIX_CACHE_KV_EVENT_SYNC_ENABLEDtrue总开关AIBRIX_PREFIX_CACHE_USE_REMOTE_TOKENIZERtrue必须启用远程 tokenizerAIBRIX_PREFIX_CACHE_TOKENIZER_TYPEremotetokenizer 类型必须为 remoteAIBRIX_PREFIX_CACHE_REMOTE_TOKENIZER_ENDPOINT非空远程 tokenizer 服务地址在 initKVEventSync 中Store 依次执行读取开关 → 校验配置 → 创建事件管理器 → 创建共享的单例SyncPrefixHashTablesyncindexer.GetSharedSyncPrefixHashTable→ 启动管理器。注意清理时cleanupKVEventSync不会关闭共享索引器单例因为它可能仍被网关路由器等其他组件使用其生命周期由全局统一管理。3.3 底层传输ZMQ 客户端与 MessagePack 编解码KV 事件的实际收发由 pkg/cache/kvcache 子包承担其 READMEkvcache/README.md描述了三个关键组件ZMQ 客户端zmq_client.go——连接 vLLM Pod 并订阅 KV 事件流具备指数退避自动重连、事件回放Replay、序号追踪检测漏事件、可配置超时与缓冲区等可靠性特性config : ZMQClientConfig{ PodName: vllm-pod-1, PodIP: 10.0.0.1, PubPort: 5557, RouterPort: 5558, PollTimeout: 100 * time.Millisecond, } client : NewZMQClient(config, handler) go client.Start(ctx)事件类型event_types.goBlockStoredEvent新 KV 块存储、BlockRemovedEvent块被移除、AllBlocksClearedEvent全部块清空对应事件处理器中block stored/removed的分发逻辑。MessagePack 编解码器msgpack_encoder.go 与 msgpack_decoder.go对事件批量序列化降低网络传输开销data, err : EncodeEventBatch(eventBatch) // 编码 events, err : DecodeEventBatch(data, modelName, podName) // 解码依赖github.com/pebbe/zmq4ZMQ Go 绑定、github.com/vmihailenco/msgpack/v5系统需安装 libzmq3 或更高版本。3.4 事件处理器与多模型部署事件处理器将收到的 KV 事件路由到合适的索引器处理块存储/移除事件并维护 Pod 重启后的一致性。事件管理器支持自动 Pod 发现与订阅、生命周期管理增/改/删以及多模型部署——同一 Pod 可为多个模型服务其 KV 事件按模型维度落入对应索引。四、请求追踪与输出预测4.1 RequestTrace单请求全链路数据RequestTrace 记录单个请求在系统中的完整轨迹请求耗时与延迟、token 数与模型信息、GPU 利用率指标。Store 以model_name - *RequestTrace的映射保存requestTrace字段并通过RequestTracker接口对外暴露。RequestTracker接口cache_api.go定义了三阶段契约AddRequestCount(ctx, requestID, modelName)路由后记录请求开始返回traceTerm追踪期标识可被多次调用实现必须保证线程安全与幂等DoneRequestCount(...)记录无 usage 信息的请求完成DoneRequestTrace(...)记录带 inputTokens / outputTokens 的请求完成。由于AddRequestCount的 ctx 可能为 nil如请求在路由完成前被取消所有实现都必须对 nil ctx 做防护。RequestTrackerRegistry允许多个 tracker 按注册顺序依次被调用便于扩展自定义追踪器。追踪数据启用后initTraceCachecache_init.go会按固定周期RequestTraceWriteInterval将窗口内轨迹批量写入 Redis 存储。4.2 OutputPredictor加权随机输出预测输出预测器output_predictor.go用于容量规划在请求到达前预测其输出 token 数帮助调度器预估 GPU 负载。其核心是SimpleOutputPredictor滑窗直方图维护最近movingWindow默认 240s与 GPU Optimizer 窗口一致见 cache_init.go内的历史分布按MovingInterval 10s滚动对数分桶输入/输出 token 分别按round(log2(tokens))分桶桶数量由maxInputTokens 1M、maxOutputTokens 1M计算cache_init.go加权随机预测Predict(inputTokens)在对应输入桶内做加权随机采样命中概率与该输出桶的历史频次成正比冷启动策略无历史数据时ColdPredictionStrategy提供四种回退——Optimistic默认预测为 1对 Profile 最友好、Random0 到MaxOutputLen4096随机、Input输出等于输入、Pessimistic直接取MaxOutputLen。五、Pod 实时负载与 GPU Profile5.1 Pod 元数据与运行中请求计数Pod结构pod.go在原生*v1.Pod之上叠加了三层缓存数据Models该 Pod 承载的模型/适配器名集合utils.Registry[string]Metrics/ModelMetricsPod 级与 Pod-模型级指标Prometheus 抓取的引擎 gauge、PromQL 结果、网关派生值等实时统计runningRequests运行中请求原子计数、completedRequests已完成请求单调计数、pendingLoadUtilization挂起负载利用率。特别值得关注的是跨网关运行中请求计数GetPodsRunningRequestscache_api.go通过一次 Redis pipeline 批量获取多个 Pod 的实时在途请求数本地原子计数作为回退供 least-request、load-balance、prefix-cache 等路由器和 in-flight 饱和过滤器使用。AdmitPodRunningRequest则提供检查并计入的原子操作避免并发调用方都读到增量前的计数而被同时放行超限。5.2 GPU Profile 缓存ProfileCache接口cache_api.go提供按 Pod 或按 Deployment 查询ModelGPUProfile的能力实现位于 model_gpu_profile.go。当启用 Profile 缓存后initProfileCache会创建pendingLoadProvider按挂起请求数定义负载的提供者并启动定时器周期刷新各 Deployment 的画像cache_init.go。六、配置总览环境变量与默认值结合 README 与源码缓存包的全部环境变量整理如下核心缓存环境变量默认值说明AIBRIX_POD_DEPLOYMENT_LABELapp.kubernetes.io/nameDeployment 识别标签见 pkg/utils/pod.goAIBRIX_POD_RAYCLUSTERFLEET_LABELorchestration.aibrix.ai/raycluster-fleet-nameRayClusterFleet 识别标签见 pkg/utils/raycluster.goKV 事件同步常量定义见 pkg/constants/kv_event_sync.go环境变量说明AIBRIX_PREFIX_CACHE_KV_EVENT_SYNC_ENABLED启用 KV 事件同步总开关AIBRIX_PREFIX_CACHE_USE_REMOTE_TOKENIZER启用远程 tokenizerKV 同步必需AIBRIX_PREFIX_CACHE_TOKENIZER_TYPEtokenizer 类型KV 同步要求remoteAIBRIX_PREFIX_CACHE_REMOTE_TOKENIZER_ENDPOINT远程 tokenizer 服务端点AIBRIX_PREFIX_CACHE_KV_EVENT_PUBLISH_ADDRZMQ 发布地址AIBRIX_PREFIX_CACHE_KV_EVENT_SUBSCRIBE_ADDRSZMQ 订阅地址性能与监控环境变量默认值说明AIBRIX_POD_METRIC_REFRESH_INTERVAL_MS见 cache_metrics.goPod 指标刷新间隔毫秒gateway-plugin 部署中显式配置示例见 config/gateway/gateway-plugin/gateway-plugin.yamlAIBRIX_MODEL_GPU_PROFILE_CACHING_FLAGtrue见 model_gpu_profile.go启用 GPU Profile 缓存AIBRIX_ZMQ_POLL_TIMEOUT100msZMQ 轮询超时AIBRIX_ZMQ_REPLAY_TIMEOUT5s事件回放请求超时AIBRIX_ZMQ_RECONNECT_INTERVAL1sZMQ 初始重连间隔生产环境可用kubectl set env动态调整例如kubectl set env deployment/aibrix-gateway-plugins -n aibrix-system AIBRIX_POD_METRIC_REFRESH_INTERVAL_MS10000七、Usage Example程序化使用缓存README 给出了完整的程序化示例——启用 KV 同步、添加 Pod、查询模型 Pod 列表并追踪请求// Create cache with KV event sync enabled os.Setenv(constants.EnvPrefixCacheKVEventSyncEnabled, true) os.Setenv(constants.EnvPrefixCacheUseRemoteTokenizer, true) store : cache.NewStore() // Add a pod - KV sync will start automatically if eligible store.AddPod(pod) // Query pods for a model pods : store.GetPodsForModel(llama-2-7b) // Track a request store.AddRequest(pod.Name, model, inputTokens) defer store.DoneRequest(pod.Name, model, outputTokens)需要说明的是生产环境通常不直接调用NewStore()而是通过cache.InitWithOptions(config, stopCh, opts)完成全局单例的初始化之后通过cache.Get()获取Cache接口未初始化时会返回cache is not initialized错误见 cache_init.go。测试环境则可用NewForTest()/InitForTest()反复重建 Store或用InitWithPods系列函数注入 Pod 与指标数据cache_init.go。八、集成点路由算法与网关插件缓存的数据最终服务于两类消费方路由算法缓存提供 Pod 可用性与负载、前缀缓存命中率、GPU 利用率指标支撑 least-request最少请求、load-balance负载均衡、prefix-cache前缀缓存优先等路由策略。KV 同步启用时路由器通过GetSyncPrefixIndexer()cache_init.go拿到同步前缀索引器做前缀匹配若返回 nil 则回退到普通索引器。网关插件gateway-plugins 服务查询缓存获取每模型的可用 Pod 列表、当前请求计数与性能预测。Store还通过MetricSubscriber注册机制向外部推送指标更新AddSubscriber见 cache_api.go。九、监控与线程安全9.1 Prometheus 指标缓存包暴露的 Prometheus 指标覆盖缓存命中/未命中率、请求队列深度、KV 事件处理速率、Pod 同步状态。kvcache子包另有事件接收/处理计数、连接状态、错误计数与处理延迟等指标见 metrics.go。9.2 线程安全设计Store 的所有操作均线程安全采用三层保障机制读写互斥锁Store.musync.RWMutex保护整体数据结构podStatsMu采用固定条带锁stripe lock同步 Pod 删除/重建周期中的计数变更避免随 Pod 键无限增长锁数量cache_init.go原子操作runningRequests、completedRequests、numRequestsTraces等计数器使用sync/atomicgatewaySnapshotCache使用atomic.Value实现快照的无锁原子交换通道驱动的事件处理podMetricsJobs/promqlJobs通道 worker 池模型解耦指标采集与主路径PromQL worker 还内置按 Pod 键去重、FIFO 顺序处理与失败回队重试的机制cache_init.go。十、测试指南缓存包拥有完善的测试体系覆盖 Store 初始化、指标刷新、运行中请求追踪、GPU Profile、KV 事件管理器校验等场景# All cache tests go test ./pkg/cache/ # KV event specific tests go test -v ./pkg/cache/ -run .*KV.* # With ZMQ support go test -tagszmq ./pkg/cache/ # kvcache subpackage tests (requires libzmq) go test -tagszmq ./pkg/cache/kvcache/相关测试用例包括 cache_test.go、cache_init_test.go、output_predictor_test.go、kv_event_manager_validation_test.go 以及 kvcache 子包的 zmq_client_test.go内含 mock publisher 的综合示例。结语pkg/cache是理解 AIBrix 路由决策机制的钥匙从 Kubernetes Informer 驱动的 Pod/模型索引到 Redis 支撑的跨网关实时计数与网关快照再到 ZMQ MessagePack 驱动的 KV 前缀缓存同步它以线程安全的单例 Store 为底座把元数据缓存、负载感知、前缀缓存路由、容量预测四件事高效地串联在一起。无论你是要扩展新的路由算法、接入自定义服务发现还是调优 KV 事件同步链路本包的接口设计Cache聚合接口、可插拔discovery.Provider、可注册的RequestTracker都预留了清晰的扩展点值得深入研读。【免费下载链接】aibrixCost-efficient and pluggable Infrastructure components for GenAI inference项目地址: https://gitcode.com/GitHub_Trending/ai/aibrix创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考