fix: Remove sagemaker_job_name from hyperparameters in TrainingStep

Payton Staub · Payton Staub · commit 298f965ab21b · 2022-02-22T15:52:00.000-08:00
diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py
@@ -301,6 +301,7 @@ def arguments(self) -> RequestType:
         )
         request_dict = self.estimator.sagemaker_session._get_train_request(**train_args)
         request_dict.pop("TrainingJobName")
+        request_dict["HyperParameters"].pop("sagemaker_job_name", None)
 
         return request_dict
 
diff --git 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 == {

Original file line number	Diff line number	Diff line change
`@@ -301,6 +301,7 @@ def arguments(self) -> RequestType:`
`301`	`301`	`)`
`302`	`302`	`request_dict = self.estimator.sagemaker_session._get_train_request(**train_args)`
`303`	`303`	`request_dict.pop("TrainingJobName")`
	`304`	`+ request_dict["HyperParameters"].pop("sagemaker_job_name", None)`
`304`	`305`
`305`	`306`	`return request_dict`
`306`	`307`
Original file line number	Diff line number	Diff line change
`@@ -399,7 +399,6 @@ def test_training_step_tensorflow(sagemaker_session):`
`399`	`399`	`name="MyTrainingStep", estimator=estimator, inputs=inputs, cache_config=cache_config`
`400`	`400`	`)`
`401`	`401`	`step_request = step.to_request()`
`402`		`- step_request["Arguments"]["HyperParameters"].pop("sagemaker_job_name", None)`
`403`	`402`	`step_request["Arguments"]["HyperParameters"].pop("sagemaker_program", None)`
`404`	`403`	`step_request["Arguments"].pop("ProfilerRuleConfigurations", None)`
`405`	`404`	`assert step_request == {`