Skip to content

feat(stepfunctions-tasks): bedrock createModelCustomizationJob integration #31913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 103 commits into from
May 23, 2025

Conversation

badmintoncryer
Copy link
Contributor

This PR was previously created and passed the community review, but the maintainer review stopped midway, and it was eventually closed. There shouldn’t be any issues with the content, so I am submitting the PR again.

Issue # (if applicable)

Closes #29042

Reason for this change

AWS stepfunctions support optimized integration with AWS bedrock.
Currently, only invokeModel is supported by CDK, but I would like createModelCustomizationJob to be supported in the same manner.

Description of changes

I've added CreatemodelCustomizationJob class.

const taskConfig = {
  baseModel: model,
  clientRequestToken: 'MyToken',
  customizationType: CustomizationType.FINE_TUNING,
  kmsKey,
  customModelName: 'MyCustomModel',
  customModelTags: [{ key: 'key1', value: 'value1' }],
  hyperParameters: {
    batchSize: '10',
  },
  jobName: 'MyCustomizationJob',
  jobTags: [{ key: 'key2', value: 'value2' }],
  outputDataS3Uri: outputBucket.s3UrlForObject(),
  trainingDataS3Uri: trainingBucket.s3UrlForObject(),
  validationDataS3Uri: [validationBucket.s3UrlForObject()],
  vpcConfig: {
    securityGroups: [new ec2.SecurityGroup(stack, 'SecurityGroup', { vpc })],
    subnets: vpc.isolatedSubnets,
  },
};

const task1 = new BedrockCreateModelCustomizationJob(stack, 'CreateModelCustomizationJob1', taskConfig);

const chain = sfn.Chain
  .start(new sfn.Pass(stack, 'Start'))
  .next(task1)
  .next(new sfn.Pass(stack, 'Done'));

new sfn.StateMachine(stack, 'StateMachine', {
  definitionBody: sfn.DefinitionBody.fromChainable(chain),
  timeout: cdk.Duration.seconds(30),
});

Description of how you validated changes

I've added both unit and integ tests.

Checklist


By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license

@aws-cdk-automation aws-cdk-automation added the pr/needs-community-review This PR needs a review from a Trusted Community Member or Core Team Member. label Feb 25, 2025
@badmintoncryer
Copy link
Contributor Author

@gracelu0 I'm really sorry for my late response. I've addressed all of your comments. Could you please reconfirm again?

/**
* The key/value pair for a tag.
*/
export interface ITag {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Avoid using the I prefix for interface names unless the interface is specifically intended to represent behavior. Reference: jsii Specification – Interfaces vs Structs

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Additionally, the interface name is too generic. It would be better to either place it within a relevant namespace or give it a more descriptive name.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've updated to CustomModelTag.

this._role = this.renderBedrockCreateModelCustomizationJobRole();
this.taskPolicies = this.renderPolicyStatements();

if (this.props.customModelKmsKey) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's important to note that addToResourcePolicy has no effect when used on a KMS key imported via fromKeyArn.

Reference: aws-cdk-lib/aws-kms/key.ts#L163.

It's recommended to check the return value of the method to ensure the policy was actually applied.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which is better to throw error or show warnings when customModelKmsKey is imported key? For now, I have implemented it to display warnings.

if (this.props.customModelKmsKey) {
      const poliyStatement = new iam.PolicyStatement({
        actions: ['kms:Decrypt', 'kms:GenerateDataKey', 'kms:DescribeKey', 'kms:CreateGrant'],
        resources: ['*'],
        principals: [new iam.ArnPrincipal(this._role.roleArn)],
        conditions: {
          StringEquals: {
            'kms:ViaService': [
              `bedrock.${Stack.of(this).region}.amazonaws.com`,
            ],
          },
        },
      });
      const result = this.props.customModelKmsKey.addToResourcePolicy(poliyStatement, true);

      // For the imported key, user must add the policy statement to the key policy
      if (result.statementAdded === false) {
        Annotations.of(this).addWarning(`You must update 'customModelKmsKey' resource policy to add the following statement:
        ${poliyStatement.toString()}`);
      }
    }

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for updating the PR. I haven’t worked with this Bedrock job much before, but here’s my take: if missing permissions on the KMS key ultimately cause the job to fail, we should throw an exception to follow the fail-fast principle. On the other hand, if the lack of permission doesn’t block the job from running, then a warning message should suffice.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure. I've updated the implementation to throw an error.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, could you help me resolve the build errors of unused imports?

@badmintoncryer
Copy link
Contributor Author

@5d Thank you for your review!! I've addressed your comments.

@5d 5d self-assigned this May 21, 2025
Copy link
Member

@5d 5d left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks for your effort on this!

Copy link
Contributor

mergify bot commented May 23, 2025

Thank you for contributing! Your pull request will be updated from main and then merged automatically (do not update manually, and be sure to allow changes to be pushed to your fork).

@aws-cdk-automation
Copy link
Collaborator

AWS CodeBuild CI Report

  • CodeBuild project: AutoBuildv2Project1C6BFA3F-wQm2hXv2jqQv
  • Commit ID: fcf5b01
  • Result: SUCCEEDED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

Copy link
Contributor

mergify bot commented May 23, 2025

Thank you for contributing! Your pull request will be updated from main and then merged automatically (do not update manually, and be sure to allow changes to be pushed to your fork).

@mergify mergify bot merged commit 0c653d3 into aws:main May 23, 2025
14 of 16 checks passed
Copy link
Contributor

Comments on closed issues and PRs are hard for our team to see.
If you need help, please open a new issue that references this one.

@github-actions github-actions bot locked as resolved and limited conversation to collaborators May 23, 2025
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
distinguished-contributor [Pilot] contributed 50+ PRs to the CDK effort/medium Medium work item – several days of effort feature-request A feature should be added or improved. needs-security-review Related to feature or issues that needs security review p2 pr/needs-community-review This PR needs a review from a Trusted Community Member or Core Team Member.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

stepfunctions-tasks: support for bedrock createModelCustomizationJob task
6 participants