Module 3: Serverless & Lambda for AI
2. Lambda Fundamentals for AI
Overview
In this capsule you'll understand the anatomy of an AI-oriented Lambda function: how the handler works, how to package Python dependencies for AI, how to manage secrets (LLM API keys), and when to use synchronous vs asynchronous invocation. It's not a generic Lambda tutorial — every line of code invokes an LLM or manages an AI workflow.
Context: The previous capsule established the landscape of serverless for AI and why Lambda has particularities when working with LLMs. Here you go from "understanding the problem" to "writing the first function." By the end, you'll have a functional Lambda that invokes OpenAI, with packaged dependencies and well-managed secrets.
Anatomy of a Lambda Handler
The handler signature
Every Lambda function in Python receives exactly two arguments:
def handler(event, context):
# event: the data that triggered the function
# context: runtime metadata (remaining time, request ID, etc.)
return response
These two arguments give you everything you need to process an invocation:
# event — Input data
# When Lambda is invoked via API Gateway:
event = {
"httpMethod": "POST",
"path": "/ask",
"headers": {
"Content-Type": "application/json",
"x-api-key": "abc123"
},
"body": "{\"prompt\": \"Explain serverless in one sentence\"}",
"queryStringParameters": None,
"pathParameters": None,
"requestContext": {
"requestId": "abc-123-def",
"stage": "prod"
}
}
# context — Runtime metadata
# context.function_name → "ai-endpoint-prod"
# context.memory_limit_in_mb → 512
# context.get_remaining_time_in_millis() → 58000 (ms remaining)
# context.aws_request_id → "abc-123-def-456"
The response format
Lambda expects a response with a specific structure when it's invoked from API Gateway:
# Valid response for API Gateway
response = {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*"
},
"body": "{\"answer\": \"Serverless runs your code without you managing servers.\"}"
}
The body is always a string (serialized JSON). If you return a dict in body, API Gateway throws an error. This is one of the most common mistakes when starting with Lambda.
Your First Lambda for AI
Complete handler that invokes OpenAI
# handler.py
import json
import os
import time
from openai import OpenAI
# Client initialized OUTSIDE the handler
# It's reused between invocations (warm starts)
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
def handler(event, context):
"""Lambda that receives a prompt and returns an LLM response."""
start_time = time.time()
try:
body = json.loads(event.get("body", "{}"))
except json.JSONDecodeError:
return _error_response(400, "Invalid JSON in request body")
prompt = body.get("prompt", "").strip()
if not prompt:
return _error_response(400, "Field 'prompt' is required and cannot be empty")
max_tokens = body.get("max_tokens", 500)
model = body.get("model", "gpt-4o-mini")
remaining_ms = context.get_remaining_time_in_millis()
# Reserve 5s for post-LLM processing
llm_timeout = (remaining_ms / 1000) - 5
try:
response = client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
max_tokens=max_tokens,
timeout=llm_timeout,
)
except Exception as e:
return _error_response(502, f"LLM API error: {str(e)}")
answer = response.choices[0].message.content
tokens_used = response.usage.total_tokens
duration_ms = int((time.time() - start_time) * 1000)
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
},
"body": json.dumps({
"answer": answer,
"model": model,
"tokens_used": tokens_used,
"duration_ms": duration_ms,
"request_id": context.aws_request_id,
}),
}
def _error_response(status_code, message):
"""Helper for consistent error responses."""
return {
"statusCode": status_code,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
},
"body": json.dumps({"error": message}),
}
Why the OpenAI client is outside the handler
# ✅ CORRECT — Client outside the handler
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
def handler(event, context):
response = client.chat.completions.create(...)
...
# ❌ INCORRECT — Client inside the handler
def handler(event, context):
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
response = client.chat.completions.create(...)
...
Lambda reuses the container between consecutive invocations (warm starts). The code outside the handler runs only once — on the cold start. The code inside the handler runs on every invocation. Initializing the OpenAI client outside means that on warm starts you don't pay the cost of creating the HTTP connection.
Cold start (first invocation):
├── Init: import openai, create client → ~500ms
└── Handler: invoke LLM, return → ~3000ms
Total: ~3500ms
Warm start (following invocations):
├── Init: SKIP (already initialized) → 0ms
└── Handler: invoke LLM, return → ~3000ms
Total: ~3000ms
This optimization seems minor but it adds up: if your function receives 1000 invocations/hour, 990+ will be warm starts that save those 500ms of initialization.
Packaging Python Dependencies
The problem
Lambda needs all the Python dependencies to be included in the deployment package. It doesn't run pip install at startup — everything must be pre-packaged. For AI, this is a challenge because the dependencies can be heavy:
Size of common dependencies for AI:
├── openai (SDK): ~5MB → Viable in zip
├── anthropic (SDK): ~8MB → Viable in zip
├── httpx: ~3MB → Viable in zip
├── langchain + deps: ~80MB → Zip limit
├── numpy: ~30MB → Requires compilation
├── pandas: ~50MB → Requires compilation
├── torch (CPU): ~800MB → Impossible in zip, needs a container
└── transformers + torch: ~2GB → Impossible in standard Lambda
Three packaging strategies
Strategy 1: Zip package (for light dependencies)
# For openai + httpx (< 50MB total)
mkdir package
pip install openai -t package/
cd package
zip -r ../deployment.zip .
cd ..
zip deployment.zip handler.py
# Result: deployment.zip (~8MB)
Zip limit: 50MB compressed, 250MB uncompressed. Enough for openai, anthropic, httpx. Not enough for langchain with all its dependencies.
Strategy 2: Lambda Layers (shared dependencies)
# Layer: reusable dependencies between functions
mkdir -p python/lib/python3.11/site-packages
pip install openai -t python/lib/python3.11/site-packages/
zip -r openai-layer.zip python/
# Publish the layer
aws lambda publish-layer-version \
--layer-name openai-sdk \
--zip-file fileb://openai-layer.zip \
--compatible-runtimes python3.11
# Attach the layer to your function
aws lambda update-function-configuration \
--function-name ai-endpoint \
--layers arn:aws:lambda:us-east-1:123456789:layer:openai-sdk:1
Advantage: the layer is shared between functions and cached by Lambda. Total limit: 250MB (layers + code, uncompressed).
Strategy 3: Container image (for heavy dependencies)
# Dockerfile
FROM public.ecr.aws/lambda/python:3.11
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY handler.py .
CMD ["handler.handler"]
# requirements.txt
openai>=1.0.0
langchain>=0.1.0
langchain-openai>=0.0.5
Advantage: up to 10GB image. No practical dependency limit. Trade-off: longer cold starts (covered in capsule 03).
When to use each strategy
| Dependencies | Size | Strategy | Example |
|---|---|---|---|
| Only openai SDK | ~5MB | Zip | API that invokes GPT-4 |
| openai + httpx + pydantic | ~15MB | Zip or Layer | API with validation |
| langchain + chromadb | ~100MB | Container | RAG endpoint |
| numpy + pandas + sklearn | ~120MB | Container | Data processing |
| torch + transformers | ~2GB | Don't use Lambda | Use SageMaker or ECS |
Environment Variables and Secrets
Environment variables in Lambda
Lambda lets you configure environment variables that your code reads with os.environ:
# In handler.py
import os
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "500"))
ENVIRONMENT = os.environ.get("ENVIRONMENT", "production")
# Configure via AWS CLI
aws lambda update-function-configuration \
--function-name ai-endpoint \
--environment "Variables={
OPENAI_API_KEY=sk-xxx,
MODEL_NAME=gpt-4o-mini,
MAX_TOKENS=500,
ENVIRONMENT=production
}"
Secrets: don't hardcode API keys
Lambda's environment variables are encrypted at rest, but they're visible in the AWS console to anyone with access to the function. For sensitive secrets (LLM API keys), you have more secure options:
Option 1: Environment variables (acceptable for most cases)
# Acceptable if you control IAM access to the function
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
Option 2: AWS Secrets Manager (production)
import boto3
import json
def get_secret(secret_name):
"""Gets a secret from AWS Secrets Manager."""
sm_client = boto3.client("secretsmanager")
response = sm_client.get_secret_value(SecretId=secret_name)
return json.loads(response["SecretString"])
# Outside the handler — runs only on cold start
secrets = get_secret("ai-endpoint/openai")
client = OpenAI(api_key=secrets["api_key"])
Option 3: AWS SSM Parameter Store (free, enough)
import boto3
def get_parameter(name):
"""Gets a parameter from SSM Parameter Store."""
ssm = boto3.client("ssm")
response = ssm.get_parameter(Name=name, WithDecryption=True)
return response["Parameter"]["Value"]
# Outside the handler
api_key = get_parameter("/ai-endpoint/openai-api-key")
client = OpenAI(api_key=api_key)
Comparison of secret options
| Option | Cost | Security | Complexity | When |
|---|---|---|---|---|
| Environment variables | Free | Medium | Low | Development, MVPs |
| SSM Parameter Store | Free (standard) | High | Medium | Production with a limited budget |
| Secrets Manager | $0.40/secret/month | Very high | Medium | Enterprise production, automatic rotation |
For this module, we use environment variables. In real production, SSM Parameter Store is the sweet spot of cost and security.
Invocation: Synchronous vs Asynchronous
Synchronous (RequestResponse)
The client waits for the response. Lambda runs, and when it finishes, it returns the result directly.
Client → API Gateway → Lambda → LLM → Lambda → API Gateway → Client
↑
The client waits here
# Synchronous invocation via CLI
aws lambda invoke \
--function-name ai-endpoint \
--invocation-type RequestResponse \
--payload '{"body": "{\"prompt\": \"What is serverless?\"}"}' \
response.json
cat response.json
# {"statusCode": 200, "body": "{\"answer\": \"Serverless runs...\"}"}
When to use it for AI: Chatbots, question-answer APIs, any case where the user waits for the response immediately.
Asynchronous (Event)
The client sends the request and receives an "OK, I'm processing it" immediately. Lambda runs in the background.
Client → API Gateway → Lambda → (202 Accepted)
↓
Processes in the background
↓
Saves the result (S3, DB, webhook)
# Asynchronous invocation via CLI
aws lambda invoke \
--function-name ai-endpoint \
--invocation-type Event \
--payload '{"body": "{\"prompt\": \"Generate a long report\"}"}' \
response.json
# Immediate response: StatusCode 202 (without waiting for the result)
When to use it for AI:
| Case | Synchronous | Asynchronous |
|---|---|---|
| Chatbot (immediate response) | ✅ | ❌ |
| Generating short summaries | ✅ | ❌ |
| Batch document processing | ❌ | ✅ |
| Generating long reports | ❌ | ✅ |
| Audio transcription (large files) | ❌ | ✅ |
| Classifying incoming emails | ❌ | ✅ |
Rule of thumb: If the user is waiting at the screen, synchronous. If the user can continue and get a notification later, asynchronous.
Async with destinations
Lambda can send the result to another service automatically:
# Configure a destination for successful async invocations
aws lambda put-function-event-invoke-config \
--function-name ai-endpoint \
--destination-config '{
"OnSuccess": {
"Destination": "arn:aws:sqs:us-east-1:123456:results-queue"
},
"OnFailure": {
"Destination": "arn:aws:sqs:us-east-1:123456:dead-letter-queue"
}
}'
SAM Template: Defining the Infrastructure
Basic template.yaml
AWS SAM (Serverless Application Model) defines your Lambda + API Gateway as code:
# template.yaml
AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: AI Endpoint — Lambda that invokes an LLM
Globals:
Function:
Timeout: 60
MemorySize: 512
Runtime: python3.11
Resources:
AIEndpointFunction:
Type: AWS::Serverless::Function
Properties:
Handler: handler.handler
CodeUri: lambda-function/
Environment:
Variables:
OPENAI_API_KEY: !Ref OpenAIApiKey
MODEL_NAME: gpt-4o-mini
MAX_TOKENS: "500"
Events:
AskAPI:
Type: Api
Properties:
Path: /ask
Method: post
Parameters:
OpenAIApiKey:
Type: String
NoEcho: true
Description: OpenAI API Key
Outputs:
ApiUrl:
Description: API endpoint URL
Value: !Sub "https://${ServerlessRestApi}.execute-api.${AWS::Region}.amazonaws.com/Prod/ask"
Local testing with SAM
# Build
sam build
# Invoke locally (without AWS)
echo '{"body": "{\"prompt\": \"What is Lambda?\"}"}' | \
sam local invoke AIEndpointFunction \
--env-vars env.json
# Local API (simulates API Gateway)
sam local start-api
# http://localhost:3000/ask available
# Test
curl -X POST http://localhost:3000/ask \
-H "Content-Type: application/json" \
-d '{"prompt": "Explain serverless"}'
// env.json — Variables for local testing
{
"AIEndpointFunction": {
"OPENAI_API_KEY": "sk-your-key-here",
"MODEL_NAME": "gpt-4o-mini",
"MAX_TOKENS": "500"
}
}
Comparison: Lambda vs FastAPI Server for AI
| Aspect | Lambda | FastAPI on a server |
|---|---|---|
| Cold start | 1-15s (first invocation) | 0s (always running) |
| Scaling | Automatic (0 to 1000+) | Manual (more instances) |
| Cost (low traffic) | ~$0 (generous free tier) | $5-40/month (always-on server) |
| Cost (high traffic) | Can be high ($$$) | Fixed (predictable) |
| Max execution time | 15 minutes | No limit |
| Max memory | 10GB | No practical limit |
| Dependencies | Manual packaging | pip install |
| Debugging | CloudWatch logs | Local logs, debugger |
| State | Stateless (each invocation independent) | Stateful if you want |
Troubleshooting
Problem 1: "ModuleNotFoundError: No module named 'openai'"
Your function can't find the dependency. The package isn't included in the deployment.
# Verify that you packaged the dependencies
unzip -l deployment.zip | grep openai
# Should show files from openai/
# If you use layers, verify that the layer is attached
aws lambda get-function-configuration \
--function-name ai-endpoint \
--query 'Layers'
Problem 2: "Task timed out after X seconds"
Lambda finished before the LLM responded.
# Check the Lambda timeout vs the LLM timeout
# Lambda timeout: 60s (configured in template.yaml)
# LLM timeout: must be < Lambda timeout
remaining = context.get_remaining_time_in_millis()
# If remaining < 10000 (10s), don't make the LLM call
Problem 3: "body must be a string, not a dict"
API Gateway expects body as a JSON string, not as a Python dict.
# ❌ INCORRECT
return {"statusCode": 200, "body": {"answer": "..."}}
# ✅ CORRECT
return {"statusCode": 200, "body": json.dumps({"answer": "..."})}
Problem 4: "CORS error in the frontend"
The Access-Control-Allow-Origin header is missing from the response.
# Include CORS headers in ALL responses (success and error)
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "POST, OPTIONS",
"Access-Control-Allow-Headers": "Content-Type, x-api-key",
},
"body": json.dumps(result),
}
Problem 5: "Environment variable OPENAI_API_KEY not found"
The variable isn't configured in the Lambda function.
# Verify the configured variables
aws lambda get-function-configuration \
--function-name ai-endpoint \
--query 'Environment.Variables'
# For local testing with SAM, create env.json
# For production, use aws lambda update-function-configuration
Hands-On Exercises
Exercise 1: Lambda with a prompt system
Modify the handler to accept a system_prompt in addition to the user's prompt. The system prompt configures the LLM's behavior (for example, "Always respond in English and in a maximum of 2 sentences").
See solution
def handler(event, context):
try:
body = json.loads(event.get("body", "{}"))
except json.JSONDecodeError:
return _error_response(400, "Invalid JSON")
prompt = body.get("prompt", "").strip()
system_prompt = body.get(
"system_prompt",
"You are a helpful assistant. Respond concisely."
)
if not prompt:
return _error_response(400, "Field 'prompt' is required")
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt},
]
try:
response = client.chat.completions.create(
model=os.environ.get("MODEL_NAME", "gpt-4o-mini"),
messages=messages,
max_tokens=int(os.environ.get("MAX_TOKENS", "500")),
)
except Exception as e:
return _error_response(502, f"LLM error: {str(e)}")
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
},
"body": json.dumps({
"answer": response.choices[0].message.content,
"tokens_used": response.usage.total_tokens,
}),
}
Exercise 2: Timeout-aware handler
Implement a handler that checks Lambda's remaining time before calling the LLM. If less than 15 seconds remain, return an error instead of making the call (which would probably fail on timeout).
See solution
MIN_REMAINING_MS = 15000
def handler(event, context):
remaining_ms = context.get_remaining_time_in_millis()
if remaining_ms < MIN_REMAINING_MS:
return _error_response(
503,
f"Insufficient time remaining: {remaining_ms}ms. "
f"Minimum required: {MIN_REMAINING_MS}ms"
)
try:
body = json.loads(event.get("body", "{}"))
except json.JSONDecodeError:
return _error_response(400, "Invalid JSON")
prompt = body.get("prompt", "").strip()
if not prompt:
return _error_response(400, "Field 'prompt' is required")
llm_timeout_s = (remaining_ms - 5000) / 1000
try:
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
max_tokens=500,
timeout=llm_timeout_s,
)
except Exception as e:
return _error_response(502, f"LLM error: {str(e)}")
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
},
"body": json.dumps({
"answer": response.choices[0].message.content,
"remaining_ms_at_start": remaining_ms,
"llm_timeout_used": llm_timeout_s,
}),
}
Exercise 3: Multi-model handler
Create a handler that supports multiple models (OpenAI and Anthropic) based on a provider parameter in the request. Use the same endpoint for both.
See solution
import json
import os
from openai import OpenAI
from anthropic import Anthropic
openai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
anthropic_client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
def _invoke_openai(prompt, max_tokens):
response = openai_client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
max_tokens=max_tokens,
)
return {
"answer": response.choices[0].message.content,
"provider": "openai",
"model": "gpt-4o-mini",
"tokens_used": response.usage.total_tokens,
}
def _invoke_anthropic(prompt, max_tokens):
response = anthropic_client.messages.create(
model="claude-3-haiku-20240307",
max_tokens=max_tokens,
messages=[{"role": "user", "content": prompt}],
)
return {
"answer": response.content[0].text,
"provider": "anthropic",
"model": "claude-3-haiku-20240307",
"tokens_used": response.usage.input_tokens + response.usage.output_tokens,
}
PROVIDERS = {
"openai": _invoke_openai,
"anthropic": _invoke_anthropic,
}
def handler(event, context):
try:
body = json.loads(event.get("body", "{}"))
except json.JSONDecodeError:
return _error_response(400, "Invalid JSON")
prompt = body.get("prompt", "").strip()
provider = body.get("provider", "openai").lower()
max_tokens = body.get("max_tokens", 500)
if not prompt:
return _error_response(400, "Field 'prompt' is required")
if provider not in PROVIDERS:
return _error_response(
400,
f"Provider '{provider}' not supported. Use: {list(PROVIDERS.keys())}"
)
try:
result = PROVIDERS[provider](prompt, max_tokens)
except Exception as e:
return _error_response(502, f"{provider} error: {str(e)}")
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
"Access-Control-Allow-Origin": "*",
},
"body": json.dumps(result),
}
Exercise 4: Package your Lambda with a Layer
Create a Lambda Layer containing the OpenAI SDK. Then create a deployment zip that only contains your handler.py. The goal is to separate code from dependencies.
See solution
# Step 1: Create the layer
mkdir -p layer/python/lib/python3.11/site-packages
pip install openai -t layer/python/lib/python3.11/site-packages/
cd layer
zip -r ../openai-layer.zip python/
cd ..
# Verify the size
ls -lh openai-layer.zip
# ~5MB expected
# Step 2: Publish the layer (on AWS)
aws lambda publish-layer-version \
--layer-name openai-sdk \
--description "OpenAI Python SDK for Lambda" \
--zip-file fileb://openai-layer.zip \
--compatible-runtimes python3.11 python3.12
# Save the ARN from the output:
# arn:aws:lambda:us-east-1:123456789:layer:openai-sdk:1
# Step 3: Package only the handler
zip deployment.zip handler.py
# ~2KB
# Step 4: Create the function with the layer
aws lambda create-function \
--function-name ai-endpoint \
--runtime python3.11 \
--handler handler.handler \
--zip-file fileb://deployment.zip \
--role arn:aws:iam::123456789:role/lambda-execution-role \
--layers arn:aws:lambda:us-east-1:123456789:layer:openai-sdk:1 \
--timeout 60 \
--memory-size 512
# Step 5: Verify
aws lambda invoke \
--function-name ai-endpoint \
--payload '{"body": "{\"prompt\": \"test\"}"}' \
output.json
The advantage: when you update your handler, you only upload ~2KB. The ~5MB layer doesn't change. This makes deployment faster and lets you share the layer between multiple functions.
Summary
- The Lambda handler receives
event(data) andcontext(runtime metadata) — everything you need to process the invocation. - Initialize clients outside the handler to reuse them on warm starts. The code outside the handler runs only once per container.
- Dependency packaging has three strategies: zip (~50MB limit), layers (shared between functions), container (up to 10GB).
- For AI with openai/anthropic SDKs, zip or layers are enough. For langchain or heavy ML libraries, you need a container.
- Environment variables are the simplest way to manage config. For secrets in production, SSM Parameter Store is the sweet spot.
- Synchronous invocation for when the user waits for a response. Asynchronous for background processing.
- SAM CLI allows local testing without an AWS account —
sam local invokeandsam local start-api.
Additional Resources
- AWS Lambda Python Handler — Official handler reference
- Lambda Layers — Lambda Layers documentation
- Lambda Environment Variables — Environment variables config
- AWS SAM CLI — Local Testing — Local testing with SAM
- Lambda Invocation Types — Sync vs Async invocation
- OpenAI Python SDK — Official OpenAI SDK
- AWS Secrets Manager vs SSM — Comparison of secret management
- Lambda Quotas — Lambda limits (size, timeout, memory)