How to Run Machine Learning Models on AWS Lambda (Container Image Guide)

AWS Lambda can serve machine learning predictions without a server to manage, and for small and medium models it is often the cheapest way to do it. It is not the right tool for everything, though: Lambda has no GPU, limited memory and a cold start that grows with the size of your model. This guide shows a working pattern (a scikit-learn model packaged in a container image), tells you which Lambda limits matter, and helps you decide when to use something else.

Updated October 2026: I rewrote this 2024 post. The original described the steps in general terms with no code, suggested pickle as the default format, and did not mention container images, which are now the practical way to ship ML dependencies on Lambda. It also pointed to AWS Batch and SageMaker without explaining when each fits. This version has code I ran locally, current limits from AWS’s documentation, and a clear checklist for choosing between Lambda and SageMaker.

Is Lambda the Right Place for Your Model?

Lambda runs your code on demand, scales automatically and bills per request and per GB-second. That suits ML workloads that are small, CPU-only, spiky and latency-tolerant. Compare your model against the limits in the Lambda quotas page:

Limit	Value	What it means for ML
Memory	128 MB to 10,240 MB	CPU power scales with memory (about one vCPU at 1,769 MB), so more memory also means faster inference
Timeout	900 seconds (15 minutes)	Fine for inference, not for training
Container image size	10 GB uncompressed	Room for PyTorch or TensorFlow plus a model
Zip deployment package	50 MB zipped, 250 MB unzipped (including layers)	Usually too small once you add NumPy, SciPy and a framework
/tmp storage	512 MB to 10,240 MB	Where you can download a model at start-up
Synchronous payload	6 MB request and response	Send images or large batches through S3 instead of the request body
GPU	Not available	Large deep-learning models need another service

• Good fit: classical models (scikit-learn, XGBoost), small neural networks, text classifiers, preprocessing, and workloads with idle periods where paying for an always-on endpoint would be wasteful.
• Poor fit: large language models, anything needing a GPU, strict single-digit-millisecond latency, or constant high traffic where an always-on instance is cheaper per request.

If your model needs a GPU or you want managed model hosting with autoscaling and monitoring, look at Amazon SageMaker AI (see our SageMaker and AWS Batch guide) instead. If you are new to the underlying idea, our serverless architecture guide covers the trade-offs.

Why a Container Image

You can package Lambda code as a zip file or as a container image. ML libraries make zip files impractical: scikit-learn with its dependencies is already large, and PyTorch or TensorFlow are far past the 250 MB unzipped limit. A container image can be up to 10 GB, so the libraries and the model file can ship together. AWS provides base images for Python (including 3.13 and 3.14) that already contain the Lambda runtime interface, which is the easiest starting point.

Putting the model inside the image has a useful side effect: there is no download step at start-up. The model is on local disk when the function initialises, so cold starts are shorter than fetching it from S3 every time. The cost is that updating the model means rebuilding and redeploying the image. If you need to swap models without redeploying, keep the model in S3 and download it to /tmp on cold start instead (shown near the end).

Step 1: Train and Save a Model

This example trains a classifier on scikit-learn’s built-in iris dataset so you can follow it without any data. Use your own model and features in practice. Save it with joblib, which is scikit-learn’s recommended way to persist models.

import joblib
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression

X, y = load_iris(return_X_y=True)
model = LogisticRegression(max_iter=500).fit(X, y)
joblib.dump({"model": model, "classes": load_iris().target_names.tolist()}, "model.joblib")
print("saved model.joblib, accuracy on training data:", round(model.score(X, y), 3))

One security rule about model files: joblib and pickle files can run arbitrary code when loaded. Only load model files you created or that come from a source you trust, and keep the S3 bucket or image repository that holds them locked down, as the scikit-learn documentation on model persistence warns. Also pin the scikit-learn version: a model saved with one version is not guaranteed to load in another, so train and serve with the same version.

Step 2: Write the Handler

import json
import os
import time

import joblib

# Runs once per execution environment (cold start), not on every request.
_start = time.perf_counter()
_artifact = joblib.load(os.path.join(os.path.dirname(__file__), "model.joblib"))
MODEL, CLASSES = _artifact["model"], _artifact["classes"]
print(f"model loaded in {time.perf_counter() - _start:.3f}s")

N_FEATURES = 4


def _response(status, body):
    return {
        "statusCode": status,
        "headers": {"content-type": "application/json"},
        "body": json.dumps(body),
    }


def handler(event, context):
    try:
        payload = json.loads(event["body"]) if isinstance(event.get("body"), str) else event
        rows = payload["instances"]
        if not rows or any(len(r) != N_FEATURES for r in rows):
            raise ValueError(f"each instance needs exactly {N_FEATURES} numbers")
        rows = [[float(v) for v in r] for r in rows]
    except (KeyError, TypeError, ValueError, json.JSONDecodeError) as err:
        return _response(400, {"error": f"bad request: {err}"})

    labels = MODEL.predict(rows)
    probs = MODEL.predict_proba(rows)
    return _response(200, {
        "predictions": [
            {"label": CLASSES[int(l)], "confidence": round(float(max(p)), 4)}
            for l, p in zip(labels, probs)
        ]
    })

The important design choices:

• The model loads at module level, outside the handler. Lambda runs that code once per execution environment, so warm invocations reuse the loaded model instead of reading it again on every request.
• Input is validated and bad requests get a 400 with a clear message rather than an exception and a 500.
• The same code works with a function URL or API Gateway (JSON in event["body"]) and with a direct invoke (the event itself is the payload).
• The response contains only what the caller needs: a label and a confidence.

Step 3: Test Locally First

Before building any image, test the function logic in plain Python:

import json
import lambda_function as lf

r = lf.handler({"instances": [[5.1, 3.5, 1.4, 0.2], [6.7, 3.0, 5.2, 2.3]]}, None)
body = json.loads(r["body"]); print(r["statusCode"], body)
assert r["statusCode"] == 200 and [p["label"] for p in body["predictions"]] == ["setosa", "virginica"]
r = lf.handler({"body": json.dumps({"instances": [[1, 2, 3]]})}, None)
assert r["statusCode"] == 400, r
r = lf.handler({"body": "not json"}, None); assert r["statusCode"] == 400
r = lf.handler({"body": json.dumps({"instances": [["a", 1, 2, 3]]})}, None); assert r["statusCode"] == 400
print("all tests passed")

python train.py
python test_handler.py
# 200 {'predictions': [{'label': 'setosa', ...}, {'label': 'virginica', ...}]}
# all tests passed

I ran this with Python 3.13, scikit-learn 1.9.1 and joblib 1.6.0, and the tests passed. They check the model’s predictions on two known samples and three kinds of bad input. They do not exercise the Lambda runtime itself; the image test below does that.

Step 4: Build the Container Image

scikit-learn==1.9.1
joblib

FROM public.ecr.aws/lambda/python:3.13

COPY requirements.txt ${LAMBDA_TASK_ROOT}
RUN pip install --no-cache-dir -r requirements.txt

COPY model.joblib lambda_function.py ${LAMBDA_TASK_ROOT}/

CMD ["lambda_function.handler"]

${LAMBDA_TASK_ROOT} is the directory where Lambda expects your code, and the CMD names the handler as file.function. Build for the architecture your function will use. The AWS docs also require --provenance=false so the image is in a format Lambda accepts:

docker buildx build --platform linux/amd64 --provenance=false -t iris-classifier:v1 .

# run it locally with the Lambda runtime emulator built into the base image
docker run --platform linux/amd64 -p 9000:8080 iris-classifier:v1

curl "http://localhost:9000/2015-03-31/functions/function/invocations" \
  -d '{"instances": [[5.1, 3.5, 1.4, 0.2]]}'

Honest note: I could not run Docker on the machine where I tested this post, so the Dockerfile and these commands follow AWS’s documented pattern for Python base images but were not built here. The handler and model code above were run. If the build fails on your machine, the base image page linked above is the reference.

Step 5: Push to Amazon ECR and Create the Function

Lambda pulls container images from Amazon ECR in the same Region. Create a repository, push the image, then create the function. Replace 111122223333 with your account ID and use your own Region:

aws ecr create-repository --repository-name iris-classifier --region us-east-1 \
  --image-scanning-configuration scanOnPush=true

aws ecr get-login-password --region us-east-1 | \
  docker login --username AWS --password-stdin 111122223333.dkr.ecr.us-east-1.amazonaws.com

docker tag iris-classifier:v1 111122223333.dkr.ecr.us-east-1.amazonaws.com/iris-classifier:v1
docker push 111122223333.dkr.ecr.us-east-1.amazonaws.com/iris-classifier:v1

aws lambda create-function \
  --function-name iris-classifier \
  --package-type Image \
  --code ImageUri=111122223333.dkr.ecr.us-east-1.amazonaws.com/iris-classifier:v1 \
  --role arn:aws:iam::111122223333:role/lambda-ex \
  --memory-size 1024 --timeout 30

The role needs only the basic execution permissions (AWSLambdaBasicExecutionRole) for logging, unless your function also reads from S3 or other services. I validated the CreateFunction, UpdateFunctionCode and ECR CreateRepository parameters against the current boto3 service model, but did not run them in a live AWS account. Invoke it with:

aws lambda invoke --function-name iris-classifier \
  --cli-binary-format raw-in-base64-out \
  --payload '{"instances": [[5.1, 3.5, 1.4, 0.2]]}' out.json && cat out.json

To deploy a new image, push it and run aws lambda update-function-code --function-name iris-classifier --image-uri <new image URI> --publish. Lambda pins the image digest when you deploy, so pushing a new image under the same tag does not update the function until you run that command.

Alternative: Keep the Model in S3

If you want to update the model without redeploying code, store it in S3 and download it once per execution environment into /tmp:

import os
import boto3
import joblib

BUCKET, KEY = os.environ["MODEL_BUCKET"], os.environ["MODEL_KEY"]
LOCAL = "/tmp/model.joblib"

if not os.path.exists(LOCAL):  # only on cold start
    boto3.client("s3").download_file(BUCKET, KEY, LOCAL)
_artifact = joblib.load(LOCAL)

Give the function role s3:GetObject on that one object path, and nothing broader. The trade-off is a longer cold start (the download time is added) in exchange for easy model updates. Versioned S3 keys (models/v7/model.joblib) let you roll back by changing an environment variable.

Performance and Cost

• Measure the cold start. The handler above prints how long the model took to load; read it in CloudWatch Logs. Large frameworks add import time on top of the model itself.
• Tune memory by measuring. More memory gives more CPU, so inference often gets faster, and sometimes cheaper in total because the function runs for less time. Try 512, 1,024 and 2,048 MB and compare duration and cost.
• Use provisioned concurrency if cold starts hurt. It keeps a set number of execution environments initialised and ready (see the provisioned concurrency docs), at an extra cost.
• Batch where you can. The handler accepts several instances at once, which is far cheaper than one invocation per row.
• Know the break-even. Lambda bills per request and duration (pricing). Steady high traffic can make an always-on service cheaper; estimate with your real request rate and duration before committing.
• Monitor. Watch duration, errors, throttles and memory use in CloudWatch, and add an alarm on errors. Log input sizes and latency, never raw personal data.

Conclusion

Lambda is a good home for small CPU-based models: package the model and its libraries in a container image, load the model outside the handler, validate input and measure cold starts. Switch to SageMaker AI or another hosted service when you need a GPU, very large models or steady heavy traffic. Start with the local test above, then deploy the same code unchanged.