Selaa lähdekoodia

add filtering to allocation api (#3155)

Signed-off-by: Alex Meijer <alexander.meijer@ibm.com>
Alex Meijer 11 kuukautta sitten
vanhempi
sitoutus
81d2b597d1
2 muutettua tiedostoa jossa 66 lisäystä ja 1 poistoa
  1. 65 0
      pkg/costmodel/aggregation.go
  2. 1 1
      pkg/costmodel/allocation.go

+ 65 - 0
pkg/costmodel/aggregation.go

@@ -15,6 +15,7 @@ import (
 	"github.com/patrickmn/go-cache"
 	prometheusClient "github.com/prometheus/client_golang/api"
 
+	"github.com/opencost/opencost/core/pkg/filter/allocation"
 	"github.com/opencost/opencost/core/pkg/log"
 	"github.com/opencost/opencost/core/pkg/opencost"
 	"github.com/opencost/opencost/core/pkg/util"
@@ -2178,6 +2179,9 @@ func (a *Accesses) ComputeAllocationHandlerSummary(w http.ResponseWriter, r *htt
 	// sums each Set in the Range, producing one Set.
 	accumulate := qp.GetBool("accumulate", false)
 
+	// Get allocation filter if provided
+	allocationFilter := qp.Get("filter", "")
+
 	// Query for AllocationSets in increments of the given step duration,
 	// appending each to the AllocationSetRange.
 	asr := opencost.NewAllocationSetRange()
@@ -2196,6 +2200,35 @@ func (a *Accesses) ComputeAllocationHandlerSummary(w http.ResponseWriter, r *htt
 		stepStart = stepEnd
 	}
 
+	// Apply allocation filter if provided
+	if allocationFilter != "" {
+		parser := allocation.NewAllocationFilterParser()
+		filterNode, err := parser.Parse(allocationFilter)
+		if err != nil {
+			WriteError(w, BadRequest(fmt.Sprintf("Invalid filter: %s", err)))
+			return
+		}
+		compiler := opencost.NewAllocationMatchCompiler(nil)
+		matcher, err := compiler.Compile(filterNode)
+		if err != nil {
+			WriteError(w, BadRequest(fmt.Sprintf("Failed to compile filter: %s", err)))
+			return
+		}
+		filteredASR := opencost.NewAllocationSetRange()
+		for _, as := range asr.Slice() {
+			filteredAS := opencost.NewAllocationSet(as.Start(), as.End())
+			for _, alloc := range as.Allocations {
+				if matcher.Matches(alloc) {
+					filteredAS.Set(alloc)
+				}
+			}
+			if filteredAS.Length() > 0 {
+				filteredASR.Append(filteredAS)
+			}
+		}
+		asr = filteredASR
+	}
+
 	// Aggregate, if requested
 	if len(aggregateBy) > 0 {
 		err = asr.AggregateBy(aggregateBy, nil)
@@ -2286,6 +2319,9 @@ func (a *Accesses) ComputeAllocationHandler(w http.ResponseWriter, r *http.Reque
 
 	shareIdle := qp.GetBool("shareIdle", false)
 
+	// Get allocation filter if provided
+	allocationFilter := qp.Get("filter", "")
+
 	asr, err := a.Model.QueryAllocation(window, resolution, step, aggregateBy, includeIdle, idleByNode, includeProportionalAssetResourceCosts, includeAggregatedMetadata, sharedLoadBalancer, accumulateBy, shareIdle)
 	if err != nil {
 		if strings.Contains(strings.ToLower(err.Error()), "bad request") {
@@ -2297,6 +2333,35 @@ func (a *Accesses) ComputeAllocationHandler(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
+	// Apply allocation filter if provided
+	if allocationFilter != "" {
+		parser := allocation.NewAllocationFilterParser()
+		filterNode, err := parser.Parse(allocationFilter)
+		if err != nil {
+			WriteError(w, BadRequest(fmt.Sprintf("Invalid filter: %s", err)))
+			return
+		}
+		compiler := opencost.NewAllocationMatchCompiler(nil)
+		matcher, err := compiler.Compile(filterNode)
+		if err != nil {
+			WriteError(w, BadRequest(fmt.Sprintf("Failed to compile filter: %s", err)))
+			return
+		}
+		filteredASR := opencost.NewAllocationSetRange()
+		for _, as := range asr.Slice() {
+			filteredAS := opencost.NewAllocationSet(as.Start(), as.End())
+			for _, alloc := range as.Allocations {
+				if matcher.Matches(alloc) {
+					filteredAS.Set(alloc)
+				}
+			}
+			if filteredAS.Length() > 0 {
+				filteredASR.Append(filteredAS)
+			}
+		}
+		asr = filteredASR
+	}
+
 	w.Write(WrapData(asr, nil))
 }
 

+ 1 - 1
pkg/costmodel/allocation.go

@@ -280,7 +280,7 @@ func (cm *CostModel) ComputeAllocation(start, end time.Time, resolution time.Dur
 				resultAlloc.RawAllocationOnly.RAMBytesUsageMax = alloc.RawAllocationOnly.RAMBytesUsageMax
 			}
 
-			if alloc.RawAllocationOnly.GPUUsageMax != nil {
+			if alloc.RawAllocationOnly.GPUUsageMax != nil && resultAlloc.RawAllocationOnly.GPUUsageMax != nil {
 				if *alloc.RawAllocationOnly.GPUUsageMax > *resultAlloc.RawAllocationOnly.GPUUsageMax {
 					resultAlloc.RawAllocationOnly.GPUUsageMax = alloc.RawAllocationOnly.GPUUsageMax
 				}