feature: add 'Domain' property to RegisterModel step (#3118)

Author: rahven14

src/sagemaker/estimator.py

@@ -1280,6 +1280,7 @@ def register(
         model_name=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
         **kwargs,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
@@ -1311,6 +1312,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during
                 deploy. For more, see the implementation docs.
@@ -1342,6 +1345,7 @@ def register(
             description,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
         )
 
     @property

src/sagemaker/huggingface/model.py

@@ -304,6 +304,7 @@ def register(
         approval_status=None,
         description=None,
         drift_check_baselines=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -331,6 +332,8 @@ def register(
                 or "PendingManualApproval". Defaults to ``PendingManualApproval``.
             description (str): Model Package description. Defaults to ``None``.
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -359,6 +362,7 @@ def register(
             approval_status,
             description,
             drift_check_baselines=drift_check_baselines,
+            domain=domain,
         )
 
     def prepare_container_def(

src/sagemaker/model.py

@@ -309,6 +309,7 @@ def register(
         drift_check_baselines=None,
         customer_metadata_properties=None,
         validation_specification=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -336,6 +337,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -365,6 +368,7 @@ def register(
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
             validation_specification=validation_specification,
+            domain=domain,
         )
         model_package = self.sagemaker_session.create_model_package_from_containers(
             **model_pkg_args

src/sagemaker/mxnet/model.py

@@ -158,6 +158,7 @@ def register(
         description=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -185,6 +186,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -214,6 +217,7 @@ def register(
             description,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
         )
 
     def prepare_container_def(

src/sagemaker/pytorch/model.py

@@ -159,6 +159,7 @@ def register(
         description=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -186,6 +187,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -215,6 +218,7 @@ def register(
             description,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
         )
 
     def prepare_container_def(

src/sagemaker/session.py

@@ -2803,6 +2803,7 @@ def create_model_package_from_containers(
         drift_check_baselines=None,
         customer_metadata_properties=None,
         validation_specification=None,
+        domain=None,
     ):
         """Get request dictionary for CreateModelPackage API.
 
@@ -2830,6 +2831,8 @@ def create_model_package_from_containers(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
         """
 
         model_pkg_request = get_create_model_package_request(
@@ -2848,6 +2851,7 @@ def create_model_package_from_containers(
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
             validation_specification=validation_specification,
+            domain=domain,
         )
 
         def submit(request):
@@ -4218,6 +4222,7 @@ def get_model_package_args(
     drift_check_baselines=None,
     customer_metadata_properties=None,
     validation_specification=None,
+    domain=None,
 ):
     """Get arguments for create_model_package method.
 
@@ -4248,6 +4253,8 @@ def get_model_package_args(
         drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
         customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
             metadata properties (default: None).
+        domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+            "MACHINE_LEARNING" (default: None).
     Returns:
         dict: A dictionary of method argument names and values.
     """
@@ -4289,6 +4296,8 @@ def get_model_package_args(
         model_package_args["customer_metadata_properties"] = customer_metadata_properties
     if validation_specification is not None:
         model_package_args["validation_specification"] = validation_specification
+    if domain is not None:
+        model_package_args["domain"] = domain
     return model_package_args
 
 
@@ -4309,6 +4318,7 @@ def get_create_model_package_request(
     drift_check_baselines=None,
     customer_metadata_properties=None,
     validation_specification=None,
+    domain=None,
 ):
     """Get request dictionary for CreateModelPackage API.
 
@@ -4337,6 +4347,8 @@ def get_create_model_package_request(
         drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
         customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
             metadata properties (default: None).
+        domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+            "MACHINE_LEARNING" (default: None).
     """
 
     if all([model_package_name, model_package_group_name]):
@@ -4362,6 +4374,8 @@ def get_create_model_package_request(
         request_dict["CustomerMetadataProperties"] = customer_metadata_properties
     if validation_specification:
         request_dict["ValidationSpecification"] = validation_specification
+    if domain is not None:
+        request_dict["Domain"] = domain
     if containers is not None:
         if not all([content_types, response_types, inference_instances, transform_instances]):
             raise ValueError(

src/sagemaker/sklearn/model.py

@@ -151,6 +151,7 @@ def register(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -175,6 +176,8 @@ def register(
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval" (default: "PendingManualApproval").
             description (str): Model Package description (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -202,6 +205,7 @@ def register(
             marketplace_cert,
             approval_status,
             description,
+            domain=domain,
         )
 
     def prepare_container_def(

src/sagemaker/tensorflow/model.py

@@ -205,6 +205,7 @@ def register(
         description=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -232,7 +233,8 @@ def register(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
-
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -262,6 +264,7 @@ def register(
             description,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
         )
 
     def deploy(

src/sagemaker/workflow/_utils.py

@@ -284,6 +284,7 @@ def __init__(
         container_def_list=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -326,6 +327,8 @@ def __init__(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
             **kwargs: additional arguments to `create_model`.
         """
         super(_RegisterModelStep, self).__init__(
@@ -356,6 +359,7 @@ def __init__(
         self.model_metrics = model_metrics
         self.drift_check_baselines = drift_check_baselines
         self.customer_metadata_properties = customer_metadata_properties
+        self.domain = domain
         self.metadata_properties = metadata_properties
         self.approval_status = approval_status
         self.image_uri = image_uri
@@ -433,6 +437,7 @@ def arguments(self) -> RequestType:
                 tags=self.tags,
                 container_def_list=self.container_def_list,
                 customer_metadata_properties=self.customer_metadata_properties,
+                domain=self.domain,
             )
 
             request_dict = get_create_model_package_request(**model_package_args)

src/sagemaker/workflow/step_collections.py

@@ -81,6 +81,7 @@ def __init__(
         model: Union[Model, PipelineModel] = None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
         **kwargs,
     ):
         """Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
@@ -122,6 +123,8 @@ def __init__(
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).
+            domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
+                "MACHINE_LEARNING" (default: None).
 
             **kwargs: additional arguments to `create_model`.
         """
@@ -241,6 +244,7 @@ def __init__(
             container_def_list=self.container_def_list,
             retry_policies=register_model_step_retry_policies,
             customer_metadata_properties=customer_metadata_properties,
+            domain=domain,
             **kwargs,
         )
         if not repack_model:

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

@@ -550,6 +550,7 @@ def test_model_registration_with_drift_check_baselines(
         ),
     )
     customer_metadata_properties = {"key1": "value1"}
+    domain = "COMPUTER_VISION"
     estimator = XGBoost(
         entry_point="training.py",
         source_dir=os.path.join(DATA_DIR, "sip"),
@@ -572,6 +573,7 @@ def test_model_registration_with_drift_check_baselines(
         model_metrics=model_metrics,
         drift_check_baselines=drift_check_baselines,
         customer_metadata_properties=customer_metadata_properties,
+        domain=domain,
     )
 
     pipeline = Pipeline(
@@ -643,6 +645,7 @@ def test_model_registration_with_drift_check_baselines(
                 == "application/json"
             )
             assert response["CustomerMetadataProperties"] == customer_metadata_properties
+            assert response["Domain"] == domain
             break
     finally:
         try:

tests/unit/test_session.py

@@ -2386,6 +2386,7 @@ def test_create_model_package_from_containers_all_args(sagemaker_session):
     approval_status = ("Approved",)
     description = "description"
     customer_metadata_properties = {"key1": "value1"}
+    domain = "COMPUTER_VISION"
     sagemaker_session.create_model_package_from_containers(
         containers=containers,
         content_types=content_types,
@@ -2400,6 +2401,7 @@ def test_create_model_package_from_containers_all_args(sagemaker_session):
         description=description,
         drift_check_baselines=drift_check_baselines,
         customer_metadata_properties=customer_metadata_properties,
+        domain=domain,
     )
     expected_args = {
         "ModelPackageName": model_package_name,
@@ -2417,6 +2419,7 @@ def test_create_model_package_from_containers_all_args(sagemaker_session):
         "ModelApprovalStatus": approval_status,
         "DriftCheckBaselines": drift_check_baselines,
         "CustomerMetadataProperties": customer_metadata_properties,
+        "Domain": domain,
     }
     sagemaker_session.sagemaker_client.create_model_package.assert_called_with(**expected_args)