Add cost tracking for blanket billing with karpenter

go.mod

@@ -4,8 +4,9 @@ go 1.21
 
 require (
 	github.com/apparentlymart/go-cidr v1.1.0
-	github.com/aws/aws-sdk-go v1.51.28
-	github.com/google/uuid v1.6.0
+	github.com/aws/aws-sdk-go v1.44.332
+	github.com/google/uuid v1.3.0
+	github.com/stretchr/testify v1.9.0
 	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/api v0.29.0
 	k8s.io/apimachinery v0.29.0
@@ -35,29 +36,32 @@ require (
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
-	github.com/google/pprof v0.0.0-20240424215950-a892ee059fd6 // indirect
-	github.com/imdario/mergo v0.3.16 // indirect
+	github.com/google/pprof v0.0.0-20211104044539-f987b9c94b31 // indirect
+	github.com/imdario/mergo v0.3.6 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.2 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
-	github.com/prometheus/client_golang v1.19.0 // indirect
-	github.com/prometheus/client_model v0.6.1 // indirect
-	github.com/prometheus/common v0.53.0 // indirect
-	github.com/prometheus/procfs v0.14.0 // indirect
-	github.com/sirupsen/logrus v1.9.3 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v1.14.0 // indirect
+	github.com/prometheus/client_model v0.3.0 // indirect
+	github.com/prometheus/common v0.37.0 // indirect
+	github.com/prometheus/procfs v0.8.0 // indirect
+	github.com/sirupsen/logrus v1.8.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	golang.org/x/net v0.24.0 // indirect
-	golang.org/x/oauth2 v0.19.0 // indirect
-	golang.org/x/sys v0.19.0 // indirect
-	golang.org/x/term v0.19.0 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
+	golang.org/x/net v0.21.0 // indirect
+	golang.org/x/oauth2 v0.10.0 // indirect
+	golang.org/x/sys v0.17.0 // indirect
+	golang.org/x/term v0.17.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	golang.org/x/time v0.5.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

hatchery/config.go

@@ -131,7 +131,9 @@ type HatcheryConfig struct {
 	Sidecar                SidecarContainer     `json:"sidecar"`
 	MoreConfigs            []AppConfigInfo      `json:"more-configs"`
 	PrismaConfig           PrismaConfig         `json:"prisma"`
+	Karpenter              bool                 `json:"karpenter"`
 	NextflowGlobalConfig   NextflowGlobalConfig `json:"nextflow-global"`
+	Developement           bool                 `json:"developement"`
 }
 
 // Config to allow for Prisma Agents

hatchery/cur.go

@@ -0,0 +1,78 @@
+package hatchery
+
+import (
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/costexplorer"
+)
+
+type costUsage struct {
+	Username  string  `json:"username"`
+	TotalCost float64 `json:"total-cost"`
+}
+
+func getCostUsageReport(username string, workflowname string) (*costUsage, error) {
+	// query cost usage report
+	// Create a Cost Explorer service client
+	sess := session.Must(session.NewSessionWithOptions(session.Options{
+		SharedConfigState: session.SharedConfigEnable,
+	}))
+	svc := costexplorer.New(sess)
+	// Build the request with date range and filter
+	// Return costs by tags
+	req := &costexplorer.GetCostAndUsageInput{
+		Metrics: []*string{
+			aws.String("UnblendedCost"),
+		},
+		TimePeriod: &costexplorer.DateInterval{
+			// 1 year ago is max
+			Start: aws.String(time.Now().AddDate(-1, 0, 0).Format("2006-01-02")),
+			// Today
+			End: aws.String(time.Now().Format("2006-01-02")),
+		},
+		Filter: &costexplorer.Expression{
+			Tags: &costexplorer.TagValues{
+				Key: aws.String("gen3username"),
+				Values: []*string{
+					aws.String(userToResourceName(username, "user")),
+				},
+			},
+		},
+		Granularity: aws.String("MONTHLY"),
+	}
+
+	if workflowname != "" {
+		req.Filter = &costexplorer.Expression{
+			Tags: &costexplorer.TagValues{
+				Key: aws.String("gen3username"),
+				Values: []*string{
+					aws.String(userToResourceName(username, "user")),
+				},
+			},
+		}
+	}
+
+	// Call Cost Explorer API
+	resp, err := svc.GetCostAndUsage(req)
+	if err != nil {
+		fmt.Println("Got error calling GetCostAndUsage:", err)
+		return nil, err
+	}
+	var total float64
+	for _, result := range resp.ResultsByTime {
+		// Get amount
+		totalAmount := result.Total["UnblendedCost"]
+		amount, _ := strconv.ParseFloat(*totalAmount.Amount, 64)
+
+		// Sum amounts
+		total += amount
+	}
+
+	ret := costUsage{Username: username, TotalCost: total}
+
+	return &ret, nil
+}

hatchery/hatchery.go

@@ -46,11 +46,46 @@ func RegisterHatchery() {
 	http.HandleFunc("/setpaymodel", setpaymodel)
 	http.HandleFunc("/resetpaymodels", resetPaymodels)
 	http.HandleFunc("/allpaymodels", allpaymodels)
+	http.HandleFunc("/cost", cost)
 
 	// ECS functions
 	http.HandleFunc("/create-ecs-cluster", createECSCluster)
 }
 
+func cost(w http.ResponseWriter, r *http.Request) {
+	// create context for http call
+	// context, cancel := context.WithTimeout(r.Context(), 30*time.Second)
+	// defer cancel()
+
+	userName := getCurrentUserName(r)
+
+	workflowname := r.URL.Query().Get("workflowname")
+	// check if workflowname is empty
+
+	// get cost usage report
+	costUsageReport, err := getCostUsageReport(userName, workflowname)
+	if err != nil {
+		Config.Logger.Print(err)
+		// Send 500 error
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// make a return object with username and cost
+	cur, err := json.Marshal(costUsageReport)
+	if err != nil {
+		Config.Logger.Print(err)
+		// Send 500 error
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	// Add header indicating this is a json response
+	w.Header().Set("Content-Type", "application/json")
+
+	// return json
+	fmt.Fprint(w, string(cur))
+}
+
 func home(w http.ResponseWriter, r *http.Request) {
 	htmlHeader := `<html>
 	<head>Gen3 Hatchery</head>
@@ -118,6 +153,8 @@ func paymodels(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	// Set header to indicate it's a json response
+	w.Header().Set("Content-Type", "application/json")
 	fmt.Fprint(w, string(out))
 }
 
@@ -142,6 +179,8 @@ func allpaymodels(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	// Set header to indicate it's a json response
+	w.Header().Set("Content-Type", "application/json")
 	fmt.Fprint(w, string(out))
 }
 
@@ -540,24 +579,27 @@ func terminate(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprintf(w, "Terminated workspace")
 	}
 
-	// Need to reset pay model only after workspace termination is completed.
-	go func() {
-		// Periodically poll for status, until it is set as "Not Found"
-		for {
-			status, err := getWorkspaceStatus(r.Context(), userName, accessToken)
-			if err != nil {
-				Config.Logger.Printf("error fetching workspace status for user %s\n err: %s", userName, err)
+	// check if dynamoDB is enabled
+	if Config.Config.PayModelsDynamodbTable != "" {
+		// Need to reset pay model only after workspace termination is completed.
+		go func() {
+			// Periodically poll for status, until it is set as "Not Found"
+			for {
+				status, err := getWorkspaceStatus(r.Context(), userName, accessToken)
+				if err != nil {
+					Config.Logger.Printf("error fetching workspace status for user %s\n err: %s", userName, err)
+				}
+				if status.Status == "Not Found" {
+					break
+				}
+				time.Sleep(5 * time.Second)
 			}
-			if status.Status == "Not Found" {
-				break
+			err = resetCurrentPaymodel(userName)
+			if err != nil {
+				Config.Logger.Printf("unable to reset current paymodel for current user %s\nerr: %s", userName, err)
 			}
-			time.Sleep(5 * time.Second)
-		}
-		err = resetCurrentPaymodel(userName)
-		if err != nil {
-			Config.Logger.Printf("unable to reset current paymodel for current user %s\nerr: %s", userName, err)
-		}
-	}()
+		}()
+	}
 }
 
 func getBearerToken(r *http.Request) string {