fix: Pass kms_key to _upload_analysis_config when provided (#4033)

Author: keerthanvasist

src/sagemaker/model_monitor/clarify_model_monitoring.py

@@ -178,7 +178,7 @@ def _create_baselining_processor(self):
         baselining_processor.base_job_name = self.base_job_name
         return baselining_processor
 
-    def _upload_analysis_config(self, analysis_config, output_s3_uri, job_definition_name):
+    def _upload_analysis_config(self, analysis_config, output_s3_uri, job_definition_name, kms_key):
         """Upload analysis config to s3://<output path>/<job name>/analysis_config.json
 
         Args:
@@ -187,6 +187,8 @@ def _upload_analysis_config(self, analysis_config, output_s3_uri, job_definition
                 Default: "s3://<default_session_bucket>/<job_name>/output"
             job_definition_name (str): Job definition name.
                 If not specified then a default one will be generated.
+            kms_key( str): The ARN of the KMS key that is used to encrypt the
+            user code file (default: None).
 
         Returns:
             str: The S3 uri of the uploaded file(s).
@@ -202,6 +204,7 @@ def _upload_analysis_config(self, analysis_config, output_s3_uri, job_definition
             json.dumps(analysis_config),
             desired_s3_uri=s3_uri,
             sagemaker_session=self.sagemaker_session,
+            kms_key=kms_key,
         )
 
     def _build_create_job_definition_request(
@@ -323,7 +326,7 @@ def _build_create_job_definition_request(
             analysis_config_uri = analysis_config
         else:
             analysis_config_uri = self._upload_analysis_config(
-                analysis_config._to_dict(), output_s3_uri, job_definition_name
+                analysis_config._to_dict(), output_s3_uri, job_definition_name, output_kms_key
             )
         app_specification["ConfigUri"] = analysis_config_uri
         app_specification["ImageUri"] = image_uri

src/sagemaker/workflow/clarify_check_step.py

@@ -435,7 +435,7 @@ def _upload_monitoring_analysis_config(self) -> str:
             job_definition_name = name_from_base(f"{_BIAS_MONITORING_CFG_BASE_NAME}-config")
 
         return self._model_monitor._upload_analysis_config(
-            analysis_config, output_s3_uri, job_definition_name
+            analysis_config, output_s3_uri, job_definition_name, self.clarify_check_config.kms_key
         )
 
     def _get_s3_base_uri_for_monitoring_analysis_config(self) -> str: