diff --git a/oidc-recipes/README.md b/oidc-recipes/README.md new file mode 100644 index 00000000..d3e7ecbc --- /dev/null +++ b/oidc-recipes/README.md @@ -0,0 +1,28 @@ +# OIDC Recipes + +These examples use the OIDC token available to running model code through [runtime OIDC](https://docs.baseten.co/organization/oidc#use-oidc-at-request-time). They show how a model can access customer-owned resources without storing long-lived credentials. + +The recipes use AWS, but the same pattern is available from other OIDC-compatible providers: + +- [AWS](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_providers_oidc.html) +- [GCP](https://docs.cloud.google.com/iam/docs/workload-identity-federation) +- [Azure](https://learn.microsoft.com/en-us/entra/workload-id/workload-identity-federation) +- [HashiCorp Vault](https://developer.hashicorp.com/vault/docs/auth/jwt#jwt-authentication) +- [Snowflake](https://docs.snowflake.com/en/user-guide/workload-identity-federation) +- [Databricks](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation) + +## Recipes + +- [`oidc-fetch-a-resource`](oidc-fetch-a-resource): Fetch an object from S3 at request time. +- [`oidc-envelope-encryption`](oidc-envelope-encryption): Protect model weights and inference payloads with customer-owned AWS KMS keys. + +## Folder structure + +Recipes contain: + +- `README.md`: What the recipe does and how to run it. +- `setup.sh`: Creates the example AWS resources and OIDC trust. +- `standard-truss/`: A regular Truss with `config.yaml` and `model/model.py`. +- `custom-base-image/`: An optional custom-server Truss. + +The envelope-encryption recipe has separate `weight-encryption/` and `payload-encryption/` flows. Fill in values marked `FILL ME` or left empty in `setup.sh` and `config.yaml`. Run the setup script first, copy its output into the Truss config, and then run `truss push` on the relevant Truss directory. diff --git a/oidc-recipes/oidc-envelope-encryption/README.md b/oidc-recipes/oidc-envelope-encryption/README.md new file mode 100644 index 00000000..19305698 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/README.md @@ -0,0 +1,243 @@ +# OIDC: Envelope Encryption +_Use customer-owned AWS KMS keys to protect model weights and inference payloads._ + +These recipes use a Baseten OIDC token to obtain short-lived AWS credentials. The running model can then ask AWS KMS to unwrap a data key without storing a long-lived AWS access key in the Truss. + +Two flows are included: + +- [Weight encryption](#weight-encryption): Decrypt model weights during startup. +- [Payload encryption](#payload-encryption): Decrypt requests and encrypt responses at inference time. + +## Terminology + +- **Root Key:** An optional top-level AWS KMS key that protects a KEK in a deeper key hierarchy. +- **KEK (Key Encryption Key):** The customer-managed AWS KMS key that wraps and unwraps the DEK. +- **DEK (Data Encryption Key):** The key that encrypts and decrypts the data. + +These recipes do not use a separate Root Key. The customer owns the KEK in their AWS account, and its key material never leaves AWS KMS. + +## How envelope encryption works + +AWS KMS is designed to protect keys, not to encrypt large data directly. Envelope encryption separates the work: + +- A generated DEK encrypts the data. +- The KEK wraps the DEK. +- The encrypted data and wrapped DEK can be stored or sent together. +- An authorized workload asks KMS to unwrap the DEK. + +Only the wrapped DEK and its encryption context are sent to KMS. The encrypted data is not. + +Both examples use a 64-byte DEK: + +- 32 bytes for AES-256-CBC encryption. +- 32 bytes for HMAC-SHA256 integrity protection. + +The HMAC covers `IV || ciphertext` and is verified before decryption. + +## Weight encryption + +This flow keeps model weights encrypted in a customer-owned S3 bucket, through Baseten Data Network (BDN), and in the model mount. The model decrypts them during `load()`. + +### Flow + +```text +Setup: + KMS GenerateDataKey + -> plaintext DEK encrypts weights.json + -> wrapped DEK is stored with weights.enc in S3 + +Runtime: + BDN mounts the encrypted bundle at /models/custom + -> Baseten OIDC token is exchanged with AWS STS + -> short-lived credentials call KMS Decrypt + -> KMS unwraps the DEK + -> the model verifies and decrypts weights.enc + -> /tmp/decrypted-weights/weights.json +``` + +The setup script uploads: + +```text +s3:///models/custom-weights/ +├── weights.enc +├── encrypted-data-key +└── envelope.json +``` + +The companion `envelope.json` stores the algorithm, IV, HMAC, and KMS encryption context. S3 metadata is not used because the BDN mount exposes files rather than the original S3 response metadata. + +### Setup + +Prerequisites: + +- AWS CLI v2 +- `jq` +- OpenSSL +- A Baseten organization and team with OIDC enabled + +Get the Baseten OIDC identifiers: + +```bash +truss whoami --show-oidc +``` + +Fill in the values at the top of [`weight-encryption/setup.sh`](weight-encryption/setup.sh), then run: + +```bash +cd weight-encryption +./setup.sh +``` + +The script creates or configures the S3 bucket, KMS key, OIDC provider, IAM role, and encrypted example weights. + +### Configure and deploy + +Copy the encrypted S3 source, role ARN, and region printed by the setup script into [`weight-encryption/standard-truss/config.yaml`](weight-encryption/standard-truss/config.yaml), then deploy: + +```bash +truss push ./weight-encryption/standard-truss +``` + +For the fictional linear weights, an input value of `4` produces `11`: + +```json +{ + "value": 4 +} +``` + +```json +{ + "value": 11.0 +} +``` + +### Custom base image + +For vLLM, SGLang, TensorRT-LLM, Triton, or another custom server, run [`weight-encryption/custom-base-image/packages/decrypt.py`](weight-encryption/custom-base-image/packages/decrypt.py) before the server command: + +```yaml +docker_server: + start_command: >- + sh -c 'python3 /packages/decrypt.py && + exec /tmp/decrypted-weights' +``` + +The included custom-base-image config shows the startup hook with vLLM. Its fictional JSON weights cannot be loaded by vLLM; replace them with a real encrypted model directory. + +### Performance cost + +This flow adds startup time for STS and KMS calls and for local decryption. Decryption time grows with the weight size. It adds no per-request encryption cost after startup. + +## Payload encryption + +This flow protects inference data beyond transport encryption. The client encrypts each request before sending it to Baseten. Customer-owned code in the model pod decrypts the request, runs inference, and encrypts the response before returning it. + +### Flow + +```text +Client: + KMS GenerateDataKey + -> plaintext DEK encrypts the request + -> wrapped DEK travels with the encrypted request + +Runtime: + Baseten OIDC token is exchanged with AWS STS + -> short-lived credentials call KMS Decrypt + -> KMS unwraps the DEK + -> the model verifies and decrypts the request + -> the model encrypts the response with the same DEK and a new IV + +Client: + retained plaintext DEK verifies and decrypts the response +``` + +The client sends a versioned JSON envelope containing the wrapped DEK, encryption context, IV, HMAC, and ciphertext. The response contains a new IV, HMAC, and ciphertext. Plaintext request and response data only exist in the client and the customer-owned model code. + +### Setup + +Prerequisites: + +- AWS CLI v2 +- `jq` +- Python with `boto3` and `cryptography` +- A Baseten organization and team with OIDC enabled + +Fill in the values at the top of [`payload-encryption/setup.sh`](payload-encryption/setup.sh), then run: + +```bash +cd payload-encryption +./setup.sh +``` + +The script creates or configures the KMS key, OIDC provider, and runtime IAM role. The runtime role can only call `kms:Decrypt` with the expected encryption context. The AWS identity running the client needs `kms:GenerateDataKey` permission for the same key. + +### Configure and deploy + +Copy the role ARN and region printed by the setup script into [`payload-encryption/standard-truss/config.yaml`](payload-encryption/standard-truss/config.yaml), then deploy: + +```bash +truss push ./payload-encryption/standard-truss +``` + +On the machine that sends inference requests, authenticate to AWS with an identity that can call `kms:GenerateDataKey` on the KMS key. Then, from the `oidc-envelope-encryption` directory, install the client dependencies and set the deployed model details: + +```bash +python -m pip install boto3 cryptography +export AWS_REGION=us-west-2 +export KMS_KEY_ARN=arn:aws:kms:us-west-2:123456789012:key/example +export BASETEN_API_KEY=YOUR_API_KEY +export MODEL_URL=https://model-example.api.baseten.co/production/predict +python ./payload-encryption/client.py +``` + +The example encrypts `{"value": 4}`. The model returns an encrypted response that the client decrypts to: + +```json +{ + "value": 9.0 +} +``` + +### Custom base image + +The custom-base-image example runs vLLM on an internal port and exposes an encryption proxy on port `8000`. The proxy decrypts each request, forwards the plaintext OpenAI request to vLLM over loopback, and encrypts the response. + +Copy the role ARN and region printed by the setup script into [`payload-encryption/custom-base-image/config.yaml`](payload-encryption/custom-base-image/config.yaml), then deploy: + +```bash +truss push ./payload-encryption/custom-base-image +``` + +Set `MODEL_URL` to the deployment URL ending in `/predict`. Do not use the OpenAI-compatible `/sync/v1` endpoint: the client sends an encrypted envelope, not a raw OpenAI request. The proxy decrypts the envelope and forwards its contents to vLLM's internal `/v1/completions` endpoint. + +```bash +export AWS_REGION=us-west-2 +export KMS_KEY_ARN=arn:aws:kms:us-west-2:123456789012:key/example +export BASETEN_API_KEY=YOUR_API_KEY +export MODEL_URL=https://model-example.api.baseten.co/deployment/example/predict +export PAYLOAD_JSON='{"model":"Qwen/Qwen2.5-0.5B-Instruct","prompt":"Hello","max_tokens":16}' +python ./payload-encryption/client.py +``` + +Replace the example model and request with the vLLM model you deploy. + +### Performance cost + +This flow adds a client-side KMS `GenerateDataKey` call and a model-side KMS `Decrypt` call to every request. It also adds request and response encryption time. The cost grows with payload size and affects request latency. + +## Security boundary + +The OIDC token is exchanged for temporary AWS credentials; it is not sent in the KMS request body. The KEK never leaves AWS KMS. Only the plaintext DEK leaves KMS after an authorized call. + +The code running in the pod is customer-owned. After KMS unwraps a DEK, that code is responsible for protecting it and discarding it after use. The client has the same responsibility for DEKs returned by `GenerateDataKey`. + +Do not log plaintext DEKs, decrypted data, OIDC tokens, or temporary AWS credentials. Use a separate KMS key or encryption context per environment, tenant, or model family when stronger isolation is needed. AWS CloudTrail records KMS operations. + +## Other use cases + +The same pattern can protect other model assets: + +- Load a customer-specific LoRA adapter only after KMS authorizes the replica. +- Decrypt a private tokenizer, prompt library, or retrieval index during startup. +- Revoke a deployed model's access without rebuilding or deleting its encrypted artifacts. diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/client.py b/oidc-recipes/oidc-envelope-encryption/payload-encryption/client.py new file mode 100644 index 00000000..77dc3444 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/client.py @@ -0,0 +1,110 @@ +import base64 +import hashlib +import hmac +import json +import os +import urllib.error +import urllib.request + +import boto3 +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + +ENVELOPE_VERSION = 1 +ENVELOPE_ALGORITHM = "AES-256-CBC-HMAC-SHA256" +ENCRYPTION_CONTEXT = {"purpose": "baseten-inference-payload"} + + +def encode(value: bytes) -> str: + return base64.b64encode(value).decode() + + +def decode_field(envelope: dict, field: str) -> bytes: + try: + return base64.b64decode(envelope[field], validate=True) + except (KeyError, ValueError) as error: + raise ValueError(f"Invalid envelope field: {field}.") from error + + +def encrypt_request(payload: dict) -> tuple[dict, bytes]: + kms = boto3.client("kms", region_name=os.environ["AWS_REGION"]) + response = kms.generate_data_key( + KeyId=os.environ["KMS_KEY_ARN"], + NumberOfBytes=64, + EncryptionContext=ENCRYPTION_CONTEXT, + ) + data_key = response["Plaintext"] + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = os.urandom(16) + plaintext = json.dumps(payload, separators=(",", ":")).encode() + + padder = padding.PKCS7(algorithms.AES.block_size).padder() + padded_plaintext = padder.update(plaintext) + padder.finalize() + encryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).encryptor() + ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize() + payload_hmac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + + return { + "version": ENVELOPE_VERSION, + "algorithm": ENVELOPE_ALGORITHM, + "encrypted_data_key": encode(response["CiphertextBlob"]), + "encryption_context": ENCRYPTION_CONTEXT, + "iv": encode(iv), + "hmac": encode(payload_hmac), + "ciphertext": encode(ciphertext), + }, data_key + + +def call_model(envelope: dict) -> dict: + request = urllib.request.Request( + os.environ["MODEL_URL"], + data=json.dumps(envelope).encode(), + headers={ + "Authorization": f"Api-Key {os.environ['BASETEN_API_KEY']}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(request) as response: + return json.load(response) + except urllib.error.HTTPError as error: + raise RuntimeError(error.read().decode()) from error + + +def decrypt_response(envelope: dict, data_key: bytes) -> dict: + if envelope.get("version") != ENVELOPE_VERSION: + raise ValueError("Unsupported payload envelope version.") + if envelope.get("algorithm") != ENVELOPE_ALGORITHM: + raise ValueError("Unsupported payload envelope algorithm.") + + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = decode_field(envelope, "iv") + ciphertext = decode_field(envelope, "ciphertext") + expected_mac = decode_field(envelope, "hmac") + + actual_mac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + if not hmac.compare_digest(actual_mac, expected_mac): + raise ValueError("Encrypted response failed integrity verification.") + + decryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).decryptor() + padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize() + unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() + plaintext = unpadder.update(padded_plaintext) + unpadder.finalize() + return json.loads(plaintext) + + +def main(): + payload = json.loads(os.environ.get("PAYLOAD_JSON", '{"value":4}')) + envelope, data_key = encrypt_request(payload) + try: + encrypted_response = call_model(envelope) + print(json.dumps(decrypt_response(encrypted_response, data_key), indent=2)) + finally: + del data_key + + +if __name__ == "__main__": + main() diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/config.yaml b/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/config.yaml new file mode 100644 index 00000000..f0478811 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/config.yaml @@ -0,0 +1,27 @@ +base_image: + image: vllm/vllm-openai:v0.8.5 +model_name: oidc-envelope-payload-encryption-vllm +docker_server: + start_command: >- + sh -c "vllm serve Qwen/Qwen2.5-0.5B-Instruct --port 8001 & + exec python3 -m uvicorn proxy:app --app-dir /packages + --host 0.0.0.0 --port 8000" + readiness_endpoint: /health + liveness_endpoint: /live + predict_endpoint: /v1/completions + server_port: 8000 +runtime: + oidc: + enabled: true +resources: + accelerator: L4 + use_gpu: true +requirements: + - boto3 + - cryptography + - fastapi + - httpx + - uvicorn +environment_variables: + AWS_ROLE_ARN: "" # e.g. arn:aws:iam::123456789012:role/BasetenOIDCPayloadRole + AWS_REGION: "" # e.g. us-west-2 diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/packages/proxy.py b/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/packages/proxy.py new file mode 100644 index 00000000..cccca418 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/custom-base-image/packages/proxy.py @@ -0,0 +1,133 @@ +import base64 +import hashlib +import hmac +import json +import os + +import boto3 +import httpx +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes +from fastapi import FastAPI, HTTPException +from fastapi.responses import Response + +ENVELOPE_VERSION = 1 +ENVELOPE_ALGORITHM = "AES-256-CBC-HMAC-SHA256" +VLLM_URL = "http://127.0.0.1:8001" + +app = FastAPI() + + +@app.get("/live") +async def live(): + return {"status": "ok"} + + +def get_kms_client(): + os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = os.environ["B10_OIDC_TOKEN_PATH"] + os.environ["AWS_ROLE_SESSION_NAME"] = "baseten-payload-decryption" + return boto3.client("kms", region_name=os.environ["AWS_REGION"]) + + +def decode_field(envelope: dict, field: str) -> bytes: + try: + return base64.b64decode(envelope[field], validate=True) + except (KeyError, ValueError) as error: + raise ValueError(f"Invalid envelope field: {field}.") from error + + +def decrypt_request(envelope: dict) -> tuple[dict, bytes]: + if envelope.get("version") != ENVELOPE_VERSION: + raise ValueError("Unsupported payload envelope version.") + if envelope.get("algorithm") != ENVELOPE_ALGORITHM: + raise ValueError("Unsupported payload envelope algorithm.") + + response = get_kms_client().decrypt( + CiphertextBlob=decode_field(envelope, "encrypted_data_key"), + EncryptionContext=envelope["encryption_context"], + ) + data_key = response["Plaintext"] + if len(data_key) != 64: + raise ValueError("Expected a 64-byte envelope data key from KMS.") + + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = decode_field(envelope, "iv") + ciphertext = decode_field(envelope, "ciphertext") + expected_mac = decode_field(envelope, "hmac") + + actual_mac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + if not hmac.compare_digest(actual_mac, expected_mac): + raise ValueError("Encrypted request failed integrity verification.") + + decryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).decryptor() + padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize() + unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() + plaintext = unpadder.update(padded_plaintext) + unpadder.finalize() + return json.loads(plaintext), data_key + + +def encrypt_response(payload: dict, data_key: bytes) -> dict: + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = os.urandom(16) + plaintext = json.dumps(payload, separators=(",", ":")).encode() + + padder = padding.PKCS7(algorithms.AES.block_size).padder() + padded_plaintext = padder.update(plaintext) + padder.finalize() + encryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).encryptor() + ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize() + payload_hmac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + + return { + "version": ENVELOPE_VERSION, + "algorithm": ENVELOPE_ALGORITHM, + "iv": base64.b64encode(iv).decode(), + "hmac": base64.b64encode(payload_hmac).decode(), + "ciphertext": base64.b64encode(ciphertext).decode(), + } + + +@app.get("/health") +async def health(): + try: + async with httpx.AsyncClient() as client: + response = await client.get(f"{VLLM_URL}/health") + response.raise_for_status() + except httpx.HTTPError as error: + raise HTTPException(status_code=503, detail="vLLM is not ready.") from error + return {"status": "ok"} + + +@app.get("/metrics") +async def metrics(): + try: + async with httpx.AsyncClient() as client: + response = await client.get(f"{VLLM_URL}/metrics") + response.raise_for_status() + except httpx.HTTPError: + return Response( + content="payload_proxy_vllm_ready 0\n", + media_type="text/plain; version=0.0.4", + ) + return Response(content=response.content, media_type="text/plain") + + +@app.post("/v1/completions") +async def completions(envelope: dict): + try: + request_payload, data_key = decrypt_request(envelope) + except (KeyError, TypeError, ValueError) as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + try: + async with httpx.AsyncClient(timeout=300) as client: + response = await client.post( + f"{VLLM_URL}/v1/completions", json=request_payload + ) + response.raise_for_status() + return encrypt_response(response.json(), data_key) + except httpx.HTTPError as error: + raise HTTPException(status_code=502, detail="vLLM request failed.") from error + finally: + del data_key diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/setup.sh b/oidc-recipes/oidc-envelope-encryption/payload-encryption/setup.sh new file mode 100755 index 00000000..54f4f20a --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/setup.sh @@ -0,0 +1,177 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ────────────────────────────────────────────── +# FILL ME: replace these with your values +# ────────────────────────────────────────────── +AWS_ACCOUNT_ID="" # Existing AWS account ID +AWS_REGION="" # AWS region for KMS (e.g. us-west-2) +KMS_ALIAS="" # KMS alias to create or reuse (must start with alias/) +ROLE_NAME="" # IAM role name to create or update +BASETEN_ORG_ID="" # From `truss whoami --show-oidc` +BASETEN_TEAM_ID="" # From `truss whoami --show-oidc` + +# ────────────────────────────────────────────── +# Helpers and configuration validation +# ────────────────────────────────────────────── +require_command() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "error: required command '$1' was not found" >&2 + exit 1 + fi +} + +require_value() { + if [[ -z "$2" ]]; then + echo "error: $1 must be set" >&2 + exit 1 + fi +} + +for command in aws jq; do + require_command "${command}" +done + +require_value AWS_ACCOUNT_ID "${AWS_ACCOUNT_ID}" +require_value AWS_REGION "${AWS_REGION}" +require_value KMS_ALIAS "${KMS_ALIAS}" +require_value ROLE_NAME "${ROLE_NAME}" +require_value BASETEN_ORG_ID "${BASETEN_ORG_ID}" +require_value BASETEN_TEAM_ID "${BASETEN_TEAM_ID}" + +if [[ "${KMS_ALIAS}" != alias/* ]]; then + echo "error: KMS_ALIAS must start with alias/" >&2 + exit 1 +fi + +# ────────────────────────────────────────────── +# 1. Authenticate the AWS CLI and verify the account +# ────────────────────────────────────────────── +if ! aws sts get-caller-identity --region "${AWS_REGION}" >/dev/null 2>&1; then + echo "No valid AWS CLI credentials were found." + echo " 1) Log in with AWS Console credentials (recommended)" + echo " 2) Configure an access key" + read -r -p "Choose an authentication method [1]: " auth_method + + case "${auth_method:-1}" in + 1) aws login --region "${AWS_REGION}" ;; + 2) aws configure ;; + *) + echo "error: invalid authentication method: ${auth_method}" >&2 + exit 1 + ;; + esac +fi + +caller_account=$(aws sts get-caller-identity \ + --region "${AWS_REGION}" \ + --query Account \ + --output text) +if [[ "${caller_account}" != "${AWS_ACCOUNT_ID}" ]]; then + echo "error: authenticated to AWS account ${caller_account}, expected ${AWS_ACCOUNT_ID}" >&2 + exit 1 +fi + +# ────────────────────────────────────────────── +# 2. Create or reuse the KMS key-encryption key +# ────────────────────────────────────────────── +if KMS_KEY_ARN=$(aws kms describe-key \ + --key-id "${KMS_ALIAS}" \ + --region "${AWS_REGION}" \ + --query KeyMetadata.Arn \ + --output text 2>/dev/null); then + echo "Using existing KMS key ${KMS_KEY_ARN}." +else + echo "Creating KMS key ${KMS_ALIAS}..." + KMS_KEY_ARN=$(aws kms create-key \ + --region "${AWS_REGION}" \ + --description "Envelope encryption key for Baseten inference payloads" \ + --query KeyMetadata.Arn \ + --output text) + aws kms create-alias \ + --region "${AWS_REGION}" \ + --alias-name "${KMS_ALIAS}" \ + --target-key-id "${KMS_KEY_ARN}" +fi + +# ────────────────────────────────────────────── +# 3. Register the Baseten OIDC identity provider +# ────────────────────────────────────────────── +OIDC_ISSUER="oidc.baseten.co" +if ! output=$(aws iam create-open-id-connect-provider \ + --url "https://${OIDC_ISSUER}" 2>&1); then + if [[ "${output}" != *"EntityAlreadyExists"* ]]; then + echo "${output}" >&2 + exit 1 + fi +fi + +# ────────────────────────────────────────────── +# 4. Create the runtime role and restrict who may assume it +# ────────────────────────────────────────────── +trust_policy=$(jq --null-input \ + --arg provider "arn:aws:iam::${AWS_ACCOUNT_ID}:oidc-provider/${OIDC_ISSUER}" \ + --arg audience "${OIDC_ISSUER}" \ + --arg subject "v=1:org=${BASETEN_ORG_ID}:team=${BASETEN_TEAM_ID}:*:type=model_container" \ + --arg audience_key "${OIDC_ISSUER}:aud" \ + --arg subject_key "${OIDC_ISSUER}:sub" \ + '{ + Version: "2012-10-17", + Statement: [{ + Effect: "Allow", + Principal: {Federated: $provider}, + Action: "sts:AssumeRoleWithWebIdentity", + Condition: { + StringEquals: {($audience_key): $audience}, + StringLike: {($subject_key): $subject} + } + }] + }') + +if aws iam get-role --role-name "${ROLE_NAME}" >/dev/null 2>&1; then + aws iam update-assume-role-policy \ + --role-name "${ROLE_NAME}" \ + --policy-document "${trust_policy}" +else + aws iam create-role \ + --role-name "${ROLE_NAME}" \ + --assume-role-policy-document "${trust_policy}" \ + --description "Baseten OIDC role for decrypting inference payloads" +fi + +# ────────────────────────────────────────────── +# 5. Grant the runtime permission to unwrap payload data keys +# ────────────────────────────────────────────── +payload_policy=$(jq --null-input \ + --arg key_arn "${KMS_KEY_ARN}" \ + '{ + Version: "2012-10-17", + Statement: [{ + Effect: "Allow", + Action: "kms:Decrypt", + Resource: $key_arn, + Condition: { + StringEquals: { + "kms:EncryptionContext:purpose": "baseten-inference-payload" + } + } + }] + }') +aws iam put-role-policy \ + --role-name "${ROLE_NAME}" \ + --policy-name "BasetenPayloadDecryptAccess" \ + --policy-document "${payload_policy}" + +# ────────────────────────────────────────────── +# Done +# ────────────────────────────────────────────── +role_arn="arn:aws:iam::${AWS_ACCOUNT_ID}:role/${ROLE_NAME}" +echo +echo "Setup complete." +echo "KMS key: ${KMS_KEY_ARN}" +echo "OIDC role: ${role_arn}" +echo +echo "Set AWS_ROLE_ARN = \"${role_arn}\" and AWS_REGION = \"${AWS_REGION}\"" +echo "in standard-truss/config.yaml or custom-base-image/config.yaml." +echo "Set KMS_KEY_ARN = \"${KMS_KEY_ARN}\" and AWS_REGION = \"${AWS_REGION}\"" +echo "when running client.py." diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/config.yaml b/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/config.yaml new file mode 100644 index 00000000..d044bf56 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/config.yaml @@ -0,0 +1,11 @@ +model_name: oidc-envelope-payload-encryption +python_version: py313 +runtime: + oidc: + enabled: true +requirements: + - boto3 + - cryptography +environment_variables: + AWS_ROLE_ARN: "" # e.g. arn:aws:iam::123456789012:role/BasetenOIDCPayloadRole + AWS_REGION: "" # e.g. us-west-2 diff --git a/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/model/model.py b/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/model/model.py new file mode 100644 index 00000000..85e8d0ab --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/payload-encryption/standard-truss/model/model.py @@ -0,0 +1,87 @@ +import base64 +import hashlib +import hmac +import json +import os + +import boto3 +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + +ENVELOPE_VERSION = 1 +ENVELOPE_ALGORITHM = "AES-256-CBC-HMAC-SHA256" + + +def get_kms_client(): + os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = os.environ["B10_OIDC_TOKEN_PATH"] + os.environ["AWS_ROLE_SESSION_NAME"] = "baseten-payload-decryption" + return boto3.client("kms", region_name=os.environ["AWS_REGION"]) + + +def decode_field(envelope: dict, field: str) -> bytes: + try: + return base64.b64decode(envelope[field], validate=True) + except (KeyError, ValueError) as error: + raise ValueError(f"Invalid envelope field: {field}.") from error + + +def decrypt_request(envelope: dict) -> tuple[dict, bytes]: + if envelope.get("version") != ENVELOPE_VERSION: + raise ValueError("Unsupported payload envelope version.") + if envelope.get("algorithm") != ENVELOPE_ALGORITHM: + raise ValueError("Unsupported payload envelope algorithm.") + + response = get_kms_client().decrypt( + CiphertextBlob=decode_field(envelope, "encrypted_data_key"), + EncryptionContext=envelope["encryption_context"], + ) + data_key = response["Plaintext"] + if len(data_key) != 64: + raise ValueError("Expected a 64-byte envelope data key from KMS.") + + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = decode_field(envelope, "iv") + ciphertext = decode_field(envelope, "ciphertext") + expected_mac = decode_field(envelope, "hmac") + + actual_mac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + if not hmac.compare_digest(actual_mac, expected_mac): + raise ValueError("Encrypted request failed integrity verification.") + + decryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).decryptor() + padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize() + unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() + plaintext = unpadder.update(padded_plaintext) + unpadder.finalize() + return json.loads(plaintext), data_key + + +def encrypt_response(payload: dict, data_key: bytes) -> dict: + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = os.urandom(16) + plaintext = json.dumps(payload, separators=(",", ":")).encode() + + padder = padding.PKCS7(algorithms.AES.block_size).padder() + padded_plaintext = padder.update(plaintext) + padder.finalize() + encryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).encryptor() + ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize() + payload_hmac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + + return { + "version": ENVELOPE_VERSION, + "algorithm": ENVELOPE_ALGORITHM, + "iv": base64.b64encode(iv).decode(), + "hmac": base64.b64encode(payload_hmac).decode(), + "ciphertext": base64.b64encode(ciphertext).decode(), + } + + +class Model: + def predict(self, model_input): + request, data_key = decrypt_request(model_input) + try: + value = float(request["value"]) + return encrypt_response({"value": value * 2 + 1}, data_key) + finally: + del data_key diff --git a/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/config.yaml b/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/config.yaml new file mode 100644 index 00000000..7f564bb4 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/config.yaml @@ -0,0 +1,30 @@ +base_image: + image: vllm/vllm-openai:v0.8.5 +model_name: oidc-envelope-weight-encryption-vllm +docker_server: + start_command: >- + sh -c "python3 /packages/decrypt.py && + exec vllm serve /tmp/decrypted-weights --port 8000" + readiness_endpoint: /health + liveness_endpoint: /health + predict_endpoint: /v1/completions + server_port: 8000 +runtime: + oidc: + enabled: true +requirements: + - boto3 + - cryptography + +# FILL ME +weights: + - source: "" + mount_location: "" + auth: + auth_method: AWS_OIDC + aws_oidc_role_arn: "" + aws_oidc_region: "" +environment_variables: + AWS_ROLE_ARN: "" + AWS_REGION: "" + DECRYPTED_WEIGHTS_PATH: "" diff --git a/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/packages/decrypt.py b/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/packages/decrypt.py new file mode 100644 index 00000000..33410554 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/weight-encryption/custom-base-image/packages/decrypt.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 + +import base64 +import hashlib +import hmac +import json +import os +from pathlib import Path + +import boto3 +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + +ENVELOPE_VERSION = 1 +ENVELOPE_ALGORITHM = "AES-256-CBC-HMAC-SHA256" +ENCRYPTED_WEIGHTS_DIR = Path("/models/custom") +ENCRYPTED_WEIGHTS_FILE = ENCRYPTED_WEIGHTS_DIR / "weights.enc" +ENCRYPTED_DATA_KEY_FILE = ENCRYPTED_WEIGHTS_DIR / "encrypted-data-key" +ENVELOPE_FILE = ENCRYPTED_WEIGHTS_DIR / "envelope.json" + + +def get_kms_client(): + os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = os.environ["B10_OIDC_TOKEN_PATH"] + os.environ["AWS_ROLE_SESSION_NAME"] = "baseten-envelope-decryption" + return boto3.client("kms", region_name=os.environ["AWS_REGION"]) + + +def read_envelope() -> dict: + envelope = json.loads(ENVELOPE_FILE.read_text()) + if envelope.get("version") != ENVELOPE_VERSION: + raise ValueError("Unsupported weights envelope version.") + if envelope.get("algorithm") != ENVELOPE_ALGORITHM: + raise ValueError("Unsupported weights envelope algorithm.") + return envelope + + +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"] + if len(data_key) != 64: + raise ValueError("Expected a 64-byte envelope data key from KMS.") + + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = base64.b64decode(envelope["iv"], validate=True) + expected_mac = base64.b64decode(envelope["hmac"], validate=True) + + actual_mac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + if not hmac.compare_digest(actual_mac, expected_mac): + raise ValueError("Encrypted weights failed integrity verification.") + + decryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).decryptor() + padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize() + unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() + plaintext = unpadder.update(padded_plaintext) + unpadder.finalize() + + output.parent.mkdir(parents=True, exist_ok=True) + temporary_path = output.with_name(f"{output.name}.tmp") + try: + temporary_path.write_bytes(plaintext) + temporary_path.chmod(0o600) + temporary_path.replace(output) + finally: + temporary_path.unlink(missing_ok=True) + return output + + +def main() -> None: + output = Path(os.environ["DECRYPTED_WEIGHTS_PATH"]) + decrypt_weights(output) + print(f"Decrypted weights to {output}") + + +if __name__ == "__main__": + main() diff --git a/oidc-recipes/oidc-envelope-encryption/weight-encryption/setup.sh b/oidc-recipes/oidc-envelope-encryption/weight-encryption/setup.sh new file mode 100755 index 00000000..8a1a796e --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/weight-encryption/setup.sh @@ -0,0 +1,317 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ────────────────────────────────────────────── +# FILL ME: replace these with your values +# ────────────────────────────────────────────── +AWS_ACCOUNT_ID="" # Existing AWS account ID +AWS_REGION="" # AWS region for S3 and KMS (e.g. us-west-2) +S3_BUCKET="" # Globally unique S3 bucket name (created or reused) +S3_PREFIX="" # S3 key prefix for the encrypted bundle (e.g. models/custom-weights) +KMS_ALIAS="" # KMS alias to create or reuse (must start with alias/) +ROLE_NAME="" # IAM role name to create or update +BASETEN_ORG_ID="" # From `truss whoami --show-oidc` +BASETEN_TEAM_ID="" # From `truss whoami --show-oidc` + +# ────────────────────────────────────────────── +# Helpers and configuration validation +# ────────────────────────────────────────────── +require_command() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "error: required command '$1' was not found" >&2 + exit 1 + fi +} + +require_value() { + if [[ -z "$2" ]]; then + echo "error: $1 must be set" >&2 + exit 1 + fi +} + +file_hex() { + od -An -tx1 "$1" | tr -d ' \n' +} + +for command in aws jq od openssl tr; do + require_command "${command}" +done + +require_value AWS_ACCOUNT_ID "${AWS_ACCOUNT_ID}" +require_value AWS_REGION "${AWS_REGION}" +require_value S3_BUCKET "${S3_BUCKET}" +require_value S3_PREFIX "${S3_PREFIX}" +require_value KMS_ALIAS "${KMS_ALIAS}" +require_value ROLE_NAME "${ROLE_NAME}" +require_value BASETEN_ORG_ID "${BASETEN_ORG_ID}" +require_value BASETEN_TEAM_ID "${BASETEN_TEAM_ID}" + +# ────────────────────────────────────────────── +# 1. Authenticate the AWS CLI and verify the account +# ────────────────────────────────────────────── +if ! aws sts get-caller-identity --region "${AWS_REGION}" >/dev/null 2>&1; then + echo "No valid AWS CLI credentials were found." + echo " 1) Log in with AWS Console credentials (recommended)" + echo " 2) Configure an access key" + read -r -p "Choose an authentication method [1]: " auth_method + + case "${auth_method:-1}" in + 1) aws login --region "${AWS_REGION}" ;; + 2) aws configure ;; + *) + echo "error: invalid authentication method: ${auth_method}" >&2 + exit 1 + ;; + esac +fi + +caller_account=$(aws sts get-caller-identity \ + --region "${AWS_REGION}" \ + --query Account \ + --output text) +if [[ "${caller_account}" != "${AWS_ACCOUNT_ID}" ]]; then + echo "error: authenticated to AWS account ${caller_account}, expected ${AWS_ACCOUNT_ID}" >&2 + exit 1 +fi + +# ────────────────────────────────────────────── +# 2. Create the bucket that will hold encrypted weights +# ────────────────────────────────────────────── +echo "Creating S3 bucket ${S3_BUCKET} in ${AWS_REGION}..." +create_bucket_args=(--bucket "${S3_BUCKET}" --region "${AWS_REGION}") +if [[ "${AWS_REGION}" != "us-east-1" ]]; then + create_bucket_args+=( + --create-bucket-configuration "LocationConstraint=${AWS_REGION}" + ) +fi + +if ! output=$(aws s3api create-bucket "${create_bucket_args[@]}" 2>&1); then + if [[ "${output}" == *"BucketAlreadyOwnedByYou"* ]]; then + echo "S3 bucket already exists and is owned by you, continuing..." + else + echo "${output}" >&2 + exit 1 + fi +fi + +# ────────────────────────────────────────────── +# 3. Create or reuse the KMS key-encryption key +# +# KMS protects the generated data key. It does not encrypt the model weights +# directly, which avoids KMS payload-size limits. +# ────────────────────────────────────────────── +if KMS_KEY_ARN=$(aws kms describe-key \ + --key-id "${KMS_ALIAS}" \ + --region "${AWS_REGION}" \ + --query KeyMetadata.Arn \ + --output text 2>/dev/null); then + echo "Using existing KMS key ${KMS_KEY_ARN}." +else + echo "Creating KMS key ${KMS_ALIAS}..." + KMS_KEY_ARN=$(aws kms create-key \ + --region "${AWS_REGION}" \ + --description "Envelope encryption key for Baseten model weights" \ + --query KeyMetadata.Arn \ + --output text) + aws kms create-alias \ + --region "${AWS_REGION}" \ + --alias-name "${KMS_ALIAS}" \ + --target-key-id "${KMS_KEY_ARN}" +fi + +# ────────────────────────────────────────────── +# 4. Create small fictional model weights +# +# All plaintext keys and weights remain in a temporary directory that is +# deleted when this script exits. +# ────────────────────────────────────────────── +work_dir=$(mktemp -d) +trap 'rm -rf "${work_dir}"' EXIT + +cat >"${work_dir}/weights.json" <<'EOF' +{ + "scale": 2.5, + "bias": 1.0, + "description": "Fictional linear-model weights for the OIDC envelope-weight-encryption recipe" +} +EOF + +# ────────────────────────────────────────────── +# 5. Ask KMS for an envelope data key +# +# KMS returns the same 64-byte key in two forms: plaintext for this one-time +# encryption operation and encrypted for storage alongside the ciphertext. +# The encryption context must match when the model later calls KMS Decrypt. +# ────────────────────────────────────────────── +encryption_context="purpose=baseten-model-weights,bucket=${S3_BUCKET}" +data_key_response=$(aws kms generate-data-key \ + --key-id "${KMS_KEY_ARN}" \ + --number-of-bytes 64 \ + --encryption-context "${encryption_context}" \ + --region "${AWS_REGION}" \ + --output json) + +jq -r '.Plaintext' <<<"${data_key_response}" \ + | tr -d '\n' \ + | openssl base64 -d -A >"${work_dir}/data-key" +jq -r '.CiphertextBlob' <<<"${data_key_response}" \ + | tr -d '\n' \ + | openssl base64 -d -A >"${work_dir}/encrypted-data-key" + +# ────────────────────────────────────────────── +# 6. Encrypt and authenticate the weights +# +# Split the data key into independent encryption and MAC keys. The HMAC covers +# IV || ciphertext, so the model verifies integrity before decrypting. +# ────────────────────────────────────────────── +dd if="${work_dir}/data-key" of="${work_dir}/encryption-key" bs=1 count=32 2>/dev/null +dd if="${work_dir}/data-key" of="${work_dir}/mac-key" bs=1 skip=32 count=32 2>/dev/null +openssl rand 16 >"${work_dir}/iv" + +openssl enc -aes-256-cbc \ + -K "$(file_hex "${work_dir}/encryption-key")" \ + -iv "$(file_hex "${work_dir}/iv")" \ + -in "${work_dir}/weights.json" \ + -out "${work_dir}/weights.enc" + +cat "${work_dir}/iv" "${work_dir}/weights.enc" >"${work_dir}/authenticated-data" +openssl dgst -sha256 \ + -mac HMAC \ + -macopt "hexkey:$(file_hex "${work_dir}/mac-key")" \ + -binary "${work_dir}/authenticated-data" >"${work_dir}/weights.hmac" + +# ────────────────────────────────────────────── +# 7. Build the envelope manifest +# +# S3 metadata is not preserved in the mounted filesystem, so the IV, HMAC, +# encryption context, and encrypted-key filename travel in this companion file. +# ────────────────────────────────────────────── +jq --null-input \ + --arg iv "$(openssl base64 -A -in "${work_dir}/iv")" \ + --arg hmac "$(openssl base64 -A -in "${work_dir}/weights.hmac")" \ + --arg bucket "${S3_BUCKET}" \ + '{ + version: 1, + algorithm: "AES-256-CBC-HMAC-SHA256", + iv: $iv, + hmac: $hmac, + encryption_context: { + purpose: "baseten-model-weights", + bucket: $bucket + } + }' >"${work_dir}/envelope.json" + +# ────────────────────────────────────────────── +# 8. Upload the encrypted payload and its companion files +# +# BDN mounts these three files at /models/custom. No plaintext weights or data +# keys are uploaded. +# ────────────────────────────────────────────── +echo "Uploading encrypted weights to s3://${S3_BUCKET}/${S3_PREFIX}/..." +for artifact in weights.enc encrypted-data-key envelope.json; do + aws s3 cp "${work_dir}/${artifact}" \ + "s3://${S3_BUCKET}/${S3_PREFIX}/${artifact}" \ + --region "${AWS_REGION}" +done + +# ────────────────────────────────────────────── +# 9. Register the Baseten OIDC identity provider +# ────────────────────────────────────────────── +OIDC_ISSUER="oidc.baseten.co" +if ! output=$(aws iam create-open-id-connect-provider \ + --url "https://${OIDC_ISSUER}" 2>&1); then + if [[ "${output}" != *"EntityAlreadyExists"* ]]; then + echo "${output}" >&2 + exit 1 + fi +fi + +# ────────────────────────────────────────────── +# 10. Create the OIDC role and restrict who may assume it +# +# BDN uses a model_build token to mirror the encrypted S3 objects. The model +# container uses a model_container token to unwrap the data key through KMS. +# ────────────────────────────────────────────── +trust_policy=$(jq --null-input \ + --arg provider "arn:aws:iam::${AWS_ACCOUNT_ID}:oidc-provider/${OIDC_ISSUER}" \ + --arg audience "${OIDC_ISSUER}" \ + --arg build_subject "v=1:org=${BASETEN_ORG_ID}:team=${BASETEN_TEAM_ID}:*:type=model_build" \ + --arg runtime_subject "v=1:org=${BASETEN_ORG_ID}:team=${BASETEN_TEAM_ID}:*:type=model_container" \ + --arg audience_key "${OIDC_ISSUER}:aud" \ + --arg subject_key "${OIDC_ISSUER}:sub" \ + '{ + Version: "2012-10-17", + Statement: [{ + Effect: "Allow", + Principal: {Federated: $provider}, + Action: "sts:AssumeRoleWithWebIdentity", + Condition: { + StringEquals: {($audience_key): $audience}, + StringLike: {($subject_key): [$build_subject, $runtime_subject]} + } + }] + }') + +if aws iam get-role --role-name "${ROLE_NAME}" >/dev/null 2>&1; then + aws iam update-assume-role-policy \ + --role-name "${ROLE_NAME}" \ + --policy-document "${trust_policy}" +else + aws iam create-role \ + --role-name "${ROLE_NAME}" \ + --assume-role-policy-document "${trust_policy}" \ + --description "Baseten OIDC role for decrypting envelope-encrypted weights" +fi + +# ────────────────────────────────────────────── +# 11. Grant build-time S3 read and runtime KMS decrypt permissions +# +# S3 access is limited to the encrypted weight prefix. KMS access is limited +# to the key that wrapped this envelope's data key. +# ────────────────────────────────────────────── +weights_policy=$(jq --null-input \ + --arg bucket_arn "arn:aws:s3:::${S3_BUCKET}" \ + --arg object_arn "arn:aws:s3:::${S3_BUCKET}/${S3_PREFIX}/*" \ + --arg prefix "${S3_PREFIX}" \ + --arg key_arn "${KMS_KEY_ARN}" \ + '{ + Version: "2012-10-17", + Statement: [ + { + Effect: "Allow", + Action: "s3:ListBucket", + Resource: $bucket_arn, + Condition: { + StringLike: {"s3:prefix": [$prefix, ($prefix + "/*")]} + } + }, + { + Effect: "Allow", + Action: "s3:GetObject", + Resource: $object_arn + }, + { + Effect: "Allow", + Action: "kms:Decrypt", + Resource: $key_arn + } + ] + }') +aws iam put-role-policy \ + --role-name "${ROLE_NAME}" \ + --policy-name "BasetenEncryptedWeightsAccess" \ + --policy-document "${weights_policy}" + +# ────────────────────────────────────────────── +# Done +# ────────────────────────────────────────────── +role_arn="arn:aws:iam::${AWS_ACCOUNT_ID}:role/${ROLE_NAME}" +echo +echo "Setup complete." +echo "Encrypted weights: s3://${S3_BUCKET}/${S3_PREFIX}" +echo "KMS key: ${KMS_KEY_ARN}" +echo "OIDC role: ${role_arn}" +echo +echo "Set AWS_ROLE_ARN = \"${role_arn}\" and AWS_REGION = \"${AWS_REGION}\"" +echo "in standard-truss/config.yaml and custom-base-image/config.yaml." diff --git a/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/config.yaml b/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/config.yaml new file mode 100644 index 00000000..7b0d99a6 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/config.yaml @@ -0,0 +1,21 @@ +model_name: oidc-envelope-weight-encryption +python_version: py313 +runtime: + oidc: + enabled: true +requirements: + - boto3 + - cryptography + +# FILL ME +weights: + - source: "" + mount_location: "" + auth: + auth_method: AWS_OIDC + aws_oidc_role_arn: "" + aws_oidc_region: "" +environment_variables: + AWS_ROLE_ARN: "" + AWS_REGION: "" + DECRYPTED_WEIGHTS_PATH: "" diff --git a/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/model/model.py b/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/model/model.py new file mode 100644 index 00000000..947a6353 --- /dev/null +++ b/oidc-recipes/oidc-envelope-encryption/weight-encryption/standard-truss/model/model.py @@ -0,0 +1,81 @@ +import base64 +import hashlib +import hmac +import json +import os +from pathlib import Path + +import boto3 +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + +ENVELOPE_VERSION = 1 +ENVELOPE_ALGORITHM = "AES-256-CBC-HMAC-SHA256" +ENCRYPTED_WEIGHTS_DIR = Path("/models/custom") +ENCRYPTED_WEIGHTS_FILE = ENCRYPTED_WEIGHTS_DIR / "weights.enc" +ENCRYPTED_DATA_KEY_FILE = ENCRYPTED_WEIGHTS_DIR / "encrypted-data-key" +ENVELOPE_FILE = ENCRYPTED_WEIGHTS_DIR / "envelope.json" + + +def get_kms_client(): + os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = os.environ["B10_OIDC_TOKEN_PATH"] + os.environ["AWS_ROLE_SESSION_NAME"] = "baseten-envelope-decryption" + return boto3.client("kms", region_name=os.environ["AWS_REGION"]) + + +def read_envelope() -> dict: + envelope = json.loads(ENVELOPE_FILE.read_text()) + if envelope.get("version") != ENVELOPE_VERSION: + raise ValueError("Unsupported weights envelope version.") + if envelope.get("algorithm") != ENVELOPE_ALGORITHM: + raise ValueError("Unsupported weights envelope algorithm.") + return envelope + + +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"] + if len(data_key) != 64: + raise ValueError("Expected a 64-byte envelope data key from KMS.") + + encryption_key = data_key[:32] + mac_key = data_key[32:] + iv = base64.b64decode(envelope["iv"], validate=True) + expected_mac = base64.b64decode(envelope["hmac"], validate=True) + + actual_mac = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest() + if not hmac.compare_digest(actual_mac, expected_mac): + raise ValueError("Encrypted weights failed integrity verification.") + + decryptor = Cipher(algorithms.AES(encryption_key), modes.CBC(iv)).decryptor() + padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize() + unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() + plaintext = unpadder.update(padded_plaintext) + unpadder.finalize() + + output.parent.mkdir(parents=True, exist_ok=True) + temporary_path = output.with_name(f"{output.name}.tmp") + try: + temporary_path.write_bytes(plaintext) + temporary_path.chmod(0o600) + temporary_path.replace(output) + finally: + temporary_path.unlink(missing_ok=True) + return output + + +class Model: + def load(self): + output = Path(os.environ["DECRYPTED_WEIGHTS_PATH"]) + weights_path = decrypt_weights(output) + self._weights = json.loads(weights_path.read_text()) + + def predict(self, model_input): + value = float(model_input["value"]) + return {"value": value * self._weights["scale"] + self._weights["bias"]} diff --git a/oidc-recipes/oidc-fetch-a-resource/README.md b/oidc-recipes/oidc-fetch-a-resource/README.md new file mode 100644 index 00000000..3181b40f --- /dev/null +++ b/oidc-recipes/oidc-fetch-a-resource/README.md @@ -0,0 +1,132 @@ +# OIDC: Fetch a Resource +_Application: Fetching text from an object in an S3 bucket_ + +This recipe demonstrates a basic and flexible pattern: fetch a resource from a remote host _without_ storing any long-term credentials. Instead, configure trust in the B10 IdP and establish a session by sending the OIDC token. + +## Setup + +1. Run `truss whoami --show-oidc` to get your Baseten organization and team IDs. +2. Fill in the values at the top of [`setup.sh`](setup.sh). +3. Run `./setup.sh`. +4. Copy the printed role, region, bucket, and key into [`standard-truss/config.yaml`](standard-truss/config.yaml). +5. Deploy with `truss push ./standard-truss`. + +At runtime, `Model.load()` maps `B10_OIDC_TOKEN_PATH` to boto3's web-identity token variable. boto3 exchanges the token for short-lived AWS credentials. Each prediction reads the configured S3 object and returns its text: + +```json +{"text": "Contents of the S3 object"} +``` + +## Other uses + +The resource contents can support many use cases. The snippets below assume `read_text()` and `write_text()` access the authenticated remote resource: + +- use extra context in an inference request (e.g. company report documents) + + ```python + def predict(self, model_input): + report = self.read_text("reports/q2-2026.txt") + return self.llm.generate(f"Report:\n{report}\n\nQuestion: {model_input['question']}") + ``` + +- write results of an inference request to the remote resource + + ```python + def predict(self, model_input): + result = self.llm.generate(model_input["prompt"]) + self.write_text(f"results/{model_input['request_id']}.txt", result) + return {"result": result} + ``` + +- pull initialization configuration without hardcoding it in your Truss (configuration can live remotely in one place and be read on every `load()` or even every `predict()`) + + ```python + def load(self): + self.settings = json.loads(self.read_text("config/production.json")) + self.model = load_model(self.settings["model_id"]) + ``` + +- load a tenant-specific prompt template from `s3://company-prompts/acme/support-agent.txt` + + ```python + def predict(self, model_input): + prompt = self.read_text(f"prompts/{model_input['tenant_id']}/support-agent.txt") + return self.llm.generate(f"{prompt}\n\nCustomer: {model_input['message']}") + ``` + +- fetch the latest product catalog before answering a recommendation request + + ```python + def predict(self, model_input): + catalog = self.read_text("catalog/current.json") + return self.llm.generate(f"Catalog: {catalog}\nRecommend: {model_input['request']}") + ``` + +- read a JSON feature-flag file that changes model behavior without redeploying the model + + ```python + def predict(self, model_input): + flags = json.loads(self.read_text("config/feature-flags.json")) + model = self.fast_model if flags["use_fast_model"] else self.quality_model + return model.generate(model_input["prompt"]) + ``` + +- retrieve a customer-specific glossary before translating industry-specific documents + + ```python + def predict(self, model_input): + glossary = self.read_text(f"glossaries/{model_input['customer_id']}.txt") + return self.translator.translate(model_input["text"], glossary=glossary) + ``` + +- load a list of blocked terms or compliance rules before generating a response + + ```python + def predict(self, model_input): + blocked_terms = self.read_text("compliance/blocked-terms.txt").splitlines() + response = self.llm.generate(model_input["prompt"]) + return {"response": redact(response, blocked_terms)} + ``` + +- fetch a small set of few-shot examples selected for a particular workflow + + ```python + def predict(self, model_input): + examples = self.read_text(f"examples/{model_input['workflow']}.jsonl") + return self.llm.generate(f"Examples:\n{examples}\n\nInput: {model_input['text']}") + ``` + +- read model routing configuration that selects an upstream model or endpoint + + ```python + def predict(self, model_input): + routes = json.loads(self.read_text("routing/models.json")) + model = self.clients[routes[model_input["task"]]] + return model.generate(model_input["prompt"]) + ``` + +- retrieve private certificates, public keys, or trust bundles needed to call another internal service + + ```python + def load(self): + trust_bundle = self.read_text("certificates/internal-ca.pem") + self.internal_client = InternalClient(ca_certificate=trust_bundle) + ``` + +- write generated transcripts, summaries, or embeddings back to a customer-owned bucket + + ```python + def predict(self, model_input): + transcript = self.transcriber.transcribe(model_input["audio"]) + self.write_text(f"transcripts/{model_input['customer_id']}.txt", transcript) + return {"transcript": transcript} + ``` + +- check for a kill-switch file that disables a workflow without requiring a new deployment + + ```python + def predict(self, model_input): + if self.read_text("controls/summarization-enabled.txt").strip() != "true": + return {"error": "Summarization is temporarily disabled"} + return {"summary": self.summarizer(model_input["text"])} + ``` diff --git a/oidc-recipes/oidc-fetch-a-resource/setup.sh b/oidc-recipes/oidc-fetch-a-resource/setup.sh new file mode 100755 index 00000000..4d14a044 --- /dev/null +++ b/oidc-recipes/oidc-fetch-a-resource/setup.sh @@ -0,0 +1,213 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ────────────────────────────────────────────── +# FILL ME: replace these with your values +# ────────────────────────────────────────────── +AWS_ACCOUNT_ID="" # Your AWS account ID (must exist prior to this script) +AWS_REGION="" # AWS region for the S3 bucket (e.g. us-west-2) +S3_BUCKET="" # S3 bucket name (to be created) +S3_KEY="" # S3 object key (like a file path) +BUCKET_TEXT="" # Text to store in the S3 object +ROLE_NAME="" # IAM role name (to be created) +BASETEN_ORG_ID="" # From `truss whoami --show-oidc` +BASETEN_TEAM_ID="" # From `truss whoami --show-oidc` + +# ────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────── +require_non_empty() { + local _name="$1" + local _desc="${2:-$1}" + local _val + eval "_val=\${${_name}-}" + if [ -z "$_val" ]; then + echo "error: ${_desc} is empty; set ${_name} in the configuration section." >&2 + exit 1 + fi +} + +require_non_empty AWS_ACCOUNT_ID "AWS account ID" +require_non_empty S3_BUCKET "S3 bucket name" +require_non_empty AWS_REGION "AWS region" +require_non_empty S3_KEY "S3 object key" +require_non_empty ROLE_NAME "IAM role name" +require_non_empty BASETEN_ORG_ID "Baseten organization ID" +require_non_empty BASETEN_TEAM_ID "Baseten team ID" + +OIDC_ISSUER="oidc.baseten.co" +OIDC_ISSUER_URL="https://${OIDC_ISSUER}" + +# ────────────────────────────────── +# 1. Authenticate the AWS CLI +# ────────────────────────────────── +if ! command -v aws >/dev/null 2>&1; then + echo "error: AWS CLI is not installed or is not on PATH." >&2 + exit 1 +fi + +if ! aws sts get-caller-identity --region "${AWS_REGION}" >/dev/null 2>&1; then + echo "No valid AWS CLI credentials were found." + echo " 1) Log in with AWS Console credentials (recommended)" + echo " 2) Configure an access key" + read -r -p "Choose an authentication method [1]: " auth_method + + case "${auth_method:-1}" in + 1) + aws login --region "${AWS_REGION}" + ;; + 2) + aws configure + ;; + *) + echo "error: invalid authentication method: ${auth_method}" >&2 + exit 1 + ;; + esac +fi + +if ! caller_account=$(aws sts get-caller-identity \ + --region "${AWS_REGION}" \ + --query Account \ + --output text); then + echo "error: AWS authentication failed." >&2 + exit 1 +fi + +if [[ "${caller_account}" != "${AWS_ACCOUNT_ID}" ]]; then + echo "error: authenticated to AWS account ${caller_account}, expected ${AWS_ACCOUNT_ID}." >&2 + exit 1 +fi + +echo "Authenticated to AWS account ${caller_account}." + +# ────────────────────────────────── +# 2. Create the S3 bucket +# ────────────────────────────────── +echo "Creating S3 bucket ${S3_BUCKET} in ${AWS_REGION}..." + +create_bucket_args=( + --bucket "${S3_BUCKET}" + --region "${AWS_REGION}" +) + +if [[ "${AWS_REGION}" != "us-east-1" ]]; then + create_bucket_args+=( + --create-bucket-configuration "LocationConstraint=${AWS_REGION}" + ) +fi + +if ! output=$(aws s3api create-bucket "${create_bucket_args[@]}" 2>&1); then + if [[ "$output" == *"BucketAlreadyOwnedByYou"* ]]; then + echo "S3 bucket already exists and is owned by you, continuing..." + else + echo "$output" >&2 + exit 1 + fi +else + echo "S3 bucket created." +fi + +# ────────────────────────────────── +# 3. Upload a text object to S3 +# ────────────────────────────────── +echo "Uploading text to s3://${S3_BUCKET}/${S3_KEY}..." +printf '%s' "${BUCKET_TEXT}" | aws s3 cp - "s3://${S3_BUCKET}/${S3_KEY}" \ + --region "${AWS_REGION}" \ + --content-type "text/plain" + +# ────────────────────────────────── +# 4. Create the OIDC identity provider +# ────────────────────────────────── +echo "Creating OIDC identity provider..." + +if ! output=$(aws iam create-open-id-connect-provider \ + --url "${OIDC_ISSUER_URL}" 2>&1); then + if [[ "$output" == *"EntityAlreadyExists"* ]]; then + echo "OIDC provider already exists, continuing..." + else + echo "$output" >&2 + exit 1 + fi +else + echo "OIDC provider created." +fi + +# ────────────────────────────────── +# 5. Create the IAM trust policy +# ────────────────────────────────── +TRUST_POLICY=$(cat < str: + resp = self._s3.get_object(Bucket=self._bucket, Key=self._key) + text = resp["Body"].read().decode("utf-8").strip() + return text + + def load(self): + self._bucket = os.environ["S3_BUCKET"] + self._key = os.environ["S3_KEY"] + self._s3 = self._get_s3_client() + + def predict(self, model_input): + text = self._read_text_from_s3() + print(f"Read from s3://{self._bucket}/{self._key}: {text}") + return {"text": text}