← Back to DevBytes

How to Use Spot Instances for LLM Batch Inference

How to Use Spot Instances for LLM Batch Inference

Running large language model (LLM) batch inference at scale can be prohibitively expensive. Spot instances — spare cloud compute capacity offered at steep discounts — can cut your GPU bill by 60–90%. The trade-off is that these instances can be reclaimed with little warning. For batch inference workloads, which are typically latency-insensitive and checkpoint-friendly, that trade-off is almost always worth it. This tutorial walks through what spot instances are, why they fit batch inference, how to architect a resilient pipeline around them, and the best practices that keep your jobs safe.

What Are Spot Instances?

Spot instances (called "Spot VMs" on Google Cloud and "Spot Instances" on AWS) are unused datacenter capacity that cloud providers sell at a significant discount. The catch: when demand for that capacity rises, the provider can interrupt or terminate your instance with as little as 30 seconds of notice. Because of this, spot pricing is dynamic and varies by region, instance type, and time of day.

For LLM workloads, the most relevant instance families are GPU-accelerated: AWS p4d, p5, g5, g6; GCP a2 and a3; and Azure ND and NC series. Discounts on these can be substantial — often 70% or more off on-demand rates.

Why Spot Instances Fit LLM Batch Inference

Batch inference has three characteristics that make it an ideal match for spot capacity:

Together, these properties mean the unreliability of spot capacity costs you time, not correctness — and the time cost is usually small compared to the dollar savings.

Architecture Overview

A robust spot-based batch inference pipeline has four components:

Step 1: Shard Your Input Data

Break your dataset into small, independent shards. Smaller shards mean less wasted work when an instance is interrupted. A good rule of thumb is 5–15 minutes of processing per shard — long enough to avoid queue overhead, short enough that re-running a shard is cheap.

import json
import os

INPUT_FILE = "prompts.jsonl"
SHARD_DIR = "shards"
SHARD_SIZE = 500  # prompts per shard

os.makedirs(SHARD_DIR, exist_ok=True)

with open(INPUT_FILE) as f:
    prompts = [json.loads(line) for line in f]

for i in range(0, len(prompts), SHARD_SIZE):
    shard = prompts[i:i + SHARD_SIZE]
    shard_id = i // SHARD_SIZE
    with open(f"{SHARD_DIR}/shard_{shard_id:05d}.jsonl", "w") as out:
        for p in shard:
            out.write(json.dumps(p) + "\n")

print(f"Wrote {(len(prompts) + SHARD_SIZE - 1) // SHARD_SIZE} shards")

Step 2: Upload Shards and Track State

Upload shards to object storage and maintain a manifest that records which shards are pending, in-progress, or complete. A simple approach is a JSON manifest file, but for production scale use a database (DynamoDB, Firestore, Postgres) so multiple workers can update it concurrently.

import boto3
import json

s3 = boto3.client("s3")
BUCKET = "my-llm-batch-bucket"
PREFIX = "shards"

manifest = {"pending": [], "in_progress": [], "complete": []}

for fname in sorted(os.listdir("shards")):
    key = f"{PREFIX}/{fname}"
    s3.upload_file(f"shards/{fname}", BUCKET, key)
    manifest["pending"].append(key)

s3.put_object(
    Bucket=BUCKET,
    Key="manifest.json",
    Body=json.dumps(manifest).encode(),
)
print("Uploaded shards and manifest.")

Step 3: Write an Interruption-Aware Worker

The worker runs on each spot instance. It loops: claim a shard, download it, run inference, upload results, mark the shard complete. Critically, it must listen for the spot interruption signal so it can stop cleanly rather than losing in-flight work.

On AWS, the interruption notice arrives via the instance metadata service at http://169.254.169.254/latest/meta-data/spot/instance-action. A 200 response means termination is imminent and includes a two-minute window.

import json
import os
import time
import urllib.request
import threading
import boto3
from transformers import AutoModelForCausalLM, AutoTokenizer

BUCKET = "my-llm-batch-bucket"
MODEL_NAME = "meta-llama/Llama-2-7b-hf"
CHECKPOINT_DIR = "/tmp/checkpoints"

os.makedirs(CHECKPOINT_DIR, exist_ok=True)
s3 = boto3.client("s3")

# --- Interruption watcher ---
INTERRUPTED = False

def watch_interruption():
    global INTERRUPTED
    url = "http://169.254.169.254/latest/meta-data/spot/instance-action"
    while not INTERRUPTED:
        try:
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=2) as resp:
                if resp.status == 200:
                    print("Spot interruption notice received!")
                    INTERRUPTED = True
                    return
        except Exception:
            pass  # 404 means no interruption scheduled yet
        time.sleep(5)

threading.Thread(target=watch_interruption, daemon=True).start()

# --- Load model once per instance ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, device_map="auto", torch_dtype="auto"
)

def claim_shard():
    """Atomically claim a pending shard. In production, use DynamoDB
    conditional writes or SQS ReceiveMessage for true atomicity."""
    manifest = json.loads(s3.get_object(Bucket=BUCKET, Key="manifest.json")["Body"].read())
    if not manifest["pending"]:
        return None
    key = manifest["pending"].pop(0)
    manifest["in_progress"].append(key)
    s3.put_object(Bucket=BUCKET, Key="manifest.json", Body=json.dumps(manifest).encode())
    return key

def mark_complete(key):
    manifest = json.loads(s3.get_object(Bucket=BUCKET, Key="manifest.json")["Body"].read())
    manifest["in_progress"].remove(key)
    manifest["complete"].append(key)
    s3.put_object(Bucket=BUCKET, Key="manifest.json", Body=json.dumps(manifest).encode())

def run_inference(prompts):
    results = []
    for p in prompts:
        if INTERRUPTED:
            break
        inputs = tokenizer(p["prompt"], return_tensors="pt").to(model.device)
        out = model.generate(**inputs, max_new_tokens=256)
        text = tokenizer.decode(out[0], skip_special_tokens=True)
        results.append({"id": p["id"], "output": text})
    return results

# --- Main loop ---
while not INTERRUPTED:
    key = claim_shard()
    if key is None:
        print("No more work. Exiting.")
        break

    local_path = f"/tmp/{os.path.basename(key)}"
    s3.download_file(BUCKET, key, local_path)

    with open(local_path) as f:
        prompts = [json.loads(line) for line in f]

    results = run_inference(prompts)

    # Upload partial or full results
    result_key = key.replace("shards/", "results/").replace(".jsonl", "_out.jsonl")
    s3.put_object(
        Bucket=BUCKET,
        Key=result_key,
        Body="\n".join(json.dumps(r) for r in results).encode(),
    )

    if INTERRUPTED:
        # Re-queue the shard so another worker can finish it
        manifest = json.loads(s3.get_object(Bucket=BUCKET, Key="manifest.json")["Body"].read())
        manifest["in_progress"].remove(key)
        manifest["pending"].append(key)
        s3.put_object(Bucket=BUCKET, Key="manifest.json", Body=json.dumps(manifest).encode())
        print(f"Interrupted. Re-queued {key} with {len(results)}/{len(prompts)} done.")
    else:
        mark_complete(key)
        print(f"Completed {key}")

print("Worker shutting down.")

Step 4: Provision a Spot Fleet

Use a spot fleet request or an autoscaling group with spot capacity to launch multiple GPU instances. Specify multiple instance types to improve the chance that capacity is available — if you only request p4d.24xlarge, you may wait a long time, but allowing p4d.24xlarge or p5.48xlarge broadens your options.

aws ec2 request-spot-fleet \
  --spot-fleet-request-config file://spot-fleet-config.json

Example spot-fleet-config.json:

{
  "SpotPrice": "12.00",
  "TargetCapacity": 4,
  "IamFleetRole": "arn:aws:iam::123456789012:role/aws-ec2-spot-fleet-tagging-role",
  "LaunchSpecifications": [
    {
      "ImageId": "ami-0abcdef1234567890",
      "InstanceType": "p4d.24xlarge",
      "SubnetId": "subnet-0abc123",
      "UserData": "IyEvYmluL2Jhc2gKZG9ja2VyIHJ1biAtLXJtIG15LWxsbS13b3JrZXI6bGF0ZXN0",
      "IamInstanceProfile": {
        "Arn": "arn:aws:iam::123456789012:instance-profile/llm-worker"
      }
    },
    {
      "ImageId": "ami-0abcdef1234567890",
      "InstanceType": "g5.48xlarge",
      "SubnetId": "subnet-0abc123",
      "UserData": "IyEvYmluL2Jhc2gKZG9ja2VyIHJ1biAtLXJtIG15LWxsbS13b3JrZXI6bGF0ZXN0",
      "IamInstanceProfile": {
        "Arn": "arn:aws:iam::123456789012:instance-profile/llm-worker"
      }
    }
  ]
}

The UserData field (base64-encoded) should start your worker container on boot. Each instance boots, loads the model, and begins pulling shards from the queue.

Step 5: Handle Interruption Gracefully

The interruption watcher in Step 3 sets a flag that the inference loop checks between prompts. When the flag is set, the worker stops generating, uploads whatever results it has, and re-queues the incomplete shard. This ensures no work is silently lost.

For even tighter guarantees, you can also register an OS signal handler for SIGTERM, which is what AWS sends before the actual termination:

import signal

def handle_sigterm(signum, frame):
    global INTERRUPTED
    INTERRUPTED = True

signal.signal(signal.SIGTERM, handle_sigterm)

Best Practices

Conclusion

Spot instances turn LLM batch inference from a budget-breaking expense into a manageable, scalable operation. The key is designing for interruption from the start: shard your work, track state durably, watch for the reclaim signal, and re-queue incomplete shards automatically. With the architecture and code patterns in this tutorial, you can run large-scale batch inference at a fraction of on-demand cost while keeping your pipeline correct and resumable. Start with a small fleet, measure your interruption and re-queue rates, and scale up once your checkpointing logic is proven.

— Ad —

Google AdSense will appear here after approval

← Back to all articles