Add cost tracking for blanket billing with karpenter

go.mod

@@ -15,6 +15,13 @@ require (
 	sigs.k8s.io/aws-iam-authenticator v0.6.29
 )
 
+replace (
+	github.com/mattn/go-sqlite3 v1.14.16 => github.com/mattn/go-sqlite3 v1.14.18
+	golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 => golang.org/x/crypto v0.1.0
+	golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550 => golang.org/x/crypto v0.1.0
+	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 => golang.org/x/crypto v0.1.0
+)
+
 require (
 	github.com/aws/aws-sdk-go-v2 v1.34.0 // indirect
 	github.com/aws/smithy-go v1.22.2 // indirect

go.sum

@@ -114,6 +114,9 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
+golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
+golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco=
 golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
 golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
 golang.org/x/oauth2 v0.25.0 h1:CY4y7XT9v0cRI9oupztF8AgiIu99L/ksR/Xp/6jrZ70=
@@ -125,12 +128,19 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
 golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
+golang.org/x/term v0.1.0/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg=
 golang.org/x/term v0.28.0/go.mod h1:Sw/lC2IAUZ92udQNf3WodGtn4k/XoLyZoh8v/8uiwek=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
 golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
@@ -139,6 +149,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
+golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.26.0 h1:v/60pFQmzmT9ExmjDv2gGIfi3OqfKoEP6I5+umXlbnQ=
 golang.org/x/tools v0.26.0/go.mod h1:TPVVj70c7JJ3WCazhD8OdXcZg/og+b9+tH/KxylGwH0=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

hatchery/config.go

@@ -1,6 +1,7 @@
 package hatchery
 
 import (
+	"github.com/aws/aws-sdk-go/service/costexplorer/costexploreriface"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 	k8sv1 "k8s.io/api/core/v1"
 
@@ -110,6 +111,10 @@ type AllPayModels struct {
 	PayModels       []PayModel `json:"all_pay_models"`
 }
 
+type CostExplorerClient struct {
+	CostExporer costexploreriface.CostExplorerAPI
+}
+
 type DbConfig struct {
 	DynamoDb dynamodbiface.DynamoDBAPI
 }
@@ -133,7 +138,11 @@ type HatcheryConfig struct {
 	Sidecar                SidecarContainer     `json:"sidecar"`
 	MoreConfigs            []AppConfigInfo      `json:"more-configs"`
 	PrismaConfig           PrismaConfig         `json:"prisma"`
+	Karpenter              bool                 `json:"karpenter"`
+	DefaultHardLimit       float32              `json:"default-hard-limit"`
+	DefaultSoftLimit       float32              `json:"default-soft-limit"`
 	NextflowGlobalConfig   NextflowGlobalConfig `json:"nextflow-global"`
+	Developement           bool                 `json:"developement"`
 }
 
 // Config to allow for Prisma Agents

hatchery/cur.go

@@ -0,0 +1,86 @@
+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"`
+}
+
+// This function will get called in the module that calls `getCostUsageRepot`
+var initializeCostExplorerClient = func() *CostExplorerClient {
+	// Create an interface to CostExplorer service client
+	sess := session.Must(session.NewSessionWithOptions(session.Options{
+		SharedConfigState: session.SharedConfigEnable,
+	}))
+	return &CostExplorerClient{
+		CostExporer: costexplorer.New(sess),
+	}
+}
+
+var getCostUsageReport = func(costexplorerclient *CostExplorerClient, username string, workflowname string) (*costUsage, error) {
+	// query cost usage report
+	// the CostExplorer service client is passed in as a parameter
+
+	// 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 := costexplorerclient.CostExporer.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/cur_test.go

@@ -0,0 +1,100 @@
+package hatchery
+
+import (
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/costexplorer"
+	"github.com/aws/aws-sdk-go/service/costexplorer/costexploreriface"
+)
+
+type MockCostOutput struct {
+	output costexplorer.GetCostAndUsageOutput
+}
+
+type CostExplorerMockClient struct {
+	costexploreriface.CostExplorerAPI
+	mockOutput *MockCostOutput
+}
+
+func (m *CostExplorerMockClient) GetCostAndUsage(input *costexplorer.GetCostAndUsageInput) (*costexplorer.GetCostAndUsageOutput, error) {
+	return &m.mockOutput.output, nil
+}
+
+func Test_GetCostUsageReport(t *testing.T) {
+	defer SetupAndTeardownTest()()
+
+	testCases := []struct {
+		name           string
+		want           *costUsage
+		mockCostOutput *MockCostOutput
+	}{
+		{
+			name: "UserHasCosts",
+			want: &costUsage{Username: "test_user", TotalCost: 100},
+			mockCostOutput: &MockCostOutput{
+				output: costexplorer.GetCostAndUsageOutput{
+					ResultsByTime: []*costexplorer.ResultByTime{
+						{
+							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")),
+							},
+							Total: map[string]*costexplorer.MetricValue{
+								"UnblendedCost": {
+									Amount: aws.String("100"),
+									Unit:   aws.String("USD"),
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name:           "NoUserCosts",
+			want:           &costUsage{Username: "test_user", TotalCost: 0},
+			mockCostOutput: &MockCostOutput{},
+		},
+	}
+
+	// Backing up original functions before mocking
+	original_getCostUsageReport := getCostUsageReport
+	defer func() {
+		// restore original functions
+		getCostUsageReport = original_getCostUsageReport
+	}()
+
+	// mock the cost explorer interface and cost report
+	costexplorerclient := initializeCostExplorerClient()
+
+	for _, testcase := range testCases {
+		t.Logf("Testing GetCostUsageReport when %s", testcase.name)
+
+		costexplorerclient.CostExporer = &CostExplorerMockClient{
+			CostExplorerAPI: nil,
+			mockOutput:      testcase.mockCostOutput,
+		}
+
+		/* Act */
+		got, err := getCostUsageReport(costexplorerclient, "test_user", "Direct Pay")
+		if nil != err {
+			t.Errorf("failed to get cost usage report, got: %v", err)
+			return
+		}
+
+		/* Assert */
+		if reflect.TypeOf(got) != reflect.TypeOf(testcase.want) {
+			t.Errorf("Return value is not correct type:\ngot: '%v'\nwant: '%v'",
+				reflect.TypeOf(got), reflect.TypeOf(testcase.want))
+		}
+		if !reflect.DeepEqual(got, testcase.want) {
+			t.Errorf("\nassertion error while testing `getCostUsageReport`: \nWant:%+v\nGot:%+v", testcase.want, got)
+		}
+
+	}
+}