Merge pull request #469 from ericzbeard/forecastv1-13

Add forecast check for SageMaker limits

go.mod

@@ -45,6 +45,7 @@ require (
 
 require (
 	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.15 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sagemaker v1.151.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect

go.sum

@@ -142,6 +142,8 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.56.1 h1:wsg9Z/vNnCmxWikfGIoOlnExtEU45
 github.com/aws/aws-sdk-go-v2/service/s3 v1.56.1/go.mod h1:8rDw3mVwmvIWWX/+LWY3PPIMZuwnQdJMCt0iVFVT3qw=
 github.com/aws/aws-sdk-go-v2/service/s3 v1.58.2 h1:sZXIzO38GZOU+O0C+INqbH7C2yALwfMWpd64tONS/NE=
 github.com/aws/aws-sdk-go-v2/service/s3 v1.58.2/go.mod h1:Lcxzg5rojyVPU/0eFwLtcyTaek/6Mtic5B1gJo7e/zE=
+github.com/aws/aws-sdk-go-v2/service/sagemaker v1.151.0 h1:zH7b/d8vOrOWdgluKEux2TAECYKhprH2eMztkpak/mI=
+github.com/aws/aws-sdk-go-v2/service/sagemaker v1.151.0/go.mod h1:lDmK3DHWV6Y6hpzeUAaXq4w+ks6fFYXdkjavIe8STCE=
 github.com/aws/aws-sdk-go-v2/service/servicequotas v1.21.9 h1:3o5zcwZYvte3CeaYpLaWafwCSkJpclPXI5KSH+lXB90=
 github.com/aws/aws-sdk-go-v2/service/servicequotas v1.21.9/go.mod h1:QZpGkzlec0TPr8CA2Td5zRUJBC5+104ib0MusH5UVfI=
 github.com/aws/aws-sdk-go-v2/service/servicequotas v1.22.1 h1:QsHvqtdy0mGzpg/A+1lZX1ilf05Vuh2rSBzNJ3f3T1I=

internal/aws/sagemaker/sagemaker.go

@@ -0,0 +1,45 @@
+package sagemaker
+
+import (
+	"context"
+
+	"github.com/aws-cloudformation/rain/internal/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sagemaker"
+)
+
+func getClient() *sagemaker.Client {
+	return sagemaker.NewFromConfig(aws.Config())
+}
+
+type NotebookInstance struct {
+	InstanceType string
+}
+
+func GetNotebookInstances() ([]NotebookInstance, error) {
+	retval := make([]NotebookInstance, 0)
+
+	client := getClient()
+	var nextToken *string
+
+	for {
+		resp, err := client.ListNotebookInstances(context.Background(),
+			&sagemaker.ListNotebookInstancesInput{
+				NextToken: nextToken,
+			})
+		if err != nil {
+			return retval, err
+		}
+
+		for _, inst := range resp.NotebookInstances {
+			retval = append(retval, NotebookInstance{InstanceType: string(inst.InstanceType)})
+		}
+
+		if resp.NextToken != nil {
+			nextToken = resp.NextToken
+		} else {
+			break
+		}
+	}
+
+	return retval, nil
+}

internal/cmd/forecast/README.md

@@ -53,6 +53,8 @@ These can be ignored with the `--ignore` argument.
 | F0015 | ELB target groups must be of type instance if they are used by an ASG          |
 | F0016 | Lambda function role exists                                                    |
 | F0017 | Lambda function role can be assumed                                            |
+| F0018 | SageMaker Notebook quota limit has not been reached                            |
+
 ## Estimates
 
 The forecast command also tries to estimate how long it thinks your stack will

internal/cmd/forecast/forecast.go

@@ -506,6 +506,7 @@ func init() {
 	forecasters["AWS::SNS::Topic"] = checkSNSTopic
 	forecasters["AWS::ElasticLoadBalancingV2::TargetGroup"] = checkELBTargetGroup
 	forecasters["AWS::Lambda::Function"] = checkLambdaFunction
+	forecasters["AWS::SageMaker::NotebookInstance"] = checkSageMakerNotebook
 
 	// Initialize estimates map
 	InitEstimates()