From 96a175e2eff54bca405b71852f2bb9f791d6aec5 Mon Sep 17 00:00:00 2001
From: Mike Yan <yiyanz@amazon.com>
Date: Tue, 19 Dec 2023 15:42:24 +0800
Subject: [PATCH] fix: update instance type for workshop

---
 src/api-deployment/ai-solution-kit-workshop-stack.ts           | 3 +++
 .../features/sagemaker/advanced-ocr-sagemaker.ts               | 2 +-
 src/api-deployment/features/sagemaker/custom-ocr-sagemaker.ts  | 2 +-
 src/api-deployment/features/sagemaker/general-nlu-sagemaker.ts | 2 +-
 4 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/api-deployment/ai-solution-kit-workshop-stack.ts b/src/api-deployment/ai-solution-kit-workshop-stack.ts
index 1ec004a0..bf261d6f 100644
--- a/src/api-deployment/ai-solution-kit-workshop-stack.ts
+++ b/src/api-deployment/ai-solution-kit-workshop-stack.ts
@@ -171,6 +171,7 @@ export class AISolutionKitStack extends Stack {
         ecrDeployment: ecrDeployment,
         updateCustomResourceProvider: updateCustomResourceProvider,
         ecrRegistry: props?.ecrRegistry,
+        instanceType: 'ml.g5.xlarge',
       });
       (advancedOCRSageMaker.nestedStackResource as CfnStack).cfnOptions.condition = cfnTemplate.getCondition('ConditionAdvancedOCRSageMaker');
       this.addOutput(cfnTemplate, api.restApiId, 'advanced-ocr-ml', 'Advanced OCR - Multilingual', 'ConditionAdvancedOCRSageMaker');
@@ -184,6 +185,7 @@ export class AISolutionKitStack extends Stack {
         ecrDeployment: ecrDeployment,
         updateCustomResourceProvider: updateCustomResourceProvider,
         ecrRegistry: props.ecrRegistry,
+        instanceType: 'ml.g5.xlarge',
       });
       (customOCRSageMakerFeatureNestedStack.nestedStackResource as CfnStack).cfnOptions.condition = cfnTemplate.getCondition('ConditionCustomOCRSageMaker');
       this.addOutput(cfnTemplate, api.restApiId, 'custom-ocr-ml', 'Custom OCR SageMaker', 'ConditionCustomOCRSageMaker');
@@ -211,6 +213,7 @@ export class AISolutionKitStack extends Stack {
         ecrDeployment: ecrDeployment,
         updateCustomResourceProvider: updateCustomResourceProvider,
         ecrRegistry: props.ecrRegistry,
+        instanceType: 'ml.g5.xlarge',
       });
       (generalNLUSageMaker.nestedStackResource as CfnStack).cfnOptions.condition = cfnTemplate.getCondition('ConditionGeneralNLUSageMaker');
       this.addOutput(cfnTemplate, api.restApiId, 'general-nlu-ml', 'General NLU SageMaker', 'ConditionGeneralNLUSageMaker');
diff --git a/src/api-deployment/features/sagemaker/advanced-ocr-sagemaker.ts b/src/api-deployment/features/sagemaker/advanced-ocr-sagemaker.ts
index 5f1466f3..5349316b 100644
--- a/src/api-deployment/features/sagemaker/advanced-ocr-sagemaker.ts
+++ b/src/api-deployment/features/sagemaker/advanced-ocr-sagemaker.ts
@@ -19,7 +19,7 @@ export class AdvancedOCRSageMakerFeatureNestedStack extends FeatureNestedStack {
       featureName: `${featureName}`,
       featureCategory: 'media',
       updateCustomResourceProvider: props.updateCustomResourceProvider,
-      sageMakerInstanceType: 'ml.g4dn.xlarge',
+      sageMakerInstanceType: props.instanceType == null ? 'ml.g4dn.xlarge' : props.instanceType,
       sageMakerEcrDeployment: props.ecrDeployment,
     });
   }
diff --git a/src/api-deployment/features/sagemaker/custom-ocr-sagemaker.ts b/src/api-deployment/features/sagemaker/custom-ocr-sagemaker.ts
index 18a6691c..3613c653 100644
--- a/src/api-deployment/features/sagemaker/custom-ocr-sagemaker.ts
+++ b/src/api-deployment/features/sagemaker/custom-ocr-sagemaker.ts
@@ -19,7 +19,7 @@ export class CustomOCRSageMakerFeatureNestedStack extends FeatureNestedStack {
       featureName: `${featureName}`,
       featureCategory: 'media',
       updateCustomResourceProvider: props.updateCustomResourceProvider,
-      sageMakerInstanceType: 'ml.g4dn.xlarge',
+      sageMakerInstanceType: props.instanceType == null ? 'ml.g4dn.xlarge' : props.instanceType,
       sageMakerEcrDeployment: props.ecrDeployment,
     });
   }
diff --git a/src/api-deployment/features/sagemaker/general-nlu-sagemaker.ts b/src/api-deployment/features/sagemaker/general-nlu-sagemaker.ts
index 7a26f8df..88c56dee 100644
--- a/src/api-deployment/features/sagemaker/general-nlu-sagemaker.ts
+++ b/src/api-deployment/features/sagemaker/general-nlu-sagemaker.ts
@@ -19,7 +19,7 @@ export class GeneralNLUSageMakerFeatureNestedStack extends FeatureNestedStack {
       featureName: `${featureName}`,
       featureCategory: 'media',
       updateCustomResourceProvider: props.updateCustomResourceProvider,
-      sageMakerInstanceType: 'ml.g4dn.xlarge',
+      sageMakerInstanceType: props.instanceType == null ? 'ml.g4dn.xlarge' : props.instanceType,
       sageMakerEcrDeployment: props.ecrDeployment,
     });
   }