Add custompricing logic to assets

pkg/costmodel/cluster.go

@@ -2,6 +2,7 @@ package costmodel
 
 import (
 	"fmt"
+	"strconv"
 	"time"
 
 	"github.com/kubecost/cost-model/pkg/util/timeutil"
@@ -169,7 +170,7 @@ func ClusterDisks(client prometheus.Client, provider cloud.Provider, duration, o
 
 	diskMap := map[string]*Disk{}
 
-	pvCosts(diskMap, resolution, resActiveMins, resPVSize, resPVCost)
+	pvCosts(diskMap, resolution, resActiveMins, resPVSize, resPVCost, provider)
 
 	for _, result := range resLocalStorageCost {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
@@ -442,23 +443,23 @@ func ClusterNodes(cp cloud.Provider, client prometheus.Client, duration, offset
 	activeDataMap := buildActiveDataMap(resActiveMins, resolution)
 
 	gpuCountMap := buildGPUCountMap(resNodeGPUCount)
+	preemptibleMap := buildPreemptibleMap(resIsSpot)
 
-	cpuCostMap, clusterAndNameToType1 := buildCPUCostMap(resNodeCPUHourlyCost)
-	ramCostMap, clusterAndNameToType2 := buildRAMCostMap(resNodeRAMHourlyCost)
-	gpuCostMap, clusterAndNameToType3 := buildGPUCostMap(resNodeGPUHourlyCost, gpuCountMap)
+	cpuCostMap, clusterAndNameToType1 := buildCPUCostMap(resNodeCPUHourlyCost, cp, preemptibleMap)
+	ramCostMap, clusterAndNameToType2 := buildRAMCostMap(resNodeRAMHourlyCost, cp, preemptibleMap)
+	gpuCostMap, clusterAndNameToType3 := buildGPUCostMap(resNodeGPUHourlyCost, gpuCountMap, cp, preemptibleMap)
 
 	clusterAndNameToTypeIntermediate := mergeTypeMaps(clusterAndNameToType1, clusterAndNameToType2)
 	clusterAndNameToType := mergeTypeMaps(clusterAndNameToTypeIntermediate, clusterAndNameToType3)
 
 	cpuCoresMap := buildCPUCoresMap(resNodeCPUCores)
-
 	ramBytesMap := buildRAMBytesMap(resNodeRAMBytes)
 
 	ramUserPctMap := buildRAMUserPctMap(resNodeRAMUserPct)
 	ramSystemPctMap := buildRAMSystemPctMap(resNodeRAMSystemPct)
 
 	cpuBreakdownMap := buildCPUBreakdownMap(resNodeCPUModeTotal)
-	preemptibleMap := buildPreemptibleMap(resIsSpot)
+
 	labelsMap := buildLabelsMap(resLabels)
 
 	costTimesMinuteAndCount(activeDataMap, cpuCostMap, cpuCoresMap)
@@ -1072,7 +1073,7 @@ func ClusterCostsOverTime(cli prometheus.Client, provider cloud.Provider, startS
 	}, nil
 }
 
-func pvCosts(diskMap map[string]*Disk, resolution time.Duration, resActiveMins, resPVSize, resPVCost []*prom.QueryResult) {
+func pvCosts(diskMap map[string]*Disk, resolution time.Duration, resActiveMins, resPVSize, resPVCost []*prom.QueryResult, cp cloud.Provider) {
 	for _, result := range resActiveMins {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
 		if err != nil {
@@ -1134,6 +1135,12 @@ func pvCosts(diskMap map[string]*Disk, resolution time.Duration, resActiveMins,
 		diskMap[key].Bytes = bytes
 	}
 
+	customPricingEnabled := cloud.CustomPricesEnabled(cp)
+	customPricingConfig, err := cp.GetConfig()
+	if err != nil {
+		log.Warningf("ClusterDisks: failed to load custom pricing: %s", err)
+	}
+
 	for _, result := range resPVCost {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
 		if err != nil {
@@ -1148,7 +1155,25 @@ func pvCosts(diskMap map[string]*Disk, resolution time.Duration, resActiveMins,
 
 		// TODO niko/assets storage class
 
-		cost := result.Values[0].Value
+		var cost float64
+
+		if customPricingEnabled && customPricingConfig != nil {
+
+			customPVCostStr := customPricingConfig.Storage
+
+			customPVCost, err := strconv.ParseFloat(customPVCostStr, 64)
+			if err != nil {
+				log.Warningf("ClusterDisks: error parsing custom PV price: %s", customPVCostStr)
+			}
+
+			cost = customPVCost
+
+		} else {
+
+			cost = result.Values[0].Value
+
+		}
+
 		key := fmt.Sprintf("%s/%s", cluster, name)
 		if _, ok := diskMap[key]; !ok {
 			diskMap[key] = &Disk{

pkg/costmodel/cluster_helpers.go

@@ -1,9 +1,11 @@
 package costmodel
 
 import (
-	"github.com/kubecost/cost-model/pkg/cloud"
+	"strconv"
 	"time"
 
+	"github.com/kubecost/cost-model/pkg/cloud"
+
 	"github.com/kubecost/cost-model/pkg/env"
 	"github.com/kubecost/cost-model/pkg/log"
 	"github.com/kubecost/cost-model/pkg/prom"
@@ -28,6 +30,8 @@ func mergeTypeMaps(clusterAndNameToType1, clusterAndNameToType2 map[nodeIdentifi
 
 func buildCPUCostMap(
 	resNodeCPUCost []*prom.QueryResult,
+	cp cloud.Provider,
+	preemptible map[NodeIdentifier]bool,
 ) (
 	map[NodeIdentifier]float64,
 	map[nodeIdentifierNoProviderID]string,
@@ -36,6 +40,12 @@ func buildCPUCostMap(
 	cpuCostMap := make(map[NodeIdentifier]float64)
 	clusterAndNameToType := make(map[nodeIdentifierNoProviderID]string)
 
+	customPricingEnabled := cloud.CustomPricesEnabled(cp)
+	customPricingConfig, err := cp.GetConfig()
+	if err != nil {
+		log.Warningf("ClusterNodes: failed to load custom pricing: %s", err)
+	}
+
 	for _, result := range resNodeCPUCost {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
 		if err != nil {
@@ -51,8 +61,6 @@ func buildCPUCostMap(
 		nodeType, _ := result.GetString("instance_type")
 		providerID, _ := result.GetString("provider_id")
 
-		cpuCost := result.Values[0].Value
-
 		key := NodeIdentifier{
 			Cluster:    cluster,
 			Name:       name,
@@ -63,8 +71,33 @@ func buildCPUCostMap(
 			Name:    name,
 		}
 
+		var cpuCost float64
+
+		if customPricingEnabled && customPricingConfig != nil {
+
+			var customCPUStr string
+			if spot, ok := preemptible[key]; ok && spot {
+				customCPUStr = customPricingConfig.SpotCPU
+			} else {
+				customCPUStr = customPricingConfig.CPU
+			}
+
+			customCPUCost, err := strconv.ParseFloat(customCPUStr, 64)
+			if err != nil {
+				log.Warningf("ClusterNodes: error parsing custom CPU price: %s", customCPUStr)
+			}
+			cpuCost = customCPUCost
+
+		} else {
+
+			cpuCost = result.Values[0].Value
+
+		}
+
 		clusterAndNameToType[keyNon] = nodeType
 
+		log.Infof("ASSET CLUSTERNODES NODE: %s WITH CPUPRICE: %f", key.Name, cpuCost)
+
 		cpuCostMap[key] = cpuCost
 	}
 
@@ -73,6 +106,8 @@ func buildCPUCostMap(
 
 func buildRAMCostMap(
 	resNodeRAMCost []*prom.QueryResult,
+	cp cloud.Provider,
+	preemptible map[NodeIdentifier]bool,
 ) (
 	map[NodeIdentifier]float64,
 	map[nodeIdentifierNoProviderID]string,
@@ -81,6 +116,12 @@ func buildRAMCostMap(
 	ramCostMap := make(map[NodeIdentifier]float64)
 	clusterAndNameToType := make(map[nodeIdentifierNoProviderID]string)
 
+	customPricingEnabled := cloud.CustomPricesEnabled(cp)
+	customPricingConfig, err := cp.GetConfig()
+	if err != nil {
+		log.Warningf("ClusterNodes: failed to load custom pricing: %s", err)
+	}
+
 	for _, result := range resNodeRAMCost {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
 		if err != nil {
@@ -96,8 +137,6 @@ func buildRAMCostMap(
 		nodeType, _ := result.GetString("instance_type")
 		providerID, _ := result.GetString("provider_id")
 
-		ramCost := result.Values[0].Value
-
 		key := NodeIdentifier{
 			Cluster:    cluster,
 			Name:       name,
@@ -108,7 +147,33 @@ func buildRAMCostMap(
 			Name:    name,
 		}
 
+		var ramCost float64
+
+		if customPricingEnabled && customPricingConfig != nil {
+
+			var customRAMStr string
+			if spot, ok := preemptible[key]; ok && spot {
+				customRAMStr = customPricingConfig.SpotRAM
+			} else {
+				customRAMStr = customPricingConfig.RAM
+			}
+
+			customRAMCost, err := strconv.ParseFloat(customRAMStr, 64)
+			if err != nil {
+				log.Warningf("ClusterNodes: error parsing custom RAM price: %s", customRAMStr)
+			}
+			ramCost = customRAMCost / 1024 / 1024 / 1024
+
+		} else {
+
+			ramCost = result.Values[0].Value
+
+		}
+
 		clusterAndNameToType[keyNon] = nodeType
+
+		log.Infof("ASSET CLUSTERNODES NODE: %s WITH RAMPRICE: %f", key.Name, ramCost)
+
 		ramCostMap[key] = ramCost
 	}
 
@@ -118,6 +183,8 @@ func buildRAMCostMap(
 func buildGPUCostMap(
 	resNodeGPUCost []*prom.QueryResult,
 	gpuCountMap map[NodeIdentifier]float64,
+	cp cloud.Provider,
+	preemptible map[NodeIdentifier]bool,
 ) (
 	map[NodeIdentifier]float64,
 	map[nodeIdentifierNoProviderID]string,
@@ -126,6 +193,12 @@ func buildGPUCostMap(
 	gpuCostMap := make(map[NodeIdentifier]float64)
 	clusterAndNameToType := make(map[nodeIdentifierNoProviderID]string)
 
+	customPricingEnabled := cloud.CustomPricesEnabled(cp)
+	customPricingConfig, err := cp.GetConfig()
+	if err != nil {
+		log.Warningf("ClusterNodes: failed to load custom pricing: %s", err)
+	}
+
 	for _, result := range resNodeGPUCost {
 		cluster, err := result.GetString(env.GetPromClusterLabel())
 		if err != nil {
@@ -141,8 +214,6 @@ func buildGPUCostMap(
 		nodeType, _ := result.GetString("instance_type")
 		providerID, _ := result.GetString("provider_id")
 
-		gpuCost := result.Values[0].Value
-
 		key := NodeIdentifier{
 			Cluster:    cluster,
 			Name:       name,
@@ -153,6 +224,29 @@ func buildGPUCostMap(
 			Name:    name,
 		}
 
+		var gpuCost float64
+
+		if customPricingEnabled && customPricingConfig != nil {
+
+			var customGPUStr string
+			if spot, ok := preemptible[key]; ok && spot {
+				customGPUStr = customPricingConfig.SpotGPU
+			} else {
+				customGPUStr = customPricingConfig.GPU
+			}
+
+			customGPUCost, err := strconv.ParseFloat(customGPUStr, 64)
+			if err != nil {
+				log.Warningf("ClusterNodes: error parsing custom GPU price: %s", customGPUStr)
+			}
+			gpuCost = customGPUCost
+
+		} else {
+
+			gpuCost = result.Values[0].Value
+
+		}
+
 		clusterAndNameToType[keyNon] = nodeType
 
 		// If gpu count is available use it to multiply gpu cost
@@ -162,6 +256,8 @@ func buildGPUCostMap(
 			gpuCostMap[key] = gpuCost
 		}
 
+		log.Infof("ASSET CLUSTERNODES NODE: %s WITH GPUPRICE: %f", key.Name, gpuCost)
+
 	}
 
 	return gpuCostMap, clusterAndNameToType

pkg/costmodel/cluster_helpers_test.go

@@ -1,12 +1,14 @@
 package costmodel
 
 import (
-	"github.com/kubecost/cost-model/pkg/prom"
-	"github.com/kubecost/cost-model/pkg/util"
 	"reflect"
 	"testing"
 	"time"
 
+	"github.com/kubecost/cost-model/pkg/cloud"
+	"github.com/kubecost/cost-model/pkg/prom"
+	"github.com/kubecost/cost-model/pkg/util"
+
 	"github.com/davecgh/go-spew/spew"
 )
 
@@ -849,7 +851,11 @@ func TestBuildGPUCostMap(t *testing.T) {
 
 	for _, testCase := range cases {
 		t.Run(testCase.name, func(t *testing.T) {
-			result, _ := buildGPUCostMap(testCase.promResult, testCase.countMap)
+			testProvider := &cloud.CustomProvider{
+				Config: cloud.NewProviderConfig("fakeFile"),
+			}
+			testPreemptible := make(map[NodeIdentifier]bool)
+			result, _ := buildGPUCostMap(testCase.promResult, testCase.countMap, testProvider, testPreemptible)
 			if !reflect.DeepEqual(result, testCase.expected) {
 				t.Errorf("buildGPUCostMap case %s failed. Got %+v but expected %+v", testCase.name, result, testCase.expected)
 			}