modified Frameworks _stage_user_code_in_s3 method and the tar_and_upload_dir method in fw_utils.py to further support s3 encryption with Framework based estimators

Author: Kessler

src/sagemaker/estimator.py

@@ -872,7 +872,8 @@ def _stage_user_code_in_s3(self):
                                   s3_key_prefix=code_s3_prefix,
                                   script=self.entry_point,
                                   directory=self.source_dir,
-                                  dependencies=self.dependencies)
+                                  dependencies=self.dependencies,
+				  output_kms_key=self.output_kms_key)
 
     def _model_source_dir(self):
         """Get the appropriate value to pass as source_dir to model constructor on deploying

src/sagemaker/fw_utils.py

@@ -136,7 +136,7 @@ def validate_source_dir(script, directory):
     return True
 
 
-def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory, dependencies=None):
+def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory, dependencies=None, output_kms_key=None):
     """Package source files and upload a compress tar file to S3. The S3 location will be
     ``s3://<bucket>/s3_key_prefix/sourcedir.tar.gz``.
 
@@ -176,8 +176,10 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory, depend
         source_files = _list_files_to_compress(script, directory) + dependencies
         tar_file = sagemaker.utils.create_tar_file(source_files,
                                                    os.path.join(tmp, _TAR_SOURCE_FILENAME))
-
-        session.resource('s3').Object(bucket, key).upload_file(tar_file)
+	if output_kms_key==None:
+	        session.resource('s3').Object(bucket, key).upload_file(tar_file)
+	else:
+		session.resource('s3').Object(bucket, key).upload_file(tar_file, ExtraArgs = {'ServerSideEncryption':'aws:kms', 'SSEKMSKeyId': output_kms_key})
     finally:
         shutil.rmtree(tmp)