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

DependenciesSizeStrategyExample
Only openai SDK~5MBZipAPI that invokes GPT-4
openai + httpx + pydantic~15MBZip or LayerAPI with validation
langchain + chromadb~100MBContainerRAG endpoint
numpy + pandas + sklearn~120MBContainerData processing
torch + transformers~2GBDon't use LambdaUse 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

OptionCostSecurityComplexityWhen
Environment variablesFreeMediumLowDevelopment, MVPs
SSM Parameter StoreFree (standard)HighMediumProduction with a limited budget
Secrets Manager$0.40/secret/monthVery highMediumEnterprise 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:

CaseSynchronousAsynchronous
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

AspectLambdaFastAPI on a server
Cold start1-15s (first invocation)0s (always running)
ScalingAutomatic (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 time15 minutesNo limit
Max memory10GBNo practical limit
DependenciesManual packagingpip install
DebuggingCloudWatch logsLocal logs, debugger
StateStateless (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) and context (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 invoke and sam local start-api.

Additional Resources

  1. AWS Lambda Python Handler — Official handler reference
  2. Lambda Layers — Lambda Layers documentation
  3. Lambda Environment Variables — Environment variables config
  4. AWS SAM CLI — Local Testing — Local testing with SAM
  5. Lambda Invocation Types — Sync vs Async invocation
  6. OpenAI Python SDK — Official OpenAI SDK
  7. AWS Secrets Manager vs SSM — Comparison of secret management
  8. Lambda Quotas — Lambda limits (size, timeout, memory)