vincbeck commented on code in PR #28024:
URL: https://github.com/apache/airflow/pull/28024#discussion_r1044731751
##########
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 guess it is question of opinion but my personal opinion is we should list
all parameters regardless of if they are being used in that class or in an
underlying class
##########
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"
Review Comment:
Agree
##########
tests/system/providers/amazon/aws/example_sagemaker.py:
##########
@@ -539,12 +564,14 @@ def delete_logs(env_id):
train_model,
await_training,
create_model,
+ register_model,
tune_model,
await_tuning,
test_model,
await_transform,
# TEST TEARDOWN
delete_ecr_repository(test_setup["ecr_repository_name"]),
+ delete_model_group(test_setup["model_package_group_name"],
register_model.output),
Review Comment:
In that case, this task will fail as well. This happens a lot in AWS system
tests with teardown tasks. Another (and more simple) example is, a S3 bucket is
created at the beginning of a system test, there is a teardown task to delete
this bucket at the end of the system test. What happens if the bucket does not
even get created because an exception is thrown before or during the creation?
The deletion fails as well. So far we are okay with this behavior
##########
tests/system/providers/amazon/aws/example_sagemaker.py:
##########
@@ -421,6 +429,14 @@ def delete_logs(env_id):
purge_logs(generated_logs)
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_model_group(group_name, model_version_arn):
+ sgmk_client = boto3.client("sagemaker")
+ # need to destroy model registered in group first
+ sgmk_client.delete_model_package(ModelPackageName=model_version_arn)
+ sgmk_client.delete_model_package_group(ModelPackageGroupName=group_name)
+
Review Comment:
No plan yet. We might want to create this operator later but again, there is
no plan as of today
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
