Merge branch 'master' into get-execution-role-fix

Author: andremoeller

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 =====
 
 * bug-fix: get_execution_role no longer fails if user can't call get_role
+* bug-fix: Session: use existing model instead of failing during ``create_model()``
 
 1.7.0
 =====

src/sagemaker/session.py

@@ -439,15 +439,24 @@ def create_model(self, name, role, primary_container):
         role = self.expand_role(role)
         primary_container = _expand_container_def(primary_container)
         LOGGER.info('Creating model with name: {}'.format(name))
-        LOGGER.debug("create_model request: {}".format({
+        LOGGER.debug('create_model request: {}'.format({
             'name': name,
             'role': role,
             'primary_container': primary_container
         }))
 
-        self.sagemaker_client.create_model(ModelName=name,
-                                           PrimaryContainer=primary_container,
-                                           ExecutionRoleArn=role)
+        try:
+            self.sagemaker_client.create_model(ModelName=name,
+                                               PrimaryContainer=primary_container,
+                                               ExecutionRoleArn=role)
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            message = e.response['Error']['Message']
+
+            if error_code == 'ValidationException' and 'Cannot create already existing model' in message:
+                LOGGER.warning('Using already existing model: {}'.format(name))
+            else:
+                raise
 
         return name
 

tests/unit/test_session.py

@@ -12,16 +12,17 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pytest
+import datetime
 import io
+import logging
+
+import pytest
 import six
+from botocore.exceptions import ClientError
 from mock import Mock, patch, call
+
 import sagemaker
 from sagemaker import s3_input, Session, get_execution_role
-import datetime
-
-from botocore.exceptions import ClientError
-
 from sagemaker.session import _tuning_job_status, _transform_job_status
 
 REGION = 'us-west-2'
@@ -515,18 +516,57 @@ def test_logs_for_job_full_lifecycle(time, cw, sagemaker_session_full_lifecycle)
                                    call(0, 'hi there #2a'), call(0, 'hi there #3')]
 
 
+MODEL_NAME = 'some-model'
+PRIMARY_CONTAINER = {
+    'Environment': {},
+    'Image': IMAGE,
+    'ModelDataUrl': 's3://sagemaker-123/output/jobname/model/model.tar.gz',
+}
+
+
+@patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER)
+def test_create_model(expand_container_def, sagemaker_session):
+    model = sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER)
+
+    assert model == MODEL_NAME
+    sagemaker_session.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE,
+                                                                       ModelName=MODEL_NAME,
+                                                                       PrimaryContainer=PRIMARY_CONTAINER)
+
+
+@patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER)
+def test_create_model_already_exists(expand_container_def, sagemaker_session, caplog):
+    error_response = {'Error': {'Code': 'ValidationException', 'Message': 'Cannot create already existing model'}}
+    exception = ClientError(error_response, 'Operation')
+    sagemaker_session.sagemaker_client.create_model.side_effect = exception
+
+    model = sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER)
+    assert model == MODEL_NAME
+
+    expected_warning = ('sagemaker', logging.WARNING, 'Using already existing model: {}'.format(MODEL_NAME))
+    assert expected_warning in caplog.record_tuples
+
+
+@patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER)
+def test_create_model_failure(expand_container_def, sagemaker_session):
+    error_message = 'this is expected'
+    sagemaker_session.sagemaker_client.create_model.side_effect = RuntimeError(error_message)
+
+    with pytest.raises(RuntimeError) as e:
+        sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER)
+
+    assert error_message in str(e)
+
+
 def test_create_model_from_job(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
     ims.create_model_from_job(JOB_NAME)
 
-    assert call(TrainingJobName='jobname') in ims.sagemaker_client.describe_training_job.call_args_list
-    ims.sagemaker_client.create_model.assert_called_with(
-        ExecutionRoleArn='arn:aws:iam::111111111111:role/ExpandedRole',
-        ModelName='jobname',
-        PrimaryContainer={
-            'Environment': {}, 'ModelDataUrl': 's3://sagemaker-123/output/jobname/model/model.tar.gz',
-            'Image': 'myimage'})
+    assert call(TrainingJobName=JOB_NAME) in ims.sagemaker_client.describe_training_job.call_args_list
+    ims.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE,
+                                                         ModelName=JOB_NAME,
+                                                         PrimaryContainer=PRIMARY_CONTAINER)
 
 
 def test_create_model_from_job_with_image(sagemaker_session):
@@ -605,7 +645,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
         Tags=tags)
 
 
-def test_wait_for_tuning_job(sagemaker_session):
+@patch('time.sleep')
+def test_wait_for_tuning_job(sleep, sagemaker_session):
     hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}
     sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
         name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc)
@@ -634,15 +675,17 @@ def test_tune_job_status_none(sagemaker_session):
     assert result is None
 
 
-def test_wait_for_transform_job_completed(sagemaker_session):
+@patch('time.sleep')
+def test_wait_for_transform_job_completed(sleep, sagemaker_session):
     transform_job_desc = {'TransformJobStatus': 'Completed'}
     sagemaker_session.sagemaker_client.describe_transform_job = Mock(
         name='describe_transform_job', return_value=transform_job_desc)
 
     assert sagemaker_session.wait_for_transform_job(JOB_NAME)['TransformJobStatus'] == 'Completed'
 
 
-def test_wait_for_transform_job_in_progress(sagemaker_session):
+@patch('time.sleep')
+def test_wait_for_transform_job_in_progress(sleep, sagemaker_session):
     transform_job_desc_in_progress = {'TransformJobStatus': 'InProgress'}
     transform_job_desc_in_completed = {'TransformJobStatus': 'Completed'}
     sagemaker_session.sagemaker_client.describe_transform_job = Mock(

tests/unit/test_tf_estimator.py

@@ -312,7 +312,8 @@ def test_run_tensorboard_locally_without_awscli_binary(time, strftime, popen, ca
 @patch('subprocess.Popen')
 @patch('time.strftime', return_value=TIMESTAMP)
 @patch('time.time', return_value=TIME)
-def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session):
+@patch('time.sleep')
+def test_run_tensorboard_locally(sleep, time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session):
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)
 
@@ -322,8 +323,7 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk
 
     popen.assert_called_with(['tensorboard', '--logdir', '/my/temp/folder', '--host', 'localhost', '--port', '6006'],
                              stderr=-1,
-                             stdout=-1
-                             )
+                             stdout=-1)
 
 
 @patch('sagemaker.tensorflow.estimator.Tensorboard._sync_directories')
@@ -335,7 +335,8 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk
 @patch('subprocess.Popen')
 @patch('time.strftime', return_value=TIMESTAMP)
 @patch('time.time', return_value=TIME)
-def test_run_tensorboard_locally_port_in_use(time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync,
+@patch('time.sleep')
+def test_run_tensorboard_locally_port_in_use(sleep, time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync,
                                              sagemaker_session):
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)