Module 3: Serverless & Lambda for AI

5. Timeout and Memory Configuration for AI

Overview

In this capsule you'll configure Lambda's timeout and memory specifically for AI workloads. It's not "set 128MB and 3 seconds" — a function that invokes GPT-4o needs a radically different configuration than one that processes a webhook. By the end, you'll know how to size Lambda for LLM inference, understand why more memory means more CPU, and use Lambda Power Tuning to find the optimal point between speed and cost.

Context: Lambda lets you configure two resources: timeout (maximum 15 minutes) and memory (128MB to 10GB). What many don't know is that CPU is allocated proportionally to memory — 1769MB gives you 1 full vCPU. This turns the memory configuration into a CPU decision, which directly affects your AI function's performance and how much you pay per invocation.

Timeout: The clock that shows no mercy

How Lambda's timeout works

Lambda runs your function until it returns a result or the timeout runs out. There are no extensions, no negotiation. If your function doesn't finish in time, Lambda kills it and returns an error.

Timeline of a Lambda invocation:
0s ──── cold start ──── init done ──── handler runs ──── response
│                                                          │
└──────────── timeout window (configured by you) ──────────┘
              If it doesn't finish here → KILLED → 504 error

Timeout configuration

# In your SAM template / serverless.yml / Terraform

# SAM template
Resources:
  AiFunction:
    Type: AWS::Serverless::Function
    Properties:
      Timeout: 60  # seconds (max: 900 = 15 min)

# Terraform
resource "aws_lambda_function" "ai_endpoint" {
  timeout = 60  # seconds
}

# AWS CLI
aws lambda update-function-configuration \
  --function-name ai-endpoint \
  --timeout 60

Why the timeout matters more for AI

A Lambda that processes a webhook takes milliseconds. A Lambda that invokes an LLM takes seconds — sometimes many:

Typical LLM API times (p95):

Operation                          Typical time
─────────────────────────────────────────────────
OpenAI gpt-4o-mini (100 tokens)   1-3s
OpenAI gpt-4o (500 tokens)        3-8s
OpenAI gpt-4o (2000 tokens)       8-20s
Anthropic Claude 3.5 (500 tokens) 2-6s
Prompt chain (3 calls)            10-30s
RAG: embed + search + generate    5-15s

Lambda's default is 3 seconds. With that timeout, most of your LLM invocations fail.

Timeout strategy for AI

Rule of thumb:
lambda_timeout = llm_client_timeout × 1.5 + overhead

Where:
- llm_client_timeout = timeout you configure in your OpenAI/Anthropic SDK
- 1.5 = margin for the SDK's internal retries
- overhead = 2-5s for cold start, parsing, logging

import os
from openai import OpenAI

LAMBDA_TIMEOUT = int(os.environ.get("LAMBDA_TIMEOUT", "60"))

# The client timeout should be LOWER than the Lambda timeout
# so your code handles the error gracefully instead of Lambda killing it
client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    timeout=LAMBDA_TIMEOUT - 10  # 10s margin for cleanup
)

def handler(event, context):
    remaining_ms = context.get_remaining_time_in_millis()

    if remaining_ms < 15000:
        return {
            "statusCode": 408,
            "body": "Insufficient time remaining for LLM call"
        }

    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": event["prompt"]}],
            max_tokens=500,
        )
        return {"statusCode": 200, "body": response.choices[0].message.content}
    except Exception as e:
        return {"statusCode": 504, "body": f"LLM timeout: {str(e)}"}

Table of recommended timeouts

Use case                     Client timeout  Lambda timeout
──────────────────────────────────────────────────────────
Simple API call (1 LLM)     30s             45s
Prompt chain (2-3 calls)    30s per call    120s
RAG pipeline                45s             90s
Batch processing            60s per item    300s
Long generation (>2K tok)   60s             90s

The `context.get_remaining_time_in_millis()` trick

Lambda gives you access to the remaining time. Use it to make smart decisions:

def handler(event, context):
    remaining = context.get_remaining_time_in_millis()

    # If little time is left, use a faster model
    if remaining < 20000:
        model = "gpt-4o-mini"
        max_tokens = 200
    else:
        model = event.get("model", "gpt-4o")
        max_tokens = event.get("max_tokens", 1000)

    response = client.chat.completions.create(
        model=model,
        messages=event["messages"],
        max_tokens=max_tokens,
    )
    return {
        "statusCode": 200,
        "body": json.dumps({
            "content": response.choices[0].message.content,
            "model_used": model,
            "time_remaining_ms": context.get_remaining_time_in_millis()
        })
    }

Memory: It's not just RAM, it's CPU

Lambda's memory-CPU relationship

Lambda doesn't let you configure CPU directly. Instead, CPU scales proportionally with memory:

Memory     CPU (approx)     Equivalence
──────────────────────────────────────────
128 MB     ~0.07 vCPU       Unusable for AI
256 MB     ~0.14 vCPU       Very slow
512 MB     ~0.29 vCPU       Minimum viable for API calls
1024 MB    ~0.58 vCPU       Adequate for 1 LLM call
1769 MB    1 vCPU           Sweet spot for AI APIs
3538 MB    2 vCPU           Parallel processing
5307 MB    3 vCPU           Local embeddings
10240 MB   6 vCPU           In-memory models (edge case)

Why this changes everything for AI

# With 128MB (default): your function has ~7% of a CPU
# JSON parsing, serialization, response handling = SLOW
# The OpenAI SDK needs CPU for TLS handshake, parsing

# With 1769MB (1 vCPU): same function, ~14x more CPU
# TLS handshake: 200ms → 15ms
# JSON parsing of a large response: 50ms → 4ms
# Total overhead: ~300ms → ~25ms

Memory configuration

# SAM template
Resources:
  AiFunction:
    Type: AWS::Serverless::Function
    Properties:
      MemorySize: 1769  # MB — 1 full vCPU

# Terraform
resource "aws_lambda_function" "ai_endpoint" {
  memory_size = 1769
}

# AWS CLI
aws lambda update-function-configuration \
  --function-name ai-endpoint \
  --memory-size 1769

Right-sizing: LLM API calls vs local processing

The optimal memory depends on WHAT your function does:

Scenario A: Lambda invokes the OpenAI API
──────────────────────────────────────
The bottleneck is NETWORK (waiting for OpenAI's response).
Your function doesn't need much CPU — just send the request and parse the response.

→ Recommended: 512MB - 1024MB
→ More memory doesn't make OpenAI respond faster

Scenario B: Lambda processes embeddings locally
──────────────────────────────────────────────────
The bottleneck is CPU (vector computation, similarity search).

→ Recommended: 3072MB - 5120MB
→ More CPU = faster processing = Lambda finishes sooner

Scenario C: Lambda does RAG (embed + search + generate)
──────────────────────────────────────────────────────────
Combination: some CPU for embeddings + network wait for the LLM.

→ Recommended: 1769MB - 3072MB
→ Balance between CPU for processing and cost per second

The paradox: more memory can be CHEAPER

Lambda charges: price × (memory_GB) × (duration_seconds)

Example with a function that takes:
- 128MB:  10 seconds → 10s × 0.125 GB = 1.25 GB-s
- 1024MB: 2 seconds  → 2s × 1.0 GB   = 2.0 GB-s   (more expensive)
- 1769MB: 1.2 seconds → 1.2s × 1.73 GB = 2.08 GB-s (more expensive)

BUT if the function is I/O bound (waiting for OpenAI):
- 128MB:  8 seconds → 8s × 0.125 GB = 1.0 GB-s
- 512MB:  7.5 seconds → 7.5s × 0.5 GB = 3.75 GB-s (more expensive)
- 1769MB: 7.2 seconds → 7.2s × 1.73 GB = 12.46 GB-s (much more expensive)

Conclusion: for I/O-bound functions (which is most AI API calls),
more memory does NOT help and DOES cost more. Measure before scaling.

Lambda Power Tuning

What it is and why to use it

Lambda Power Tuning is an open-source tool that runs your function with different memory configurations and shows you the optimal point between cost and speed.

Typical Power Tuning result:

Memory    Duration   Cost          Ratio
────────────────────────────────────────────
128 MB    12.3s      $0.0000253    Slow and cheap
256 MB    8.1s       $0.0000333    Better
512 MB    4.2s       $0.0000344    Diminishing returns begin
1024 MB   2.8s       $0.0000459    Faster but more expensive
1769 MB   2.5s       $0.0000710    CPU doesn't help (I/O bound)
3072 MB   2.4s       $0.0001178    Waste

How to use it

# Step 1: Deploy the State Machine (once)
# Use the SAR (Serverless Application Repository)
aws serverlessrepo create-cloud-formation-change-set \
  --application-id arn:aws:serverlessrepo:us-east-1:451282441545:applications/aws-lambda-power-tuning \
  --stack-name lambda-power-tuning

# Step 2: Run the tuning
aws stepfunctions start-execution \
  --state-machine-arn arn:aws:states:us-east-1:ACCOUNT:stateMachine:powerTuningStateMachine \
  --input '{
    "lambdaARN": "arn:aws:lambda:us-east-1:ACCOUNT:function:ai-endpoint",
    "powerValues": [128, 256, 512, 1024, 1769, 3072],
    "num": 20,
    "payload": "{\"prompt\": \"Explain serverless in one sentence\"}"
  }'

Interpreting results for AI

What you'll typically see in a function that makes API calls to an LLM:

Speed vs Memory for an AI API call:

Duration (s)
│
12 ┤  ●
│
8  ┤     ●
│
4  ┤        ●
│
3  ┤           ●──●──●──●    ← Plateau: I/O bound
│
└──┬──┬──┬──┬──┬──┬──┬─ Memory (MB)
  128 256 512 1K 1.7K 3K 5K

After ~512MB-1024MB, more memory doesn't reduce
the duration because the bottleneck is the network.

Practical conclusion

For functions that call LLM APIs (the most common case in this guide):

Recommended configuration:
├── Memory: 512MB - 1024MB
├── Timeout: 45-60s
├── Reason: I/O bound, more CPU doesn't help
└── Exception: if you do heavy JSON processing, bump to 1769MB

For functions that process data locally (embeddings, transformations):

Recommended configuration:
├── Memory: 1769MB - 3072MB
├── Timeout: 60-120s
├── Reason: CPU bound, more memory = more CPU = faster
└── Exception: if the data fits in 512MB of RAM, don't bump up for CPU

Complete Configuration: AI Endpoint Example

# handler.py — Lambda handler with optimized timeout and memory
import json
import os
import logging
import time
from openai import OpenAI

logger = logging.getLogger()
logger.setLevel(logging.INFO)

LAMBDA_TIMEOUT = int(os.environ.get("LAMBDA_TIMEOUT", "60"))
MIN_REMAINING_MS = 10000

client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    timeout=LAMBDA_TIMEOUT - 10,
    max_retries=1,
)

def handler(event, context):
    start_time = time.time()

    remaining_ms = context.get_remaining_time_in_millis()
    if remaining_ms < MIN_REMAINING_MS:
        logger.warning(f"Insufficient time: {remaining_ms}ms remaining")
        return {
            "statusCode": 408,
            "body": json.dumps({"error": "Insufficient time for LLM call"})
        }

    try:
        body = json.loads(event.get("body", "{}"))
        prompt = body.get("prompt", "")
        max_tokens = min(body.get("max_tokens", 500), 2000)
    except (json.JSONDecodeError, AttributeError):
        return {
            "statusCode": 400,
            "body": json.dumps({"error": "Invalid request body"})
        }

    if not prompt:
        return {
            "statusCode": 400,
            "body": json.dumps({"error": "prompt is required"})
        }

    try:
        response = client.chat.completions.create(
            model=os.environ.get("MODEL_NAME", "gpt-4o-mini"),
            messages=[
                {"role": "system", "content": "Respond concisely and helpfully."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=max_tokens,
        )

        duration_ms = round((time.time() - start_time) * 1000)
        answer = response.choices[0].message.content
        tokens = response.usage.total_tokens

        logger.info(f"LLM call completed in {duration_ms}ms, {tokens} tokens")

        return {
            "statusCode": 200,
            "headers": {"Content-Type": "application/json"},
            "body": json.dumps({
                "answer": answer,
                "tokens_used": tokens,
                "duration_ms": duration_ms,
                "remaining_ms": context.get_remaining_time_in_millis()
            })
        }

    except Exception as e:
        logger.error(f"LLM error after {round((time.time() - start_time) * 1000)}ms: {e}")
        return {
            "statusCode": 502,
            "body": json.dumps({"error": f"LLM call failed: {str(e)}"})
        }

# template.yaml (SAM) — Configuration optimized for AI
AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31

Globals:
  Function:
    Runtime: python3.11
    Architectures:
      - arm64  # Graviton: 20% cheaper, similar performance for I/O bound

Resources:
  AiEndpoint:
    Type: AWS::Serverless::Function
    Properties:
      Handler: handler.handler
      CodeUri: ./src/
      MemorySize: 768  # Sweet spot for API calls: enough CPU without waste
      Timeout: 60
      Environment:
        Variables:
          OPENAI_API_KEY: !Ref OpenAiApiKey
          MODEL_NAME: gpt-4o-mini
          LAMBDA_TIMEOUT: "60"
      Events:
        ApiEvent:
          Type: Api
          Properties:
            Path: /ask
            Method: post

Parameters:
  OpenAiApiKey:
    Type: String
    NoEcho: true

Cold Start and Its Relationship with Memory

More memory also reduces the cold start — initialization has more CPU available:

Typical cold start for an AI function (with the openai SDK):

Memory      Cold start    Warm invocation
──────────────────────────────────────────
128 MB      ~8-12s        ~3-5s
512 MB      ~3-5s         ~2-4s
1024 MB     ~1.5-3s       ~2-3s
1769 MB     ~1-2s         ~2-3s
3072 MB     ~0.8-1.5s     ~2-3s

The cold start improves with more memory (more CPU for imports).
The warm invocation barely changes (I/O bound waiting for OpenAI).

# Optimize imports to reduce the cold start
# Imports run OUTSIDE the handler (during init)

# ✅ Import only what's needed
from openai import OpenAI
import json
import os

# ❌ Don't import heavy libraries you don't need
# import pandas  # +500ms to the cold start
# import numpy   # +300ms to the cold start
# import langchain  # +800ms to the cold start — import only the specific modules

Monitoring Timeout and Memory

CloudWatch Metrics

# Automatic metrics that Lambda sends to CloudWatch:
# - Duration: execution time (ms)
# - Max Memory Used: memory actually used (MB)
# - Timeouts: invocations that exceeded the timeout

# Query metrics with the AWS CLI
aws cloudwatch get-metric-statistics \
  --namespace AWS/Lambda \
  --metric-name Duration \
  --dimensions Name=FunctionName,Value=ai-endpoint \
  --start-time 2026-03-01T00:00:00Z \
  --end-time 2026-03-08T00:00:00Z \
  --period 3600 \
  --statistics Average Maximum p99

Custom metrics for AI

import json
import time
import os
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    timeout=50,
)

def handler(event, context):
    start = time.time()
    body = json.loads(event.get("body", "{}"))

    llm_start = time.time()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": body["prompt"]}],
        max_tokens=500,
    )
    llm_duration = time.time() - llm_start

    total_duration = time.time() - start
    memory_used = context.memory_limit_in_mb

    # Structured logging for CloudWatch Insights
    print(json.dumps({
        "metric": "llm_call",
        "llm_duration_ms": round(llm_duration * 1000),
        "total_duration_ms": round(total_duration * 1000),
        "overhead_ms": round((total_duration - llm_duration) * 1000),
        "tokens": response.usage.total_tokens,
        "memory_configured_mb": memory_used,
        "remaining_ms": context.get_remaining_time_in_millis(),
    }))

    return {
        "statusCode": 200,
        "body": json.dumps({"answer": response.choices[0].message.content})
    }

# CloudWatch Insights query to analyze performance
fields @timestamp, metric, llm_duration_ms, overhead_ms, tokens, memory_configured_mb
| filter metric = "llm_call"
| stats avg(llm_duration_ms) as avg_llm,
        max(llm_duration_ms) as p100_llm,
        avg(overhead_ms) as avg_overhead
  by bin(1h)

ARM64 Architecture (Graviton)

Lambda offers two architectures: x86_64 and arm64 (Graviton). For AI functions that call APIs:

x86_64 vs arm64 for AI API calls:

Aspect            x86_64       arm64 (Graviton)
──────────────────────────────────────────────────
Price             $0.0000166667/GB-s   $0.0000133334/GB-s
Difference        Base         20% cheaper
Cold start        Similar      Similar (sometimes 5% better)
Compatibility     Everything   Almost everything (except x86 binaries)
OpenAI SDK        ✅            ✅
Anthropic SDK     ✅            ✅
numpy/scipy       ✅            ✅ (arm64 wheels available)

Recommendation: use arm64 whenever you don't have dependencies
that require specific x86 binaries.

Troubleshooting

Problem 1: "Task timed out after X seconds"

# The function exceeded its timeout
# Diagnosis:
aws logs filter-log-events \
  --log-group-name /aws/lambda/ai-endpoint \
  --filter-pattern "Task timed out"

# Solutions:
# 1. Increase the Lambda timeout
# 2. Lower the client timeout to handle the error in your code
# 3. Reduce max_tokens for shorter responses
# 4. Use a faster model (gpt-4o-mini vs gpt-4o)

Problem 2: "Runtime exited with error: signal: killed" (OOM)

# The function exceeded its memory limit
# Lambda kills it immediately (Out of Memory)

# Diagnosis: look for "Max Memory Used" near the limit
aws cloudwatch get-metric-statistics \
  --namespace AWS/Lambda \
  --metric-name "Max Memory Used" \
  --dimensions Name=FunctionName,Value=ai-endpoint \
  --period 300 --statistics Maximum

# If Max Memory Used > 90% of MemorySize → increase the memory
# For AI API calls, you rarely need more than 512MB of real RAM
# If OOM happens, you're probably loading something large into memory

Problem 3: "The function is slow but doesn't time out"

# Possible causes:
# 1. Memory too low → insufficient CPU for the TLS handshake
#    Solution: bump to 512MB minimum

# 2. Cold start — first invocation after inactivity
#    Solution: provisioned concurrency or keep-warm

# 3. The LLM is slow (not your Lambda)
#    Diagnosis: check llm_duration_ms vs overhead_ms in the logs
#    If llm_duration_ms is >90% of the total → the bottleneck is OpenAI, not Lambda

Problem 4: "Costs higher than expected"

# Diagnosis: check real vs configured duration
aws logs filter-log-events \
  --log-group-name /aws/lambda/ai-endpoint \
  --filter-pattern "REPORT" \
  --limit 20

# Each REPORT line shows:
# Duration: 2345.67 ms  Billed Duration: 2400 ms  Memory Size: 1769 MB  Max Memory Used: 180 MB
#                                                   ^^^^^^^^^^^^^^^^^     ^^^^^^^^^^^^^^^^^^^^^^
#                                                   CONFIGURED memory     REAL memory used

# If Max Memory Used << Memory Size → you're wasting money
# Lower the memory until Max Memory Used is ~70% of Memory Size

Hands-On Exercises

Exercise 1: Configure adaptive timeout

Write a handler that uses context.get_remaining_time_in_millis() to decide: if more than 30s is left, use gpt-4o with 1000 tokens; if 15-30s is left, use gpt-4o-mini with 300 tokens; if less than 15s is left, return an error without calling the LLM.

See solution

import json
import os
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], timeout=50)

def handler(event, context):
    remaining_ms = context.get_remaining_time_in_millis()

    if remaining_ms < 15000:
        return {
            "statusCode": 408,
            "body": json.dumps({
                "error": "Not enough time for LLM call",
                "remaining_ms": remaining_ms
            })
        }

    if remaining_ms > 30000:
        model = "gpt-4o"
        max_tokens = 1000
    else:
        model = "gpt-4o-mini"
        max_tokens = 300

    body = json.loads(event.get("body", "{}"))
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": body["prompt"]}],
        max_tokens=max_tokens,
    )

    return {
        "statusCode": 200,
        "body": json.dumps({
            "answer": response.choices[0].message.content,
            "model_used": model,
            "remaining_ms": context.get_remaining_time_in_millis()
        })
    }

Exercise 2: Calculate the optimal memory

Your AI function has this CloudWatch data:

With 256MB: average duration 8.2s, max memory used 120MB
With 512MB: average duration 4.1s, max memory used 135MB
With 1024MB: average duration 3.8s, max memory used 142MB
With 1769MB: average duration 3.6s, max memory used 148MB

Calculate the cost per invocation for each configuration and determine which is optimal. Price: $0.0000166667 per GB-second.

See solution

configs = [
    {"memory_mb": 256,  "duration_s": 8.2, "max_used_mb": 120},
    {"memory_mb": 512,  "duration_s": 4.1, "max_used_mb": 135},
    {"memory_mb": 1024, "duration_s": 3.8, "max_used_mb": 142},
    {"memory_mb": 1769, "duration_s": 3.6, "max_used_mb": 148},
]

price_per_gb_second = 0.0000166667

print(f"{'Memory':>8} {'Duration':>10} {'GB-s':>8} {'Cost/inv':>12} {'Used%':>8}")
print("─" * 52)

for c in configs:
    gb_seconds = (c["memory_mb"] / 1024) * c["duration_s"]
    cost = gb_seconds * price_per_gb_second
    used_pct = (c["max_used_mb"] / c["memory_mb"]) * 100
    print(f"{c['memory_mb']:>6}MB {c['duration_s']:>8.1f}s {gb_seconds:>8.2f} ${cost:>11.7f} {used_pct:>6.1f}%")

# Result:
#   Memory   Duration     GB-s     Cost/inv    Used%
# ────────────────────────────────────────────────────
#   256MB       8.2s     2.05  $0.0000342    46.9%
#   512MB       4.1s     2.05  $0.0000342    26.4%
#  1024MB       3.8s     3.80  $0.0000633    13.9%
#  1769MB       3.6s     6.23  $0.0001038     8.4%
#
# 256MB and 512MB cost the same ($0.0000342), but 512MB is 2x faster.
# The duration barely improves from 512MB to 1769MB (I/O bound).
# → Optimal: 512MB — same cost as 256MB, twice as fast.
# → 1024MB+ only if you need to shave off another 0.3-0.5s and it's worth the 85% cost increase.

Exercise 3: Structured logging for monitoring

Add structured logging to your handler that emits: llm_duration_ms, total_duration_ms, overhead_ms, tokens_used, model, memory_configured_mb, and cold_start (boolean). The cold start is detected with a global variable that initializes to True and changes to False after the first invocation.

See solution

import json
import os
import time
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], timeout=50)

IS_COLD_START = True

def handler(event, context):
    global IS_COLD_START
    was_cold = IS_COLD_START
    IS_COLD_START = False

    start = time.time()
    body = json.loads(event.get("body", "{}"))

    llm_start = time.time()
    response = client.chat.completions.create(
        model=os.environ.get("MODEL_NAME", "gpt-4o-mini"),
        messages=[{"role": "user", "content": body["prompt"]}],
        max_tokens=body.get("max_tokens", 500),
    )
    llm_end = time.time()

    total_end = time.time()

    llm_ms = round((llm_end - llm_start) * 1000)
    total_ms = round((total_end - start) * 1000)

    print(json.dumps({
        "event": "llm_invocation",
        "cold_start": was_cold,
        "llm_duration_ms": llm_ms,
        "total_duration_ms": total_ms,
        "overhead_ms": total_ms - llm_ms,
        "tokens_used": response.usage.total_tokens,
        "model": response.model,
        "memory_configured_mb": int(context.memory_limit_in_mb),
    }))

    return {
        "statusCode": 200,
        "body": json.dumps({
            "answer": response.choices[0].message.content,
            "cold_start": was_cold,
        })
    }

Exercise 4: Timeout safety net with fallback

Implement a handler that, if the main LLM call (gpt-4o) takes more than 20 seconds, cancels it and retries with gpt-4o-mini using a simplified prompt. Use asyncio with wait_for or threading with a timeout.

See solution

import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], timeout=25)
executor = ThreadPoolExecutor(max_workers=1)

def call_llm(model, prompt, max_tokens):
    return client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
    )

def handler(event, context):
    body = json.loads(event.get("body", "{}"))
    prompt = body["prompt"]

    model_used = "gpt-4o"
    fallback_used = False

    try:
        future = executor.submit(call_llm, "gpt-4o", prompt, 1000)
        response = future.result(timeout=20)
    except TimeoutError:
        model_used = "gpt-4o-mini"
        fallback_used = True
        simplified = f"Answer briefly: {prompt}"
        response = call_llm("gpt-4o-mini", simplified, 300)
    except Exception as e:
        return {
            "statusCode": 502,
            "body": json.dumps({"error": str(e)})
        }

    return {
        "statusCode": 200,
        "body": json.dumps({
            "answer": response.choices[0].message.content,
            "model_used": model_used,
            "fallback_used": fallback_used,
            "tokens": response.usage.total_tokens,
        })
    }

Summary

Lambda's default timeout (3s) is insufficient for AI. Configure 45-60s minimum for LLM API calls.
The OpenAI/Anthropic client timeout should be LOWER than the Lambda timeout — so your code handles the error, not Lambda.
context.get_remaining_time_in_millis() lets you make smart decisions (fast vs slow model, abort if there's no time).
More memory = more CPU (proportional). 1769MB = 1 full vCPU.
For AI API calls (I/O bound), 512-1024MB is the sweet spot. More memory doesn't make OpenAI respond faster.
For local processing (CPU bound), 1769-3072MB. More CPU = faster = shorter duration.
Lambda Power Tuning gives you real data for your function — don't guess, measure.
arm64 (Graviton) is 20% cheaper with similar performance for AI workloads.
Monitor Memory Used vs Memory Configured. If you use 150MB of 1769MB configured, you're throwing money away.

Additional Resources

Lambda Memory and CPU — Official documentation on memory configuration
Lambda Power Tuning — Open-source tool to optimize memory
Lambda Timeout Best Practices — Timeout configuration
Lambda Graviton2 (arm64) — ARM architecture for Lambda
CloudWatch Insights for Lambda — Queries to analyze performance
OpenAI API Latency — Latency optimization in OpenAI calls
Lambda Pricing — Current pricing model
AWS Well-Architected Serverless Lens — Serverless best practices

← Previous4. Cold Starts in AI Next →6. API Gateway Integration

5. Timeout and Memory Configuration for AI

Overview

Timeout: The clock that shows no mercy

How Lambda's timeout works

Timeout configuration

Why the timeout matters more for AI

Timeout strategy for AI

Table of recommended timeouts

The context.get_remaining_time_in_millis() trick

Memory: It's not just RAM, it's CPU

Lambda's memory-CPU relationship

Why this changes everything for AI

Memory configuration

Right-sizing: LLM API calls vs local processing

The paradox: more memory can be CHEAPER

Lambda Power Tuning

What it is and why to use it

How to use it

Interpreting results for AI

Practical conclusion

Complete Configuration: AI Endpoint Example

Cold Start and Its Relationship with Memory

Monitoring Timeout and Memory

CloudWatch Metrics

Custom metrics for AI

ARM64 Architecture (Graviton)

Troubleshooting

Problem 1: "Task timed out after X seconds"

Problem 2: "Runtime exited with error: signal: killed" (OOM)

Problem 3: "The function is slow but doesn't time out"

Problem 4: "Costs higher than expected"

Hands-On Exercises

Exercise 1: Configure adaptive timeout

Exercise 2: Calculate the optimal memory

Exercise 3: Structured logging for monitoring

Exercise 4: Timeout safety net with fallback

Summary

Additional Resources

The `context.get_remaining_time_in_millis()` trick