feature: adding 'Domain' property to RegisterModel step

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/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.
 
@@ -4362,6 +4372,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/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

@@ -280,6 +280,7 @@ def __init__(
         container_def_list=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        domain=None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -321,6 +322,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__(
@@ -351,6 +354,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
@@ -428,6 +432,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

@@ -78,6 +78,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.
@@ -118,6 +119,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`.
         """
@@ -236,6 +239,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)