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
3535from sagemaker .ai_registry .dataset import DataSet
3636from sagemaker .train .common import TrainingType
3737from 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