From 0a917e71ea6aabbfaff5bd19c8b74e8efebc2657 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 15 Oct 2026 17:12:00 -0700 Subject: [PATCH] test(train): patch mocks at the lookup site so unit tests stop calling AWS Five tests in sagemaker-train's unit suite issued real AWS requests. They pass in CI only because the build role carries credentials, so the mocks were silently ineffective rather than visibly broken. Root cause in three shapes: 1. Patch targeting the defining module instead of the lookup site. base_trainer does `from ...finetune_utils import get_recipe_s3_uri` at module level, so patching finetune_utils.get_recipe_s3_uri never rebinds the reference base_trainer already holds and the real function ran, calling DescribeHubContent. Same for get_training_image, and for TrainDefaults, which handed back a real Session so download_file issued a live HeadObject. Only two of the four finetune_utils patches in that block were affected: _render_recipe_placeholders and _get_smtj_override_spec are imported inside _train_serverful_smtj, so that import re-executes under patch and finetune_utils is the correct target for them. Those are left as they are. 2. Missing mock. base_evaluator imports verify_evaluation_caller_permissions in the same statement as resolve_and_validate_role, but only the latter was ever patched, so _get_aws_execution_context reached STS GetCallerIdentity and IAM. 3. Incomplete mock. test_model_trainer_no_from_recipe_raises patched TrainDefaults.get_role for exactly this reason, but STS was also reached by a second route: model_post_init builds a default OutputDataConfig, which calls session.default_bucket(). The session is stubbed rather than get_output_data_config so that the defaulting logic stays under test. Verified with deliberately invalid credentials plus a transport-level recorder that counts swallowed calls: - the three files go from 5 failed / 97 passed to 102 passed - the full sagemaker-train unit suite goes from 19 to 14 failures; the remaining 14 are a local botocore shape mismatch, not these tests - zero outbound AWS requests from the three files Coverage moves by one line, 27402 to 27401, with no real loss. The 12 lines no longer covered are all sagemaker-core session_helper AWS plumbing (default_bucket, generate_default_sagemaker_bucket_name, botocore_resolver, sts_regional_endpoint) that these tests only covered by genuinely calling STS; all 12 are covered by sagemaker-core's own unit suite. Offsetting that, 11 lines become newly covered in base_trainer, base_evaluator, benchmark_evaluator, custom_scorer_evaluator and model_trainer, because the tests no longer abort partway through on an AWS error. No production code changes. --- .../evaluate/test_mtrl_evaluator_handshake.py | 18 +++++++++++++- .../tests/unit/train/test_sft_trainer.py | 15 +++++++++--- .../train/test_trainer_recipe_integration.py | 24 ++++++++++++++++--- 3 files changed, 50 insertions(+), 7 deletions(-) diff --git a/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator_handshake.py b/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator_handshake.py index f593372a6d..329a451ac5 100644 --- a/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator_handshake.py +++ b/sagemaker-train/tests/unit/train/evaluate/test_mtrl_evaluator_handshake.py @@ -292,6 +292,11 @@ def test_custom_scorer_evaluator_rejects_mtrl_trainer_without_job(self, mock_mlf class TestEvaluateSubmissionWithMTRLTrainer: """Test that evaluate() can be called successfully when model is an MTRLTrainer.""" + # _get_aws_execution_context calls verify_evaluation_caller_permissions, + # which reaches STS GetCallerIdentity and IAM. base_evaluator imports it in + # the same statement as resolve_and_validate_role, so it needs its own patch + # in that namespace; without it this test resolves a real Session. + @patch("sagemaker.train.evaluate.base_evaluator.verify_evaluation_caller_permissions") @patch( "sagemaker.train.evaluate.base_evaluator.resolve_and_validate_role", side_effect=lambda provided_role, **kwargs: provided_role, @@ -307,7 +312,14 @@ class TestEvaluateSubmissionWithMTRLTrainer: new_callable=PropertyMock, ) def test_benchmark_evaluate_submission_with_mtrl_trainer( - self, mock_hyperparams, mock_start, mock_artifact, mock_resolve_mp, mock_mlflow, mock_role + self, + mock_hyperparams, + mock_start, + mock_artifact, + mock_resolve_mp, + mock_mlflow, + mock_role, + mock_verify_permissions, ): """BenchMarkEvaluator.evaluate() should successfully submit when model is MTRLTrainer.""" from sagemaker.train.evaluate import BenchMarkEvaluator, get_benchmarks @@ -349,6 +361,9 @@ def test_benchmark_evaluate_submission_with_mtrl_trainer( assert execution.arn is not None mock_start.assert_called_once() + # See the note on the benchmark test above: verify_evaluation_caller_permissions + # needs its own patch in base_evaluator's namespace or this reaches STS. + @patch("sagemaker.train.evaluate.base_evaluator.verify_evaluation_caller_permissions") @patch("sagemaker.train.evaluate.custom_scorer_evaluator.validate_data_path_exists") @patch( "sagemaker.train.evaluate.base_evaluator.resolve_and_validate_role", @@ -373,6 +388,7 @@ def test_custom_scorer_evaluate_submission_with_mtrl_trainer( mock_mlflow, mock_role, mock_validate_data, + mock_verify_permissions, ): """CustomScorerEvaluator.evaluate() should successfully submit when model is MTRLTrainer.""" from sagemaker.train.evaluate import CustomScorerEvaluator diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py index b25da58322..61fda75129 100644 --- a/sagemaker-train/tests/unit/train/test_sft_trainer.py +++ b/sagemaker-train/tests/unit/train/test_sft_trainer.py @@ -1565,13 +1565,22 @@ def _run_smtj_and_capture_input_config(self, model_name): mock_session.sagemaker_config = {} with ( - patch("sagemaker.train.defaults.TrainDefaults") as mock_defaults, + # These three are imported into base_trainer at module level, so + # they must be patched in base_trainer's namespace -- patching them + # where they are defined does not rebind the reference base_trainer + # already holds. Left unpatched there, TrainDefaults hands back a + # real Session (so download_file issues a live HeadObject) and + # get_recipe_s3_uri calls DescribeHubContent for real. The two + # finetune_utils patches below are correct as written, because + # _train_serverful_smtj imports those names in its own body, so the + # import re-executes while the patch is active. + patch("sagemaker.train.base_trainer.TrainDefaults") as mock_defaults, patch( - "sagemaker.train.common_utils.finetune_utils.get_recipe_s3_uri", + "sagemaker.train.base_trainer.get_recipe_s3_uri", return_value="s3://bucket/recipe.yaml", ), patch( - "sagemaker.train.common_utils.finetune_utils.get_training_image", + "sagemaker.train.base_trainer.get_training_image", return_value="img:latest", ), patch( diff --git a/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py b/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py index b2af01ea1f..3b0a3026d4 100644 --- a/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py +++ b/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py @@ -552,14 +552,32 @@ def test_model_trainer_from_recipe_get_resolved(self, tmp_path): def test_model_trainer_no_from_recipe_raises(self): """ModelTrainer not created via from_recipe raises AttributeError.""" + from sagemaker.core.helper.session_helper import Session from sagemaker.train.model_trainer import ModelTrainer from sagemaker.train.configs import Compute # Construction resolves an execution role via STS; mock it so this unit # test (which only exercises get_resolved_recipe) runs without AWS creds. - with patch( - "sagemaker.train.defaults.TrainDefaults.get_role", - return_value="arn:aws:iam::123456789012:role/test-role", + # It also builds a default OutputDataConfig, which calls + # session.default_bucket() and so reaches STS by a second route, so the + # session is stubbed too. Stubbing the session rather than + # get_output_data_config keeps that defaulting logic under test. + # spec=Session is required because ModelTrainer validates the field as a + # Session instance; default_bucket_prefix is set on instances, not the + # class, so it has to be assigned explicitly. + mock_session = MagicMock(spec=Session) + mock_session.default_bucket.return_value = "test-bucket" + mock_session.default_bucket_prefix = None + + with ( + patch( + "sagemaker.train.defaults.TrainDefaults.get_role", + return_value="arn:aws:iam::123456789012:role/test-role", + ), + patch( + "sagemaker.train.defaults.TrainDefaults.get_sagemaker_session", + return_value=mock_session, + ), ): trainer = ModelTrainer( training_image="123456789012.dkr.ecr.us-east-1.amazonaws.com/test:latest",