Skip to content

Commit 12e81df

Browse files
author
Payton Staub
committed
Use correct boto model name for RegisterModelStep properties
1 parent 9fc5755 commit 12e81df

File tree

2 files changed

+13
-4
lines changed

2 files changed

+13
-4
lines changed

src/sagemaker/workflow/_utils.py

+2-4
Original file line numberDiff line numberDiff line change
@@ -284,9 +284,7 @@ def __init__(
284284
self.kwargs = kwargs
285285
self.container_def_list = container_def_list
286286

287-
self._properties = Properties(
288-
path=f"Steps.{name}", shape_name="DescribeModelPackageResponse"
289-
)
287+
self._properties = Properties(path=f"Steps.{name}", shape_name="DescribeModelPackageOutput")
290288

291289
@property
292290
def arguments(self) -> RequestType:
@@ -362,5 +360,5 @@ def arguments(self) -> RequestType:
362360

363361
@property
364362
def properties(self):
365-
"""A Properties object representing the DescribeTrainingJobResponse data model."""
363+
"""A Properties object representing the DescribeModelPackageOutput data model."""
366364
return self._properties

tests/unit/sagemaker/workflow/test_properties.py

+11
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,14 @@ def test_properties_tuning_job():
6767
assert prop.TrainingJobSummaries[0].TrainingJobName.expr == {
6868
"Get": "Steps.MyStep.TrainingJobSummaries[0].TrainingJobName"
6969
}
70+
71+
72+
def test_properties_describe_model_package_output():
73+
prop = Properties("Steps.MyStep", "DescribeModelPackageOutput")
74+
some_prop_names = ["ModelPackageName", "ModelPackageGroupName", "ModelPackageArn"]
75+
for name in some_prop_names:
76+
assert name in prop.__dict__.keys()
77+
assert prop.ModelPackageName.expr == {"Get": "Steps.MyStep.ModelPackageName"}
78+
assert prop.ValidationSpecification.ValidationRole.expr == {
79+
"Get": "Steps.MyStep.ValidationSpecification.ValidationRole"
80+
}

0 commit comments

Comments
 (0)