[GitHub] [airflow] vandonr-amz commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version

[GitBox <gi...@apache.org>]

vandonr-amz commented on code in PR #28024:
URL: https://github.com/apache/airflow/pull/28024#discussion_r1046290741


##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -750,3 +751,99 @@ def execute(self, context: Context) -> Any:
         sagemaker_hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
         sagemaker_hook.delete_model(model_name=self.config["ModelName"])
         self.log.info("Model %s deleted successfully.", self.config["ModelName"])
+
+
+class ApprovalStatus(Enum):
+    """Approval statuses for a Sagemaker Model Package."""
+
+    APPROVED = "Approved"
+    REJECTED = "Rejected"
+    PENDING_MANUAL_APPROVAL = "PendingManualApproval"
+
+
+class SageMakerRegisterModelVersionOperator(SageMakerBaseOperator):
+    """
+    Registers an Amazon SageMaker model by creating a model version that specifies the model group to which it
+    belongs. Will create the model group if it does not exist already.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SageMakerRegisterModelVersionOperator`
+
+    :param image_uri: The Amazon EC2 Container Registry (Amazon ECR) path where inference code is stored.
+    :param model_url: The Amazon S3 path where the model artifacts (the trained weights of the model), which
+        result from model training, are stored. This path must point to a single gzip compressed tar archive
+        (.tar.gz suffix).
+    :param package_group_name: The name of the model package group that the model is going to be registered
+        to. Will be created if it doesn't already exist.
+    :param package_group_desc: Description of the model package group, if it was to be created (optional).
+    :param package_desc: Description of the model package (optional).
+    :param model_approval: Approval status of the model package. Defaults to PendingManualApproval
+    :param aws_conn_id: The AWS connection ID to use.
+    :param config: Can contain extra parameters for the boto call to create_model_package, and/or overrides
+        for any parameter defined above. For a complete list of available parameters, see
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_model_package

Review Comment:
   I think it makes sense to document it here because it's not only used for configuring the base operator, but also to override values for stuff specific to this operator (and also, not having to piece documentation together from different places is nice from the user perspective I think ?)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org