-
-
Notifications
You must be signed in to change notification settings - Fork 358
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support deletion of SageMaker Notebook Instances (#332)
* add nuking of sagemaker notebook instance * update README.md * go fmt * add sagemaker notebook filtering * Fix build failure on config Co-authored-by: Jan <[email protected]>
- Loading branch information
1 parent
044f1f8
commit 84e51f2
Showing
15 changed files
with
538 additions
and
21 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| package aws | ||
|
|
||
| import ( | ||
| "time" | ||
|
|
||
| awsgo "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/aws/aws-sdk-go/aws/session" | ||
| "github.com/aws/aws-sdk-go/service/sagemaker" | ||
| "github.com/gruntwork-io/cloud-nuke/config" | ||
| "github.com/gruntwork-io/cloud-nuke/logging" | ||
| "github.com/gruntwork-io/go-commons/errors" | ||
| ) | ||
|
|
||
| func getAllNotebookInstances(session *session.Session, excludeAfter time.Time, configObj config.Config) ([]*string, error) { | ||
| svc := sagemaker.New(session) | ||
|
|
||
| result, err := svc.ListNotebookInstances(&sagemaker.ListNotebookInstancesInput{}) | ||
|
|
||
| if err != nil { | ||
| return nil, errors.WithStackTrace(err) | ||
| } | ||
|
|
||
| var names []*string | ||
|
|
||
| for _, notebook := range result.NotebookInstances { | ||
| if notebook.CreationTime == nil { | ||
| continue | ||
| } | ||
| if !excludeAfter.After(awsgo.TimeValue(notebook.CreationTime)) { | ||
| continue | ||
| } | ||
| if !config.ShouldInclude(awsgo.StringValue(notebook.NotebookInstanceName), configObj.S3.IncludeRule.NamesRegExp, configObj.S3.ExcludeRule.NamesRegExp){ | ||
| continue | ||
| } | ||
| names = append(names, notebook.NotebookInstanceName) | ||
| } | ||
|
|
||
|
|
||
| return names, nil | ||
| } | ||
|
|
||
| func nukeAllNotebookInstances(session *session.Session, names []*string) error { | ||
| svc := sagemaker.New(session) | ||
|
|
||
| if len(names) == 0 { | ||
| logging.Logger.Infof("No Sagemaker Notebook Instance to nuke in region %s", *session.Config.Region) | ||
| return nil | ||
| } | ||
|
|
||
| logging.Logger.Infof("Deleting all Sagemaker Notebook Instances in region %s", *session.Config.Region) | ||
| deletedNames := []*string{} | ||
|
|
||
| for _, name := range names { | ||
| params := &sagemaker.DeleteNotebookInstanceInput{ | ||
| NotebookInstanceName: name, | ||
| } | ||
|
|
||
| _, err := svc.StopNotebookInstance(&sagemaker.StopNotebookInstanceInput{ | ||
| NotebookInstanceName: name, | ||
| }) | ||
| if err != nil { | ||
| logging.Logger.Errorf("[Failed] %s: %s", *name, err) | ||
| } | ||
|
|
||
| err = svc.WaitUntilNotebookInstanceStopped(&sagemaker.DescribeNotebookInstanceInput{ | ||
| NotebookInstanceName: name, | ||
| }) | ||
|
|
||
| if err != nil { | ||
| logging.Logger.Errorf("[Failed] %s", err) | ||
| } | ||
|
|
||
| _, err = svc.DeleteNotebookInstance(params) | ||
|
|
||
| if err != nil { | ||
| logging.Logger.Errorf("[Failed] %s: %s", *name, err) | ||
| } else { | ||
| deletedNames = append(deletedNames, name) | ||
| logging.Logger.Infof("Deleted Sagemaker Notebook Instance: %s", awsgo.StringValue(name)) | ||
| } | ||
| } | ||
|
|
||
| if len(deletedNames) > 0 { | ||
| for _, name := range deletedNames { | ||
|
|
||
| err := svc.WaitUntilNotebookInstanceDeleted(&sagemaker.DescribeNotebookInstanceInput{ | ||
| NotebookInstanceName: name, | ||
| }) | ||
|
|
||
| if err != nil { | ||
| logging.Logger.Errorf("[Failed] %s", err) | ||
| return errors.WithStackTrace(err) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| logging.Logger.Infof("[OK] %d Sagemaker Notebook Instance(s) deleted in %s", len(deletedNames), *session.Config.Region) | ||
| return nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| package aws | ||
|
|
||
| import ( | ||
| "strings" | ||
| "testing" | ||
| "time" | ||
|
|
||
| awsgo "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/aws/aws-sdk-go/aws/session" | ||
| "github.com/aws/aws-sdk-go/service/sagemaker" | ||
| "github.com/gruntwork-io/cloud-nuke/config" | ||
| "github.com/gruntwork-io/cloud-nuke/logging" | ||
| "github.com/gruntwork-io/cloud-nuke/util" | ||
| "github.com/gruntwork-io/go-commons/errors" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| // There's a built-in function WaitUntilDBInstanceAvailable but | ||
| // the times that it was tested, it wasn't returning anything so we'll leave with the | ||
| // custom one. | ||
|
|
||
| func waitUntilNotebookInstanceCreated(svc *sagemaker.SageMaker, name *string) error { | ||
| input := &sagemaker.DescribeNotebookInstanceInput{ | ||
| NotebookInstanceName: name, | ||
| } | ||
|
|
||
| for i := 0; i < 600; i++ { | ||
| instance, err := svc.DescribeNotebookInstance(input) | ||
| status := instance.NotebookInstanceStatus | ||
|
|
||
| if awsgo.StringValue(status) != "Pending" { | ||
| return nil | ||
| } | ||
|
|
||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| time.Sleep(1 * time.Second) | ||
| logging.Logger.Debug("Waiting for SageMaker Notebook Instance to be created") | ||
| } | ||
|
|
||
| return SageMakerNotebookInstanceDeleteError{name: *name} | ||
| } | ||
|
|
||
| func createTestNotebookInstance(t *testing.T, session *session.Session, name string, roleArn string) { | ||
| svc := sagemaker.New(session) | ||
|
|
||
| params := &sagemaker.CreateNotebookInstanceInput{ | ||
| InstanceType: awsgo.String("ml.t2.medium"), | ||
| NotebookInstanceName: awsgo.String(name), | ||
| RoleArn: awsgo.String(roleArn), | ||
| } | ||
|
|
||
| _, err := svc.CreateNotebookInstance(params) | ||
| require.NoError(t, err) | ||
|
|
||
| waitUntilNotebookInstanceCreated(svc, &name) | ||
| } | ||
|
|
||
| func TestNukeNotebookInstance(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| region, err := getRandomRegion() | ||
|
|
||
| require.NoError(t, errors.WithStackTrace(err)) | ||
|
|
||
| session, err := session.NewSessionWithOptions( | ||
| session.Options{ | ||
| SharedConfigState: session.SharedConfigEnable, | ||
| Config: awsgo.Config{ | ||
| Region: awsgo.String(region), | ||
| }, | ||
| }, | ||
| ) | ||
|
|
||
| notebookName := "cloud-nuke-test-" + util.UniqueID() | ||
| excludeAfter := time.Now().Add(1 * time.Hour) | ||
|
|
||
| role := createNotebookRole(t, session, notebookName+"-role") | ||
| defer deleteNotebookRole(session, role) | ||
|
|
||
| createTestNotebookInstance(t, session, notebookName, *role.Arn) | ||
|
|
||
| defer func() { | ||
| nukeAllNotebookInstances(session, []*string{¬ebookName}) | ||
|
|
||
| notebookNames, _ := getAllNotebookInstances(session, excludeAfter, config.Config{}) | ||
|
|
||
| assert.NotContains(t, awsgo.StringValueSlice(notebookNames), strings.ToLower(notebookName)) | ||
| }() | ||
|
|
||
| instances, err := getAllNotebookInstances(session, excludeAfter, config.Config{}) | ||
|
|
||
| if err != nil { | ||
| assert.Failf(t, "Unable to fetch list of SageMaker Notebook Instances", errors.WithStackTrace(err).Error()) | ||
| } | ||
|
|
||
| assert.Contains(t, awsgo.StringValueSlice(instances), notebookName) | ||
|
|
||
| } |
Oops, something went wrong.