Add OIDC Recipes - #585
Conversation
|
|
24291a3 to
f7770a6
Compare
| AWS KMS is designed to protect encryption keys, not to encrypt large model files directly. KMS `Encrypt` and `Decrypt` operations have small payload limits, while model weights may be many gigabytes. | ||
|
|
||
| Envelope encryption separates the work: | ||
|
|
||
| - A random **data key** encrypts the large payload locally. | ||
| - A KMS **key-encryption key** encrypts, or wraps, that data key. | ||
| - The encrypted payload and encrypted data key can be stored together safely. | ||
| - A workload must be authorized by KMS before it can recover the plaintext data key. |
There was a problem hiding this comment.
"Why envelope encryption?" should answer why the customer wants this, which is not really what you're saying here. The goal is to protect the weights from exfiltration or misuse. Envelope encryption accomplishes this two ways:
- Customer owns the key lifecycle, including access management. They can revoke access at any time, rendering the encrypted weights unusable. And they get third party access logs (from AWS) tracking all reads of the DEK.
- Minimizes exposure of decrypted weights to only the model pod, dramatically reducing attack surface area
There was a problem hiding this comment.
That's fair in the context of what if does for a B10 customer. How does this sound?
----
Envelope encryption protects proprietary model weights from exfiltration and misuse in two ways:
- Customer-controlled access and key lifecycle. The customer owns the KMS key and its access policy, and can revoke access at any time. Without permission to decrypt the data-encryption key (DEK), a copy of the encrypted weights is unusable. AWS CloudTrail also provides an independent audit trail of every KMS
Decryptrequest used to recover the DEK. - Minimal plaintext exposure. Storage, transfer, and mirroring systems handle only encrypted artifacts. The weights are decrypted only inside the authorized model pod, significantly reducing the systems and identities that can access the plaintext weights.
| -> /tmp/decrypted-weights/weights.json | ||
| ``` | ||
|
|
||
| The OIDC token is used to obtain temporary AWS credentials; it is not sent in the KMS request body. The encrypted model weights are never sent to KMS. Only the small encrypted data key and its encryption context are sent to KMS, and the returned plaintext data key is used inside the model container. |
There was a problem hiding this comment.
Only the small encrypted data key and its encryption context are sent to KMS
The DEK originates from KSM, it's not sent to KMS, correct?
There was a problem hiding this comment.
No, I think this is the process:
- The encrypted weights and encrypted DEK are stored/mounted together, e.g. by BDN.
- At runtime, the model container uses OIDC to assume an AWS role / obtain temporary AWS credentials.
- Then the container sends the encrypted DEK plus encryption context to kms:Decrypt.
- KMS returns the plaintext DEK.
- The container uses that plaintext DEK locally to decrypt the model weights.
This is what happens in code:
def decrypt_weights(output: Path) -> Path:
envelope = read_envelope()
encrypted_data_key = ENCRYPTED_DATA_KEY_FILE.read_bytes()
ciphertext = ENCRYPTED_WEIGHTS_FILE.read_bytes()
response = get_kms_client().decrypt(
CiphertextBlob=encrypted_data_key,
EncryptionContext=envelope["encryption_context"],
)
data_key = response["Plaintext"]
We get the encrypted_data_key sent (e.g. by BDN), send it to KMS, and KMS sends us back the decrypted version.
| ## Other use cases | ||
|
|
||
| The same pattern applies whenever a model must process sensitive payloads but should not carry a long-lived decryption credential. The snippets below focus on business logic and assume `decrypt_envelope()` and `encrypt_envelope()` implement the same authenticated envelope format. | ||
|
|
||
| - Decrypt an inference request supplied by a customer application | ||
|
|
||
| ```python | ||
| def predict(self, model_input): | ||
| request = json.loads(self.decrypt_envelope(model_input["encrypted_request"])) | ||
| return self.model.generate(request["prompt"]) | ||
| ``` |
There was a problem hiding this comment.
IMO it's worth splitting out envelope payload encryption as a separate use case e2e since the setup and inference client code are separate, and the motivations behind it are separate too (some customers care more about weight security, others more about payload security). Maybe too much to take on right now, but as-is, it's a bit burried
| @@ -0,0 +1,315 @@ | |||
| #!/usr/bin/env bash | |||
There was a problem hiding this comment.
Would this be clearer and simpler in python with boto3?
There was a problem hiding this comment.
Maybe, but I riffed off a one-shot script Daniel already had in the docs that was a shell script. I think either works fine. Since the shell examples are already there, I think it's easier to just keep it.
| @@ -0,0 +1,269 @@ | |||
| # OIDC: Envelope Encryption | |||
There was a problem hiding this comment.
A few comments on this README:
1.Add a terminology section at the top. It may look like basic info, but it's worth defining up front:
Root Key
KEK (Key Encryption Key)
DEK (Data Encryption Key)
Please also state explicitly that the customer owns the Root Key and KEK in their own AWS account, and that neither ever leaves AWS.
- Call out the performance cost. Customers need to be aware of the overhead this flow introduces.
3.Clarify the ownership boundary inside the pod. The code running in the pod is customer-owned. Once the DEK is unwrapped, handling the plaintext DEK — and discarding it after use — is the customer's responsibility.
|
|
||
| Envelope encryption separates the work: | ||
|
|
||
| - A random **data key** encrypts the large payload locally. |
There was a problem hiding this comment.
use professional word like wrap or unwrap, instead of random or locally
4 participants
Context
( See https://basetenlabs.slack.com/archives/C06CZ3RSXRU/p1787240030385759)
Supporting the runtime OIDC release, we want to add a home for recipes. Since runtime OIDC is a primitive we are handing to customers, a flexible git repo folder like this one makes sense. Me (@MarkGotesman) and @pat-baseten aligned on a Zoom call that this was a fair direction.
Content
Adding today:
AWS KMS).