Skip to content

Commit 85299be

Browse files
committed
fix: correct test imports and mock setup for sequence_length tests
1 parent 8032756 commit 85299be

1 file changed

Lines changed: 10 additions & 13 deletions

File tree

‎sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@
3030
_validate_s3_path_exists,
3131
_parse_context_length
3232
)
33-
from sagemaker.core.resources import ModelPackage
34-
from sagemaker.core.utils.utils import Unassigned, ModelPackageGroup
33+
from sagemaker.core.resources import ModelPackage, ModelPackageGroup
34+
from sagemaker.core.utils.utils import Unassigned
3535
from sagemaker.ai_registry.dataset import DataSet
3636
from sagemaker.train.common import TrainingType
3737
from sagemaker.train.configs import InputData
@@ -467,7 +467,6 @@ def test__convert_input_data_to_channels(self):
467467
def test__validate_eula_for_gated_model_with_model_package(self):
468468
"""Test EULA validation returns True for ModelPackage input"""
469469
from sagemaker.core.resources import ModelPackage
470-
from sagemaker.core.utils.utils import Unassigned
471470
model_package = Mock(spec=ModelPackage)
472471

473472
result = _validate_eula_for_gated_model(model_package, False, True)
@@ -898,10 +897,14 @@ def test__parse_context_length_with_empty(self):
898897
assert _parse_context_length("") == 0
899898

900899
@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
901-
@patch('boto3.client')
902-
def test__get_fine_tuning_options_filters_by_sequence_length(self, mock_boto_client, mock_get_hub_content):
900+
def test__get_fine_tuning_options_filters_by_sequence_length(self, mock_get_hub_content):
903901
mock_session = Mock()
904902
mock_session.boto_session.region_name = "us-east-1"
903+
mock_s3 = Mock()
904+
mock_s3.get_object.return_value = {
905+
"Body": Mock(read=Mock(return_value=b'{"max_length": {"default": 32768}}'))
906+
}
907+
mock_session.boto_session.client.return_value = mock_s3
905908

906909
mock_get_hub_content.return_value = {
907910
'hub_content_arn': "arn:aws:sagemaker:us-east-1:123456789012:model/test-model",
@@ -926,19 +929,13 @@ def test__get_fine_tuning_options_filters_by_sequence_length(self, mock_boto_cli
926929
}
927930
}
928931

929-
mock_s3_client = Mock()
930-
mock_boto_client.return_value = mock_s3_client
931-
mock_s3_client.get_object.return_value = {
932-
"Body": Mock(read=Mock(return_value=b'{"max_length": {"default": 32768}}'))
933-
}
934-
935932
result = _get_fine_tuning_options_and_model_arn("test-model", "SFT", "LORA", mock_session, sequence_length="8K")
936933

937934
if result is not None:
938935
options, model_arn, is_gated_model = result
939936
# Should pick the 32K recipe (smallest >= 8K)
940-
mock_s3_client.get_object.assert_called_once()
941-
call_args = mock_s3_client.get_object.call_args[1]
937+
mock_s3.get_object.assert_called_once()
938+
call_args = mock_s3.get_object.call_args[1]
942939
assert "params-32k" in call_args["Key"]
943940

944941
@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')

0 commit comments

Comments
 (0)