10.2k6 Script for LLM Endpoint Load Testing

// k6-llm-load-test.js
import http from 'k6/http';
import { check, sleep } from 'k6';
import { Trend, Counter, Rate } from 'k6/metrics';

// Custom LLM-specific metrics
const ttft = new Trend('time_to_first_token', true);
const totalGenTime = new Trend('total_generation_time', true);
const tokensPerSecond = new Trend('tokens_per_second');
const rateLimitHits = new Counter('rate_limit_hits');
const timeoutRate = new Rate('timeout_rate');

export const options = {
  scenarios: {
    // Simulate gradual ramp to find the breaking point
    ramp_to_limit: {
      executor: 'ramping-arrival-rate',
      startRate: 1,
      timeUnit: '1s',
      preAllocatedVUs: 50,
      maxVUs: 200,
      stages: [
        { duration: '2m', target: 5 },    // 5 req/s
        { duration: '3m', target: 10 },   // 10 req/s
        { duration: '3m', target: 20 },   // 20 req/s -- likely hits rate limits
        { duration: '2m', target: 5 },    // cool down
      ],
    },
  },
  thresholds: {
    time_to_first_token: ['p(95)<3000'],     // TTFT under 3s for 95th pctile
    total_generation_time: ['p(95)<15000'],   // Total gen under 15s
    tokens_per_second: ['avg>30'],            // At least 30 tok/s average
    timeout_rate: ['rate<0.05'],              // Under 5% timeouts
  },
};

// Varied prompts to simulate realistic usage
const prompts = [
  "Summarize the key differences between REST and GraphQL in 3 sentences.",
  "Write a Python function that validates an email address using regex.",
  "Explain the CAP theorem to a junior developer.",
  "Generate a SQL query to find the top 10 customers by revenue last quarter.",
  "What are the SOLID principles? Give a one-line explanation of each.",
];

export default function () {
  const prompt = prompts[Math.floor(Math.random() * prompts.length)];

  const payload = JSON.stringify({
    model: "gpt-5.5",
    messages: [{ role: "user", content: prompt }],
    max_tokens: 256,
    stream: false,
  });

  const startTime = Date.now();

  const res = http.post(
    'https://api.example.com/v1/chat/completions',
    payload,
    {
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${__ENV.LLM_API_KEY}`,
      },
      timeout: '30s',
    }
  );

  const elapsed = Date.now() - startTime;

  // Track rate limit responses separately
  if (res.status === 429) {
    rateLimitHits.add(1);
    console.warn(`Rate limited at ${new Date().toISOString()}`);
    sleep(5); // back off to avoid cascading rate limits
    return;
  }

  // Track timeouts
  timeoutRate.add(res.status === 0 || elapsed > 29000);

  if (res.status === 200) {
    const body = JSON.parse(res.body);
    const completionTokens = body.usage?.completion_tokens || 0;
    const totalTime = elapsed / 1000; // seconds

    // Approximate TTFT (for accurate TTFT, use streaming with k6 WebSocket)
    ttft.add(elapsed * 0.15); // rough heuristic for non-streaming
    totalGenTime.add(elapsed);

    if (totalTime > 0 && completionTokens > 0) {
      tokensPerSecond.add(completionTokens / totalTime);
    }

    check(res, {
      'status is 200': (r) => r.status === 200,
      'response has content': () => body.choices?.[0]?.message?.content?.length > 0,
      'under token budget': () => body.usage?.total_tokens < 1000,
    });
  }

  sleep(Math.random() * 2 + 0.5); // think time between requests
}