Skip to content

Commit

Permalink
feat(bedrock): add bedrock custom models, customization jobs, provisi…
Browse files Browse the repository at this point in the history
…oned throughput and logging
  • Loading branch information
npellegrin committed Nov 18, 2024
1 parent 674525d commit 2153da1
Show file tree
Hide file tree
Showing 4 changed files with 353 additions and 0 deletions.
89 changes: 89 additions & 0 deletions resources/bedrock-custom-models.go
Original file line number Diff line number Diff line change
@@ -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(&registry.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)
}
105 changes: 105 additions & 0 deletions resources/bedrock-model-customization-jobs.go
Original file line number Diff line number Diff line change
@@ -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(&registry.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
}
65 changes: 65 additions & 0 deletions resources/bedrock-model-invocation-logging-configuration.go
Original file line number Diff line number Diff line change
@@ -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(&registry.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()
}
94 changes: 94 additions & 0 deletions resources/bedrock-provisioned-model-throughputs.go
Original file line number Diff line number Diff line change
@@ -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(&registry.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)
}

0 comments on commit 2153da1

Please sign in to comment.