You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2022/12/01 01:01:39 UTC
[GitHub] [airflow] vandonr-amz opened a new pull request, #28024: Add AWS SageMaker operator to register a model's version
vandonr-amz opened a new pull request, #28024:
URL: https://github.com/apache/airflow/pull/28024
Add a SageMaker operator that allows registering a model version.
One would think that this takes a model and a version as input, but it's not quite as simple. Registering a model package is actually more like pushing a new version to prod.
I tried to make this as user-friendly as possible by expliciting the required parameters, and adding some of the most common ones, and the rest of the many possible parameters can be specified using kwargs.
The system test can be helpful in understanding what kind of inputs this expects.
Here is the official doc : https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry-version.html
And here is some more context around model versions and package groups: https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry.html
--
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
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
vandonr-amz commented on code in PR #28024:
URL: https://github.com/apache/airflow/pull/28024#discussion_r1047675077
##########
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:
come to think of it, maybe I should split that, and have a config for the base operator on one side, and the extra params for mine in an other dict
--
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
[GitHub] [airflow] eladkal merged pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
eladkal merged PR #28024:
URL: https://github.com/apache/airflow/pull/28024
--
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
[GitHub] [airflow] eladkal commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
eladkal commented on code in PR #28024:
URL: https://github.com/apache/airflow/pull/28024#discussion_r1038809101
##########
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:
I think it's best to have the enum in `utils/sagemaker.py` if not mistaken that is also how we handle other enums
##########
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:
Why are we mentioning here parameter of `SageMakerBaseOperator`?
Shouldn't that description be in the parent class?
##########
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:
What do we expect to be deleted in this phase when register_model raised exception (not the happy path)?
##########
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:
Is there intention to create operator to do that?
--
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
[GitHub] [airflow] uranusjr commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
uranusjr commented on code in PR #28024:
URL: https://github.com/apache/airflow/pull/28024#discussion_r1036743714
##########
airflow/providers/amazon/aws/hooks/sagemaker.py:
##########
@@ -1010,3 +1010,35 @@ def delete_model(self, model_name: str):
except Exception as general_error:
self.log.error("Failed to delete model, error: %s", general_error)
raise
+
+ def create_model_package_group(self, package_group_name: str, package_group_desc: str = "") -> bool:
+ """
+ Creates a Model Package Group if it does not already exist
+
+ :param package_group_name: Name of the model package group to create if not already present.
+ :param package_group_desc: Description of the model package group, if it was to be created (optional).
+
+ :return: True if the model package group was created, False if it already existed.
+ """
+ try:
+ res = self.conn.create_model_package_group(
+ ModelPackageGroupName=package_group_name,
+ ModelPackageGroupDescription=package_group_desc,
+ )
+ self.log.info(
+ "Created new Model Package Group with name %s (ARN: %s)",
+ package_group_name,
+ res["ModelPackageGroupArn"],
+ )
+ return True
+ except ClientError as e:
+ # ValidationException can also happen if the package group name contains invalid char,
+ # so we have to look at the error message too
+ if e.response["Error"]["Code"] == "ValidationException" and e.response["Error"][
+ "Message"
+ ].startswith("Model Package Group already exists"):
+ self.log.info(e.response["Error"]["Message"]) # log msg only so it doesn't look like an error
Review Comment:
```suggestion
self.log.info("%s", e.response["Error"]["Message"]) # log msg only so it doesn't look like an error
```
Don’t pass a variable to the logger directly.
--
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
[GitHub] [airflow] vandonr-amz commented on pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
vandonr-amz commented on PR #28024:
URL: https://github.com/apache/airflow/pull/28024#issuecomment-1355533568
hey @eladkal do you think there are more things to change on this PR ?
if not could you drop your `change requested` review ? 🙏
--
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
[GitHub] [airflow] eladkal commented on pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
eladkal commented on PR #28024:
URL: https://github.com/apache/airflow/pull/28024#issuecomment-1355539041
I'll review this week
--
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
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by 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
[GitHub] [airflow] vincbeck commented on a diff in pull request #28024: Add AWS SageMaker operator to register a model's version
Posted by GitBox <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org