From fb577ce3d19ef6adc56ca032cffa84da057b22af Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Wed, 15 Jul 2026 12:09:59 -0700 Subject: [PATCH 1/2] cleanup: use a single logger in base_trainer --- sagemaker-train/src/sagemaker/train/base_trainer.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index da94cc0982..c17199a611 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/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 ) From 706e73e7c9dd8172fdb2a4ab63a014d3ae021cfb Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Wed, 15 Jul 2026 13:42:39 -0700 Subject: [PATCH 2/2] fix: update recipe resolver unit tests --- .../train/test_trainer_recipe_integration.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) 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 bb873a6b8b..97a754b24d 100644 --- a/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py +++ b/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"), \