fix: add description parameter for RegisterModelStep

src/sagemaker/workflow/_utils.py

@@ -220,6 +220,7 @@ def __init__(
         approval_status="PendingManualApproval",
         image_uri=None,
         compile_model_family=None,
+        description=None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -246,6 +247,7 @@ def __init__(
                 Estimator's training container image will be used (default: None).
             compile_model_family (str): Instance family for compiled model, if specified, a compiled
                 model will be used (default: None).
+            description (str): Model Package description (default: None).
             **kwargs: additional arguments to `create_model`.
         """
         super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL)
@@ -261,6 +263,7 @@ def __init__(
         self.approval_status = approval_status
         self.image_uri = image_uri
         self.compile_model_family = compile_model_family
+        self.description = description
         self.kwargs = kwargs
 
         self._properties = Properties(
@@ -314,6 +317,7 @@ def arguments(self) -> RequestType:
             model_metrics=self.model_metrics,
             metadata_properties=self.metadata_properties,
             approval_status=self.approval_status,
+            description=self.description,
         )
         request_dict = model.sagemaker_session._get_create_model_package_request(
             **model_package_args

src/sagemaker/workflow/step_collections.py

@@ -65,6 +65,7 @@ def __init__(
         approval_status=None,
         image_uri=None,
         compile_model_family=None,
+        description=None,
         **kwargs,
     ):
         """Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
@@ -89,6 +90,7 @@ def __init__(
                 Estimator's training container image is used (default: None).
             compile_model_family (str): The instance family for the compiled model. If
                 specified, a compiled model is used (default: None).
+            description (str): Model Package description (default: None).
             **kwargs: additional arguments to `create_model`.
         """
         steps: List[Step] = []
@@ -120,6 +122,7 @@ def __init__(
             approval_status=approval_status,
             image_uri=image_uri,
             compile_model_family=compile_model_family,
+            description=description,
             **kwargs,
         )
         steps.append(register_model_step)

tests/integ/test_workflow.py

@@ -393,6 +393,7 @@ def test_conditional_pytorch_training_model_registration(
         response_types=["*"],
         inference_instances=["*"],
         transform_instances=["*"],
+        description="test-description",
     )
 
     model = Model(

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -143,6 +143,7 @@ def test_register_model(estimator, model_metrics):
         model_package_group_name="mpg",
         model_metrics=model_metrics,
         approval_status="Approved",
+        description="description",
     )
     assert ordered(register_model.request_dicts()) == ordered(
         [
@@ -168,6 +169,7 @@ def test_register_model(estimator, model_metrics):
                             },
                         },
                     },
+                    "ModelPackageDescription": "description",
                     "ModelPackageGroupName": "mpg",
                 },
             },