Add addition mapping

get_image_uri -> sagemaker.image_uris.retrieve()

Author: brightsparc

tests/integ/test_state_machine_definition.py

@@ -16,7 +16,7 @@
 import json
 
 from sagemaker.utils import unique_name_from_base
-from sagemaker.amazon.amazon_estimator import get_image_uri
+from sagemaker.image_uris import retrieve 
 from stepfunctions import steps
 from stepfunctions.workflow import Workflow
 from tests.integ.utils import state_machine_delete_wait
@@ -25,7 +25,7 @@
 def training_job_parameters(sagemaker_session, sagemaker_role_arn, record_set_fixture):
     parameters = { 
         "AlgorithmSpecification": { 
-            "TrainingImage": get_image_uri(sagemaker_session.boto_session.region_name, 'pca'),
+            "TrainingImage": retrieve(region=sagemaker_session.boto_session.region_name, framework='pca'),
             "TrainingInputMode": "File"
         },
         "OutputDataConfig": { 

tests/integ/test_training_pipeline_estimators.py

@@ -25,7 +25,7 @@
 
 # import Sagemaker
 from sagemaker.amazon.pca import PCA
-from sagemaker.amazon.amazon_estimator import get_image_uri
+from sagemaker.image_uris import retrieve 
 
 # import StepFunctions
 from stepfunctions.template.pipeline import TrainingPipeline
@@ -105,7 +105,7 @@ def test_pca_estimator(sfn_client, sagemaker_session, sagemaker_role_arn, sfn_ro
         job_name = workflow_execution_info['name']
         s3_manifest_uri = inputs.s3_data
         status = 'SUCCEEDED'
-        estimator_image_uri = get_image_uri(sagemaker_session.boto_region_name, 'pca')
+        estimator_image_uri = retrieve(region=sagemaker_session.boto_region_name,  framework='pca')
 
         execution_info = sfn_client.describe_execution(executionArn=execution_arn)
         execution_info['input'] = json.loads(execution_info['input'])