Run Batch Inference at Scale

{"prompt": "Summarize this review: 'Great product, fast shipping.'", "id": "r001"}
{"prompt": "Summarize this review: 'Arrived broken. Returning.'", "id": "r002"}

runcrate instances create --name batch-job --gpu H100 --template ubuntu-inference
runcrate instances status batch-job

runcrate cp ./prompts.jsonl batch-job:/workspace/prompts.jsonl

runcrate ssh batch-job -- "pip install vllm"

# batch_infer.py
import json
from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-3.1-70B-Instruct", max_model_len=4096)
params = SamplingParams(max_tokens=256, temperature=0.1)

with open("/workspace/prompts.jsonl") as f:
    items = [json.loads(line) for line in f]

prompts = [item["prompt"] for item in items]
outputs = llm.generate(prompts, params)

with open("/workspace/results.jsonl", "w") as f:
    for item, output in zip(items, outputs):
        f.write(json.dumps({
            "id": item["id"],
            "response": output.outputs[0].text,
        }) + "\n")

print(f"Processed {len(items)} prompts.")

runcrate cp ./batch_infer.py batch-job:/workspace/batch_infer.py
runcrate ssh batch-job -- "cd /workspace && python batch_infer.py"

runcrate cp batch-job:/workspace/results.jsonl ./results.jsonl
runcrate instances delete batch-job

runcrate ssh batch-job -- nvidia-smi
runcrate ssh batch-job -- "wc -l /workspace/results.jsonl"