Skip to content

Commit 60a051c

Browse files
mohamedzeidan2021Mohamed Zeidan
andauthored
test(serve): stop unit tests resolving IAM roles over the network (#6496)
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^..5a65c55 ("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. Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>
1 parent 4f486d7 commit 60a051c

2 files changed

Lines changed: 97 additions & 0 deletions

File tree

‎sagemaker-core/tests/unit/helper/test_iam_role_resolver.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
_role_trusts_service,
2525
_trusted_services_in_document,
2626
_expected_trust_services,
27+
_config_path_for_role_type,
28+
_resolve_config_default_role,
2729
)
2830

2931

@@ -1391,3 +1393,27 @@ class TestBackwardCompatibleExceptions:
13911393

13921394
def test_role_auto_creation_error_importable(self):
13931395
assert issubclass(RoleAutoCreationError, Exception)
1396+
1397+
1398+
class TestConfigDefaultRole:
1399+
"""Config-default lookup is only consulted for role types that have a path."""
1400+
1401+
@pytest.mark.parametrize("role_type", ["training", "feature_store"])
1402+
def test_role_types_with_a_config_path(self, role_type):
1403+
assert _config_path_for_role_type(role_type) is not None
1404+
1405+
@pytest.mark.parametrize("role_type", ["serving", "bedrock", "not_a_role_type"])
1406+
def test_role_types_without_a_config_path(self, role_type):
1407+
assert _config_path_for_role_type(role_type) is None
1408+
1409+
@pytest.mark.parametrize("role_type", ["serving", "bedrock"])
1410+
def test_returns_none_without_reading_config(self, role_type):
1411+
"""No config path means no config to read, so resolution short-circuits.
1412+
1413+
The caller then falls back to caller-identity resolution. Asserting that
1414+
resolve_value_from_config is never consulted keeps the short-circuit from
1415+
regressing into a config read for role types that have no default.
1416+
"""
1417+
with patch("sagemaker.core.common_utils.resolve_value_from_config") as mock_resolve:
1418+
assert _resolve_config_default_role(role_type) is None
1419+
mock_resolve.assert_not_called()
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Shared fixtures for the sagemaker-serve unit tests.
14+
15+
Keeps role resolution off the network. ``ModelBuilder.__post_init__`` resolves a
16+
default execution role when none is supplied, and ``resolve_and_validate_role``
17+
does that by calling ``sts:GetCallerIdentity`` followed by a paginated
18+
``iam:SimulatePrincipalPolicy``. Any unit test that builds a ModelBuilder without
19+
passing ``role_arn`` therefore makes live IAM calls, which only succeed because
20+
CI credentials happen to be present.
21+
22+
Those calls are the dominant source of CI flakiness for this package: IAM
23+
throttles on a shared, account-wide quota, so builds fail with
24+
``Throttling ... SimulatePrincipalPolicy ... reached max retries`` on whichever
25+
test happened to be running. The failures are not reproducible -- each build runs
26+
the suite three times (py310/py311/py312) and these failures almost always hit
27+
only one of the three.
28+
"""
29+
from __future__ import annotations
30+
31+
from contextlib import ExitStack
32+
from unittest.mock import patch
33+
34+
import pytest
35+
36+
# Stands in for "the caller's own identity role", which the real resolver obtains
37+
# from sts:GetCallerIdentity.
38+
DEFAULT_UNIT_TEST_ROLE = "arn:aws:iam::123456789012:role/UnitTestExecutionRole"
39+
40+
# resolve_and_validate_role has to be patched once per importing module, not once
41+
# at its definition site: both modules below bind it with a module-level
42+
# ``from ... import resolve_and_validate_role``, so each holds its own reference
43+
# and patching sagemaker.core.helper.iam_role_resolver would not rebind either.
44+
_RESOLVER_LOOKUP_SITES = (
45+
"sagemaker.serve.model_builder.resolve_and_validate_role",
46+
"sagemaker.serve.bedrock_model_builder.resolve_and_validate_role",
47+
)
48+
49+
50+
def _resolve_role_offline(provided_role=None, *, role_arn=None, **kwargs):
51+
"""Mirror resolve_and_validate_role's contract without touching IAM.
52+
53+
The real function returns the role it was given and only falls back to the
54+
caller's identity when none was supplied. Preserving that passthrough matters:
55+
four of its five call sites in this package pass a real role, and tests assert
56+
on the value that comes back.
57+
"""
58+
return provided_role or role_arn or DEFAULT_UNIT_TEST_ROLE
59+
60+
61+
@pytest.fixture(autouse=True)
62+
def stub_iam_role_resolution():
63+
"""Resolve execution roles locally for every unit test in this package.
64+
65+
Tests that patch ``resolve_and_validate_role`` themselves still win, because
66+
their patch is applied inside this one and so rebinds the same attribute last.
67+
"""
68+
with ExitStack() as stack:
69+
for target in _RESOLVER_LOOKUP_SITES:
70+
stack.enter_context(patch(target, side_effect=_resolve_role_offline))
71+
yield

0 commit comments

Comments
 (0)