Bläddra i källkod

support multiple GPU labels

Signed-off-by: Ajay Tripathy <4tripathy@gmail.com>
Ajay Tripathy 3 år sedan
förälder
incheckning
f797a15da9
3 ändrade filer med 35 tillägg och 7 borttagningar
  1. 1 1
      configs/pricing_schema.csv
  2. 9 6
      pkg/cloud/csvprovider.go
  3. 25 0
      test/cloud_test.go

+ 1 - 1
configs/pricing_schema.csv

@@ -1,4 +1,4 @@
 EndTimestamp,InstanceID,Region,AssetClass,InstanceIDField,InstanceType,MarketPriceHourly,Version
 2019-04-17 23:34:22 UTC,gke-standard-cluster-1-pool-1-91dc432d-cg69,,node,metadata.name,,0.1337,
 2019-04-17 23:34:22 UTC,Quadro_RTX_4000,,gpu,nvidia.com/gpu_type,,0.75,
-
+2019-04-17 23:34:22 UTC,Quadro_RTX_4001,,gpu,gpu.nvidia.com/class,,0.80,

+ 9 - 6
pkg/cloud/csvprovider.go

@@ -34,7 +34,7 @@ type CSVProvider struct {
 	PricingPV               map[string]*price
 	PVMapField              string
 	GPUClassPricing         map[string]*price
-	GPUMapField             string
+	GPUMapFields            []string
 	UsesRegion              bool
 	DownloadPricingDataLock sync.RWMutex
 }
@@ -62,6 +62,7 @@ func (c *CSVProvider) DownloadPricingData() error {
 	nodeclasscount := make(map[string]float64)
 	pvpricing := make(map[string]*price)
 	gpupricing := make(map[string]*price)
+	c.GPUMapFields = make([]string, 0, 1)
 	header, err := csvutil.Header(price{}, "csv")
 	if err != nil {
 		return err
@@ -171,7 +172,7 @@ func (c *CSVProvider) DownloadPricingData() error {
 			c.NodeMapField = p.InstanceIDField
 		} else if p.AssetClass == "gpu" {
 			gpupricing[key] = &p
-			c.GPUMapField = strings.ToLower(p.InstanceIDField)
+			c.GPUMapFields = append(c.GPUMapFields, strings.ToLower(p.InstanceIDField))
 		} else {
 			log.Infof("Unrecognized asset class %s, defaulting to node", p.AssetClass)
 			pricing[key] = &p
@@ -193,7 +194,7 @@ func (c *CSVProvider) DownloadPricingData() error {
 type csvKey struct {
 	Labels     map[string]string
 	ProviderID string
-	GPULabel   string
+	GPULabel   []string
 	GPU        int64
 }
 
@@ -210,8 +211,10 @@ func (k *csvKey) GPUCount() int {
 }
 
 func (k *csvKey) GPUType() string {
-	if val, ok := k.Labels[k.GPULabel]; ok {
-		return val
+	for _, label := range k.GPULabel {
+		if val, ok := k.Labels[label]; ok {
+			return val
+		}
 	}
 	return ""
 }
@@ -346,7 +349,7 @@ func (c *CSVProvider) GetKey(l map[string]string, n *v1.Node) Key {
 	return &csvKey{
 		ProviderID: id,
 		Labels:     l,
-		GPULabel:   c.GPUMapField,
+		GPULabel:   c.GPUMapFields,
 		GPU:        gpuCount,
 	}
 }

+ 25 - 0
test/cloud_test.go

@@ -143,6 +143,15 @@ func TestNodePriceFromCSVWithGPU(t *testing.T) {
 	n.Status.Capacity = v1.ResourceList{"nvidia.com/gpu": *resource.NewScaledQuantity(2, 0)}
 	wantPrice := "1.633700"
 
+	n2 := &v1.Node{}
+	n2.Spec.ProviderID = providerIDWant
+	n2.Name = nameWant
+	n2.Labels = make(map[string]string)
+	n2.Labels["foo"] = labelFooWant
+	n2.Labels["gpu.nvidia.com/class"] = "Quadro_RTX_4001"
+	n2.Status.Capacity = v1.ResourceList{"nvidia.com/gpu": *resource.NewScaledQuantity(2, 0)}
+	wantPrice2 := "1.733700"
+
 	c := &cloud.CSVProvider{
 		CSVLocation: "../configs/pricing_schema.csv",
 		CustomProvider: &cloud.CustomProvider{
@@ -167,6 +176,22 @@ func TestNodePriceFromCSVWithGPU(t *testing.T) {
 
 	}
 
+	k2 := c.GetKey(n2.Labels, n2)
+	resN2, err := c.NodePricing(k2)
+	if err != nil {
+		t.Errorf("Error in NodePricing: %s", err.Error())
+	} else {
+		gotGPU := resN2.GPU
+		gotPrice := resN2.Cost
+		if gotGPU != wantGPU {
+			t.Errorf("Wanted gpu count '%s' got gpu count '%s'", wantGPU, gotGPU)
+		}
+		if gotPrice != wantPrice2 {
+			t.Errorf("Wanted price '%s' got price '%s'", wantPrice2, gotPrice)
+		}
+
+	}
+
 }
 
 func TestNodePriceFromCSV(t *testing.T) {