change: Replaced generic ValueError with custom subclass when reporting unexpected resource status

Author: arodiss

src/sagemaker/exceptions.py

@@ -0,0 +1,21 @@
+# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+
+class UnexpectedStatusException(ValueError):
+    """Raised when resource status is not expected and thus not allowed for further execution"""
+    def __init__(self, message, allowed_statuses, actual_status):
+        self.allowed_statuses = allowed_statuses
+        self.actual_status = actual_status
+        super(UnexpectedStatusException, self).__init__(message)

src/sagemaker/session.py

@@ -34,6 +34,7 @@
     secondary_training_status_changed,
     secondary_training_status_message,
 )
+from sagemaker import exceptions
 
 LOGGER = logging.getLogger("sagemaker")
 
@@ -206,7 +207,7 @@ def default_bucket(self):
                     Bucket=default_bucket, CreateBucketConfiguration={"LocationConstraint": region}
                 )
 
-            LOGGER.info("Created S3 bucket: %s", default_bucket)
+            LOGGER.info("Created S3 bucket: {}".format(default_bucket))
         except ClientError as e:
             error_code = e.response["Error"]["Code"]
             message = e.response["Error"]["Message"]
@@ -343,8 +344,8 @@ def train(  # noqa: C901
         if encrypt_inter_container_traffic:
             train_request["EnableInterContainerTrafficEncryption"] = encrypt_inter_container_traffic
 
-        LOGGER.info("Creating training-job with name: %s", job_name)
-        LOGGER.debug("train request: %s", json.dumps(train_request, indent=4))
+        LOGGER.info("Creating training-job with name: {}".format(job_name))
+        LOGGER.debug("train request: {}".format(json.dumps(train_request, indent=4)))
         self.sagemaker_client.create_training_job(**train_request)
 
     def compile_model(
@@ -379,7 +380,7 @@ def compile_model(
         if tags is not None:
             compilation_job_request["Tags"] = tags
 
-        LOGGER.info("Creating compilation-job with name: %s", job_name)
+        LOGGER.info("Creating compilation-job with name: {}".format(job_name))
         self.sagemaker_client.create_compilation_job(**compilation_job_request)
 
     def tune(
@@ -521,8 +522,8 @@ def tune(
         if encrypt_inter_container_traffic:
             tune_request["TrainingJobDefinition"]["EnableInterContainerTrafficEncryption"] = True
 
-        LOGGER.info("Creating hyperparameter tuning job with name: %s", job_name)
-        LOGGER.debug("tune request: %s", json.dumps(tune_request, indent=4))
+        LOGGER.info("Creating hyperparameter tuning job with name: {}".format(job_name))
+        LOGGER.debug("tune request: {}".format(json.dumps(tune_request, indent=4)))
         self.sagemaker_client.create_hyper_parameter_tuning_job(**tune_request)
 
     def stop_tuning_job(self, name):
@@ -535,17 +536,18 @@ def stop_tuning_job(self, name):
             ClientError: If an error occurs while trying to stop the hyperparameter tuning job.
         """
         try:
-            LOGGER.info("Stopping tuning job: %s", name)
+            LOGGER.info("Stopping tuning job: {}".format(name))
             self.sagemaker_client.stop_hyper_parameter_tuning_job(HyperParameterTuningJobName=name)
         except ClientError as e:
             error_code = e.response["Error"]["Code"]
             # allow to pass if the job already stopped
             if error_code == "ValidationException":
-                LOGGER.info("Tuning job: %s is already stopped or not running.", name)
+                LOGGER.info("Tuning job: {} is already stopped or not running.".format(name))
             else:
                 LOGGER.error(
-                    "Error occurred while attempting to stop tuning job: %s. Please try again.",
-                    name,
+                    "Error occurred while attempting to stop tuning job: {}. Please try again.".format(
+                        name
+                    )
                 )
                 raise
 
@@ -607,8 +609,8 @@ def transform(
         if data_processing is not None:
             transform_request["DataProcessing"] = data_processing
 
-        LOGGER.info("Creating transform job with name: %s", job_name)
-        LOGGER.debug("Transform request: %s", json.dumps(transform_request, indent=4))
+        LOGGER.info("Creating transform job with name: {}".format(job_name))
+        LOGGER.debug("Transform request: {}".format(json.dumps(transform_request, indent=4)))
         self.sagemaker_client.create_transform_job(**transform_request)
 
     def create_model(
@@ -680,8 +682,8 @@ def create_model(
         if enable_network_isolation:
             create_model_request["EnableNetworkIsolation"] = True
 
-        LOGGER.info("Creating model with name: %s", name)
-        LOGGER.debug("CreateModel request: %s", json.dumps(create_model_request, indent=4))
+        LOGGER.info("Creating model with name: {}".format(name))
+        LOGGER.debug("CreateModel request: {}".format(json.dumps(create_model_request, indent=4)))
 
         try:
             self.sagemaker_client.create_model(**create_model_request)
@@ -693,7 +695,7 @@ def create_model(
                 error_code == "ValidationException"
                 and "Cannot create already existing model" in message
             ):
-                LOGGER.warning("Using already existing model: %s", name)
+                LOGGER.warning("Using already existing model: {}".format(name))
             else:
                 raise
 
@@ -764,14 +766,14 @@ def create_model_package_from_algorithm(self, name, description, algorithm_arn,
             },
         }
         try:
-            LOGGER.info("Creating model package with name: %s", name)
+            LOGGER.info("Creating model package with name: {}".format(name))
             self.sagemaker_client.create_model_package(**request)
         except ClientError as e:
             error_code = e.response["Error"]["Code"]
             message = e.response["Error"]["Message"]
 
             if error_code == "ValidationException" and "ModelPackage already exists" in message:
-                LOGGER.warning("Using already existing model package: %s", name)
+                LOGGER.warning("Using already existing model package: {}".format(name))
             else:
                 raise
 
@@ -792,10 +794,10 @@ def wait_for_model_package(self, model_package_name, poll=5):
 
         if status != "Completed":
             reason = desc.get("FailureReason", None)
-            raise ValueError(
-                "Error creating model package {}: {} Reason: {}".format(
-                    model_package_name, status, reason
-                )
+            raise exceptions.UnexpectedStatusException(
+                message="Error creating model package {}: {} Reason: {}".format(model_package_name, status, reason),
+                allowed_statuses=["Completed"],
+                actual_status=status
             )
         return desc
 
@@ -832,7 +834,7 @@ def create_endpoint_config(
         Returns:
             str: Name of the endpoint point configuration created.
         """
-        LOGGER.info("Creating endpoint-config with name %s", name)
+        LOGGER.info("Creating endpoint-config with name {}".format(name))
 
         tags = tags or []
 
@@ -871,7 +873,7 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
         Returns:
             str: Name of the Amazon SageMaker ``Endpoint`` created.
         """
-        LOGGER.info("Creating endpoint with name %s", endpoint_name)
+        LOGGER.info("Creating endpoint with name {}".format(endpoint_name))
 
         tags = tags or []
 
@@ -914,7 +916,7 @@ def delete_endpoint(self, endpoint_name):
         Args:
             endpoint_name (str): Name of the Amazon SageMaker ``Endpoint`` to delete.
         """
-        LOGGER.info("Deleting endpoint with name: %s", endpoint_name)
+        LOGGER.info("Deleting endpoint with name: {}".format(endpoint_name))
         self.sagemaker_client.delete_endpoint(EndpointName=endpoint_name)
 
     def delete_endpoint_config(self, endpoint_config_name):
@@ -923,7 +925,7 @@ def delete_endpoint_config(self, endpoint_config_name):
         Args:
             endpoint_config_name (str): Name of the Amazon SageMaker endpoint configuration to delete.
         """
-        LOGGER.info("Deleting endpoint configuration with name: %s", endpoint_config_name)
+        LOGGER.info("Deleting endpoint configuration with name: {}".format(endpoint_config_name))
         self.sagemaker_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
 
     def delete_model(self, model_name):
@@ -933,7 +935,7 @@ def delete_model(self, model_name):
             model_name (str): Name of the Amazon SageMaker model to delete.
 
         """
-        LOGGER.info("Deleting model with name: %s", model_name)
+        LOGGER.info("Deleting model with name: {}".format(model_name))
         self.sagemaker_client.delete_model(ModelName=model_name)
 
     def wait_for_job(self, job, poll=5):
@@ -947,7 +949,7 @@ def wait_for_job(self, job, poll=5):
             (dict): Return value from the ``DescribeTrainingJob`` API.
 
         Raises:
-            ValueError: If the training job fails.
+            exceptions.UnexpectedStatusException: If the training job fails.
         """
         desc = _wait_until_training_done(
             lambda last_desc: _train_done(self.sagemaker_client, job, last_desc), None, poll
@@ -966,7 +968,7 @@ def wait_for_compilation_job(self, job, poll=5):
             (dict): Return value from the ``DescribeCompilationJob`` API.
 
         Raises:
-            ValueError: If the compilation job fails.
+            exceptions.UnexpectedStatusException: If the compilation job fails.
         """
         desc = _wait_until(lambda: _compilation_job_status(self.sagemaker_client, job), poll)
         self._check_job_status(job, desc, "CompilationJobStatus")
@@ -983,7 +985,7 @@ def wait_for_tuning_job(self, job, poll=5):
             (dict): Return value from the ``DescribeHyperParameterTuningJob`` API.
 
         Raises:
-            ValueError: If the hyperparameter tuning job fails.
+            exceptions.UnexpectedStatusException: If the hyperparameter tuning job fails.
         """
         desc = _wait_until(lambda: _tuning_job_status(self.sagemaker_client, job), poll)
         self._check_job_status(job, desc, "HyperParameterTuningJobStatus")
@@ -1000,23 +1002,23 @@ def wait_for_transform_job(self, job, poll=5):
             (dict): Return value from the ``DescribeTransformJob`` API.
 
         Raises:
-            ValueError: If the transform job fails.
+            exceptions.UnexpectedStatusException: If the transform job fails.
         """
         desc = _wait_until(lambda: _transform_job_status(self.sagemaker_client, job), poll)
         self._check_job_status(job, desc, "TransformJobStatus")
         return desc
 
     def _check_job_status(self, job, desc, status_key_name):
         """Check to see if the job completed successfully and, if not, construct and
-        raise a ValueError.
+        raise a exceptions.UnexpectedStatusException.
 
         Args:
             job (str): The name of the job to check.
             desc (dict[str, str]): The result of ``describe_training_job()``.
             status_key_name (str): Status key name to check for.
 
         Raises:
-            ValueError: If the training job fails.
+            exceptions.UnexpectedStatusException: If the training job fails.
         """
         status = desc[status_key_name]
         # If the status is capital case, then convert it to Camel case
@@ -1025,7 +1027,11 @@ def _check_job_status(self, job, desc, status_key_name):
         if status != "Completed" and status != "Stopped":
             reason = desc.get("FailureReason", "(No reason provided)")
             job_type = status_key_name.replace("JobStatus", " job")
-            raise ValueError("Error for {} {}: {} Reason: {}".format(job_type, job, status, reason))
+            raise exceptions.UnexpectedStatusException(
+                message="Error for {} {}: {} Reason: {}".format(job_type, job, status, reason),
+                allowed_statuses=["Completed", "Stopped"],
+                actual_status=status
+            )
 
     def wait_for_endpoint(self, endpoint, poll=5):
         """Wait for an Amazon SageMaker endpoint deployment to complete.
@@ -1042,8 +1048,10 @@ def wait_for_endpoint(self, endpoint, poll=5):
 
         if status != "InService":
             reason = desc.get("FailureReason", None)
-            raise ValueError(
-                "Error hosting endpoint {}: {} Reason: {}".format(endpoint, status, reason)
+            raise exceptions.UnexpectedStatusException(
+                message="Error hosting endpoint {}: {} Reason: {}".format(endpoint, status, reason),
+                allowed_statuses=["InService"],
+                actual_status=status
             )
         return desc
 
@@ -1257,8 +1265,9 @@ def get_caller_identity_arn(self):
             role = self.boto_session.client("iam").get_role(RoleName=role_name)["Role"]["Arn"]
         except ClientError:
             LOGGER.warning(
-                "Couldn't call 'get_role' to get Role ARN from role name %s to get Role path.",
-                role_name,
+                "Couldn't call 'get_role' to get Role ARN from role name {} to get Role path.".format(
+                    role_name
+                )
             )
 
         return role
@@ -1276,7 +1285,7 @@ def logs_for_job(  # noqa: C901 - suppress complexity warning for this method
             poll (int): The interval in seconds between polling for new log entries and job completion (default: 5).
 
         Raises:
-            ValueError: If waiting and the training job fails.
+            exceptions.UnexpectedStatusException: If waiting and the training job fails.
         """
 
         description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)

tests/unit/test_exception_on_bad_status.py

@@ -0,0 +1,93 @@
+# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import pytest
+from mock import Mock, MagicMock
+import sagemaker
+
+EXPANDED_ROLE = 'arn:aws:iam::111111111111:role/ExpandedRole'
+REGION = 'us-west-2'
+MODEL_PACKAGE_NAME = 'my_model_package'
+JOB_NAME = 'my_job_name'
+ENDPOINT_NAME = 'the_point_of_end'
+
+
+def get_sagemaker_session(returns_status):
+    boto_mock = Mock(name='boto_session', region_name=REGION)
+    client_mock = Mock()
+    client_mock.describe_model_package = MagicMock(return_value={'ModelPackageStatus': returns_status})
+    client_mock.describe_endpoint = MagicMock(return_value={'EndpointStatus': returns_status})
+    ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=client_mock)
+    ims.expand_role = Mock(return_value=EXPANDED_ROLE)
+    return ims
+
+
+def test_does_not_raise_when_successfully_created_package():
+    try:
+        sagemaker_session = get_sagemaker_session(returns_status='Completed')
+        sagemaker_session.wait_for_model_package(MODEL_PACKAGE_NAME)
+    except sagemaker.exceptions.UnexpectedStatusException:
+        pytest.fail("UnexpectedStatusException was thrown while it should not")
+
+
+def test_raise_when_failed_created_package():
+    try:
+        sagemaker_session = get_sagemaker_session(returns_status='EnRoute')
+        sagemaker_session.wait_for_model_package(MODEL_PACKAGE_NAME)
+        assert False, 'sagemaker.exceptions.UnexpectedStatusException should have been raised but was not'
+    except Exception as e:
+        assert type(e) == sagemaker.exceptions.UnexpectedStatusException
+        assert e.actual_status == 'EnRoute'
+        assert 'Completed' in e.allowed_statuses
+
+
+def test_does_not_raise_when_correct_job_status():
+    try:
+        job = Mock()
+        sagemaker_session = get_sagemaker_session(returns_status='Stopped')
+        sagemaker_session._check_job_status(job, {'TransformationJobStatus': 'Stopped'}, 'TransformationJobStatus')
+    except sagemaker.exceptions.UnexpectedStatusException:
+        pytest.fail("UnexpectedStatusException was thrown while it should not")
+
+
+def test_does_raise_when_incorrect_job_status():
+    try:
+        job = Mock()
+        sagemaker_session = get_sagemaker_session(returns_status='Failed')
+        sagemaker_session._check_job_status(job, {'TransformationJobStatus': 'Failed'}, 'TransformationJobStatus')
+        assert False, 'sagemaker.exceptions.UnexpectedStatusException should have been raised but was not'
+    except Exception as e:
+        assert type(e) == sagemaker.exceptions.UnexpectedStatusException
+        assert e.actual_status == 'Failed'
+        assert 'Completed' in e.allowed_statuses
+        assert 'Stopped' in e.allowed_statuses
+
+
+def test_does_not_raise_when_successfully_deployed_endpoint():
+    try:
+        sagemaker_session = get_sagemaker_session(returns_status='InService')
+        sagemaker_session.wait_for_endpoint(ENDPOINT_NAME)
+    except sagemaker.exceptions.UnexpectedStatusException:
+        pytest.fail("UnexpectedStatusException was thrown while it should not")
+
+
+def test_raise_when_failed_to_deploy_endpoint():
+    try:
+        sagemaker_session = get_sagemaker_session(returns_status='Failed')
+        assert sagemaker_session.wait_for_endpoint(ENDPOINT_NAME)
+        assert False, 'sagemaker.exceptions.UnexpectedStatusException should have been raised but was not'
+    except Exception as e:
+        assert type(e) == sagemaker.exceptions.UnexpectedStatusException
+        assert e.actual_status == 'Failed'
+        assert 'InService' in e.allowed_statuses