Merge pull request aws#314 from yangaws/pyspark

Use get_image_uri for pyspark_mnist_customer_estimator notebook

Author: djarpin

sagemaker-spark/pyspark_mnist/pyspark_mnist_custom_estimator.ipynb

@@ -138,26 +138,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
     "from sagemaker_pyspark import SageMakerEstimator\n",
     "from sagemaker_pyspark.transformation.deserializers import KMeansProtobufResponseRowDeserializer\n",
     "from sagemaker_pyspark.transformation.serializers import ProtobufRequestRowSerializer\n",
     "from sagemaker_pyspark import IAMRole\n",
     "from sagemaker_pyspark import RandomNamePolicyFactory\n",
     "from sagemaker_pyspark import EndpointCreationPolicy\n",
     "\n",
-    "# Use the Docker Registry paths for the K-Means containers from each region\n",
-    "kmeans_ecr_region_path = {\n",
-    "    'us-west-2': '174872318107.dkr.ecr.us-west-2.amazonaws.com',\n",
-    "    'us-east-1': '382416733822.dkr.ecr.us-east-1.amazonaws.com',\n",
-    "    'us-east-2': '404615174143.dkr.ecr.us-east-2.amazonaws.com',\n",
-    "    'eu-west-1': '438346466558.dkr.ecr.eu-west-1.amazonaws.com'\n",
-    "    'ap-northeast-1': '351501993468.dkr.ecr.ap-northeast-1.amazonaws.com',\n",
-    "    'ap-northeast-2': '835164637446.dkr.ecr.ap-northeast-2.amazonaws.com'}\n",
-    "\n",
     "# Create an Estimator from scratch\n",
     "estimator = SageMakerEstimator(\n",
-    "    trainingImage = kmeans_ecr_region_path[region] + \"/kmeans:1\", # Training image \n",
-    "    modelImage = kmeans_ecr_region_path[region] + \"/kmeans:1\", # Model image\n",
+    "    trainingImage = get_image_uri(region, 'kmeans'), # Training image \n",
+    "    modelImage = get_image_uri(region, 'kmeans'), # Model image\n",
     "    requestRowSerializer = ProtobufRequestRowSerializer(),\n",
     "    responseRowDeserializer = KMeansProtobufResponseRowDeserializer(),\n",
     "    hyperParameters = {\"k\": \"10\", \"feature_dim\": \"784\"}, # Set parameters for K-Means\n",