From 120badd70203194f073ecfbd610d5f2eb13a0f23 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Fri, 16 Oct 2026 12:32:08 -0700 Subject: [PATCH] test(serve): stop unit tests resolving IAM roles over the network ModelBuilder.__post_init__ resolves a default execution role when none is supplied, and resolve_and_validate_role does that with sts:GetCallerIdentity followed by a paginated iam:SimulatePrincipalPolicy. Every unit test that builds a ModelBuilder without role_arn therefore makes live IAM calls, which succeed only because CI credentials are present. IAM throttles on an account-wide quota, so those builds fail with Throttling ... SimulatePrincipalPolicy (reached max retries: 4): Rate exceeded on whichever test happened to be running. Over the 184 decided builds of sagemaker-python-sdk-ci-health-unit-test-v3 with SUBMODULE=sagemaker-serve since 2026-09-23, throttling accounts for 37 of 42 failures, and 100% of it (37/37 builds, 90/90 per-env instances) comes from tests/unit/test_model_builder_methods.py -- the one remaining file whose tests construct ModelBuilder without a role and without stubbing the resolver. All 23 of its unprotected tests have thrown it in production. The failures are not reproducible: each build runs the suite three times (py310/py311/py312) and these hit only one of the three 89% of the time, while genuine failures hit all three every time. This generalises the fixture 5432663^..5a65c559 ("fix PR check tests", #6272) added to test_model_builder_coverage_boost.py, which drove that file from 83% of throttling occurrences to zero. test_bedrock_model_builder.py and test_ai_inference_recommender/test_model_builder_recommendations.py carry the same fixture and have never thrown a throttle in 531 builds. Moving it into tests/unit/conftest.py covers the package rather than one file at a time. Two details worth noting: - The stub passes the provided role through rather than returning a fixed value, matching the real contract. Four of the five call sites pass a real role, and tests assert on what comes back. - resolve_and_validate_role is patched once per importing module. model_builder and bedrock_model_builder each bind it with a module-level `from ... import`, so each holds its own reference and patching the definition site in sagemaker.core.helper.iam_role_resolver would rebind neither. Tests that patch the resolver themselves still win, since their patch is applied inside the fixture's and so rebinds the attribute last. Verified with deliberately invalid credentials plus a transport-level request recorder (with a positive control, so a silent failure to record cannot be read as absence of traffic): - serve unit suite: 67 failed -> 44 failed, 2087 -> 2110 passed - newly broken tests: 0 - tests issuing STS calls: 24 -> 1 (the remaining one is the telemetry daemon thread, attributed outside any test, and is unrelated) - STS/IAM errors in the run: 0 The 44 remaining failures are all s3:GetObject against the JumpStart cache. Those succeed in CI, so they produce no CI failures and are left for a separate change. Coverage is held whole. serve/model_builder.py gains 5 covered lines because 23 tests now run to completion. 18 lines of sagemaker-core's iam_role_resolver stop being covered by serve, having previously been reached only by making real IAM calls; 17 were already covered by that module's own tests, and the new TestConfigDefaultRole covers the last one (the early return for role types with no config path, which serve was exercising incidentally via role_type="serving"). No production code changes. --- .../unit/helper/test_iam_role_resolver.py | 26 +++++++ sagemaker-serve/tests/unit/conftest.py | 71 +++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 sagemaker-serve/tests/unit/conftest.py diff --git a/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py b/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py index 8d0962b4f3..c726050e1a 100644 --- a/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py +++ b/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py @@ -24,6 +24,8 @@ _role_trusts_service, _trusted_services_in_document, _expected_trust_services, + _config_path_for_role_type, + _resolve_config_default_role, ) @@ -1391,3 +1393,27 @@ class TestBackwardCompatibleExceptions: def test_role_auto_creation_error_importable(self): assert issubclass(RoleAutoCreationError, Exception) + + +class TestConfigDefaultRole: + """Config-default lookup is only consulted for role types that have a path.""" + + @pytest.mark.parametrize("role_type", ["training", "feature_store"]) + def test_role_types_with_a_config_path(self, role_type): + assert _config_path_for_role_type(role_type) is not None + + @pytest.mark.parametrize("role_type", ["serving", "bedrock", "not_a_role_type"]) + def test_role_types_without_a_config_path(self, role_type): + assert _config_path_for_role_type(role_type) is None + + @pytest.mark.parametrize("role_type", ["serving", "bedrock"]) + def test_returns_none_without_reading_config(self, role_type): + """No config path means no config to read, so resolution short-circuits. + + The caller then falls back to caller-identity resolution. Asserting that + resolve_value_from_config is never consulted keeps the short-circuit from + regressing into a config read for role types that have no default. + """ + with patch("sagemaker.core.common_utils.resolve_value_from_config") as mock_resolve: + assert _resolve_config_default_role(role_type) is None + mock_resolve.assert_not_called() diff --git a/sagemaker-serve/tests/unit/conftest.py b/sagemaker-serve/tests/unit/conftest.py new file mode 100644 index 0000000000..ab164faa4b --- /dev/null +++ b/sagemaker-serve/tests/unit/conftest.py @@ -0,0 +1,71 @@ +# 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 fixtures for the sagemaker-serve unit tests. + +Keeps role resolution off the network. ``ModelBuilder.__post_init__`` resolves a +default execution role when none is supplied, and ``resolve_and_validate_role`` +does that by calling ``sts:GetCallerIdentity`` followed by a paginated +``iam:SimulatePrincipalPolicy``. Any unit test that builds a ModelBuilder without +passing ``role_arn`` therefore makes live IAM calls, which only succeed because +CI credentials happen to be present. + +Those calls are the dominant source of CI flakiness for this package: IAM +throttles on a shared, account-wide quota, so builds fail with +``Throttling ... SimulatePrincipalPolicy ... reached max retries`` on whichever +test happened to be running. The failures are not reproducible -- each build runs +the suite three times (py310/py311/py312) and these failures almost always hit +only one of the three. +""" +from __future__ import annotations + +from contextlib import ExitStack +from unittest.mock import patch + +import pytest + +# Stands in for "the caller's own identity role", which the real resolver obtains +# from sts:GetCallerIdentity. +DEFAULT_UNIT_TEST_ROLE = "arn:aws:iam::123456789012:role/UnitTestExecutionRole" + +# resolve_and_validate_role has to be patched once per importing module, not once +# at its definition site: both modules below bind it with a module-level +# ``from ... import resolve_and_validate_role``, so each holds its own reference +# and patching sagemaker.core.helper.iam_role_resolver would not rebind either. +_RESOLVER_LOOKUP_SITES = ( + "sagemaker.serve.model_builder.resolve_and_validate_role", + "sagemaker.serve.bedrock_model_builder.resolve_and_validate_role", +) + + +def _resolve_role_offline(provided_role=None, *, role_arn=None, **kwargs): + """Mirror resolve_and_validate_role's contract without touching IAM. + + The real function returns the role it was given and only falls back to the + caller's identity when none was supplied. Preserving that passthrough matters: + four of its five call sites in this package pass a real role, and tests assert + on the value that comes back. + """ + return provided_role or role_arn or DEFAULT_UNIT_TEST_ROLE + + +@pytest.fixture(autouse=True) +def stub_iam_role_resolution(): + """Resolve execution roles locally for every unit test in this package. + + Tests that patch ``resolve_and_validate_role`` themselves still win, because + their patch is applied inside this one and so rebinds the same attribute last. + """ + with ExitStack() as stack: + for target in _RESOLVER_LOOKUP_SITES: + stack.enter_context(patch(target, side_effect=_resolve_role_offline)) + yield