@@ -815,7 +815,135 @@ def test_train_with_pipeline_session_does_not_launch_job(
815815
816816 result = trainer .train ()
817817
818+ from sagemaker .core .workflow .pipeline_context import _StepArguments
819+ # @runnable_by_pipeline intercepts and returns _StepArguments
820+ assert isinstance (result , _StepArguments )
821+ assert result .caller_name == "train"
822+ assert result .func is not None
823+ assert result .func_args [0 ] is trainer
818824 mock_training_job_create .assert_not_called ()
819- pipeline_session ._intercept_create_request .assert_called_once ()
820- assert pipeline_session ._intercept_create_request .call_args [0 ][2 ] == "train"
821- assert result == step_args
825+
826+ @patch ('sagemaker.train.dpo_trainer._create_model_package_config' )
827+ @patch ('sagemaker.train.dpo_trainer._create_mlflow_config' )
828+ @patch ('sagemaker.train.dpo_trainer._create_output_config' )
829+ @patch ('sagemaker.train.dpo_trainer._create_serverless_config' )
830+ @patch ('sagemaker.train.dpo_trainer._convert_input_data_to_channels' )
831+ @patch ('sagemaker.train.dpo_trainer._create_input_data_config' )
832+ @patch ('sagemaker.train.dpo_trainer._get_unique_name' )
833+ @patch ('sagemaker.train.dpo_trainer.TrainDefaults.get_role' )
834+ @patch ('sagemaker.train.dpo_trainer.TrainDefaults.get_sagemaker_session' )
835+ @patch ('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group' )
836+ @patch ('sagemaker.train.dpo_trainer._get_fine_tuning_options_and_model_arn' )
837+ @patch ('sagemaker.train.dpo_trainer._resolve_model_and_name' )
838+ @patch ('sagemaker.train.common_utils.finetune_utils._get_beta_session' )
839+ @patch ('sagemaker.core.resources.TrainingJob.create' )
840+ def test_train_pipeline_session_produces_valid_step_arguments (
841+ self , mock_training_job_create , mock_beta_session , mock_resolve_model ,
842+ mock_finetuning_options , mock_validate_group , mock_get_session , mock_get_role ,
843+ mock_unique_name , mock_input_config , mock_convert_channels ,
844+ mock_serverless_config , mock_output_config , mock_mlflow_config , mock_model_package_config ,
845+ ):
846+ """TrainingStep.arguments produces valid PascalCase dict."""
847+ from sagemaker .train .dpo_trainer import DPOTrainer
848+ from sagemaker .core .workflow .pipeline_context import PipelineSession , _StepArguments
849+ from sagemaker .mlops .workflow .steps import TrainingStep
850+
851+ pipeline_session = PipelineSession .__new__ (PipelineSession )
852+ pipeline_session ._context = None
853+ pipeline_session .boto_session = Mock ()
854+ pipeline_session .boto_session .region_name = "us-west-2"
855+ mock_get_session .return_value = pipeline_session
856+
857+ mock_resolve_model .return_value = ("test-model" , "test" )
858+ mock_hyperparams = Mock ()
859+ mock_hyperparams .to_dict .return_value = {"lr" : "0.001" }
860+ mock_hyperparams ._specs = {"lr" : {"type" : "string" }}
861+ mock_hyperparams ._user_set = set ()
862+ mock_finetuning_options .return_value = (mock_hyperparams , "arn:model" , False )
863+ mock_validate_group .return_value = "grp"
864+ mock_get_role .return_value = "arn:aws:iam::123:role/Role"
865+ mock_unique_name .return_value = "test-job"
866+ mock_input_config .return_value = [{"DataSource" : {"S3DataSource" : {"S3Uri" : "s3://data" }}}]
867+ mock_convert_channels .return_value = [{"ChannelName" : "train" }]
868+ mock_serverless_config .return_value = {"BaseModelArn" : "arn:model" , "JobType" : "FineTuning" }
869+ mock_output_config .return_value = {"S3OutputPath" : "s3://output" }
870+ mock_mlflow_config .return_value = None
871+ mock_model_package_config .return_value = None
872+ mock_beta_session .return_value = pipeline_session
873+
874+ trainer = DPOTrainer (model = "test-model" , training_dataset = "s3://bucket/data" , model_package_group = "grp" , sagemaker_session = pipeline_session )
875+ trainer ._model_arn = "arn:model"
876+ trainer ._model_name = "test-model"
877+ trainer .accept_eula = True
878+ trainer .hyperparameters = mock_hyperparams
879+
880+ result = trainer .train ()
881+ step = TrainingStep (name = "TestDPO" , step_args = result )
882+ arguments = step .arguments
883+
884+ assert isinstance (arguments , dict )
885+ assert "session" not in arguments
886+ assert "region" not in arguments
887+ tags = arguments .get ("Tags" , [])
888+ for t in tags :
889+ assert "Key" in t and "Value" in t
890+
891+ @patch ('sagemaker.train.dpo_trainer._create_model_package_config' )
892+ @patch ('sagemaker.train.dpo_trainer._create_mlflow_config' )
893+ @patch ('sagemaker.train.dpo_trainer._create_output_config' )
894+ @patch ('sagemaker.train.dpo_trainer._create_serverless_config' )
895+ @patch ('sagemaker.train.dpo_trainer._convert_input_data_to_channels' )
896+ @patch ('sagemaker.train.dpo_trainer._create_input_data_config' )
897+ @patch ('sagemaker.train.dpo_trainer._get_unique_name' )
898+ @patch ('sagemaker.train.dpo_trainer.TrainDefaults.get_role' )
899+ @patch ('sagemaker.train.dpo_trainer.TrainDefaults.get_sagemaker_session' )
900+ @patch ('sagemaker.train.dpo_trainer._validate_and_resolve_model_package_group' )
901+ @patch ('sagemaker.train.dpo_trainer._get_fine_tuning_options_and_model_arn' )
902+ @patch ('sagemaker.train.dpo_trainer._resolve_model_and_name' )
903+ @patch ('sagemaker.train.common_utils.finetune_utils._get_beta_session' )
904+ @patch ('sagemaker.train.common_utils.data_utils.validate_data_path_exists' )
905+ @patch ('sagemaker.core.resources.TrainingJob.create' )
906+ def test_train_without_pipeline_session_launches_job (
907+ self , mock_training_job_create , mock_validate_path , mock_beta_session ,
908+ mock_resolve_model , mock_finetuning_options , mock_validate_group ,
909+ mock_get_session , mock_get_role , mock_unique_name , mock_input_config ,
910+ mock_convert_channels , mock_serverless_config , mock_output_config ,
911+ mock_mlflow_config , mock_model_package_config ,
912+ ):
913+ """Regular Session launches job normally."""
914+ from sagemaker .train .dpo_trainer import DPOTrainer
915+
916+ regular_session = Mock ()
917+ regular_session .boto_session = Mock ()
918+ regular_session .boto_session .region_name = "us-west-2"
919+ regular_session .sagemaker_config = {}
920+ mock_get_session .return_value = regular_session
921+
922+ mock_resolve_model .return_value = ("test-model" , "test" )
923+ mock_hyperparams = Mock ()
924+ mock_hyperparams .to_dict .return_value = {"lr" : "0.001" }
925+ mock_hyperparams ._specs = {"lr" : {"type" : "string" }}
926+ mock_hyperparams ._user_set = set ()
927+ mock_finetuning_options .return_value = (mock_hyperparams , "arn:model" , False )
928+ mock_validate_group .return_value = "grp"
929+ mock_get_role .return_value = "arn:aws:iam::123:role/Role"
930+ mock_unique_name .return_value = "test-job"
931+ mock_input_config .return_value = {}
932+ mock_convert_channels .return_value = []
933+ mock_serverless_config .return_value = {"BaseModelArn" : "arn:model" }
934+ mock_output_config .return_value = {"S3OutputPath" : "s3://output" }
935+ mock_mlflow_config .return_value = None
936+ mock_model_package_config .return_value = None
937+ mock_beta_session .return_value = regular_session
938+ mock_training_job = Mock ()
939+ mock_training_job_create .return_value = mock_training_job
940+
941+ trainer = DPOTrainer (model = "test-model" , training_dataset = "s3://bucket/data" , model_package_group = "grp" , sagemaker_session = regular_session )
942+ trainer ._model_arn = "arn:model"
943+ trainer ._model_name = "test-model"
944+ trainer .accept_eula = True
945+ trainer .hyperparameters = mock_hyperparams
946+
947+ result = trainer .train (wait = False )
948+ mock_training_job_create .assert_called_once ()
949+ assert result == mock_training_job
0 commit comments