Implement retrieval of prices from 3rd party source in a runtime

Author: dm3ch

pkg/providers/pricing/pricing.go

@@ -16,19 +16,24 @@ package pricing
 
 import (
 	"context"
+	"encoding/csv"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"sync"
+	"time"
 
 	"k8s.io/klog/v2"
 )
 
 const (
 	initialOnDemandPricesFile = "initial-on-demand-prices.json"
+	pricingCSVURL             = "https://gcloud-compute.com/machine-types-regions.csv"
 )
 
 type Provider interface {
@@ -46,12 +51,16 @@ type DefaultProvider struct {
 	muSpot                      sync.RWMutex
 	region                      string
 	prices                      map[string]map[string]float64
+	spotPrices                  map[string]map[string]float64
+	lastUpdated                 time.Time
 }
 
 func NewDefaultProvider(ctx context.Context, region string) *DefaultProvider {
 	p := &DefaultProvider{
-		region: region,
-		prices: make(map[string]map[string]float64),
+		region:      region,
+		prices:      make(map[string]map[string]float64),
+		spotPrices:  make(map[string]map[string]float64),
+		lastUpdated: time.Now(),
 	}
 
 	if err := p.loadInitialPrices(); err != nil {
@@ -113,7 +122,6 @@ func (p *DefaultProvider) OnDemandPrice(instanceType string) (float64, bool) {
 
 	regionPrices, ok := p.prices[p.region]
 	if !ok {
-		fmt.Println(p.prices)
 		return 0, false
 	}
 
@@ -122,16 +130,124 @@ func (p *DefaultProvider) OnDemandPrice(instanceType string) (float64, bool) {
 }
 
 func (p *DefaultProvider) SpotPrice(instanceType string, zone string) (float64, bool) {
-	// Currently, we don't have spot price information
-	return 0, false
+	p.muSpot.RLock()
+	defer p.muSpot.RUnlock()
+
+	regionPrices, ok := p.spotPrices[p.region]
+	if !ok {
+		return 0, false
+	}
+
+	price, ok := regionPrices[instanceType]
+	return price, ok
 }
 
-func (p *DefaultProvider) UpdateOnDemandPricing(ctx context.Context) error {
-	// For now, we only use static pricing data
+type priceUpdateConfig struct {
+	priceColumn string
+	targetMap   *map[string]map[string]float64
+	mu          *sync.RWMutex
+}
+
+func (p *DefaultProvider) updatePricing(ctx context.Context, config priceUpdateConfig) error {
+	config.mu.Lock()
+	defer config.mu.Unlock()
+
+	// Download the CSV file
+	resp, err := http.Get(pricingCSVURL)
+	if err != nil {
+		return fmt.Errorf("downloading CSV: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// Read the CSV data
+	reader := csv.NewReader(resp.Body)
+	reader.Comma = ','
+
+	// Read header
+	header, err := reader.Read()
+	if err != nil {
+		return fmt.Errorf("reading CSV header: %w", err)
+	}
+
+	// Find the required column indices
+	priceColIndex := -1
+	regionColIndex := -1
+	machineTypeColIndex := -1
+
+	for i, col := range header {
+		switch col {
+		case config.priceColumn:
+			priceColIndex = i
+		case "region":
+			regionColIndex = i
+		case "name":
+			machineTypeColIndex = i
+		}
+	}
+
+	if priceColIndex == -1 || regionColIndex == -1 || machineTypeColIndex == -1 {
+		return fmt.Errorf("could not find required columns in CSV")
+	}
+
+	// Process the data
+	newPrices := make(map[string]map[string]float64)
+
+	for {
+		record, err := reader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			klog.Errorf("Error reading CSV record: %v", err)
+			continue
+		}
+
+		// Get machine type, region and price
+		machineType := record[machineTypeColIndex]
+		region := record[regionColIndex]
+		priceStr := record[priceColIndex]
+
+		// Parse price
+		price, err := strconv.ParseFloat(priceStr, 64)
+		if err != nil {
+			klog.Errorf("Error parsing price for %s in region %s: %v", machineType, region, err)
+			continue
+		}
+
+		// Initialize region map if it doesn't exist
+		if _, exists := newPrices[region]; !exists {
+			newPrices[region] = make(map[string]float64)
+		}
+
+		// Store the price
+		newPrices[region][machineType] = price
+	}
+
+	// Check if we got any prices
+	if len(newPrices) == 0 {
+		return fmt.Errorf("no prices found during price update for %s pricing", config.priceColumn)
+	}
+
+	// Update the prices
+	*config.targetMap = newPrices
+	p.lastUpdated = time.Now()
+
+	klog.V(2).Infof("Updated prices for %d regions", len(newPrices))
 	return nil
 }
 
+func (p *DefaultProvider) UpdateOnDemandPricing(ctx context.Context) error {
+	return p.updatePricing(ctx, priceUpdateConfig{
+		priceColumn: "hour",
+		targetMap:   &p.prices,
+		mu:          &p.muOnDemand,
+	})
+}
+
 func (p *DefaultProvider) UpdateSpotPricing(ctx context.Context) error {
-	// Currently, we don't have spot price information
-	return nil
+	return p.updatePricing(ctx, priceUpdateConfig{
+		priceColumn: "hourSpot",
+		targetMap:   &p.spotPrices,
+		mu:          &p.muSpot,
+	})
 }

pkg/providers/pricing/pricing_test.go

@@ -23,49 +23,124 @@ import (
 	"github.com/cloudpilot-ai/karpenter-provider-gcp/pkg/providers/pricing"
 )
 
-func TestDefaultProvider_OnDemandPrice(t *testing.T) {
+var testInstanceTypes = []string{
+	"e2-standard-32",
+	"n2-standard-16",
+	"c2-standard-8",
+	"m1-ultramem-40",
+	"a2-highgpu-1g",
+}
+
+func TestPricingProvider_InitialPrices(t *testing.T) {
+	// Initialize the provider with europe-west4 region
+	provider := pricing.NewDefaultProvider(context.Background(), "europe-west4")
+
+	// Test getting prices for various instance types from initial file
+	for _, instanceType := range testInstanceTypes {
+		_, found := provider.OnDemandPrice(instanceType)
+		if !found {
+			t.Errorf("Failed to find initial on-demand price for %s", instanceType)
+			continue
+		}
+	}
+
+	// Get all instance types from initial file
+	types := provider.InstanceTypes()
+	if len(types) == 0 {
+		t.Fatal("No instance types found in initial prices")
+	}
+
+	// Check if all test instance types are in the initial list
+	for _, instanceType := range testInstanceTypes {
+		found := false
+		for _, t := range types {
+			if t == instanceType {
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Errorf("%s not found in initial instance types list", instanceType)
+		}
+	}
+}
+
+func TestPricingProvider_OnDemandPrice(t *testing.T) {
 	// Initialize the provider with europe-west4 region
 	provider := pricing.NewDefaultProvider(context.Background(), "europe-west4")
-	provider.UpdateOnDemandPricing(context.Background())
 
-	// Test getting price for e2-standard-32
-	price, found := provider.OnDemandPrice("e2-standard-32")
-	if !found {
-		t.Fatal("Failed to find price for e2-standard-32")
+	// Test price update
+	if err := provider.UpdateOnDemandPricing(context.Background()); err != nil {
+		t.Fatalf("Failed to update on-demand pricing: %v", err)
 	}
 
-	t.Logf("On-demand price for e2-standard-32 in europe-west4: $%.4f per hour", price)
+	// Test getting prices for various instance types
+	for _, instanceType := range testInstanceTypes {
+		_, found := provider.OnDemandPrice(instanceType)
+		if !found {
+			t.Errorf("Failed to find on-demand price for %s", instanceType)
+			continue
+		}
+	}
 
 	// Test getting price for a non-existent instance type
-	_, found = provider.OnDemandPrice("non-existent-type")
+	_, found := provider.OnDemandPrice("non-existent-type")
 	if found {
 		t.Error("Expected to not find price for non-existent instance type")
 	}
 }
 
-func TestDefaultProvider_InstanceTypes(t *testing.T) {
+func TestPricingProvider_SpotPrice(t *testing.T) {
 	// Initialize the provider with europe-west4 region
 	provider := pricing.NewDefaultProvider(context.Background(), "europe-west4")
-	provider.UpdateOnDemandPricing(context.Background())
+
+	// Test price update
+	if err := provider.UpdateSpotPricing(context.Background()); err != nil {
+		t.Fatalf("Failed to update spot pricing: %v", err)
+	}
+
+	// Test getting spot prices for various instance types
+	for _, instanceType := range testInstanceTypes {
+		_, found := provider.SpotPrice(instanceType, "europe-west4-a")
+		if !found {
+			t.Errorf("Failed to find spot price for %s", instanceType)
+			continue
+		}
+	}
+
+	// Test getting price for a non-existent instance type
+	_, found := provider.SpotPrice("non-existent-type", "europe-west4-a")
+	if found {
+		t.Error("Expected to not find price for non-existent instance type")
+	}
+}
+
+func TestPricingProvider_InstanceTypes(t *testing.T) {
+	// Initialize the provider with europe-west4 region
+	provider := pricing.NewDefaultProvider(context.Background(), "europe-west4")
+
+	// Test price update
+	if err := provider.UpdateOnDemandPricing(context.Background()); err != nil {
+		t.Fatalf("Failed to update on-demand pricing: %v", err)
+	}
 
 	// Get all instance types
 	types := provider.InstanceTypes()
 	if len(types) == 0 {
 		t.Fatal("No instance types found")
 	}
 
-	t.Logf("Found %d instance types in europe-west4", len(types))
-
-	// Check if e2-standard-32 is in the list
-	found := false
-	for _, instanceType := range types {
-		if instanceType == "e2-standard-32" {
-			found = true
-			break
+	// Check if all test instance types are in the list
+	for _, instanceType := range testInstanceTypes {
+		found := false
+		for _, t := range types {
+			if t == instanceType {
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Errorf("%s not found in instance types list", instanceType)
 		}
-	}
-
-	if !found {
-		t.Error("e2-standard-32 not found in instance types list")
 	}
 }