Nrt support for existing algorithms in spark sdk (#50)

* Add support for NRT region for existing algorithms in Spark SDK.

* updated version to 1.1.0

Author: yijiezh

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.1.0
+=====
+
+* spark/pyspark: Enable NRT region support for spark SDK
+
 1.0.5
 =====
 

sagemaker-pyspark-sdk/setup.py

@@ -7,7 +7,7 @@
 from setuptools import setup
 
 
-VERSION = "1.0.5"
+VERSION = "1.1.0"
 
 TEMP_PATH = "deps"
 JARS_TARGET = os.path.join(TEMP_PATH, "jars")

sagemaker-spark-sdk/build.sbt

@@ -19,7 +19,7 @@ scalaVersion := "2.11.7"
 // to change the version of spark add -DSPARK_VERSION=2.x.x when running sbt
 // for example: "sbt -DSPARK_VERSION=2.1.1 clean compile test doc package"
 val sparkVersion = System.getProperty("SPARK_VERSION", "2.2.0")
-version := "spark_" + sparkVersion + "-1.0.5"
+version := "spark_" + sparkVersion + "-1.1.0"
 
 lazy val SageMakerSpark = (project in file("."))
 

sagemaker-spark-sdk/src/main/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/ImageURIProvider.scala

@@ -36,23 +36,26 @@ private[algorithms] object SagerMakerRegionAccountMaps {
     Regions.EU_WEST_1.getName -> "438346466558",
     Regions.US_EAST_1.getName -> "382416733822",
     Regions.US_EAST_2.getName -> "404615174143",
-    Regions.US_WEST_2.getName -> "174872318107"
+    Regions.US_WEST_2.getName -> "174872318107",
+    Regions.AP_NORTHEAST_1.getName -> "351501993468"
   )
 
   // For LDA
   val LDAAccountMap: Map[String, String] = Map(
     Regions.EU_WEST_1.getName -> "999678624901",
     Regions.US_EAST_1.getName -> "766337827248",
     Regions.US_EAST_2.getName -> "999911452149",
-    Regions.US_WEST_2.getName -> "266724342769"
+    Regions.US_WEST_2.getName -> "266724342769",
+    Regions.AP_NORTHEAST_1.getName -> "258307448986"
   )
 
   // For XGBoost
   val ApplicationsAccountMap: Map[String, String] = Map(
     Regions.EU_WEST_1.getName -> "685385470294",
     Regions.US_EAST_1.getName -> "811284229777",
     Regions.US_EAST_2.getName -> "825641698319",
-    Regions.US_WEST_2.getName -> "433757028032"
+    Regions.US_WEST_2.getName -> "433757028032",
+    Regions.AP_NORTHEAST_1.getName -> "501404015308"
   )
 }
 

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/FactorizationMachinesSageMakerEstimatorTests.scala

@@ -76,6 +76,11 @@ class FactorizationMachinesSageMakerEstimatorTests extends FlatSpec with Mockito
       createFactorizationMachinesBinaryClassifier(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "174872318107.dkr.ecr.us-west-2.amazonaws.com/factorization-machines:1")
+
+    val estimatorAPNorthEast1 =
+      createFactorizationMachinesBinaryClassifier(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/factorization-machines:1")
   }
 
   it should "use the correct defaults for regressor" in {
@@ -102,6 +107,11 @@ class FactorizationMachinesSageMakerEstimatorTests extends FlatSpec with Mockito
     val estimatorUSWest2 = createFactorizationMachinesRegressor(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "174872318107.dkr.ecr.us-west-2.amazonaws.com/factorization-machines:1")
+
+    val estimatorAPNorthEast1 =
+      createFactorizationMachinesRegressor(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/factorization-machines:1")
   }
 
   it should "setFeatureDim" in {

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/KMeansSageMakerEstimatorTests.scala

@@ -68,6 +68,10 @@ class KMeansSageMakerEstimatorTests extends FlatSpec with Matchers with MockitoS
     val estimatorUSWest2 = createKMeansEstimator(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage
       == "174872318107.dkr.ecr.us-west-2.amazonaws.com/kmeans:1")
+
+    val estimatorAPNorthEast1 = createKMeansEstimator(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage
+      == "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/kmeans:1")
   }
 
   it should "setK" in {

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/LDASageMakerEstimatorTests.scala

@@ -66,6 +66,11 @@ class LDASageMakerEstimatorTests extends FlatSpec with MockitoSugar {
       createLDAEstimator(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "266724342769.dkr.ecr.us-west-2.amazonaws.com/lda:1")
+
+    val estimatorAPNorthEast1 =
+      createLDAEstimator(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "258307448986.dkr.ecr.ap-northeast-1.amazonaws.com/lda:1")
   }
 
   it should "setFeatureDim" in {

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/LinearLearnerSageMakerEstimatorTests.scala

@@ -71,6 +71,11 @@ class LinearLearnerSageMakerEstimatorTests extends FlatSpec with MockitoSugar {
     val estimatorUSWest2 = createLinearLearnerBinaryClassifier(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "174872318107.dkr.ecr.us-west-2.amazonaws.com/linear-learner:1")
+
+    val estimatorAPNorthEast1 =
+      createLinearLearnerBinaryClassifier(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/linear-learner:1")
   }
 
   it should "use the correct defaults for regressor" in {
@@ -97,6 +102,11 @@ class LinearLearnerSageMakerEstimatorTests extends FlatSpec with MockitoSugar {
     val estimatorUSWest2 = createLinearLearnerRegressor(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "174872318107.dkr.ecr.us-west-2.amazonaws.com/linear-learner:1")
+
+    val estimatorAPNorthEast1 =
+      createLinearLearnerRegressor(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/linear-learner:1")
   }
 
   it should "setFeatureDim" in {

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/PCASageMakerEstimatorTests.scala

@@ -48,6 +48,10 @@ class PCASageMakerEstimatorTests extends FlatSpec with MockitoSugar {
 
     val estimatorUSWest2 = createPCAEstimator(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage == "174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1")
+
+    val estimatorAPNorthEast1 = createPCAEstimator(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/pca:1")
   }
 
   it should "use the correct defaults" in {

sagemaker-spark-sdk/src/test/scala/com/amazonaws/services/sagemaker/sparksdk/algorithms/XGBoostSageMakerEstimatorTests.scala

@@ -67,6 +67,10 @@ class XGBoostSageMakerEstimatorTests extends FlatSpec with Matchers with Mockito
     val estimatorUSWest2 = createXGBoostEstimator(region = Regions.US_WEST_2.getName)
     assert(estimatorUSWest2.trainingImage ==
       "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1")
+
+    val estimatorAPNorthEast1 = createXGBoostEstimator(region = Regions.AP_NORTHEAST_1.getName)
+    assert(estimatorAPNorthEast1.trainingImage ==
+      "501404015308.dkr.ecr.ap-northeast-1.amazonaws.com/xgboost:1")
   }
 
   it should "setBooster" in {