TF Framework support for encrypting code upload

[gnbk-aws]

System Information

• Framework (e.g. TensorFlow) / Algorithm (e.g. KMeans): TensorFlow Framework
• Framework Version: 1.10,1.11
• Python Version: 3.x
• CPU or GPU: GPU
• Python SDK Version: latest
• Are you using a custom image: no

Describe the problem

A large data science org has a requirement for encrypted S3 buckets. BlazingText works effectively and pushes data to encrypted S3 buckets. With the same parameters as BlazingText (role, session, subnets, security groups, and output_path) plus a few others for the Framework specifically (code_location, output_kms_key), I am getting a putObject error (access denied).

Minimal repro / logs

The problem is specifically in the process _prepare_for_training -> _stage_user_code_in_s3 -> tar_and_upload_dir -> object_upload_file. It looks like the Boto3 command here is not using the KMS key, even when one is provided to the framework. Here's the specific line:

sagemaker-python-sdk/src/sagemaker/fw_utils.py

Line 180 in ddfc301

session.resource('s3').Object(bucket, key).upload_file(tar_file)

The upload_file command could have extra ServerSideEncryption and SSEMKSKeyId args for the S3 Transfer Manager.

• Exact command to reproduce:

estimator = TensorFlow(entry_point = 'model_file.py', role = role, output_path = 's3://encryptedbucket/key', code_location = 's3://encryptedbucket/code_key', hyperparameters = hyperparameter_dict, training_steps = training_steps, evaluation_steps = evaluation_steps, train_instance_count = 1, train_instance_type = 'ml.p3.2xlarge', sagemaker_session = sagemaker_session, subnets = ['subnet'], security_group_ids = ['securitygroup'], output_kms_ket = 'alias/aws/s3', train_volume_kms_key = 'alias/aws/s3')

estimator.fit(inputs = 's3://encryptedbucket/training_key')

Please let me know if you need any further information--thank you!

[ChoiByungWook]

Hello @gnbk-aws,

Thank you for the suggestion. I think this change would require us to modify at the very least the Estimator class to take in an additional parameter for the s3_kms_key: https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/estimator.py#L51

I'll relay this to the team as a feature request, however I am unable to provide any ETA.

[gnbk-aws]

Hi @ChoiByungWook,

Thanks for the prompt response!

I noticed that the estimator has a "outout_kms_key" parameter used for encrypting training results.

sagemaker-python-sdk/src/sagemaker/estimator.py

Line 80 in bf6805f

output_kms_key (str): Optional. KMS key ID for encrypting the training output (default: None).

I believe this is how estimators provide encrypted training outputs to S3, and it could use the same for staging the user code. From _stage_user_code_in_s3 if you pass the existing KMS key param to tar_and_upload_dir, then we might have a quick fix.

Edit: I forked the repo and committed these edits to be a bit more explicit. It's a draft, of course, and I haven't tested it yet but the commit is here: gnbk-aws/sagemaker-python-sdk@74d2451

[mvsusp]

Hi @gnbk-aws ,

Thanks for your code change. Unfortunately, for that change to work, it would be necessary to send the kms to the training container as well. Another possible issue is that the output bucket and the code bucket are different with potentially different kms keys.

We have your feature request in the roadmap.

Thanks for using SageMaker!

[gnbk-aws]

Hi @mvsusp ,

That makes sense. Since the buckets are distinct, you'd need a new framework param for the kms key of the upload_code bucket. And I see how the training containers need that key in order to pull the model.py file from S3, so you'd need to send it to the training containers with the other framework params like here:

sagemaker-python-sdk/src/sagemaker/estimator.py

Line 852 in 8b33a30

self._hyperparameters[DIR_PARAM_NAME] = code_dir

I'm not sure what would need to be modified within the training container from there. Perhaps I'll look around for that later.

Thanks for following-through here :)

[mvsusp]

Hi @gnbk-aws,

The code changes to support server side encryption is already pushed to the master branch. We are expected to release a new version of the SDK today or tomorrow.

[gnbk-aws]

This is great news--thanks @mvsusp !

[mvsusp]

Hi @gnbk-aws,

We just released sagemaker 1.18.5 including Server Side Encryption support.

Thanks for using SageMaker.