Merge pull request #16135 from hashicorp/t-gov-sagemaker-useds

tests/resource/sagemaker_endpoint: Make region agnostic
hashicorp · Dec 8, 2020 · e51f827 · e51f827
2 parents b30a0f4 + 4d98411
commit e51f827
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 10 deletions.
diff --git a/aws/data_source_aws_sagemaker_prebuilt_ecr_image.go b/aws/data_source_aws_sagemaker_prebuilt_ecr_image.go
@@ -47,6 +47,10 @@ const (
 	sageMakerRepositoryScikitLearn = "sagemaker-scikit-learn"
 	// SageMaker Library Spark ML
 	sageMakerRepositorySparkML = "sagemaker-sparkml-serving"
+	// SageMaker Library TensorFlow Serving
+	sageMakerRepositoryTensorFlowServing = "sagemaker-tensorflow-serving"
+	// SageMaker Library TensorFlow Serving EIA
+	sageMakerRepositoryTensorFlowServingEIA = "sagemaker-tensorflow-serving-eia"
 	// SageMaker Repo MXNet Inference
 	sageMakerRepositoryMXNetInference = "mxnet-inference"
 	// SageMaker Repo MXNet Inference EIA
@@ -204,6 +208,7 @@ var sageMakerPrebuiltECRImageIDByRegion_SparkML = map[string]string{
 }
 
 // https://github.com/aws/deep-learning-containers/blob/master/available_images.md
+// https://github.com/aws/sagemaker-tensorflow-serving-container
 var sageMakerPrebuiltECRImageIDByRegion_DeepLearning = map[string]string{
 	endpoints.ApEast1RegionID:      "871362719292",
 	endpoints.ApNortheast1RegionID: "763104351884",
@@ -223,10 +228,35 @@ var sageMakerPrebuiltECRImageIDByRegion_DeepLearning = map[string]string{
 	endpoints.SaEast1RegionID:      "763104351884",
 	endpoints.UsEast1RegionID:      "763104351884",
 	endpoints.UsEast2RegionID:      "763104351884",
+	endpoints.UsIsoEast1RegionID:   "886529160074",
 	endpoints.UsWest1RegionID:      "763104351884",
 	endpoints.UsWest2RegionID:      "763104351884",
 }
 
+// https://github.com/aws/sagemaker-tensorflow-serving-container
+var sageMakerPrebuiltECRImageIDByRegion_TensorFlowServing = map[string]string{
+	endpoints.ApEast1RegionID:      "057415533634",
+	endpoints.ApNortheast1RegionID: "520713654638",
+	endpoints.ApNortheast2RegionID: "520713654638",
+	endpoints.ApSouth1RegionID:     "520713654638",
+	endpoints.ApSoutheast1RegionID: "520713654638",
+	endpoints.ApSoutheast2RegionID: "520713654638",
+	endpoints.CaCentral1RegionID:   "520713654638",
+	endpoints.CnNorth1RegionID:     "520713654638",
+	endpoints.CnNorthwest1RegionID: "520713654638",
+	endpoints.EuCentral1RegionID:   "520713654638",
+	endpoints.EuNorth1RegionID:     "520713654638",
+	endpoints.EuWest1RegionID:      "520713654638",
+	endpoints.EuWest2RegionID:      "520713654638",
+	endpoints.EuWest3RegionID:      "520713654638",
+	endpoints.MeSouth1RegionID:     "724002660598",
+	endpoints.SaEast1RegionID:      "520713654638",
+	endpoints.UsEast1RegionID:      "520713654638",
+	endpoints.UsEast2RegionID:      "520713654638",
+	endpoints.UsWest1RegionID:      "520713654638",
+	endpoints.UsWest2RegionID:      "520713654638",
+}
+
 func dataSourceAwsSageMakerPrebuiltECRImage() *schema.Resource {
 	return &schema.Resource{
 		Read: dataSourceAwsSageMakerPrebuiltECRImageRead,
@@ -244,25 +274,27 @@ func dataSourceAwsSageMakerPrebuiltECRImage() *schema.Resource {
 					sageMakerRepositoryKNearestNeighbor,
 					sageMakerRepositoryLDA,
 					sageMakerRepositoryLinearLearner,
+					sageMakerRepositoryMXNetInference,
+					sageMakerRepositoryMXNetInferenceEIA,
+					sageMakerRepositoryMXNetTraining,
 					sageMakerRepositoryNeuralTopicModel,
 					sageMakerRepositoryObject2Vec,
 					sageMakerRepositoryObjectDetection,
 					sageMakerRepositoryPCA,
+					sageMakerRepositoryPyTorchInference,
+					sageMakerRepositoryPyTorchInferenceEIA,
+					sageMakerRepositoryPyTorchTraining,
 					sageMakerRepositoryRandomCutForest,
+					sageMakerRepositoryScikitLearn,
 					sageMakerRepositorySemanticSegmentation,
 					sageMakerRepositorySeq2Seq,
-					sageMakerRepositoryXGBoost,
-					sageMakerRepositoryScikitLearn,
 					sageMakerRepositorySparkML,
-					sageMakerRepositoryMXNetInference,
-					sageMakerRepositoryMXNetInferenceEIA,
-					sageMakerRepositoryMXNetTraining,
-					sageMakerRepositoryPyTorchInference,
-					sageMakerRepositoryPyTorchInferenceEIA,
-					sageMakerRepositoryPyTorchTraining,
 					sageMakerRepositoryTensorFlowInference,
 					sageMakerRepositoryTensorFlowInferenceEIA,
+					sageMakerRepositoryTensorFlowServing,
+					sageMakerRepositoryTensorFlowServingEIA,
 					sageMakerRepositoryTensorFlowTraining,
+					sageMakerRepositoryXGBoost,
 				}, false),
 			},
 
@@ -324,6 +356,8 @@ func dataSourceAwsSageMakerPrebuiltECRImageRead(d *schema.ResourceData, meta int
 		id = sageMakerPrebuiltECRImageIDByRegion_XGBoost[region]
 	case sageMakerRepositoryScikitLearn, sageMakerRepositorySparkML:
 		id = sageMakerPrebuiltECRImageIDByRegion_SparkML[region]
+	case sageMakerRepositoryTensorFlowServing, sageMakerRepositoryTensorFlowServingEIA:
+		id = sageMakerPrebuiltECRImageIDByRegion_TensorFlowServing[region]
 	case sageMakerRepositoryMXNetInference,
 		sageMakerRepositoryMXNetInferenceEIA,
 		sageMakerRepositoryMXNetTraining,

diff --git a/aws/resource_aws_sagemaker_endpoint_test.go b/aws/resource_aws_sagemaker_endpoint_test.go
@@ -225,13 +225,15 @@ data "aws_iam_policy_document" "access" {
   }
 }
 
+data "aws_partition" "current" {}
+
 data "aws_iam_policy_document" "assume_role" {
   statement {
     actions = ["sts:AssumeRole"]
 
     principals {
       type        = "Service"
-      identifiers = ["sagemaker.amazonaws.com"]
+      identifiers = ["sagemaker.${data.aws_partition.current.dns_suffix}"]
     }
   }
 }
@@ -258,12 +260,17 @@ resource "aws_s3_bucket_object" "test" {
   source = "test-fixtures/sagemaker-tensorflow-serving-test-model.tar.gz"
 }
 
+data "aws_sagemaker_prebuilt_ecr_image" "test" {
+  repository_name = "sagemaker-tensorflow-serving"
+  image_tag       = "1.12-cpu"
+}
+
 resource "aws_sagemaker_model" "test" {
   name               = %[1]q
   execution_role_arn = aws_iam_role.test.arn
 
   primary_container {
-    image          = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow-serving:1.12-cpu"
+    image          = data.aws_sagemaker_prebuilt_ecr_image.test.registry_path
     model_data_url = "https://${aws_s3_bucket.test.bucket_regional_domain_name}/${aws_s3_bucket_object.test.key}"
   }