diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py index 99f3444f23..1ef63ef915 100644 --- a/src/sagemaker/workflow/steps.py +++ b/src/sagemaker/workflow/steps.py @@ -301,6 +301,8 @@ def arguments(self) -> RequestType: ) request_dict = self.estimator.sagemaker_session._get_train_request(**train_args) request_dict.pop("TrainingJobName") + if "HyperParameters" in request_dict: + request_dict["HyperParameters"].pop("sagemaker_job_name", None) return request_dict diff --git a/tests/unit/sagemaker/workflow/test_step_collections.py b/tests/unit/sagemaker/workflow/test_step_collections.py index d2f1f07059..ea810796f4 100644 --- a/tests/unit/sagemaker/workflow/test_step_collections.py +++ b/tests/unit/sagemaker/workflow/test_step_collections.py @@ -457,7 +457,8 @@ def test_register_model_with_model_repack_with_estimator( assert len(request_dict["DependsOn"]) == 1 assert request_dict["DependsOn"][0] == "TestStep" arguments = request_dict["Arguments"] - repacker_job_name = arguments["HyperParameters"]["sagemaker_job_name"] + assert BUCKET in arguments["HyperParameters"]["sagemaker_submit_directory"] + arguments["HyperParameters"].pop("sagemaker_submit_directory") assert ordered(arguments) == ordered( { "AlgorithmSpecification": { @@ -472,12 +473,8 @@ def test_register_model_with_model_repack_with_estimator( "inference_script": '"dummy_script.py"', "dependencies": f'"{dummy_requirements}"', "model_archive": '"model.tar.gz"', - "sagemaker_submit_directory": '"s3://{}/{}/source/sourcedir.tar.gz"'.format( - BUCKET, repacker_job_name.replace('"', "") - ), "sagemaker_program": '"_repack_model.py"', "sagemaker_container_log_level": "20", - "sagemaker_job_name": repacker_job_name, "sagemaker_region": f'"{REGION}"', "source_dir": "null", }, @@ -585,7 +582,8 @@ def test_register_model_with_model_repack_with_model(model, model_metrics, drift assert len(request_dict["DependsOn"]) == 1 assert request_dict["DependsOn"][0] == "TestStep" arguments = request_dict["Arguments"] - repacker_job_name = arguments["HyperParameters"]["sagemaker_job_name"] + assert BUCKET in arguments["HyperParameters"]["sagemaker_submit_directory"] + arguments["HyperParameters"].pop("sagemaker_submit_directory") assert ordered(arguments) == ordered( { "AlgorithmSpecification": { @@ -599,12 +597,8 @@ def test_register_model_with_model_repack_with_model(model, model_metrics, drift "HyperParameters": { "inference_script": '"dummy_script.py"', "model_archive": '"model.tar.gz"', - "sagemaker_submit_directory": '"s3://{}/{}/source/sourcedir.tar.gz"'.format( - BUCKET, repacker_job_name.replace('"', "") - ), "sagemaker_program": '"_repack_model.py"', "sagemaker_container_log_level": "20", - "sagemaker_job_name": repacker_job_name, "sagemaker_region": f'"{REGION}"', "dependencies": "null", "source_dir": "null", @@ -717,7 +711,8 @@ def test_register_model_with_model_repack_with_pipeline_model( assert len(request_dict["DependsOn"]) == 1 assert request_dict["DependsOn"][0] == "TestStep" arguments = request_dict["Arguments"] - repacker_job_name = arguments["HyperParameters"]["sagemaker_job_name"] + assert BUCKET in arguments["HyperParameters"]["sagemaker_submit_directory"] + arguments["HyperParameters"].pop("sagemaker_submit_directory") assert ordered(arguments) == ordered( { "AlgorithmSpecification": { @@ -732,12 +727,8 @@ def test_register_model_with_model_repack_with_pipeline_model( "dependencies": "null", "inference_script": '"dummy_script.py"', "model_archive": '"model.tar.gz"', - "sagemaker_submit_directory": '"s3://{}/{}/source/sourcedir.tar.gz"'.format( - BUCKET, repacker_job_name.replace('"', "") - ), "sagemaker_program": '"_repack_model.py"', "sagemaker_container_log_level": "20", - "sagemaker_job_name": repacker_job_name, "sagemaker_region": f'"{REGION}"', "source_dir": "null", }, @@ -917,7 +908,6 @@ def test_estimator_transformer_with_model_repack_with_estimator(estimator): arguments = request_dict["Arguments"] # pop out the dynamic generated fields arguments["HyperParameters"].pop("sagemaker_submit_directory") - arguments["HyperParameters"].pop("sagemaker_job_name") assert arguments == { "AlgorithmSpecification": { "TrainingInputMode": "File", diff --git a/tests/unit/sagemaker/workflow/test_steps.py b/tests/unit/sagemaker/workflow/test_steps.py index 674c715617..fd3bd7d0b9 100644 --- a/tests/unit/sagemaker/workflow/test_steps.py +++ b/tests/unit/sagemaker/workflow/test_steps.py @@ -399,7 +399,6 @@ def test_training_step_tensorflow(sagemaker_session): name="MyTrainingStep", estimator=estimator, inputs=inputs, cache_config=cache_config ) step_request = step.to_request() - step_request["Arguments"]["HyperParameters"].pop("sagemaker_job_name", None) step_request["Arguments"]["HyperParameters"].pop("sagemaker_program", None) step_request["Arguments"].pop("ProfilerRuleConfigurations", None) assert step_request == {