fix: enable kms support for repack_model (#1061)

* fix: enable kms support for repack_model

Currently repack_model doesn't accept a kms key. This change
added a kms_key argument to the fucntion. In addition repack_model
will always use the output_kms_key inside the Estimator if it's set.

Author: icywang86rui

src/sagemaker/amazon/kmeans.py

@@ -148,7 +148,7 @@ def __init__(
         self.center_factor = center_factor
         self.eval_metrics = eval_metrics
 
-    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT, **kwargs):
         """Return a :class:`~sagemaker.amazon.kmeans.KMeansModel` referencing
         the latest s3 model data produced by this Estimator.
 
@@ -158,12 +158,14 @@ def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
                 Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            **kwargs: Additional kwargs passed to the KMeansModel constructor.
         """
         return KMeansModel(
             self.model_data,
             self.role,
             self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            **kwargs
         )
 
     def _prepare_for_training(self, records, mini_batch_size=5000, job_name=None):

src/sagemaker/amazon/lda.py

@@ -122,7 +122,7 @@ def __init__(
         self.max_iterations = max_iterations
         self.tol = tol
 
-    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT, **kwargs):
         """Return a :class:`~sagemaker.amazon.LDAModel` referencing the latest
         s3 model data produced by this Estimator.
 
@@ -132,12 +132,14 @@ def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
                 Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            **kwargs: Additional kwargs passed to the LDAModel constructor.
         """
         return LDAModel(
             self.model_data,
             self.role,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            **kwargs
         )
 
     def _prepare_for_training(  # pylint: disable=signature-differs

src/sagemaker/amazon/linear_learner.py

@@ -373,7 +373,7 @@ def __init__(
                 "value greater than 2."
             )
 
-    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
+    def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT, **kwargs):
         """Return a :class:`~sagemaker.amazon.LinearLearnerModel` referencing
         the latest s3 model data produced by this Estimator.
 
@@ -382,12 +382,14 @@ def create_model(self, vpc_config_override=VPC_CONFIG_DEFAULT):
                 the model. Default: use subnets and security groups from this Estimator.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            **kwargs: Additional kwargs passed to the LinearLearnerModel constructor.
         """
         return LinearLearnerModel(
             self.model_data,
             self.role,
             self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            **kwargs
         )
 
     def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):

src/sagemaker/chainer/estimator.py

@@ -162,6 +162,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Create a SageMaker ``ChainerModel`` object that can be deployed to an
         ``Endpoint``.
@@ -186,6 +187,7 @@ def create_model(
             dependencies (list[str]): A list of paths to directories (absolute or relative) with
                 any additional libraries that will be exported to the container.
                 If not specified, the dependencies from training are used.
+            **kwargs: Additional kwargs passed to the ChainerModel constructor.
 
         Returns:
             sagemaker.chainer.model.ChainerModel: A SageMaker ``ChainerModel``

src/sagemaker/estimator.py

@@ -547,6 +547,7 @@ def deploy(
                 )
             model = self._compiled_models[family]
         else:
+            kwargs["model_kms_key"] = self.output_kms_key
             model = self.create_model(**kwargs)
         model.name = model_name
         return model.deploy(
@@ -734,7 +735,9 @@ def transformer(
             model_name = self._current_job_name
         else:
             model_name = self.latest_training_job.name
-            model = self.create_model(vpc_config_override=vpc_config_override)
+            model = self.create_model(
+                vpc_config_override=vpc_config_override, model_kms_key=self.output_kms_key
+            )
 
             # not all create_model() implementations have the same kwargs
             model.name = model_name
@@ -1716,6 +1719,7 @@ def transformer(
                 model_server_workers=model_server_workers,
                 entry_point=entry_point,
                 vpc_config_override=vpc_config_override,
+                model_kms_key=self.output_kms_key,
             )
             model._create_sagemaker_model(instance_type, tags=tags)
 

src/sagemaker/model.py

@@ -79,6 +79,7 @@ def __init__(
         vpc_config=None,
         sagemaker_session=None,
         enable_network_isolation=False,
+        model_kms_key=None,
     ):
         """Initialize an SageMaker ``Model``.
 
@@ -114,6 +115,8 @@ def __init__(
                 network isolation in the endpoint, isolating the model
                 container. No inbound or outbound network calls can be made to
                 or from the model container.
+            model_kms_key (str): KMS key ARN used to encrypt the repacked
+                model archive file if the model is repacked
         """
         self.model_data = model_data
         self.image = image
@@ -127,6 +130,7 @@ def __init__(
         self.endpoint_name = None
         self._is_compiled_model = False
         self._enable_network_isolation = enable_network_isolation
+        self.model_kms_key = model_kms_key
 
     def prepare_container_def(
         self, instance_type, accelerator_type=None
@@ -799,6 +803,7 @@ def _upload_code(self, key_prefix, repack=False):
                 model_uri=self.model_data,
                 repacked_model_uri=repacked_model_data,
                 sagemaker_session=self.sagemaker_session,
+                kms_key=self.model_kms_key,
             )
 
             self.repacked_model_data = repacked_model_data

src/sagemaker/mxnet/estimator.py

@@ -141,6 +141,7 @@ def create_model(
         source_dir=None,
         dependencies=None,
         image_name=None,
+        **kwargs
     ):
         """Create a SageMaker ``MXNetModel`` object that can be deployed to an
         ``Endpoint``.
@@ -171,6 +172,7 @@ def create_model(
                 Examples:
                     123.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0
                     custom-image:latest.
+            **kwargs: Additional kwargs passed to the MXNetModel constructor.
 
         Returns:
             sagemaker.mxnet.model.MXNetModel: A SageMaker ``MXNetModel`` object.

src/sagemaker/pytorch/estimator.py

@@ -115,6 +115,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Create a SageMaker ``PyTorchModel`` object that can be deployed to an
         ``Endpoint``.
@@ -139,6 +140,7 @@ def create_model(
             dependencies (list[str]): A list of paths to directories (absolute or relative) with
                 any additional libraries that will be exported to the container.
                 If not specified, the dependencies from training are used.
+            **kwargs: Additional kwargs passed to the PyTorchModel constructor.
 
         Returns:
             sagemaker.pytorch.model.PyTorchModel: A SageMaker ``PyTorchModel``

src/sagemaker/rl/estimator.py

@@ -163,6 +163,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Create a SageMaker ``RLEstimatorModel`` object that can be deployed
         to an Endpoint.
@@ -189,6 +190,7 @@ def create_model(
                 folders will be copied to SageMaker in the same folder where the
                 entry_point is copied. If the ```source_dir``` points to S3,
                 code will be uploaded and the S3 location will be used instead.
+            **kwargs: Additional kwargs passed to the FrameworkModel constructor.
 
         Returns:
             sagemaker.model.FrameworkModel: Depending on input parameters returns

src/sagemaker/tensorflow/estimator.py

@@ -504,6 +504,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Create a ``Model`` object that can be used for creating SageMaker model entities,
         deploying to a SageMaker endpoint, or starting SageMaker Batch Transform jobs.
@@ -537,6 +538,8 @@ def create_model(
                 If not specified and ``endpoint_type`` is 'tensorflow-serving', ``dependencies`` is
                 set to ``None``.
                 If ``endpoint_type`` is also ``None``, then the dependencies from training are used.
+            **kwargs: Additional kwargs passed to ``sagemaker.tensorflow.serving.Model`` constructor
+                and ``sagemaker.tensorflow.model.TensorFlowModel`` constructor.
 
         Returns:
             sagemaker.tensorflow.model.TensorFlowModel or sagemaker.tensorflow.serving.Model: A
@@ -552,6 +555,7 @@ def create_model(
                 entry_point=entry_point,
                 source_dir=source_dir,
                 dependencies=dependencies,
+                **kwargs
             )
 
         return self._create_default_model(
@@ -561,6 +565,7 @@ def create_model(
             entry_point=entry_point,
             source_dir=source_dir,
             dependencies=dependencies,
+            **kwargs
         )
 
     def _create_tfs_model(
@@ -570,6 +575,7 @@ def _create_tfs_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Placeholder docstring"""
         return Model(
@@ -585,6 +591,7 @@ def _create_tfs_model(
             source_dir=source_dir,
             dependencies=dependencies,
             enable_network_isolation=self.enable_network_isolation(),
+            **kwargs
         )
 
     def _create_default_model(
@@ -595,6 +602,7 @@ def _create_default_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
+        **kwargs
     ):
         """Placeholder docstring"""
         return TensorFlowModel(
@@ -615,6 +623,7 @@ def _create_default_model(
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=dependencies or self.dependencies,
             enable_network_isolation=self.enable_network_isolation(),
+            **kwargs
         )
 
     def hyperparameters(self):