Skip to content

Commit 1db06c5

Browse files
mujtaba1747Syed Jafri
andauthored
Add unit test to prevent future regression of preset reward function (#6182)
* fix: rever preset reward function deletion from hyperparams dict * testing: add unit test to prevent future regression of preset_reward_function --------- Co-authored-by: Syed Jafri <syedjfr@amazon.com>
1 parent 10ada9b commit 1db06c5

1 file changed

Lines changed: 26 additions & 0 deletions

File tree

‎sagemaker-train/tests/unit/train/test_rlvr_trainer.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,32 @@ def test_train_without_datasets_raises_error(self, mock_finetuning_options, mock
200200
with pytest.raises(Exception):
201201
trainer.train(wait=False)
202202

203+
@patch('sagemaker.train.common_utils.finetune_utils._get_beta_session')
204+
@patch('sagemaker.train.rlvr_trainer._validate_and_resolve_model_package_group')
205+
@patch('sagemaker.train.rlvr_trainer._get_fine_tuning_options_and_model_arn')
206+
def test_train_raises_when_no_reward_signal(self, mock_finetuning_options, mock_validate_group, mock_get_session):
207+
"""Test train() raises ValueError when no reward signal is configured.
208+
209+
Neither custom_reward_function nor the preset_reward_function hyperparameter
210+
is set, so the guard in train() must raise. Using Mock(spec=[]) ensures
211+
getattr(hyperparameters, "preset_reward_function", None) returns None rather
212+
than an auto-created (truthy) Mock attribute.
213+
"""
214+
mock_validate_group.return_value = "test-group"
215+
mock_get_session.return_value = Mock()
216+
mock_hyperparams = Mock(spec=["to_dict"]) # no preset_reward_function attr
217+
mock_hyperparams.to_dict.return_value = {}
218+
mock_finetuning_options.return_value = (mock_hyperparams, "model-arn", False)
219+
220+
trainer = RLVRTrainer(
221+
model="test-model",
222+
model_package_group="test-group",
223+
training_dataset="s3://bucket/train",
224+
) # no custom_reward_function
225+
226+
with pytest.raises(ValueError, match="requires a reward signal"):
227+
trainer.train(wait=False)
228+
203229
@patch('sagemaker.train.common_utils.finetune_utils._get_beta_session')
204230
@patch('sagemaker.train.common_utils.finetune_utils._resolve_model_name')
205231
@patch('sagemaker.train.rlvr_trainer._get_fine_tuning_options_and_model_arn')

0 commit comments

Comments
 (0)