From 2153da1f3cbe1f851d73c52bd15c86a32c2bda29 Mon Sep 17 00:00:00 2001 From: Nicolas Pellegrin Date: Fri, 3 May 2024 20:22:41 +0200 Subject: [PATCH] feat(bedrock): add bedrock custom models, customization jobs, provisioned throughput and logging --- resources/bedrock-custom-models.go | 89 +++++++++++++++ resources/bedrock-model-customization-jobs.go | 105 ++++++++++++++++++ ...-model-invocation-logging-configuration.go | 65 +++++++++++ .../bedrock-provisioned-model-throughputs.go | 94 ++++++++++++++++ 4 files changed, 353 insertions(+) create mode 100644 resources/bedrock-custom-models.go create mode 100644 resources/bedrock-model-customization-jobs.go create mode 100644 resources/bedrock-model-invocation-logging-configuration.go create mode 100644 resources/bedrock-provisioned-model-throughputs.go diff --git a/resources/bedrock-custom-models.go b/resources/bedrock-custom-models.go new file mode 100644 index 00000000..8e9b4cf4 --- /dev/null +++ b/resources/bedrock-custom-models.go @@ -0,0 +1,89 @@ +package resources + +import ( + "context" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrock" + + "github.com/ekristen/libnuke/pkg/registry" + "github.com/ekristen/libnuke/pkg/resource" + "github.com/ekristen/libnuke/pkg/types" + + "github.com/ekristen/aws-nuke/v3/pkg/nuke" +) + +const BedrockCustomModelResource = "BedrockCustomModel" + +func init() { + registry.Register(®istry.Registration{ + Name: BedrockCustomModelResource, + Scope: nuke.Account, + Lister: &BedrockCustomModelLister{}, + }) +} + +type BedrockCustomModelLister struct{} + +func (l *BedrockCustomModelLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) { + opts := o.(*nuke.ListerOpts) + + svc := bedrock.New(opts.Session) + resources := make([]resource.Resource, 0) + + params := &bedrock.ListCustomModelsInput{ + MaxResults: aws.Int64(30), + } + + for { + resp, err := svc.ListCustomModels(params) + if err != nil { + return nil, err + } + + for _, modelSummary := range resp.ModelSummaries { + tagResp, err := svc.ListTagsForResource( + &bedrock.ListTagsForResourceInput{ + ResourceARN: modelSummary.ModelArn, + }) + if err != nil { + return nil, err + } + resources = append(resources, &BedrockCustomModel{ + svc: svc, + Name: modelSummary.ModelName, + Tags: tagResp.Tags, + }) + } + + if resp.NextToken == nil { + break + } + + params.NextToken = resp.NextToken + } + + return resources, nil +} + +type BedrockCustomModel struct { + svc *bedrock.Bedrock + Name *string + Tags []*bedrock.Tag +} + +func (r *BedrockCustomModel) Remove(_ context.Context) error { + _, err := r.svc.DeleteCustomModel(&bedrock.DeleteCustomModelInput{ + ModelIdentifier: r.Name, + }) + + return err +} + +func (r *BedrockCustomModel) String() string { + return *r.Name +} + +func (r *BedrockCustomModel) Properties() types.Properties { + return types.NewPropertiesFromStruct(r) +} diff --git a/resources/bedrock-model-customization-jobs.go b/resources/bedrock-model-customization-jobs.go new file mode 100644 index 00000000..8d21ad36 --- /dev/null +++ b/resources/bedrock-model-customization-jobs.go @@ -0,0 +1,105 @@ +package resources + +import ( + "context" + + "fmt" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrock" + + "github.com/ekristen/libnuke/pkg/registry" + "github.com/ekristen/libnuke/pkg/resource" + "github.com/ekristen/libnuke/pkg/types" + + "github.com/ekristen/aws-nuke/v3/pkg/nuke" +) + +const BedrockModelCustomizationJobResource = "BedrockModelCustomizationJob" + +func init() { + registry.Register(®istry.Registration{ + Name: BedrockModelCustomizationJobResource, + Scope: nuke.Account, + Lister: &BedrockModelCustomizationJobLister{}, + }) +} + +type BedrockModelCustomizationJobLister struct{} + +func (l *BedrockModelCustomizationJobLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) { + opts := o.(*nuke.ListerOpts) + + svc := bedrock.New(opts.Session) + resources := make([]resource.Resource, 0) + + params := &bedrock.ListModelCustomizationJobsInput{ + MaxResults: aws.Int64(30), + } + + for { + resp, err := svc.ListModelCustomizationJobs(params) + if err != nil { + return nil, err + } + + for _, modelCustomizationJobSummary := range resp.ModelCustomizationJobSummaries { + tagResp, err := svc.ListTagsForResource( + &bedrock.ListTagsForResourceInput{ + ResourceARN: modelCustomizationJobSummary.JobArn, + }) + if err != nil { + return nil, err + } + resources = append(resources, &BedrockModelCustomizationJob{ + svc: svc, + Arn: modelCustomizationJobSummary.JobArn, + JobName: modelCustomizationJobSummary.JobName, + ModelName: modelCustomizationJobSummary.CustomModelName, + Status: modelCustomizationJobSummary.Status, + Tags: tagResp.Tags, + }) + } + + if resp.NextToken == nil { + break + } + + params.NextToken = resp.NextToken + } + + return resources, nil +} + +type BedrockModelCustomizationJob struct { + svc *bedrock.Bedrock + Arn *string + ModelName *string + JobName *string + Status *string + Tags []*bedrock.Tag +} + +func (r *BedrockModelCustomizationJob) Remove(_ context.Context) error { + _, err := r.svc.StopModelCustomizationJob(&bedrock.StopModelCustomizationJobInput{ + JobIdentifier: r.Arn, + }) + + return err +} + +func (r *BedrockModelCustomizationJob) String() string { + return *r.JobName +} + +func (r *BedrockModelCustomizationJob) Properties() types.Properties { + return types.NewPropertiesFromStruct(r) +} + +func (r *BedrockModelCustomizationJob) Filter() error { + if *r.Status != bedrock.ModelCustomizationJobStatusInProgress { + // May be completed, failed, stopping or stopped + return fmt.Errorf("already stopped") + } + return nil +} diff --git a/resources/bedrock-model-invocation-logging-configuration.go b/resources/bedrock-model-invocation-logging-configuration.go new file mode 100644 index 00000000..10ef6e28 --- /dev/null +++ b/resources/bedrock-model-invocation-logging-configuration.go @@ -0,0 +1,65 @@ +package resources + +import ( + "context" + + "github.com/aws/aws-sdk-go/service/bedrock" + + "github.com/ekristen/libnuke/pkg/registry" + "github.com/ekristen/libnuke/pkg/resource" + "github.com/ekristen/libnuke/pkg/types" + + "github.com/ekristen/aws-nuke/v3/pkg/awsutil" + "github.com/ekristen/aws-nuke/v3/pkg/nuke" +) + +const BedrockModelInvocationLoggingConfigurationResource = "BedrockModelInvocationLoggingConfiguration" + +func init() { + registry.Register(®istry.Registration{ + Name: BedrockModelInvocationLoggingConfigurationResource, + Scope: nuke.Account, + Lister: &BedrockModelInvocationLoggingConfigurationLister{}, + }) +} + +type BedrockModelInvocationLoggingConfigurationLister struct{} + +func (l *BedrockModelInvocationLoggingConfigurationLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) { + opts := o.(*nuke.ListerOpts) + + svc := bedrock.New(opts.Session) + resources := make([]resource.Resource, 0) + + // There is nothing to "list" because there is only one logging configuration per account and deleting it require no params + resp, err := svc.GetModelInvocationLoggingConfiguration(&bedrock.GetModelInvocationLoggingConfigurationInput{}) + if err != nil { + return nil, err + } + + if resp != nil && resp.LoggingConfig != nil { + resources = append(resources, &BedrockModelInvocationLoggingConfiguration{ + svc: svc, + }) + } + + return resources, nil +} + +type BedrockModelInvocationLoggingConfiguration struct { + svc *bedrock.Bedrock +} + +func (r *BedrockModelInvocationLoggingConfiguration) Remove(_ context.Context) error { + _, err := r.svc.DeleteModelInvocationLoggingConfiguration(&bedrock.DeleteModelInvocationLoggingConfigurationInput{}) + + return err +} + +func (r *BedrockModelInvocationLoggingConfiguration) String() string { + return awsutil.Default +} + +func (r *BedrockModelInvocationLoggingConfiguration) Properties() types.Properties { + return types.NewProperties() +} diff --git a/resources/bedrock-provisioned-model-throughputs.go b/resources/bedrock-provisioned-model-throughputs.go new file mode 100644 index 00000000..58b66efd --- /dev/null +++ b/resources/bedrock-provisioned-model-throughputs.go @@ -0,0 +1,94 @@ +package resources + +import ( + "context" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrock" + + "github.com/ekristen/libnuke/pkg/registry" + "github.com/ekristen/libnuke/pkg/resource" + "github.com/ekristen/libnuke/pkg/types" + + "github.com/ekristen/aws-nuke/v3/pkg/nuke" +) + +const BedrockProvisionedModelThroughputResource = "BedrockProvisionedModelThroughput" + +func init() { + registry.Register(®istry.Registration{ + Name: BedrockProvisionedModelThroughputResource, + Scope: nuke.Account, + Lister: &BedrockProvisionedModelThroughputLister{}, + }) +} + +type BedrockProvisionedModelThroughputLister struct{} + +func (l *BedrockProvisionedModelThroughputLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) { + opts := o.(*nuke.ListerOpts) + + svc := bedrock.New(opts.Session) + resources := make([]resource.Resource, 0) + + params := &bedrock.ListProvisionedModelThroughputsInput{ + MaxResults: aws.Int64(30), + } + + for { + resp, err := svc.ListProvisionedModelThroughputs(params) + if err != nil { + return nil, err + } + + for _, provisionedModelSummary := range resp.ProvisionedModelSummaries { + tagResp, err := svc.ListTagsForResource( + &bedrock.ListTagsForResourceInput{ + ResourceARN: provisionedModelSummary.ProvisionedModelArn, + }) + if err != nil { + return nil, err + } + + resources = append(resources, &BedrockProvisionedModelThroughput{ + svc: svc, + Arn: provisionedModelSummary.ProvisionedModelArn, + Name: provisionedModelSummary.ProvisionedModelName, + Status: provisionedModelSummary.Status, + Tags: tagResp.Tags, + }) + } + + if resp.NextToken == nil { + break + } + + params.NextToken = resp.NextToken + } + + return resources, nil +} + +type BedrockProvisionedModelThroughput struct { + svc *bedrock.Bedrock + Arn *string + Name *string + Status *string + Tags []*bedrock.Tag +} + +func (r *BedrockProvisionedModelThroughput) Remove(_ context.Context) error { + _, err := r.svc.DeleteProvisionedModelThroughput(&bedrock.DeleteProvisionedModelThroughputInput{ + ProvisionedModelId: r.Arn, + }) + + return err +} + +func (r *BedrockProvisionedModelThroughput) String() string { + return *r.Name +} + +func (r *BedrockProvisionedModelThroughput) Properties() types.Properties { + return types.NewPropertiesFromStruct(r) +}