fix: job wait logic to wait for tag propagation

Author: mufaddal-rohawala

src/sagemaker/session.py

@@ -42,6 +42,7 @@
     secondary_training_status_changed,
     secondary_training_status_message,
     sts_regional_endpoint,
+    retries,
 )
 from sagemaker import exceptions
 from sagemaker.session_settings import SessionSettings
@@ -5812,19 +5813,56 @@ def _deploy_done(sagemaker_client, endpoint_name):
 
 def _wait_until_training_done(callable_fn, desc, poll=5):
     """Placeholder docstring"""
-    job_desc, finished = callable_fn(desc)
+    elapsed_time = 0
+    finished = None
+    job_desc = desc
     while not finished:
-        time.sleep(poll)
-        job_desc, finished = callable_fn(job_desc)
+        try:
+            elapsed_time += poll
+            time.sleep(poll)
+            job_desc, finished = callable_fn(job_desc)
+        except botocore.exceptions.ClientError as err:
+            # For initial 5 mins we accept/pass AccessDeniedException.
+            # The reason is to await tag propagation to avoid false AccessDenied claims for an access
+            # policy based on resource tags, The caveat here is for true AccessDenied cases the routine
+            # will fail after 5 mins
+            if err.response["Error"]["Code"] == "AccessDeniedException" and elapsed_time <= 300:
+                LOGGER.warning(
+                    "Received AccessDeniedException. This could mean the IAM role does not "
+                    "have the resource permissions, in which case please add resource access "
+                    "and retry. For cases where the role has tag based resource policy, "
+                    "continuing to wait for tag propagation.."
+                )
+                continue
+            else:
+                raise err
     return job_desc
 
 
 def _wait_until(callable_fn, poll=5):
     """Placeholder docstring"""
-    result = callable_fn()
+    elapsed_time = 0
+    result = None
     while result is None:
-        time.sleep(poll)
-        result = callable_fn()
+        try:
+            elapsed_time += poll
+            time.sleep(poll)
+            result = callable_fn()
+        except botocore.exceptions.ClientError as err:
+            # For initial 5 mins we accept/pass AccessDeniedException.
+            # The reason is to await tag propagation to avoid false AccessDenied claims for an
+            # access policy based on resource tags, The caveat here is for true AccessDenied
+            # cases the routine will fail after 5 mins
+            if err.response["Error"]["Code"] == "AccessDeniedException" and elapsed_time <= 300:
+                LOGGER.warning(
+                    "Received AccessDeniedException. This could mean the IAM role does not "
+                    "have the resource permissions, in which case please add resource access "
+                    "and retry. For cases where the role has tag based resource policy, "
+                    "continuing to wait for tag propagation.."
+                )
+                continue
+            else:
+                raise err
     return result
 
 

src/sagemaker/workflow/pipeline.py

@@ -246,10 +246,7 @@ def upsert(
         except ClientError as ce:
             error_code = ce.response["Error"]["Code"]
             error_message = ce.response["Error"]["Message"]
-            if not (
-                    error_code == "ValidationException"
-                    and "already exists" in error_message
-            ):
+            if not (error_code == "ValidationException" and "already exists" in error_message):
                 raise ce
             # already exists
             response = self.update(role_arn, description)

tests/unit/common.py

@@ -0,0 +1,38 @@
+# Copyright 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
+
+
+from botocore.exceptions import ClientError
+
+
+def _raise_unexpected_client_error(**kwargs):
+    response = {
+        "Error": {"Code": "ValidationException", "Message": "Name does not satisfy expression."}
+    }
+    raise ClientError(error_response=response, operation_name="foo")
+
+
+def _raise_does_not_exist_client_error(**kwargs):
+    response = {"Error": {"Code": "ValidationException", "Message": "Could not find entity."}}
+    raise ClientError(error_response=response, operation_name="foo")
+
+
+def _raise_does_already_exists_client_error(**kwargs):
+    response = {"Error": {"Code": "ValidationException", "Message": "Resource already exists."}}
+    raise ClientError(error_response=response, operation_name="foo")
+
+
+def _raise_access_denied_client_error(**kwargs):
+    response = {"Error": {"Code": "AccessDeniedException", "Message": "Could not access entity."}}
+    raise ClientError(error_response=response, operation_name="foo")

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -36,7 +36,6 @@
 from botocore.exceptions import ClientError
 
 
-
 @pytest.fixture
 def role_arn():
     return "arn:role"
@@ -207,8 +206,10 @@ def _raise_does_already_exists_client_error(**kwargs):
     pipeline.upsert(role_arn=role_arn, tags=tags)
 
     sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
-        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
-        Tags=tags
+        PipelineName="MyPipeline",
+        RoleArn=role_arn,
+        PipelineDefinition=pipeline.definition(),
+        Tags=tags,
     )
 
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_once_with(
@@ -223,6 +224,7 @@ def _raise_does_already_exists_client_error(**kwargs):
         ResourceArn="mock_pipeline_arn", Tags=tags
     )
 
+
 def test_pipeline_upsert_create_unexpected_failure(sagemaker_session_mock, role_arn):
 
     # case 2: unexpected failure on create
@@ -258,16 +260,17 @@ def _raise_unexpected_client_error(**kwargs):
     with pytest.raises(ClientError):
         pipeline.upsert(role_arn=role_arn, tags=tags)
 
-
-
     sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
-        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
-        Tags=tags
+        PipelineName="MyPipeline",
+        RoleArn=role_arn,
+        PipelineDefinition=pipeline.definition(),
+        Tags=tags,
     )
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_not_called()
     sagemaker_session_mock.sagemaker_client.list_tags.assert_not_called()
     sagemaker_session_mock.sagemaker_client.add_tags.assert_not_called()
 
+
 def test_pipeline_upsert_resourse_doesnt_exist(sagemaker_session_mock, role_arn):
 
     # case 3: resource does not exist
@@ -295,11 +298,13 @@ def test_pipeline_upsert_resourse_doesnt_exist(sagemaker_session_mock, role_arn)
     try:
         pipeline.upsert(role_arn=role_arn, tags=tags)
     except ClientError:
-        assert False, f"Unexpected ClientError raised"
+        assert False, "Unexpected ClientError raised"
 
     sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
-        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
-        Tags=tags
+        PipelineName="MyPipeline",
+        RoleArn=role_arn,
+        PipelineDefinition=pipeline.definition(),
+        Tags=tags,
     )
 
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_not_called()

tests/unit/test_endpoint_from_model_data.py

@@ -17,6 +17,11 @@
 from mock import MagicMock, Mock
 from mock import patch
 
+from .common import (
+    _raise_unexpected_client_error,
+    _raise_does_already_exists_client_error,
+    _raise_does_not_exist_client_error,
+)
 import sagemaker
 
 ENDPOINT_NAME = "myendpoint"
@@ -126,7 +131,7 @@ def test_model_and_endpoint_config_exist(name_from_image_mock, sagemaker_session
             wait=False,
         )
     except ClientError:
-        assert False, f"Unexpected ClientError raised for resource already exists scenario"
+        assert False, "Unexpected ClientError raised for resource already exists scenario"
 
     sagemaker_session.create_model.assert_called_once_with(
         name=NAME_FROM_IMAGE,
@@ -208,20 +213,3 @@ def test_entity_doesnt_exist():
 def test_describe_failure():
     with pytest.raises(ClientError):
         sagemaker.session._deployment_entity_exists(_raise_unexpected_client_error)
-
-
-def _raise_unexpected_client_error(**kwargs):
-    response = {
-        "Error": {"Code": "ValidationException", "Message": "Name does not satisfy expression."}
-    }
-    raise ClientError(error_response=response, operation_name="foo")
-
-
-def _raise_does_not_exist_client_error(**kwargs):
-    response = {"Error": {"Code": "ValidationException", "Message": "Could not find entity."}}
-    raise ClientError(error_response=response, operation_name="foo")
-
-
-def _raise_does_already_exists_client_error(**kwargs):
-    response = {"Error": {"Code": "ValidationException", "Message": "Resource already exists."}}
-    raise ClientError(error_response=response, operation_name="foo")

tests/unit/test_session.py

@@ -22,13 +22,16 @@
 from botocore.exceptions import ClientError
 from mock import ANY, MagicMock, Mock, patch, call, mock_open
 
+from .common import _raise_unexpected_client_error
 import sagemaker
 from sagemaker import TrainingInput, Session, get_execution_role, exceptions
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.session import (
     _tuning_job_status,
     _transform_job_status,
     _train_done,
+    _wait_until,
+    _wait_until_training_done,
     NOTEBOOK_METADATA_FILE,
 )
 from sagemaker.tuner import WarmStartConfig, WarmStartTypes
@@ -2342,6 +2345,81 @@ def test_train_done_in_progress(sagemaker_session):
     assert training_finished is False
 
 
+@patch("time.sleep", return_value=None)
+def test_wait_until_training_done_raises_other_exception(patched_sleep):
+    response = {"Error": {"Code": "ValidationException", "Message": "Could not access entity."}}
+    mock_func = Mock(
+        name="describe_training_job",
+        side_effect=ClientError(error_response=response, operation_name="foo"),
+    )
+    desc = "dummy"
+    with pytest.raises(ClientError) as error:
+        _wait_until_training_done(mock_func, desc)
+
+    mock_func.assert_called_once()
+    assert "ValidationException" in str(error)
+
+
+@patch("time.sleep", return_value=None)
+def test_wait_until_training_done_tag_propagation(patched_sleep):
+    response = {"Error": {"Code": "AccessDeniedException", "Message": "Could not access entity."}}
+    side_effect_iter = [ClientError(error_response=response, operation_name="foo")] * 3
+    side_effect_iter.append(("result", "result"))
+    mock_func = Mock(name="describe_training_job", side_effect=side_effect_iter)
+    desc = "dummy"
+    result = _wait_until_training_done(mock_func, desc)
+    assert result == "result"
+    assert mock_func.call_count == 4
+
+
+@patch("time.sleep", return_value=None)
+def test_wait_until_training_done_fail_access_denied_after_5_mins(patched_sleep):
+    response = {"Error": {"Code": "AccessDeniedException", "Message": "Could not access entity."}}
+    side_effect_iter = [ClientError(error_response=response, operation_name="foo")] * 70
+    mock_func = Mock(name="describe_training_job", side_effect=side_effect_iter)
+    desc = "dummy"
+    with pytest.raises(ClientError) as error:
+        _wait_until_training_done(mock_func, desc)
+
+    # mock_func should be retried 300(elapsed time)/5(default poll delay) = 60 times
+    assert mock_func.call_count == 61
+    assert "AccessDeniedException" in str(error)
+
+
+@patch("time.sleep", return_value=None)
+def test_wait_until_raises_other_exception(patched_sleep):
+    mock_func = Mock(name="describe_training_job", side_effect=_raise_unexpected_client_error)
+    with pytest.raises(ClientError) as error:
+        _wait_until(mock_func)
+
+    mock_func.assert_called_once()
+    assert "ValidationException" in str(error)
+
+
+@patch("time.sleep", return_value=None)
+def test_wait_until_tag_propagation(patched_sleep):
+    response = {"Error": {"Code": "AccessDeniedException", "Message": "Could not access entity."}}
+    side_effect_iter = [ClientError(error_response=response, operation_name="foo")] * 3
+    side_effect_iter.append("result")
+    mock_func = Mock(name="describe_training_job", side_effect=side_effect_iter)
+    result = _wait_until(mock_func)
+    assert result == "result"
+    assert mock_func.call_count == 4
+
+
+@patch("time.sleep", return_value=None)
+def test_wait_until_fail_access_denied_after_5_mins(patched_sleep):
+    response = {"Error": {"Code": "AccessDeniedException", "Message": "Could not access entity."}}
+    side_effect_iter = [ClientError(error_response=response, operation_name="foo")] * 70
+    mock_func = Mock(name="describe_training_job", side_effect=side_effect_iter)
+    with pytest.raises(ClientError) as error:
+        _wait_until(mock_func)
+
+    # mock_func should be retried 300(elapsed time)/5(default poll delay) = 60 times
+    assert mock_func.call_count == 61
+    assert "AccessDeniedException" in str(error)
+
+
 DEFAULT_EXPECTED_AUTO_ML_JOB_ARGS = {
     "AutoMLJobName": JOB_NAME,
     "InputDataConfig": [