K6 Integration for Loadtesting

[burakemirsezen]

This branch adds a cluster-native loadtesting workflow to SML.

It introduces sml loadtest commands for launching a model and running k6 against it, running k6 against an already-running model, running against an external OpenAI-compatible URL, and batch-testing multiple model configs from YAML. k6 runs as its own SLURM job inside a container on the cluster, so load generation does not happen locally.

[robmsmt]

Initially I was thinking something like:

set -euo pipefail

: "${LOADTEST_SERVER_URL:?export LOADTEST_SERVER_URL}"
: "${LOADTEST_API_KEY:?export LOADTEST_API_KEY}"
LOADTEST_PROMPTS_FILE="${LOADTEST_PROMPTS_FILE:-/capstor/store/cscs/swissai/infra01/loadtest/prompts.json}"
LOADTEST_SCENARIO="${LOADTEST_SCENARIO:-throughput}"

MODEL="/capstor/store/cscs/swissai/infra01/hf_models/models/swiss-ai/Apertus-8B-Instruct-2509"
ENV="src/swiss_ai_model_launch/assets/envs/sglang.toml"
TIME="04:00:00"
SUFFIX="$(whoami)"

SERVED_1="swiss-ai/Apertus-8B-Instruct-2509-tp4-${SUFFIX}"
SERVED_2="swiss-ai/Apertus-8B-Instruct-2509-tp1-dp4-${SUFFIX}"

# --- Experiment 1: TP=4 ---
OUT_1=$(sml advanced \
  --firecrest-system clariden \
  --partition normal \
  --slurm-nodes 1 \
  --slurm-time "$TIME" \
  --serving-framework sglang \
  --slurm-environment "$ENV" \
  --framework-args "--model-path $MODEL \
    --served-model-name $SERVED_1 \
    --host 0.0.0.0 --port 8080 \
    --tp-size 4 \
    --enable-metrics")
echo "$OUT_1"
JOB_1=$(echo "$OUT_1" | grep "Job submitted:" | awk '{print $3}')

echo "Waiting for $SERVED_1 to be healthy..."
until curl -fsS -H "Authorization: Bearer $LOADTEST_API_KEY" "$LOADTEST_SERVER_URL/v1/models" | grep -q "$SERVED_1"; do
  sleep 30
done

sml loadtest run \
  --firecrest-system clariden \
  --partition normal \
  --loadtest-server-url "$LOADTEST_SERVER_URL" \
  --loadtest-api-key "$LOADTEST_API_KEY" \
  --loadtest-model "$SERVED_1" \
  --loadtest-scenario "$LOADTEST_SCENARIO" \
  --loadtest-prompts-file "$LOADTEST_PROMPTS_FILE" \
  --no-wait-until-healthy

scancel "$JOB_1"

# --- Experiment 2: TP=1, DP=4 ---
OUT_2=$(sml advanced \
  --firecrest-system clariden \
  --partition normal \
  --slurm-nodes 1 \
  --slurm-time "$TIME" \
  --serving-framework sglang \
  --slurm-environment "$ENV" \
  --framework-args "--model-path $MODEL \
    --served-model-name $SERVED_2 \
    --host 0.0.0.0 --port 8080 \
    --tp-size 1 --dp-size 4 \
    --enable-metrics")
echo "$OUT_2"
JOB_2=$(echo "$OUT_2" | grep "Job submitted:" | awk '{print $3}')

echo "Waiting for $SERVED_2 to be healthy..."
until curl -fsS -H "Authorization: Bearer $LOADTEST_API_KEY" "$LOADTEST_SERVER_URL/v1/models" | grep -q "$SERVED_2"; do
  sleep 30
done

sml loadtest run \
  --firecrest-system clariden \
  --partition normal \
  --loadtest-server-url "$LOADTEST_SERVER_URL" \
  --loadtest-api-key "$LOADTEST_API_KEY" \
  --loadtest-model "$SERVED_2" \
  --loadtest-scenario "$LOADTEST_SCENARIO" \
  --loadtest-prompts-file "$LOADTEST_PROMPTS_FILE" \
  --no-wait-until-healthy

scancel "$JOB_2"

the sml loadtest advanced section has a clear advantage with --cancel-after-loadtest ... what is not clear to me is that we've added some complexity here and now have to maintain another set of flags in (def _add_loadtest_arguments) and if it's worth that trade off? Could we not import these so don't repeat the flags?

[burakemirsezen]

Initially I was thinking something like:

set -euo pipefail

: "${LOADTEST_SERVER_URL:?export LOADTEST_SERVER_URL}"
: "${LOADTEST_API_KEY:?export LOADTEST_API_KEY}"
LOADTEST_PROMPTS_FILE="${LOADTEST_PROMPTS_FILE:-/capstor/store/cscs/swissai/infra01/loadtest/prompts.json}"
LOADTEST_SCENARIO="${LOADTEST_SCENARIO:-throughput}"

MODEL="/capstor/store/cscs/swissai/infra01/hf_models/models/swiss-ai/Apertus-8B-Instruct-2509"
ENV="src/swiss_ai_model_launch/assets/envs/sglang.toml"
TIME="04:00:00"
SUFFIX="$(whoami)"

SERVED_1="swiss-ai/Apertus-8B-Instruct-2509-tp4-${SUFFIX}"
SERVED_2="swiss-ai/Apertus-8B-Instruct-2509-tp1-dp4-${SUFFIX}"

# --- Experiment 1: TP=4 ---
OUT_1=$(sml advanced \
  --firecrest-system clariden \
  --partition normal \
  --slurm-nodes 1 \
  --slurm-time "$TIME" \
  --serving-framework sglang \
  --slurm-environment "$ENV" \
  --framework-args "--model-path $MODEL \
    --served-model-name $SERVED_1 \
    --host 0.0.0.0 --port 8080 \
    --tp-size 4 \
    --enable-metrics")
echo "$OUT_1"
JOB_1=$(echo "$OUT_1" | grep "Job submitted:" | awk '{print $3}')

echo "Waiting for $SERVED_1 to be healthy..."
until curl -fsS -H "Authorization: Bearer $LOADTEST_API_KEY" "$LOADTEST_SERVER_URL/v1/models" | grep -q "$SERVED_1"; do
  sleep 30
done

sml loadtest run \
  --firecrest-system clariden \
  --partition normal \
  --loadtest-server-url "$LOADTEST_SERVER_URL" \
  --loadtest-api-key "$LOADTEST_API_KEY" \
  --loadtest-model "$SERVED_1" \
  --loadtest-scenario "$LOADTEST_SCENARIO" \
  --loadtest-prompts-file "$LOADTEST_PROMPTS_FILE" \
  --no-wait-until-healthy

scancel "$JOB_1"

# --- Experiment 2: TP=1, DP=4 ---
OUT_2=$(sml advanced \
  --firecrest-system clariden \
  --partition normal \
  --slurm-nodes 1 \
  --slurm-time "$TIME" \
  --serving-framework sglang \
  --slurm-environment "$ENV" \
  --framework-args "--model-path $MODEL \
    --served-model-name $SERVED_2 \
    --host 0.0.0.0 --port 8080 \
    --tp-size 1 --dp-size 4 \
    --enable-metrics")
echo "$OUT_2"
JOB_2=$(echo "$OUT_2" | grep "Job submitted:" | awk '{print $3}')

echo "Waiting for $SERVED_2 to be healthy..."
until curl -fsS -H "Authorization: Bearer $LOADTEST_API_KEY" "$LOADTEST_SERVER_URL/v1/models" | grep -q "$SERVED_2"; do
  sleep 30
done

sml loadtest run \
  --firecrest-system clariden \
  --partition normal \
  --loadtest-server-url "$LOADTEST_SERVER_URL" \
  --loadtest-api-key "$LOADTEST_API_KEY" \
  --loadtest-model "$SERVED_2" \
  --loadtest-scenario "$LOADTEST_SCENARIO" \
  --loadtest-prompts-file "$LOADTEST_PROMPTS_FILE" \
  --no-wait-until-healthy

scancel "$JOB_2"

the sml loadtest advanced section has a clear advantage with --cancel-after-loadtest ... what is not clear to me is that we've added some complexity here and now have to maintain another set of flags in (def _add_loadtest_arguments) and if it's worth that trade off? Could we not import these so don't repeat the flags?

Should I just get rid of the loadtest batch then?

[burakemirsezen]

I removed batch submitting, made it so that the parser uses advanced arguments underneath, did some cleanup and a lot of debugging. Also pinned the k6 version. It is in a working state now, we can improve it later on.

[burakemirsezen]

I added prometheus remote write for k6 and added/removed some flags for loadtest.

This should be more or less good to go now. Lmk if there are any issues or things that we want to add.

[sonarqubecloud]

Quality Gate failed

Failed conditions
61.8% Coverage on New Code (required ≥ 80%)

See analysis details on SonarQube Cloud

[robmsmt]

looks good, nearly there... just some docs updates to change?

[robmsmt]

LGTM

[sonarqubecloud]

Quality Gate passed

Issues
7 New issues
0 Accepted issues

Measures
0 Security Hotspots
86.4% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

can't see what is blocking suggest we move forwards with this