Skip to content

Commit 5a65c55

Browse files
authored
fix PR check tests (#6272)
* test commit * test PR checks * Fix flaky unit tests * Revert "test PR checks" This reverts commit 39c127a. * Revert "test commit" This reverts commit 52c765f. * Update unit tests * update integ test to avoid failure from conflicting updates
1 parent 4ac6376 commit 5a65c55

6 files changed

Lines changed: 142 additions & 15 deletions

File tree

‎sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
from dataclasses import dataclass
99
import tempfile
1010

11+
import pytest
12+
1113
from sagemaker.serve.model_builder import ModelBuilder
1214
from sagemaker.serve.mode.function_pointers import Mode
1315
from sagemaker.serve.utils.types import ModelServer
@@ -16,6 +18,26 @@
1618
from sagemaker.core.inference_config import AsyncInferenceConfig
1719
from botocore.exceptions import ClientError
1820

21+
TEST_ROLE_ARN = "arn:aws:iam::123456789012:role/SageMakerRole"
22+
23+
24+
@pytest.fixture(autouse=True)
25+
def stub_role_resolution():
26+
"""Keep ModelBuilder construction offline.
27+
28+
``ModelBuilder.__post_init__`` auto-resolves a serving role when no
29+
``role_arn`` is given, which calls sts:GetCallerIdentity and the paginated
30+
iam:SimulatePrincipalPolicy. Tests here construct ``ModelBuilder`` without a
31+
role, so unpatched they issue live IAM calls and fail on throttling
32+
(SimulatePrincipalPolicy "Rate exceeded") rather than on the behavior under
33+
test. Tests that patch the resolver themselves still override this.
34+
"""
35+
with patch(
36+
"sagemaker.serve.model_builder.resolve_and_validate_role",
37+
side_effect=lambda provided_role=None, **kwargs: provided_role or TEST_ROLE_ARN,
38+
):
39+
yield
40+
1941

2042
class TestModelBuilderInit(unittest.TestCase):
2143
"""Test ModelBuilder initialization."""

‎sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,31 @@
5959
EVALUATION_TIMEOUT_SECONDS = 7200 # 2 hours
6060
POLL_INTERVAL_SECONDS = 30
6161

62+
# SageMaker's message when UpdatePipeline loses an optimistic-concurrency race.
63+
_PIPELINE_CONFLICT_MESSAGE = "has been modified since your last read"
64+
65+
66+
def _evaluate_or_skip_on_pipeline_conflict(evaluator):
67+
"""Start an evaluation, skipping if a concurrent evaluation won the pipeline race.
68+
69+
All evaluations of one eval type share a single SageMaker pipeline, so two
70+
evaluations starting at the same time both call UpdatePipeline and SageMaker
71+
rejects the loser with a conflict. ``evaluate()`` reports that as a Failed
72+
execution with no ARN. The tests in this module run on separate xdist workers
73+
and can overlap, so hitting it says nothing about the code under test --
74+
give up on this pipeline update and move on rather than failing.
75+
"""
76+
execution = evaluator.evaluate()
77+
78+
failure_reason = getattr(execution.status, "failure_reason", None) or ""
79+
if execution.arn is None and _PIPELINE_CONFLICT_MESSAGE in failure_reason:
80+
pytest.skip(
81+
f"A concurrent evaluation modified the shared evaluation pipeline: "
82+
f"{failure_reason}"
83+
)
84+
85+
return execution
86+
6287

6388
def _prefix_has_content(s3_client, bucket_name: str, prefix: str) -> bool:
6489
"""Check if an S3 prefix has any objects."""
@@ -136,7 +161,7 @@ def test_inspect_ai_bedrock_evaluation(
136161
)
137162

138163
logger.info("Starting InspectAI evaluation with Bedrock inference...")
139-
execution = evaluator.evaluate()
164+
execution = _evaluate_or_skip_on_pipeline_conflict(evaluator)
140165

141166
assert execution is not None
142167
assert execution.arn is not None
@@ -200,7 +225,7 @@ def test_inspect_ai_upload_benchmarks(
200225
)
201226

202227
logger.info("Starting evaluation with pre-existing benchmarks...")
203-
execution = evaluator2.evaluate()
228+
execution = _evaluate_or_skip_on_pipeline_conflict(evaluator2)
204229

205230
assert execution is not None
206231
assert execution.arn is not None

‎sagemaker-train/tests/unit/ai_registry/test_dataset.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -296,10 +296,27 @@ def mock_exists(path):
296296
assert document['DatasetS3Prefix'] == 'path/to/dataset.jsonl'
297297
assert document['DatasetS3Bucket'] == 'test-bucket'
298298

299+
@patch('sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name', return_value="test-hub")
300+
@patch('sagemaker.ai_registry.dataset._get_default_bucket', return_value="test-bucket")
301+
@patch('sagemaker.train.defaults.TrainDefaults.get_role', return_value="arn:aws:iam::123456789012:role/SageMakerRole")
302+
@patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session')
303+
@patch('sagemaker.ai_registry.dataset._get_current_domain_id', return_value=None)
304+
@patch('sagemaker.ai_registry.dataset.Session')
299305
@patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_file')
300306
@patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_format')
301307
@patch('sagemaker.ai_registry.dataset.AIRHub')
302-
def test_create_with_local_file(self, mock_air_hub, mock_validate_format, mock_validate_file):
308+
def test_create_with_local_file(
309+
self,
310+
mock_air_hub,
311+
mock_validate_format,
312+
mock_validate_file,
313+
mock_session,
314+
mock_get_domain_id,
315+
mock_get_session,
316+
mock_get_role,
317+
mock_default_bucket,
318+
mock_get_hub_name,
319+
):
303320
mock_air_hub.upload_to_s3.return_value = "s3://bucket/path"
304321
mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"}
305322
mock_air_hub.describe_hub_content.return_value = {

‎sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,23 @@ def sample_dataset_file():
4646

4747
class TestDataSetDomainId:
4848
"""Test domain-id is added to SearchKeywords when available."""
49-
49+
50+
@pytest.fixture(autouse=True)
51+
def stub_account_lookups(self):
52+
"""Keep these unit tests offline.
53+
54+
Uploading a local dataset derives the default bucket, and constructing the entity
55+
derives the hub name — both call STS ``GetCallerIdentity``. ``AIRHub`` is patched on
56+
the base-entity module because ``AIRHubEntity.__init__`` resolves the hub name
57+
through its own import, not the one patched on ``dataset``.
58+
"""
59+
with patch(
60+
'sagemaker.ai_registry.dataset._get_default_bucket', return_value='test-bucket'
61+
), patch(
62+
'sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name', return_value='test-hub'
63+
):
64+
yield
65+
5066
@patch('sagemaker.core.helper.session_helper.Session')
5167
@patch('sagemaker.ai_registry.dataset._get_current_domain_id')
5268
@patch('sagemaker.ai_registry.dataset.AIRHub')

‎sagemaker-train/tests/unit/ai_registry/test_evaluator.py‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,43 @@
2323
)
2424

2525

26+
DUMMY_ROLE = "arn:aws:iam::123456789012:role/SageMakerRole"
27+
28+
2629
def _keywords_from_import_call(mock_air_hub):
2730
"""Extract the SearchKeyword strings passed to import_hub_content."""
2831
tags = mock_air_hub.import_hub_content.call_args.kwargs["tags"]
2932
return [f"{tag[0]}:{tag[1]}" for tag in tags]
3033

3134

3235
class TestEvaluator:
36+
@pytest.fixture(autouse=True)
37+
def stub_aws_resolution(self):
38+
"""Keep these unit tests offline.
39+
40+
``Evaluator.create`` builds a default ``Session``, auto-detects the Studio domain
41+
ID (STS ``GetCallerIdentity``) and resolves/validates an execution role (IAM
42+
``SimulatePrincipalPolicy``); constructing the entity also derives the hub name
43+
from the caller's account. Unmocked, those reach real AWS and make the suite slow
44+
and flaky — CI hit ``Throttling: Rate exceeded`` on ``SimulatePrincipalPolicy``.
45+
46+
``AIRHub`` is patched on the base-entity module as well as on ``evaluator``, since
47+
``AIRHubEntity.__init__`` resolves the hub name through its own import.
48+
"""
49+
session = MagicMock()
50+
with patch("sagemaker.ai_registry.evaluator.Session", return_value=session), patch(
51+
"sagemaker.ai_registry.evaluator._get_current_domain_id", return_value=None
52+
), patch(
53+
"sagemaker.train.defaults.TrainDefaults.get_sagemaker_session",
54+
return_value=session,
55+
), patch(
56+
"sagemaker.train.defaults.TrainDefaults.get_role", return_value=DUMMY_ROLE
57+
), patch(
58+
"sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name",
59+
return_value="test-hub",
60+
):
61+
yield
62+
3363
@patch('sagemaker.ai_registry.evaluator.AIRHub')
3464
def test_create_with_lambda_arn(self, mock_air_hub):
3565
mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"}
@@ -52,9 +82,10 @@ def test_create_with_lambda_arn(self, mock_air_hub):
5282
assert evaluator.method == EvaluatorMethod.LAMBDA
5383
mock_air_hub.import_hub_content.assert_called_once()
5484

85+
@patch('sagemaker.ai_registry.evaluator._get_default_bucket', return_value="test-bucket")
5586
@patch('sagemaker.ai_registry.evaluator.boto3')
5687
@patch('sagemaker.ai_registry.evaluator.AIRHub')
57-
def test_create_with_byoc(self, mock_air_hub, mock_boto3):
88+
def test_create_with_byoc(self, mock_air_hub, mock_boto3, mock_default_bucket):
5889
mock_lambda_client = MagicMock()
5990
mock_boto3.client.return_value = mock_lambda_client
6091
mock_lambda_client.create_function.return_value = {"FunctionArn": "lambda-arn"}

‎sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,18 +18,38 @@
1818

1919
class TestEvaluatorDomainId:
2020
"""Test domain-id is added to SearchKeywords when available."""
21-
22-
@patch('sagemaker.core.helper.session_helper.Session')
21+
22+
@pytest.fixture(autouse=True)
23+
def stub_aws_resolution(self):
24+
"""Keep these unit tests offline.
25+
26+
``Evaluator.create`` builds a default ``Session``, resolves/validates an execution
27+
role via IAM ``SimulatePrincipalPolicy``, and derives the hub name from the
28+
caller's account. Unmocked, those reach real AWS and made CI fail with
29+
``Throttling: Rate exceeded``.
30+
"""
31+
session = Mock()
32+
with patch("sagemaker.ai_registry.evaluator.Session", return_value=session), patch(
33+
"sagemaker.train.defaults.TrainDefaults.get_sagemaker_session",
34+
return_value=session,
35+
), patch(
36+
"sagemaker.train.defaults.TrainDefaults.get_role",
37+
return_value="arn:aws:iam::123456789012:role/test-role",
38+
), patch(
39+
"sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name",
40+
return_value="test-hub",
41+
):
42+
yield
43+
2344
@patch('sagemaker.ai_registry.evaluator._get_current_domain_id')
2445
@patch('sagemaker.ai_registry.evaluator.AIRHub')
2546
def test_domain_id_added_when_available(
26-
self, mock_air_hub, mock_get_domain_id, mock_session
47+
self, mock_air_hub, mock_get_domain_id
2748
):
2849
"""Test that domain-id is added to tags when available."""
2950
# Setup mocks
3051
mock_domain_id = "d-test123456"
3152
mock_get_domain_id.return_value = mock_domain_id
32-
mock_session.return_value = Mock()
3353

3454
# Mock AIRHub methods
3555
mock_air_hub.import_hub_content = Mock()
@@ -62,16 +82,14 @@ def test_domain_id_added_when_available(
6282
# Verify domain-id is in tags
6383
assert any(tag[0] == '@domain' and tag[1] == mock_domain_id for tag in tags)
6484

65-
@patch('sagemaker.core.helper.session_helper.Session')
6685
@patch('sagemaker.ai_registry.evaluator._get_current_domain_id')
6786
@patch('sagemaker.ai_registry.evaluator.AIRHub')
6887
def test_domain_id_not_added_when_unavailable(
69-
self, mock_air_hub, mock_get_domain_id, mock_session
88+
self, mock_air_hub, mock_get_domain_id
7089
):
7190
"""Test that domain-id is not added when unavailable (non-Studio)."""
7291
# Setup mocks - domain_id returns None
7392
mock_get_domain_id.return_value = None
74-
mock_session.return_value = Mock()
7593

7694
# Mock AIRHub methods
7795
mock_air_hub.import_hub_content = Mock()
@@ -104,18 +122,16 @@ def test_domain_id_not_added_when_unavailable(
104122
# Verify domain-id is NOT in tags
105123
assert not any(tag[0] == '@domain' for tag in tags)
106124

107-
@patch('sagemaker.core.helper.session_helper.Session')
108125
@patch('sagemaker.ai_registry.evaluator._get_current_domain_id')
109126
@patch('sagemaker.ai_registry.evaluator.AIRHub')
110127
def test_explicit_domain_id_used_without_auto_detection(
111-
self, mock_air_hub, mock_get_domain_id, mock_session
128+
self, mock_air_hub, mock_get_domain_id
112129
):
113130
"""An explicit domain_id is tagged and auto-detection is not invoked.
114131
115132
Covers the non-Studio case (P467494019) where the domain cannot be inferred and
116133
must be supplied by the caller.
117134
"""
118-
mock_session.return_value = Mock()
119135
mock_air_hub.import_hub_content = Mock()
120136
mock_air_hub.describe_hub_content = Mock(return_value={
121137
'HubContentName': 'test-evaluator',

0 commit comments

Comments
 (0)