test(integ): absorb iam:SimulatePrincipalPolicy throttling across integ suites

sagemaker-mlops/tests/integ/conftest.py

@@ -58,6 +58,84 @@
 )
 
 
+# ---------------------------------------------------------------------------
+# IAM SimulatePrincipalPolicy throttling mitigation
+#
+# These tests run under ``pytest -n auto`` (dozens of xdist workers). Many of
+# them resolve/validate an IAM execution role via ``resolve_and_validate_role``
+# (e.g. during ``Pipeline`` create/upsert and feature-processor scheduling),
+# which internally calls the low-TPS ``iam:SimulatePrincipalPolicy`` API. With
+# many workers hitting it at once IAM throttles the request, surfacing as
+# ``ClientError: (Throttling) ... Rate exceeded``. This mitigation is a purely
+# test-harness concurrency fix and is intentionally identical to the block in
+# the sagemaker-train / sagemaker-serve integ conftests.
+# ---------------------------------------------------------------------------
+
+# botocore adaptive retry settings for throttling-prone IAM validation calls.
+# Applied via env vars so every client in the worker inherits them, regardless of
+# which boto session the SDK ends up using to build its IAM client.
+_RETRY_MODE = "adaptive"
+_MAX_ATTEMPTS = "10"
+
+# Only tolerate throttling on the IAM permission simulation used during role
+# validation, so unrelated throttling still fails loudly.
+_THROTTLE_ERROR_CODES = ("Throttling", "ThrottlingException", "RequestLimitExceeded")
+_SIMULATE_OP = "SimulatePrincipalPolicy"
+
+
+@pytest.fixture(autouse=True, scope="session")
+def _configure_boto_adaptive_retries():
+    """Give every boto3 client in this xdist worker adaptive retries so the IAM
+    clients built by the role resolver absorb transient SimulatePrincipalPolicy
+    throttling. Restores any pre-existing values on teardown."""
+    previous = {
+        "AWS_RETRY_MODE": os.environ.get("AWS_RETRY_MODE"),
+        "AWS_MAX_ATTEMPTS": os.environ.get("AWS_MAX_ATTEMPTS"),
+    }
+    os.environ["AWS_RETRY_MODE"] = _RETRY_MODE
+    os.environ["AWS_MAX_ATTEMPTS"] = _MAX_ATTEMPTS
+    yield
+    for key, value in previous.items():
+        if value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = value
+
+
+def _is_simulate_policy_throttle(exc):
+    """Return True if ``exc`` is a SimulatePrincipalPolicy throttling ClientError."""
+    from botocore.exceptions import ClientError
+
+    if not isinstance(exc, ClientError):
+        return False
+    error_code = exc.response.get("Error", {}).get("Code", "")
+    operation = getattr(exc, "operation_name", "") or ""
+    return error_code in _THROTTLE_ERROR_CODES and operation == _SIMULATE_OP
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_makereport(item, call):
+    """Skip (rather than fail) on residual SimulatePrincipalPolicy throttling that
+    survives the adaptive retries under heavy concurrency. Applies to setup and
+    call phases, since role resolution frequently happens in fixture setup."""
+    outcome = yield
+    report = outcome.get_result()
+
+    if report.when not in ("setup", "call") or not report.failed:
+        return
+
+    exc = call.excinfo.value if call.excinfo else None
+    # Walk the exception chain so a wrapped ClientError is still detected.
+    while exc is not None:
+        if _is_simulate_policy_throttle(exc):
+            report.outcome = "skipped"
+            report.wasxfail = (
+                "iam:SimulatePrincipalPolicy throttled under concurrent test load"
+            )
+            break
+        exc = exc.__cause__ or exc.__context__
+
+
 # ---------------------------------------------------------------------------
 # CLI options
 # ---------------------------------------------------------------------------

sagemaker-serve/tests/integ/conftest.py

@@ -18,28 +18,34 @@
 ``ClientError: (Throttling) ... Rate exceeded`` and failing the build.
 
 This is purely a test-harness concurrency problem, so the mitigation lives here in
-the test layer rather than in SDK source:
-
-* ``_configure_default_boto_retries`` (autouse) — approach "A": set an adaptive
-  retry policy on boto3's *default* session. The role resolver, when no explicit
-  session is passed, falls back to ``sagemaker.core...Session()`` whose boto
-  session is ``boto3.DEFAULT_SESSION`` (see session_helper._initialize). Clients
-  created from it therefore inherit these retries, letting the internal IAM calls
-  ride out transient throttling without any source change.
-
-* ``pytest_runtest_makereport`` — approach "C": a belt-and-suspenders fallback. If
-  a residual ``SimulatePrincipalPolicy`` throttling error still escapes after the
-  adaptive retries, convert the failure into an xfail so a transient rate limit
-  never reds the build. Genuine failures are untouched.
+the test layer rather than in SDK source. It is intentionally identical to the
+block in ``sagemaker-train`` / ``sagemaker-mlops`` integ conftests:
+
+* ``_configure_boto_adaptive_retries`` (autouse) — set adaptive retries via the
+  ``AWS_RETRY_MODE`` / ``AWS_MAX_ATTEMPTS`` environment variables. Env vars apply
+  to *every* boto3 client created in the worker, so the internal IAM calls ride
+  out transient throttling whether the resolver falls back to the default session
+  or builds its client from an explicitly-passed ``Session`` (several serve tests
+  pass their own session, whose IAM client would otherwise carry botocore's
+  default 4-attempt retry policy). ``adaptive`` mode also adds client-side rate
+  limiting to smooth bursts.
+
+* ``pytest_runtest_makereport`` — a belt-and-suspenders fallback. If a residual
+  ``SimulatePrincipalPolicy`` throttling error still escapes after the adaptive
+  retries, convert the failure into a skip so a transient rate limit never reds
+  the build. Genuine failures are untouched.
 """
 from __future__ import absolute_import
 
-import boto3
+import os
+
 import pytest
-from botocore.config import Config
 
-# Adaptive retry budget for throttling-prone IAM validation calls.
-_ADAPTIVE_RETRY_CONFIG = Config(retries={"max_attempts": 10, "mode": "adaptive"})
+# botocore adaptive retry settings for throttling-prone IAM validation calls.
+# Applied via env vars so every client in the worker inherits them, regardless of
+# which boto session the SDK ends up using to build its IAM client.
+_RETRY_MODE = "adaptive"
+_MAX_ATTEMPTS = "10"
 
 # Only tolerate throttling on the IAM permission simulation used during role
 # validation, so unrelated throttling still fails loudly.
@@ -48,12 +54,22 @@
 
 
 @pytest.fixture(autouse=True, scope="session")
-def _configure_default_boto_retries():
-    """Give boto3's default session adaptive retries so IAM clients built by the
-    role resolver absorb transient SimulatePrincipalPolicy throttling."""
-    boto3.setup_default_session()
-    boto3.DEFAULT_SESSION._session.set_default_client_config(_ADAPTIVE_RETRY_CONFIG)
+def _configure_boto_adaptive_retries():
+    """Give every boto3 client in this xdist worker adaptive retries so the IAM
+    clients built by the role resolver absorb transient SimulatePrincipalPolicy
+    throttling. Restores any pre-existing values on teardown."""
+    previous = {
+        "AWS_RETRY_MODE": os.environ.get("AWS_RETRY_MODE"),
+        "AWS_MAX_ATTEMPTS": os.environ.get("AWS_MAX_ATTEMPTS"),
+    }
+    os.environ["AWS_RETRY_MODE"] = _RETRY_MODE
+    os.environ["AWS_MAX_ATTEMPTS"] = _MAX_ATTEMPTS
     yield
+    for key, value in previous.items():
+        if value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = value
 
 
 def _is_simulate_policy_throttle(exc):
@@ -70,11 +86,12 @@ def _is_simulate_policy_throttle(exc):
 @pytest.hookimpl(hookwrapper=True)
 def pytest_runtest_makereport(item, call):
     """Skip (rather than fail) on residual SimulatePrincipalPolicy throttling that
-    survives the adaptive retries under heavy concurrency."""
+    survives the adaptive retries under heavy concurrency. Applies to setup and
+    call phases, since role resolution frequently happens in fixture setup."""
     outcome = yield
     report = outcome.get_result()
 
-    if report.when != "call" or not report.failed:
+    if report.when not in ("setup", "call") or not report.failed:
         return
 
     exc = call.excinfo.value if call.excinfo else None

sagemaker-train/tests/integ/conftest.py

@@ -0,0 +1,108 @@
+# 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.
+"""Shared pytest configuration for sagemaker-train integration tests.
+
+These tests run under ``pytest -n auto`` (dozens of xdist workers). Many of them
+resolve/validate an IAM execution role via ``TrainDefaults.get_role`` ->
+``resolve_and_validate_role``, which internally calls the low-TPS
+``iam:SimulatePrincipalPolicy`` API. With many workers hitting it at once IAM
+throttles the request, surfacing as ``ClientError: (Throttling) ... Rate
+exceeded`` and failing the test during setup.
+
+This is purely a test-harness concurrency problem, so the mitigation lives here
+in the test layer rather than in SDK source. It is intentionally identical to the
+block in ``sagemaker-serve`` / ``sagemaker-mlops`` integ conftests:
+
+* ``_configure_boto_adaptive_retries`` (autouse) — set adaptive retries via the
+  ``AWS_RETRY_MODE`` / ``AWS_MAX_ATTEMPTS`` environment variables. Unlike setting
+  a retry ``Config`` on a single boto session, env vars apply to *every* boto3
+  client created in the worker — including the IAM clients the role resolver
+  builds from an explicitly-passed ``Session`` (which carries no retry config of
+  its own). ``adaptive`` mode adds client-side rate limiting so bursts of
+  ``SimulatePrincipalPolicy`` calls ride out transient throttling.
+
+* ``pytest_runtest_makereport`` — a belt-and-suspenders fallback. If a residual
+  ``SimulatePrincipalPolicy`` throttling error still escapes after the adaptive
+  retries, convert the failure into a skip so a transient rate limit never reds
+  the build. Genuine failures (and throttling on any other operation) are
+  untouched.
+"""
+from __future__ import absolute_import
+
+import os
+
+import pytest
+
+# botocore adaptive retry settings for throttling-prone IAM validation calls.
+# Applied via env vars so every client in the worker inherits them, regardless of
+# which boto session the SDK ends up using to build its IAM client.
+_RETRY_MODE = "adaptive"
+_MAX_ATTEMPTS = "10"
+
+# Only tolerate throttling on the IAM permission simulation used during role
+# validation, so unrelated throttling still fails loudly.
+_THROTTLE_ERROR_CODES = ("Throttling", "ThrottlingException", "RequestLimitExceeded")
+_SIMULATE_OP = "SimulatePrincipalPolicy"
+
+
+@pytest.fixture(autouse=True, scope="session")
+def _configure_boto_adaptive_retries():
+    """Give every boto3 client in this xdist worker adaptive retries so the IAM
+    clients built by the role resolver absorb transient SimulatePrincipalPolicy
+    throttling. Restores any pre-existing values on teardown."""
+    previous = {
+        "AWS_RETRY_MODE": os.environ.get("AWS_RETRY_MODE"),
+        "AWS_MAX_ATTEMPTS": os.environ.get("AWS_MAX_ATTEMPTS"),
+    }
+    os.environ["AWS_RETRY_MODE"] = _RETRY_MODE
+    os.environ["AWS_MAX_ATTEMPTS"] = _MAX_ATTEMPTS
+    yield
+    for key, value in previous.items():
+        if value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = value
+
+
+def _is_simulate_policy_throttle(exc):
+    """Return True if ``exc`` is a SimulatePrincipalPolicy throttling ClientError."""
+    from botocore.exceptions import ClientError
+
+    if not isinstance(exc, ClientError):
+        return False
+    error_code = exc.response.get("Error", {}).get("Code", "")
+    operation = getattr(exc, "operation_name", "") or ""
+    return error_code in _THROTTLE_ERROR_CODES and operation == _SIMULATE_OP
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_makereport(item, call):
+    """Skip (rather than fail) on residual SimulatePrincipalPolicy throttling that
+    survives the adaptive retries under heavy concurrency. Applies to setup and
+    call phases, since role resolution frequently happens in fixture setup."""
+    outcome = yield
+    report = outcome.get_result()
+
+    if report.when not in ("setup", "call") or not report.failed:
+        return
+
+    exc = call.excinfo.value if call.excinfo else None
+    # Walk the exception chain so a wrapped ClientError is still detected.
+    while exc is not None:
+        if _is_simulate_policy_throttle(exc):
+            report.outcome = "skipped"
+            report.wasxfail = (
+                "iam:SimulatePrincipalPolicy throttled under concurrent test load"
+            )
+            break
+        exc = exc.__cause__ or exc.__context__