feat(Sagemaker): allow deploying llama 3.1 from Sagemaker Jumpstart (#…

…692) * feat(Sagemaker): add Jumpstart support for new regions and add gatedBucket into consideration. --------- Signed-off-by: arief hidayat <[email protected]> Co-authored-by: Alain Krok <[email protected]>
awslabs · Sep 19, 2024 · 6ab07ac · 6ab07ac
1 parent 9e91372
commit 6ab07ac
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 3 deletions.
diff --git a/apidocs/interfaces/IJumpStartModelSpec.md b/apidocs/interfaces/IJumpStartModelSpec.md
@@ -30,6 +30,12 @@
 
 ***
 
+### gatedBucket
+
+> **gatedBucket**: `boolean`
+
+***
+
 ### instanceAliases?
 
 > `optional` **instanceAliases**: [`IInstanceAliase`](IInstanceAliase.md)[]

diff --git a/...tterns/gen-ai/aws-model-deployment-sagemaker/code-generation/generate-jumpstart-models.ts b/...tterns/gen-ai/aws-model-deployment-sagemaker/code-generation/generate-jumpstart-models.ts
@@ -34,6 +34,7 @@ interface JumpStartModelSpec {
   hosting_artifact_key?: string;
   hosting_script_key?: string;
   hosting_prepacked_artifact_key?: string;
+  gated_bucket: boolean;
   hosting_eula_key?: string;
   inference_environment_variables: {
     name: string;
@@ -113,6 +114,7 @@ export async function download_data() {
         hosting_script_key,
         hosting_artifact_key,
         hosting_prepacked_artifact_key,
+        gated_bucket,
         inference_environment_variables,
         hosting_instance_type_variants,
         hosting_eula_key,
@@ -149,6 +151,7 @@ export async function download_data() {
         hosting_artifact_key,
         hosting_script_key,
         hosting_prepacked_artifact_key,
+        gated_bucket,
         inference_environment_variables,
         hosting_instance_type_variants,
         hosting_eula_key,
@@ -220,6 +223,7 @@ function generateCode() {
         instanceTypes: specSource.supported_inference_instance_types,
         modelPackageArns: specSource.hosting_model_package_arns,
         prepackedArtifactKey: specSource.hosting_prepacked_artifact_key,
+        gatedBucket: specSource.gated_bucket,
         artifactKey: specSource.hosting_artifact_key,
         environment,
         instanceAliases: instanceAliasesArr,
@@ -270,6 +274,7 @@ export interface IJumpStartModelSpec {
   instanceTypes: string[];
   modelPackageArns?: { [region: string]: string };
   prepackedArtifactKey?: string;
+  gatedBucket: boolean;
   artifactKey?: string;
   environment: { [key: string]: string | number | boolean };
   instanceAliases?: IInstanceAliase[];

diff --git a/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-model.ts b/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-model.ts
@@ -31,6 +31,7 @@ export interface IJumpStartModelSpec {
   instanceTypes: string[];
   modelPackageArns?: { [region: string]: string };
   prepackedArtifactKey?: string;
+  gatedBucket: boolean;
   artifactKey?: string;
   environment: { [key: string]: string | number | boolean };
   instanceAliases?: IInstanceAliase[];

diff --git a/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-models.json b/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-models.json
diff --git a/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-sagemaker-endpoint.ts b/src/patterns/gen-ai/aws-model-deployment-sagemaker/jumpstart-sagemaker-endpoint.ts
@@ -183,6 +183,10 @@ export class JumpStartSageMakerEndpoint extends SageMakerEndpointBase {
       ...this.environment,
     };
 
+    if (environment.SAGEMAKER_SUBMIT_DIRECTORY) {
+      delete environment.SAGEMAKER_SUBMIT_DIRECTORY;
+    }
+
     return environment;
   }
 
@@ -195,7 +199,8 @@ export class JumpStartSageMakerEndpoint extends SageMakerEndpointBase {
     vpcConfig: sagemaker.CfnModel.VpcConfigProperty | undefined,
   ) {
     const key = this.spec.prepackedArtifactKey ?? this.spec.artifactKey;
-    const bucket = JumpStartConstants.JUMPSTART_LAUNCHED_REGIONS[this.region]?.contentBucket;
+    const bucket = this.spec.gatedBucket ? JumpStartConstants.JUMPSTART_LAUNCHED_REGIONS[this.region]?.gatedContentBucket :
+      JumpStartConstants.JUMPSTART_LAUNCHED_REGIONS[this.region]?.contentBucket;
     if (!bucket) {
       throw new Error(`JumpStart is not available in the region ${this.region}.`);
     }
@@ -224,7 +229,7 @@ export class JumpStartSageMakerEndpoint extends SageMakerEndpointBase {
       executionRoleArn: this.role.roleArn,
       enableNetworkIsolation: true,
       primaryContainer: isArtifactCompressed ? {
-      // True: Artifact is a tarball
+        // True: Artifact is a tarball
         image,
         modelDataUrl: modelArtifactUrl,
         environment,
@@ -252,6 +257,18 @@ export class JumpStartSageMakerEndpoint extends SageMakerEndpointBase {
           key: 'modelVersion',
           value: this.spec.version,
         },
+        {
+          key: 'sagemaker-studio:jumpstart-model-id',
+          value: this.spec.modelId,
+        },
+        {
+          key: 'sagemaker-studio:jumpstart-model-version',
+          value: this.spec.version,
+        },
+        {
+          key: 'sagemaker-studio:jumpstart-hub-name',
+          value: 'SageMakerPublicHub',
+        },
       ],
       vpcConfig: vpcConfig,
     });

diff --git a/src/patterns/gen-ai/aws-model-deployment-sagemaker/private/jumpstart-constants.ts b/src/patterns/gen-ai/aws-model-deployment-sagemaker/private/jumpstart-constants.ts
@@ -44,6 +44,10 @@ export abstract class JumpStartConstants {
       contentBucket: 'jumpstart-cache-prod-eu-north-1',
       gatedContentBucket: 'jumpstart-private-cache-prod-eu-north-1',
     },
+    'me-central-1': {
+      contentBucket: 'jumpstart-cache-prod-me-central-1',
+      gatedContentBucket: 'jumpstart-private-cache-prod-me-central-1',
+    },
     'me-south-1': {
       contentBucket: 'jumpstart-cache-prod-me-south-1',
       gatedContentBucket: 'jumpstart-private-cache-prod-me-south-1',
@@ -72,6 +76,10 @@ export abstract class JumpStartConstants {
       contentBucket: 'jumpstart-cache-prod-ap-northeast-2',
       gatedContentBucket: 'jumpstart-private-cache-prod-ap-northeast-2',
     },
+    'ap-northeast-3': {
+      contentBucket: 'jumpstart-cache-prod-ap-northeast-3',
+      gatedContentBucket: 'jumpstart-private-cache-prod-ap-northeast-3',
+    },
     'eu-west-2': {
       contentBucket: 'jumpstart-cache-prod-eu-west-2',
       gatedContentBucket: 'jumpstart-private-cache-prod-eu-west-2',
@@ -96,13 +104,29 @@ export abstract class JumpStartConstants {
       contentBucket: 'jumpstart-cache-prod-ap-southeast-2',
       gatedContentBucket: 'jumpstart-private-cache-prod-ap-southeast-2',
     },
+    'ap-southeast-3': {
+      contentBucket: 'jumpstart-cache-prod-ap-southeast-3',
+      gatedContentBucket: 'jumpstart-private-cache-prod-ap-southeast-3',
+    },
     'ca-central-1': {
       contentBucket: 'jumpstart-cache-prod-ca-central-1',
       gatedContentBucket: 'jumpstart-private-cache-prod-ca-central-1',
     },
     'cn-north-1': {
       contentBucket: 'jumpstart-cache-prod-cn-north-1',
     },
+    'il-central-1': {
+      contentBucket: 'jumpstart-cache-prod-il-central-1',
+      gatedContentBucket: 'jumpstart-private-cache-prod-il-central-1',
+    },
+    'us-gov-east-1': {
+      contentBucket: 'jumpstart-cache-prod-us-gov-east-1',
+      gatedContentBucket: 'jumpstart-private-cache-prod-us-gov-east-1',
+    },
+    'us-gov-west-1': {
+      contentBucket: 'jumpstart-cache-prod-us-gov-west-1',
+      gatedContentBucket: 'jumpstart-private-cache-prod-us-gov-west-1',
+    },
   };
 
   public static JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = 'models_manifest.json';