Skip to content

Commit 91553e5

Browse files
vinayak-kukrejavinayak-kukreja
authored
feat(sagemaker): support dlc images in sagemaker model (#25018)
CDK currently does not support [deep learning container images for AWS Sagemaker models](https://docs.aws.amazon.com/sagemaker/latest/dg/pre-built-containers-frameworks-deep-learning.html). Looks like this functionality was proposed as part of a [PR](https://github.com/aws/aws-cdk/pull/17399/files#diff-356f35099770f68f4ceee2e63d34aad8729b0a9be6c933a0c05e999be7374685R98-R145) but did not make it to the RFC. This PR is adding that functionality and tests for it. ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
1 parent 4a04f92 commit 91553e5

File tree

11 files changed

+751
-87
lines changed

11 files changed

+751
-87
lines changed

‎packages/@aws-cdk/aws-sagemaker/README.md

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,19 @@ const repository = ecr.Repository.fromRepositoryName(this, 'Repository', 'repo')
107107
const image = sagemaker.ContainerImage.fromEcrRepository(repository, 'tag');
108108
```
109109

110+
#### DLC Image
111+
112+
Reference a deep learning container image:
113+
114+
```typescript
115+
import * as sagemaker from '@aws-cdk/aws-sagemaker-alpha';
116+
117+
const repositoryName = 'huggingface-pytorch-training';
118+
const tag = '1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04';
119+
120+
const image = sagemaker.ContainerImage.fromDlc(repositoryName, tag);
121+
```
122+
110123
### Model Artifacts
111124

112125
If you choose to decouple your model artifacts from your inference code (as is natural given

‎packages/@aws-cdk/aws-sagemaker/lib/container-image.ts

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ import * as assets from 'aws-cdk-lib/aws-ecr-assets';
33
import { Construct } from 'constructs';
44
import { Model } from './model';
55
import { hashcode } from './private/util';
6+
import { FactName } from 'aws-cdk-lib/region-info';
7+
import { Stack } from 'aws-cdk-lib/core';
68

79
/**
810
* The configuration for creating a container image.
@@ -38,6 +40,13 @@ export abstract class ContainerImage {
3840
return new AssetImage(directory, options);
3941
}
4042

43+
/**
44+
* Reference an AWS Deep Learning Container image
45+
*/
46+
public static fromDlc(repositoryName: string, tag: string, accountId?: string): ContainerImage {
47+
return new DlcEcrImage(repositoryName, tag, accountId);
48+
}
49+
4150
/**
4251
* Called when the image is used by a Model
4352
*/
@@ -81,3 +90,27 @@ class AssetImage extends ContainerImage {
8190
};
8291
}
8392
}
93+
94+
class DlcEcrImage extends ContainerImage {
95+
96+
constructor(private readonly repositoryName: string, private readonly tag: string, private readonly accountId?: string) {
97+
super();
98+
}
99+
100+
public bind(scope: Construct, model: Model): ContainerImageConfig {
101+
const accountId = this.accountId ?? Stack.of(scope).regionalFact(FactName.DLC_REPOSITORY_ACCOUNT);
102+
103+
const repository = ecr.Repository.fromRepositoryAttributes(scope, 'DlcRepository', {
104+
repositoryName: this.repositoryName,
105+
repositoryArn: ecr.Repository.arnForLocalRepository(
106+
this.repositoryName,
107+
scope,
108+
accountId,
109+
),
110+
});
111+
112+
repository.grantPull(model);
113+
114+
return { imageName: `${repository.repositoryUri}:${this.tag}` };
115+
}
116+
}

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/aws-cdk-sagemaker-model.assets.json

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"version": "21.0.0",
2+
"version": "31.0.0",
33
"files": {
44
"126d48fa0e32fbef5078b9d88658b35ad29d4291eb86675a64c75fa4f1338916": {
55
"source": {
@@ -14,15 +14,15 @@
1414
}
1515
}
1616
},
17-
"eebf956b730b436212d72ff616d7ab38ac8a7c128b71f54c4e9d1f0d29b36604": {
17+
"be0577003620b4cc844de880dc7d01b90719400dea1cc12c03d93a3ac8e451a3": {
1818
"source": {
1919
"path": "aws-cdk-sagemaker-model.template.json",
2020
"packaging": "file"
2121
},
2222
"destinations": {
2323
"current_account-current_region": {
2424
"bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
25-
"objectKey": "eebf956b730b436212d72ff616d7ab38ac8a7c128b71f54c4e9d1f0d29b36604.json",
25+
"objectKey": "be0577003620b4cc844de880dc7d01b90719400dea1cc12c03d93a3ac8e451a3.json",
2626
"assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
2727
}
2828
}

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/aws-cdk-sagemaker-model.template.json

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,198 @@
770770
"InferencePipelineModelRoleDefaultPolicy2AF0CDCF",
771771
"InferencePipelineModelRole6A99C5B3"
772772
]
773+
},
774+
"HuggingFaceModelRoleDA17DA00": {
775+
"Type": "AWS::IAM::Role",
776+
"Properties": {
777+
"AssumeRolePolicyDocument": {
778+
"Statement": [
779+
{
780+
"Action": "sts:AssumeRole",
781+
"Effect": "Allow",
782+
"Principal": {
783+
"Service": "sagemaker.amazonaws.com"
784+
}
785+
}
786+
],
787+
"Version": "2012-10-17"
788+
},
789+
"ManagedPolicyArns": [
790+
{
791+
"Fn::Join": [
792+
"",
793+
[
794+
"arn:",
795+
{
796+
"Ref": "AWS::Partition"
797+
},
798+
":iam::aws:policy/AmazonSageMakerFullAccess"
799+
]
800+
]
801+
}
802+
]
803+
}
804+
},
805+
"HuggingFaceModelRoleDefaultPolicy50587D35": {
806+
"Type": "AWS::IAM::Policy",
807+
"Properties": {
808+
"PolicyDocument": {
809+
"Statement": [
810+
{
811+
"Action": [
812+
"ecr:BatchCheckLayerAvailability",
813+
"ecr:BatchGetImage",
814+
"ecr:GetDownloadUrlForLayer"
815+
],
816+
"Effect": "Allow",
817+
"Resource": {
818+
"Fn::Join": [
819+
"",
820+
[
821+
"arn:",
822+
{
823+
"Ref": "AWS::Partition"
824+
},
825+
":ecr:",
826+
{
827+
"Ref": "AWS::Region"
828+
},
829+
":",
830+
{
831+
"Fn::FindInMap": [
832+
"DlcRepositoryAccountMap",
833+
{
834+
"Ref": "AWS::Region"
835+
},
836+
"value"
837+
]
838+
},
839+
":repository/huggingface-pytorch-training"
840+
]
841+
]
842+
}
843+
},
844+
{
845+
"Action": "ecr:GetAuthorizationToken",
846+
"Effect": "Allow",
847+
"Resource": "*"
848+
}
849+
],
850+
"Version": "2012-10-17"
851+
},
852+
"PolicyName": "HuggingFaceModelRoleDefaultPolicy50587D35",
853+
"Roles": [
854+
{
855+
"Ref": "HuggingFaceModelRoleDA17DA00"
856+
}
857+
]
858+
}
859+
},
860+
"HuggingFaceModel8E13DAA5": {
861+
"Type": "AWS::SageMaker::Model",
862+
"Properties": {
863+
"ExecutionRoleArn": {
864+
"Fn::GetAtt": [
865+
"HuggingFaceModelRoleDA17DA00",
866+
"Arn"
867+
]
868+
},
869+
"PrimaryContainer": {
870+
"Image": {
871+
"Fn::Join": [
872+
"",
873+
[
874+
{
875+
"Fn::FindInMap": [
876+
"DlcRepositoryAccountMap",
877+
{
878+
"Ref": "AWS::Region"
879+
},
880+
"value"
881+
]
882+
},
883+
".dkr.ecr.",
884+
{
885+
"Ref": "AWS::Region"
886+
},
887+
".",
888+
{
889+
"Ref": "AWS::URLSuffix"
890+
},
891+
"/huggingface-pytorch-training:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04"
892+
]
893+
]
894+
}
895+
}
896+
},
897+
"DependsOn": [
898+
"HuggingFaceModelRoleDefaultPolicy50587D35",
899+
"HuggingFaceModelRoleDA17DA00"
900+
]
901+
}
902+
},
903+
"Mappings": {
904+
"DlcRepositoryAccountMap": {
905+
"ap-east-1": {
906+
"value": "871362719292"
907+
},
908+
"ap-northeast-1": {
909+
"value": "763104351884"
910+
},
911+
"ap-northeast-2": {
912+
"value": "763104351884"
913+
},
914+
"ap-south-1": {
915+
"value": "763104351884"
916+
},
917+
"ap-southeast-1": {
918+
"value": "763104351884"
919+
},
920+
"ap-southeast-2": {
921+
"value": "763104351884"
922+
},
923+
"ca-central-1": {
924+
"value": "763104351884"
925+
},
926+
"cn-north-1": {
927+
"value": "727897471807"
928+
},
929+
"cn-northwest-1": {
930+
"value": "727897471807"
931+
},
932+
"eu-central-1": {
933+
"value": "763104351884"
934+
},
935+
"eu-north-1": {
936+
"value": "763104351884"
937+
},
938+
"eu-west-1": {
939+
"value": "763104351884"
940+
},
941+
"eu-west-2": {
942+
"value": "763104351884"
943+
},
944+
"eu-west-3": {
945+
"value": "763104351884"
946+
},
947+
"me-south-1": {
948+
"value": "217643126080"
949+
},
950+
"sa-east-1": {
951+
"value": "763104351884"
952+
},
953+
"us-east-1": {
954+
"value": "763104351884"
955+
},
956+
"us-east-2": {
957+
"value": "763104351884"
958+
},
959+
"us-west-1": {
960+
"value": "763104351884"
961+
},
962+
"us-west-2": {
963+
"value": "763104351884"
964+
}
773965
}
774966
},
775967
"Parameters": {

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/cdk.out

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
{"version":"21.0.0"}
1+
{"version":"31.0.0"}

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/integ.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"version": "21.0.0",
2+
"version": "31.0.0",
33
"testCases": {
44
"integtest-model/DefaultTest": {
55
"stacks": [

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/integtestmodelDefaultTestDeployAssertCF40BD53.assets.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"version": "21.0.0",
2+
"version": "31.0.0",
33
"files": {
44
"21fbb51d7b23f6a6c262b46a9caee79d744a3ac019fd45422d988b96d44b2a22": {
55
"source": {

‎packages/@aws-cdk/aws-sagemaker/test/integ.model.js.snapshot/manifest.json

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"version": "21.0.0",
2+
"version": "31.0.0",
33
"artifacts": {
44
"aws-cdk-sagemaker-model.assets": {
55
"type": "cdk:asset-manifest",
@@ -17,7 +17,7 @@
1717
"validateOnSynth": false,
1818
"assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}",
1919
"cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}",
20-
"stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/eebf956b730b436212d72ff616d7ab38ac8a7c128b71f54c4e9d1f0d29b36604.json",
20+
"stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/be0577003620b4cc844de880dc7d01b90719400dea1cc12c03d93a3ac8e451a3.json",
2121
"requiresBootstrapStackVersion": 6,
2222
"bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version",
2323
"additionalDependencies": [
@@ -213,6 +213,30 @@
213213
"data": "InferencePipelineModel3564F141"
214214
}
215215
],
216+
"/aws-cdk-sagemaker-model/HuggingFaceModel/Role/Resource": [
217+
{
218+
"type": "aws:cdk:logicalId",
219+
"data": "HuggingFaceModelRoleDA17DA00"
220+
}
221+
],
222+
"/aws-cdk-sagemaker-model/HuggingFaceModel/Role/DefaultPolicy/Resource": [
223+
{
224+
"type": "aws:cdk:logicalId",
225+
"data": "HuggingFaceModelRoleDefaultPolicy50587D35"
226+
}
227+
],
228+
"/aws-cdk-sagemaker-model/HuggingFaceModel/Model": [
229+
{
230+
"type": "aws:cdk:logicalId",
231+
"data": "HuggingFaceModel8E13DAA5"
232+
}
233+
],
234+
"/aws-cdk-sagemaker-model/DlcRepositoryAccountMap": [
235+
{
236+
"type": "aws:cdk:logicalId",
237+
"data": "DlcRepositoryAccountMap"
238+
}
239+
],
216240
"/aws-cdk-sagemaker-model/BootstrapVersion": [
217241
{
218242
"type": "aws:cdk:logicalId",
@@ -224,6 +248,33 @@
224248
"type": "aws:cdk:logicalId",
225249
"data": "CheckBootstrapVersion"
226250
}
251+
],
252+
"HuggingFaceModelRole8CE687FE": [
253+
{
254+
"type": "aws:cdk:logicalId",
255+
"data": "HuggingFaceModelRole8CE687FE",
256+
"trace": [
257+
"!!DESTRUCTIVE_CHANGES: WILL_DESTROY"
258+
]
259+
}
260+
],
261+
"HuggingFaceModelRoleDefaultPolicy365BBC22": [
262+
{
263+
"type": "aws:cdk:logicalId",
264+
"data": "HuggingFaceModelRoleDefaultPolicy365BBC22",
265+
"trace": [
266+
"!!DESTRUCTIVE_CHANGES: WILL_DESTROY"
267+
]
268+
}
269+
],
270+
"HuggingFaceModelD53E0910": [
271+
{
272+
"type": "aws:cdk:logicalId",
273+
"data": "HuggingFaceModelD53E0910",
274+
"trace": [
275+
"!!DESTRUCTIVE_CHANGES: WILL_DESTROY"
276+
]
277+
}
227278
]
228279
},
229280
"displayName": "aws-cdk-sagemaker-model"

0 commit comments

Comments
 (0)