test(train): patch mocks at the lookup site so unit tests stop calling AWS (#6491)

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.

Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>

Author: mohamedzeidan2021

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

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(

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",