diff --git a/sagemaker-mlops/tests/integ/conftest.py b/sagemaker-mlops/tests/integ/conftest.py index dee3642cb9..a0981c9b5f 100644 --- a/sagemaker-mlops/tests/integ/conftest.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/sagemaker-serve/tests/integ/conftest.py b/sagemaker-serve/tests/integ/conftest.py index 4170a1a934..6ffc1a65f1 100644 --- a/sagemaker-serve/tests/integ/conftest.py +++ b/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 diff --git a/sagemaker-train/tests/integ/conftest.py b/sagemaker-train/tests/integ/conftest.py new file mode 100644 index 0000000000..87c861092a --- /dev/null +++ b/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__