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