Fix sm-train unit tests + use single logger in base trainer (#6030)

* cleanup: use a single logger in base_trainer

* fix: update recipe resolver unit tests

Author: mujtaba1747

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -34,6 +34,9 @@
 from sagemaker.train.defaults import TrainDefaults
 from sagemaker.train.utils import _get_unique_name
 
+logger = logging.getLogger(__name__)
+
+
 class BaseTrainer(ABC):
     """Abstract base class for all SageMaker training workflows.
 
@@ -156,7 +159,6 @@ def _fetch_full_recipe_template(self) -> Optional[Dict[str, Any]]:
 
         Returns None if the template can't be fetched (fallback to synthetic template).
         """
-        logger = logging.getLogger(__name__)
         frt = getattr(self.hyperparameters, '_full_recipe_template', None) if hasattr(self, 'hyperparameters') else None
         if isinstance(frt, dict):
             return frt
@@ -392,8 +394,6 @@ def _train_serverful_smtj(self, training_dataset=None, validation_dataset=None,
         )
         from sagemaker.train.defaults import TrainDefaults
 
-        logger = logging.getLogger(__name__)
-
         sagemaker_session = TrainDefaults.get_sagemaker_session(
             sagemaker_session=self.sagemaker_session
         )
@@ -848,8 +848,6 @@ def _train_hyperpod(self, training_dataset=None, validation_dataset=None,
         (SFT, DPO, RLVR).
         """
 
-        logger = logging.getLogger(__name__)
-
         sagemaker_session = TrainDefaults.get_sagemaker_session(
             sagemaker_session=self.sagemaker_session
         )

sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py

@@ -611,6 +611,9 @@ def test_non_spec_keys_flow_into_train_hyperparameters(
             overrides={"training_config": {"sequence_length": 8192}},
         )
 
+        # For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
+        trainer.compute = MagicMock()
+
         with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
              patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
              patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
@@ -640,6 +643,54 @@ def test_non_spec_keys_flow_into_train_hyperparameters(
             assert "warmup_ratio" in final_hp
             assert final_hp["warmup_ratio"] == "0.1"
 
+    @patch("sagemaker.train.sft_trainer._validate_eula_for_gated_model", return_value=False)
+    @patch("sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn")
+    @patch("sagemaker.train.sft_trainer._validate_and_resolve_model_package_group", return_value="my-group")
+    @patch("sagemaker.train.sft_trainer._resolve_model_and_name", return_value=("model_obj", "nova-lite-v2"))
+    def test_serverless_non_spec_keys_dont_flow_into_train_hyperparameters(
+        self, mock_resolve, mock_validate_group, mock_get_options, mock_eula,
+        mock_hyperparams_with_full_template
+    ):
+        """Non-spec keys from full template are included in final training hyperparameters."""
+        mock_get_options.return_value = (mock_hyperparams_with_full_template, "model-arn", False)
+
+        from sagemaker.train.sft_trainer import SFTTrainer
+
+        trainer = SFTTrainer(
+            model="nova-lite-v2",
+            model_package_group="my-group",
+            training_dataset="s3://bucket/train.jsonl",
+            overrides={"training_config": {"sequence_length": 8191}},
+        )
+
+        with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
+             patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
+             patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
+             patch("sagemaker.train.sft_trainer._convert_input_data_to_channels", return_value=[]), \
+             patch("sagemaker.train.sft_trainer._create_output_config", return_value=MagicMock()), \
+             patch("sagemaker.train.sft_trainer._create_serverless_config", return_value=MagicMock()), \
+             patch("sagemaker.train.sft_trainer._create_mlflow_config", return_value=None), \
+             patch("sagemaker.train.sft_trainer._create_model_package_config", return_value=None), \
+             patch("sagemaker.train.sft_trainer._validate_hyperparameter_values"), \
+             patch("sagemaker.train.sft_trainer._get_jumpstart_tags", return_value=[]):
+
+            mock_session = MagicMock()
+            mock_session.boto_session.region_name = "us-west-2"
+            mock_defaults.get_sagemaker_session.return_value = mock_session
+            mock_defaults.get_role.return_value = "arn:aws:iam::123:role/Role"
+            mock_tj.create.return_value = MagicMock()
+
+            trainer.train(wait=False)
+
+            call_kwargs = mock_tj.create.call_args[1]
+            final_hp = call_kwargs["hyper_parameters"]
+
+            # Non-spec key is now in final hyperparameters
+            assert "sequence_length" in final_hp
+            assert final_hp["sequence_length"] == "8191"
+            # Other full template keys also present
+            assert "warmup_ratio" not in final_hp
+
     @patch("sagemaker.train.sft_trainer._validate_eula_for_gated_model", return_value=False)
     @patch("sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn")
     @patch("sagemaker.train.sft_trainer._validate_and_resolve_model_package_group", return_value="my-group")
@@ -660,6 +711,9 @@ def test_nested_keys_flow_into_train_hyperparameters(
             overrides={"training_config": {"lr_scheduler": {"warmup_steps": 30}}},
         )
 
+        # For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
+        trainer.compute = MagicMock()
+
         with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
              patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
              patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
@@ -725,6 +779,9 @@ def test_deeply_nested_peft_keys_flow_into_hyperparameters(
             overrides={"training_config": {"peft": {"lora_tuning": {"alpha": 128}}}},
         )
 
+        # For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
+        trainer.compute = MagicMock()
+        
         with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
              patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
              patch("sagemaker.train.sft_trainer._create_input_data_config"), \