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'])