collector.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659
  1. package inferencecost
  2. import (
  3. "context"
  4. "fmt"
  5. "strings"
  6. "time"
  7. "github.com/opencost/opencost/core/pkg/filter/allocation"
  8. "github.com/opencost/opencost/core/pkg/filter/ops"
  9. "github.com/opencost/opencost/core/pkg/log"
  10. "github.com/opencost/opencost/core/pkg/opencost"
  11. "github.com/opencost/opencost/core/pkg/source"
  12. )
  13. // AllocationQuerier is the subset of the cost model needed to fetch per-model
  14. // infrastructure costs. Abstracted as an interface for testability.
  15. type AllocationQuerier interface {
  16. // ComputeAllocation returns an AllocationSet for the given time window.
  17. ComputeAllocation(start, end time.Time) (*opencost.AllocationSet, error)
  18. }
  19. // Collector gathers per-model infrastructure costs from the OpenCost allocation
  20. // layer and token/timing/cache metrics from the data source.
  21. type Collector struct {
  22. allocationQuerier AllocationQuerier
  23. metricsQuerier source.MetricsQuerier
  24. config *Config
  25. }
  26. // NewCollector creates a Collector that uses the provided MetricsQuerier for
  27. // inference metrics.
  28. func NewCollector(config *Config, querier AllocationQuerier, metricsQuerier source.MetricsQuerier) (*Collector, error) {
  29. return &Collector{
  30. allocationQuerier: querier,
  31. metricsQuerier: metricsQuerier,
  32. config: config,
  33. }, nil
  34. }
  35. // CollectMetrics queries all data sources and returns one InferenceCost per
  36. // model/namespace combination. start and end define the time window to query;
  37. // the caller is responsible for choosing appropriate boundaries (e.g. the
  38. // runner uses now-interval..now; the API uses the request window).
  39. func (c *Collector) CollectMetrics(ctx context.Context, start, end time.Time) ([]*InferenceCost, error) {
  40. // --- Infrastructure costs from OpenCost allocation layer ---
  41. allocationCosts, err := c.queryAllocationCosts(ctx, start, end)
  42. if err != nil {
  43. return nil, fmt.Errorf("failed to query allocation costs: %w", err)
  44. }
  45. log.Infof("InferenceCost: collected allocation costs for %d model/namespace combinations", len(allocationCosts))
  46. // --- Token metrics from data source ---
  47. // Query all metrics concurrently using Futures
  48. promptTokensFuture := c.metricsQuerier.QueryInferencePromptTokens(start, end)
  49. generationTokensFuture := c.metricsQuerier.QueryInferenceGenerationTokens(start, end)
  50. inputTimeFuture := c.metricsQuerier.QueryInferenceInputProcessingTime(start, end)
  51. outputTimeFuture := c.metricsQuerier.QueryInferenceOutputProcessingTime(start, end)
  52. cachedTokensFuture := c.metricsQuerier.QueryInferenceCachedTokens(start, end)
  53. cacheConfigFuture := c.metricsQuerier.QueryInferenceCacheConfig(end)
  54. // Await required metrics (prompt and generation tokens)
  55. promptTokensResults, err := promptTokensFuture.Await()
  56. if err != nil {
  57. return nil, fmt.Errorf("failed to query prompt tokens: %w", err)
  58. }
  59. promptTokens := mergeTokenResults(promptTokensResults)
  60. generationTokensResults, err := generationTokensFuture.Await()
  61. if err != nil {
  62. return nil, fmt.Errorf("failed to query generation tokens: %w", err)
  63. }
  64. generationTokens := mergeTokenResults(generationTokensResults)
  65. // --- Timing metrics (optional — degrade gracefully) ---
  66. inputProcessingTime := make(map[string]float64)
  67. if inputTimeResults, err := inputTimeFuture.Await(); err != nil {
  68. log.Warnf("InferenceCost: failed to query input processing time (will use multiplier fallback): %v", err)
  69. } else {
  70. inputProcessingTime = mergeProcessingTimeResults(inputTimeResults)
  71. }
  72. outputProcessingTime := make(map[string]float64)
  73. if outputTimeResults, err := outputTimeFuture.Await(); err != nil {
  74. log.Warnf("InferenceCost: failed to query output processing time (will use multiplier fallback): %v", err)
  75. } else {
  76. outputProcessingTime = mergeProcessingTimeResults(outputTimeResults)
  77. }
  78. // --- KV cache hits (optional — degrade gracefully) ---
  79. cachedTokens := make(map[string]float64)
  80. if cachedTokensResults, err := cachedTokensFuture.Await(); err != nil {
  81. log.Warnf("InferenceCost: failed to query KV cache hits (cacheSavingsFraction will be zero): %v", err)
  82. } else {
  83. cachedTokens = mergeTokenResults(cachedTokensResults)
  84. }
  85. // --- KV cache config (prefix caching enabled flag only) ---
  86. cacheConfigs := make(map[string]*cacheConfig)
  87. if cacheConfigResults, err := cacheConfigFuture.Await(); err != nil {
  88. log.Warnf("InferenceCost: failed to query cache config (prefix_caching_off detection disabled): %v", err)
  89. } else {
  90. cacheConfigs = mergeCacheConfigResults(cacheConfigResults)
  91. }
  92. return c.combineMetrics(allocationCosts, promptTokens, generationTokens,
  93. inputProcessingTime, outputProcessingTime, cachedTokens, cacheConfigs, start, end), nil
  94. }
  95. // cacheConfig holds per-model KV cache configuration from vllm:cache_config_info.
  96. type cacheConfig struct {
  97. prefixCachingEnabled bool
  98. }
  99. // allocationResult holds the two cost figures derived from one Allocation.
  100. type allocationResult struct {
  101. allocationTotalCost float64
  102. usageTotalCost float64
  103. namespace string
  104. cluster string
  105. pod string
  106. controller string
  107. controllerKind string
  108. container string
  109. }
  110. // queryAllocationCosts calls the OpenCost allocation layer twice:
  111. // once with idle sharing (for allocation costs) and once without (for usage costs).
  112. // This ensures allocation costs reconcile to the bill while usage costs reflect
  113. // only active compute without idle or waste.
  114. // This approach was chosen rather than doing a single call and deducting idle and optionally shared, so that
  115. // core logic is not duplicated. A performance penalty is paid though.
  116. func (c *Collector) queryAllocationCosts(ctx context.Context, start, end time.Time) (map[string]*allocationResult, error) {
  117. // Query 1: Allocation costs with idle sharing (reconciles to bill)
  118. allocationCosts, err := c.queryAllocationCostsWithIdle(ctx, start, end)
  119. if err != nil {
  120. return nil, fmt.Errorf("failed to query allocation costs with idle: %w", err)
  121. }
  122. // Query 2: Usage costs without idle sharing (active compute only)
  123. usageCosts, err := c.queryAllocationCostsWithoutIdle(ctx, start, end)
  124. if err != nil {
  125. return nil, fmt.Errorf("failed to query usage costs without idle: %w", err)
  126. }
  127. // Merge results: allocation costs from first query, usage costs from second
  128. results := make(map[string]*allocationResult)
  129. for key, allocResult := range allocationCosts {
  130. if allocResult == nil {
  131. continue
  132. }
  133. // Copy the full allocationResult so we retain pod/controller/container metadata
  134. copied := *allocResult
  135. copied.usageTotalCost = 0 // Will be filled from usageCosts
  136. results[key] = &copied
  137. }
  138. // Fill in usage costs from the second query
  139. for key, usageResult := range usageCosts {
  140. if result, exists := results[key]; exists {
  141. result.usageTotalCost = usageResult.usageTotalCost
  142. } else {
  143. // Model exists in usage query but not allocation query (shouldn't happen)
  144. log.Warnf("InferenceCost: model %s has usage cost but no allocation cost", key)
  145. results[key] = usageResult
  146. }
  147. }
  148. // Log the differences
  149. for key, result := range results {
  150. modelName, namespace := parseKey(key)
  151. if result.allocationTotalCost > 0 {
  152. log.Debugf("InferenceCost: model=%s ns=%s alloc=$%.4f usage=$%.4f (%.1f%% of alloc)",
  153. modelName, namespace, result.allocationTotalCost, result.usageTotalCost,
  154. (result.usageTotalCost/result.allocationTotalCost)*100)
  155. }
  156. }
  157. return results, nil
  158. }
  159. // queryAllocationCostsWithIdle queries allocations with idle sharing enabled.
  160. func (c *Collector) queryAllocationCostsWithIdle(ctx context.Context, start, end time.Time) (map[string]*allocationResult, error) {
  161. as, err := c.allocationQuerier.ComputeAllocation(start, end)
  162. if err != nil {
  163. return nil, err
  164. }
  165. // Create a filter to match shared infrastructure allocations by label
  166. // This ensures allocations with the shared infra label are moved to shareSet
  167. // and distributed among other allocations, rather than aggregating into __unallocated__
  168. shareFilter := ops.Eq(
  169. ops.WithKey(allocation.FieldLabel, c.config.SharedInfraLabel),
  170. c.config.SharedInfraLabelValue,
  171. )
  172. opts := &opencost.AllocationAggregationOptions{
  173. ShareIdle: opencost.ShareWeighted,
  174. ShareSplit: opencost.ShareWeighted,
  175. Share: shareFilter,
  176. SharedLabels: map[string][]string{c.config.SharedInfraLabel: {c.config.SharedInfraLabelValue}},
  177. }
  178. aggregateBy := []string{"label:" + c.config.ModelLabel}
  179. if err := as.AggregateBy(aggregateBy, opts); err != nil {
  180. return nil, fmt.Errorf("AggregateBy label:%s: %w", c.config.ModelLabel, err)
  181. }
  182. return c.extractAllocationResults(as, true)
  183. }
  184. // queryAllocationCostsWithoutIdle queries allocations without idle or shared
  185. // infrastructure cost sharing. Usage costs reflect active compute only.
  186. func (c *Collector) queryAllocationCostsWithoutIdle(ctx context.Context, start, end time.Time) (map[string]*allocationResult, error) {
  187. as, err := c.allocationQuerier.ComputeAllocation(start, end)
  188. if err != nil {
  189. return nil, err
  190. }
  191. // Create a filter to match shared infrastructure allocations by label
  192. // Even though we're not sharing costs (ShareSplit: ShareNone), we still need
  193. // the Share filter to identify and separate shared infra allocations from
  194. // regular allocations, preventing them from aggregating into __unallocated__
  195. shareFilter := ops.Eq(
  196. ops.WithKey(allocation.FieldLabel, c.config.SharedInfraLabel),
  197. c.config.SharedInfraLabelValue,
  198. )
  199. opts := &opencost.AllocationAggregationOptions{
  200. ShareIdle: opencost.ShareNone,
  201. ShareSplit: opencost.ShareNone,
  202. Share: shareFilter,
  203. SharedLabels: map[string][]string{c.config.SharedInfraLabel: {c.config.SharedInfraLabelValue}},
  204. }
  205. aggregateBy := []string{"label:" + c.config.ModelLabel}
  206. if err := as.AggregateBy(aggregateBy, opts); err != nil {
  207. return nil, fmt.Errorf("AggregateBy label:%s: %w", c.config.ModelLabel, err)
  208. }
  209. return c.extractAllocationResults(as, false)
  210. }
  211. // extractAllocationResults extracts cost data from an AllocationSet.
  212. func (c *Collector) extractAllocationResults(as *opencost.AllocationSet, isAllocationCost bool) (map[string]*allocationResult, error) {
  213. results := make(map[string]*allocationResult)
  214. for name, alloc := range as.Allocations {
  215. if alloc == nil {
  216. continue
  217. }
  218. // Skip the synthetic __idle__ and __unallocated__ entries.
  219. if strings.HasPrefix(name, "__") {
  220. continue
  221. }
  222. modelName := extractModelName(alloc, c.config.ModelLabel)
  223. if modelName == "" {
  224. continue
  225. }
  226. namespace := ""
  227. cluster := ""
  228. pod := ""
  229. controller := ""
  230. controllerKind := ""
  231. container := ""
  232. if alloc.Properties != nil {
  233. namespace = alloc.Properties.Namespace
  234. cluster = alloc.Properties.Cluster
  235. pod = alloc.Properties.Pod
  236. controller = alloc.Properties.Controller
  237. controllerKind = alloc.Properties.ControllerKind
  238. container = alloc.Properties.Container
  239. }
  240. key := modelNamespaceKey(modelName, namespace)
  241. // Accumulate costs for the same model/namespace key
  242. existing, exists := results[key]
  243. if !exists {
  244. existing = &allocationResult{
  245. namespace: namespace,
  246. cluster: cluster,
  247. pod: pod,
  248. controller: controller,
  249. controllerKind: controllerKind,
  250. container: container,
  251. }
  252. results[key] = existing
  253. }
  254. if isAllocationCost {
  255. // For allocation cost: use TotalCost() which includes idle and shared
  256. existing.allocationTotalCost += alloc.TotalCost()
  257. } else {
  258. // For usage cost: use TotalCost() from the ShareNone query (no idle),
  259. // then scale GPU, CPU, and RAM by their actual utilisation when available.
  260. // This ensures costBasis=usage reflects actual resource consumption rather
  261. // than the full reservation cost.
  262. //
  263. // Resources intentionally left unscaled:
  264. // Network — already billed by actual bytes transferred, no reservation to remove.
  265. // PV — billed by provisioned capacity; no IO utilisation metric available.
  266. // LB — billed per hour of existence; no per-request utilisation signal.
  267. // External — opaque cloud billing pass-through; no usage signal attached.
  268. cost := alloc.TotalCost()
  269. // GPU: scale by SM duty cycle (GPUUsageAverage ∈ [0,1]).
  270. // GPUHours (and therefore GPUCost) always reflects the full reservation;
  271. // GPUUsageAverage is the fraction of time the GPU cores were active.
  272. // Any non-nil value is clamped to [0,1] so that zero utilisation
  273. // correctly reduces cost to $0 and out-of-range values are handled
  274. // deterministically rather than silently ignored.
  275. if alloc.GPUAllocation != nil && alloc.GPUAllocation.GPUUsageAverage != nil {
  276. util := *alloc.GPUAllocation.GPUUsageAverage
  277. if util < 0 {
  278. util = 0
  279. } else if util > 1 {
  280. util = 1
  281. }
  282. scaledGPUCost := alloc.GPUTotalCost() * util
  283. cost = cost - alloc.GPUTotalCost() + scaledGPUCost
  284. log.Debugf("InferenceCost usage: GPU scaled model=%s ns=%s orig=$%.4f scaled=$%.4f util=%.1f%%",
  285. modelName, namespace, alloc.GPUTotalCost(), scaledGPUCost, util*100)
  286. }
  287. // CPU: scale by core utilisation ratio (usage / request).
  288. // CPUCoreRequestAverage and CPUCoreUsageAverage are plain float64 (not pointers).
  289. if alloc.CPUCoreRequestAverage > 0 &&
  290. alloc.CPUCoreUsageAverage > 0 &&
  291. alloc.CPUCoreUsageAverage < alloc.CPUCoreRequestAverage {
  292. cpuUtil := alloc.CPUCoreUsageAverage / alloc.CPUCoreRequestAverage
  293. scaledCPUCost := alloc.CPUTotalCost() * cpuUtil
  294. cost = cost - alloc.CPUTotalCost() + scaledCPUCost
  295. log.Debugf("InferenceCost usage: CPU scaled model=%s ns=%s orig=$%.4f scaled=$%.4f util=%.1f%%",
  296. modelName, namespace, alloc.CPUTotalCost(), scaledCPUCost, cpuUtil*100)
  297. }
  298. // RAM: scale by byte utilisation ratio (usage / request).
  299. // RAMBytesRequestAverage and RAMBytesUsageAverage are plain float64 (not pointers).
  300. if alloc.RAMBytesRequestAverage > 0 &&
  301. alloc.RAMBytesUsageAverage > 0 &&
  302. alloc.RAMBytesUsageAverage < alloc.RAMBytesRequestAverage {
  303. ramUtil := alloc.RAMBytesUsageAverage / alloc.RAMBytesRequestAverage
  304. scaledRAMCost := alloc.RAMTotalCost() * ramUtil
  305. cost = cost - alloc.RAMTotalCost() + scaledRAMCost
  306. log.Debugf("InferenceCost usage: RAM scaled model=%s ns=%s orig=$%.4f scaled=$%.4f util=%.1f%%",
  307. modelName, namespace, alloc.RAMTotalCost(), scaledRAMCost, ramUtil*100)
  308. }
  309. existing.usageTotalCost += cost
  310. }
  311. // When aggregating multiple allocations, preserve the first non-empty values
  312. // for pod, controller, and container. This provides representative values
  313. // when costs are aggregated across multiple pods/containers.
  314. if existing.pod == "" && pod != "" {
  315. existing.pod = pod
  316. }
  317. if existing.controller == "" && controller != "" {
  318. existing.controller = controller
  319. }
  320. if existing.controllerKind == "" && controllerKind != "" {
  321. existing.controllerKind = controllerKind
  322. }
  323. if existing.container == "" && container != "" {
  324. existing.container = container
  325. }
  326. }
  327. return results, nil
  328. }
  329. // extractModelName extracts the model name from the allocation name or label.
  330. // After AggregateBy("label:<key>"), the allocation Name is the label value.
  331. func extractModelName(alloc *opencost.Allocation, _ string) string {
  332. if alloc == nil {
  333. return ""
  334. }
  335. // AggregateBy sets the Name to the label value.
  336. return alloc.Name
  337. }
  338. // canonicalModelName normalizes a model name by stripping any org/vendor prefix
  339. // before the last "/".
  340. // Examples:
  341. // - "MiniMaxAI/MiniMax-M2.7" -> "MiniMax-M2.7"
  342. // - "google/gemma-4-31B" -> "gemma-4-31B"
  343. func canonicalModelName(modelName string) string {
  344. if idx := strings.LastIndex(modelName, "/"); idx >= 0 {
  345. return modelName[idx+1:]
  346. }
  347. return modelName
  348. }
  349. // reconcileTokenKeys re-keys entries only when there is a confirmed mismatch
  350. // between the metric key and the allocation-backed model key for the same
  351. // namespace.
  352. //
  353. // Two common mismatch examples:
  354. // 1. Fully-qualified vLLM model name vs short allocation label:
  355. // "google/gemma-4-31B:llm-d-pic" -> "gemma-4-31B:llm-d-pic"
  356. // 2. Fully-qualified vLLM model name vs short allocation label with a
  357. // different vendor/org prefix:
  358. // "MiniMaxAI/MiniMax-M2.7:llm-d-pic" -> "MiniMax-M2.7:llm-d-pic"
  359. //
  360. // Exact matches are preserved. Keys with no matching allocation-backed target
  361. // are also preserved unchanged. A warning is logged for every remapped key so
  362. // the mismatch is auditable.
  363. //
  364. // Returns both the reconciled map and a set of keys that were remapped (to be
  365. // excluded later).
  366. func reconcileTokenKeys(tokens map[string]float64, allocCosts map[string]*allocationResult) (map[string]float64, map[string]struct{}) {
  367. // Build index: normalizedShortName:namespace -> allocKey, preferring
  368. // allocation keys that are already in short-name form.
  369. shortIndex := make(map[string]string, len(allocCosts))
  370. for allocKey := range allocCosts {
  371. modelName, namespace := parseKey(allocKey)
  372. shortName := canonicalModelName(modelName)
  373. shortKey := modelNamespaceKey(shortName, namespace)
  374. if existing, found := shortIndex[shortKey]; found {
  375. existingModelName, _ := parseKey(existing)
  376. if existingModelName == shortName {
  377. continue
  378. }
  379. }
  380. shortIndex[shortKey] = allocKey
  381. }
  382. out := make(map[string]float64, len(tokens))
  383. remappedKeys := make(map[string]struct{})
  384. for k, v := range tokens {
  385. modelName, namespace := parseKey(k)
  386. shortName := canonicalModelName(modelName)
  387. shortKey := modelNamespaceKey(shortName, namespace)
  388. if allocKey, found := shortIndex[shortKey]; found {
  389. if k != allocKey {
  390. log.Warnf("InferenceCost: remapping metric key %q → %q (model-name mismatch with allocation label)", k, allocKey)
  391. out[allocKey] += v
  392. remappedKeys[k] = struct{}{}
  393. continue
  394. }
  395. }
  396. out[k] = v
  397. }
  398. return out, remappedKeys
  399. }
  400. // reconcileCacheConfigKeys re-keys a cacheConfig map the same way reconcileTokenKeys
  401. // does for float64 maps — handling fully-qualified vs short model name mismatches.
  402. func reconcileCacheConfigKeys(configs map[string]*cacheConfig, allocCosts map[string]*allocationResult) (map[string]*cacheConfig, map[string]struct{}) {
  403. shortIndex := make(map[string]string, len(allocCosts))
  404. for allocKey := range allocCosts {
  405. modelName, namespace := parseKey(allocKey)
  406. shortName := canonicalModelName(modelName)
  407. shortKey := modelNamespaceKey(shortName, namespace)
  408. if _, exists := shortIndex[shortKey]; !exists {
  409. shortIndex[shortKey] = allocKey
  410. }
  411. }
  412. out := make(map[string]*cacheConfig, len(configs))
  413. remappedKeys := make(map[string]struct{})
  414. for k, v := range configs {
  415. modelName, namespace := parseKey(k)
  416. shortName := canonicalModelName(modelName)
  417. shortKey := modelNamespaceKey(shortName, namespace)
  418. if allocKey, found := shortIndex[shortKey]; found {
  419. if k != allocKey {
  420. log.Warnf("InferenceCost: remapping cache config key %q → %q (model-name mismatch with allocation label)", k, allocKey)
  421. out[allocKey] = v
  422. remappedKeys[k] = struct{}{}
  423. continue
  424. }
  425. }
  426. out[k] = v
  427. }
  428. return out, remappedKeys
  429. }
  430. // combineMetrics joins all data sources into InferenceCost structs.
  431. func (c *Collector) combineMetrics(
  432. allocCosts map[string]*allocationResult,
  433. promptTokens, generationTokens,
  434. inputProcessingTime, outputProcessingTime,
  435. cachedTokens map[string]float64,
  436. cacheConfigs map[string]*cacheConfig,
  437. start, end time.Time,
  438. ) []*InferenceCost {
  439. // Reconcile token map keys against allocation keys to handle the case where
  440. // vLLM reports a fully-qualified model name (e.g. "org/model") but the K8s
  441. // pod label uses only the short name ("model"). Re-keying fires only when a
  442. // mismatch is detected; keys that already match are left unchanged.
  443. // Track which keys were remapped so we can exclude them from final results.
  444. var remappedKeys map[string]struct{}
  445. promptTokens, remappedKeys = reconcileTokenKeys(promptTokens, allocCosts)
  446. var remapped map[string]struct{}
  447. generationTokens, remapped = reconcileTokenKeys(generationTokens, allocCosts)
  448. for k := range remapped {
  449. remappedKeys[k] = struct{}{}
  450. }
  451. inputProcessingTime, remapped = reconcileTokenKeys(inputProcessingTime, allocCosts)
  452. for k := range remapped {
  453. remappedKeys[k] = struct{}{}
  454. }
  455. outputProcessingTime, remapped = reconcileTokenKeys(outputProcessingTime, allocCosts)
  456. for k := range remapped {
  457. remappedKeys[k] = struct{}{}
  458. }
  459. cachedTokens, remapped = reconcileTokenKeys(cachedTokens, allocCosts)
  460. for k := range remapped {
  461. remappedKeys[k] = struct{}{}
  462. }
  463. cacheConfigs, remapped = reconcileCacheConfigKeys(cacheConfigs, allocCosts)
  464. for k := range remapped {
  465. remappedKeys[k] = struct{}{}
  466. }
  467. // Union of all keys across sources.
  468. // Include timing/cache maps as well so models that only appear in those
  469. // sources are not dropped before cost calculation.
  470. keys := make(map[string]struct{})
  471. for k := range allocCosts {
  472. keys[k] = struct{}{}
  473. }
  474. for k := range promptTokens {
  475. keys[k] = struct{}{}
  476. }
  477. for k := range generationTokens {
  478. keys[k] = struct{}{}
  479. }
  480. for k := range inputProcessingTime {
  481. keys[k] = struct{}{}
  482. }
  483. for k := range outputProcessingTime {
  484. keys[k] = struct{}{}
  485. }
  486. for k := range cachedTokens {
  487. keys[k] = struct{}{}
  488. }
  489. for k := range cacheConfigs {
  490. keys[k] = struct{}{}
  491. }
  492. results := make([]*InferenceCost, 0, len(keys))
  493. for key := range keys {
  494. // Skip keys that were remapped to avoid duplicate series
  495. if _, wasRemapped := remappedKeys[key]; wasRemapped {
  496. continue
  497. }
  498. modelName, namespace := parseKey(key)
  499. cfg := cacheConfigs[key]
  500. var prefixCachingEnabled, cacheConfigKnown bool
  501. if cfg != nil {
  502. prefixCachingEnabled = cfg.prefixCachingEnabled
  503. cacheConfigKnown = true
  504. }
  505. ic := &InferenceCost{
  506. Properties: InferenceCostProperties{
  507. ModelName: modelName,
  508. Namespace: namespace,
  509. WorkloadType: "inference",
  510. },
  511. PromptTokens: promptTokens[key],
  512. GenerationTokens: generationTokens[key],
  513. InputProcessingTime: inputProcessingTime[key],
  514. OutputProcessingTime: outputProcessingTime[key],
  515. CachedTokens: cachedTokens[key],
  516. PrefixCachingEnabled: prefixCachingEnabled,
  517. CacheConfigKnown: cacheConfigKnown,
  518. Timestamp: end,
  519. }
  520. ic.Window.Start = start
  521. ic.Window.End = end
  522. if ar, ok := allocCosts[key]; ok {
  523. ic.AllocationTotalCost = ar.allocationTotalCost
  524. ic.UsageTotalCost = ar.usageTotalCost
  525. ic.Properties.Cluster = ar.cluster
  526. ic.Properties.Pod = ar.pod
  527. ic.Properties.Controller = ar.controller
  528. ic.Properties.ControllerKind = ar.controllerKind
  529. ic.Properties.Container = ar.container
  530. if namespace == "" {
  531. ic.Properties.Namespace = ar.namespace
  532. }
  533. }
  534. ic.TotalTokens = ic.PromptTokens + ic.GenerationTokens
  535. ic.EffectiveInputTokens = ic.PromptTokens - ic.CachedTokens
  536. if ic.EffectiveInputTokens < 0 {
  537. ic.EffectiveInputTokens = 0
  538. }
  539. results = append(results, ic)
  540. }
  541. return results
  542. }
  543. func modelNamespaceKey(modelName, namespace string) string {
  544. return modelName + ":" + namespace
  545. }
  546. func parseKey(key string) (modelName, namespace string) {
  547. idx := strings.IndexByte(key, ':')
  548. if idx < 0 {
  549. return key, "unknown"
  550. }
  551. return key[:idx], key[idx+1:]
  552. }
  553. // mergeTokenResults merges multiple InferenceTokensResult into a single map
  554. func mergeTokenResults(results []*source.InferenceTokensResult) map[string]float64 {
  555. merged := make(map[string]float64)
  556. for _, result := range results {
  557. for k, v := range result.Values {
  558. merged[k] = v
  559. }
  560. }
  561. return merged
  562. }
  563. // mergeProcessingTimeResults merges multiple InferenceProcessingTimeResult into a single map
  564. func mergeProcessingTimeResults(results []*source.InferenceProcessingTimeResult) map[string]float64 {
  565. merged := make(map[string]float64)
  566. for _, result := range results {
  567. for k, v := range result.Values {
  568. merged[k] = v
  569. }
  570. }
  571. return merged
  572. }
  573. // mergeCacheConfigResults merges multiple InferenceCacheConfigResult into a single map
  574. func mergeCacheConfigResults(results []*source.InferenceCacheConfigResult) map[string]*cacheConfig {
  575. merged := make(map[string]*cacheConfig)
  576. for _, result := range results {
  577. for k, v := range result.Configs {
  578. merged[k] = &cacheConfig{prefixCachingEnabled: v.PrefixCachingEnabled}
  579. }
  580. }
  581. return merged
  582. }