Posted on
• Artificial Intelligence

Artificial Intelligence Load Balancing

Author
  • User
    linuxbash
    Posts by this author
    Posts by this author

Artificial Intelligence Load Balancing on Linux (with Bash-friendly steps)

Ever watched one GPU melt while the others idle, or one model server queue explode while the rest nap? That’s the cost of poor load distribution. AI workloads are spiky, heavy, and unpredictable. Smart load balancing turns that chaos into consistent latency, higher throughput, and fewer 3 a.m. incidents.

In this guide, you’ll:

  • Understand why AI traffic needs different balancing than web apps.

  • Stand up a production‑ready HAProxy layer for model serving.

  • Add GPU/CPU-aware adaptive weights with a tiny Bash script.

  • Test the setup with a local mock model server.

  • See NGINX as an alternative.

All commands are Linux-first and copy/paste friendly.

Why load balancing for AI is different

  • Uneven, bursty latency: Token-by-token generation, batch size, and dynamic model paths cause wide request times. Averages lie; p95/p99 drive UX.

  • Resource pinning: GPUs are discrete and fill up fast. One hot card can tank tail latency if you don’t route around it.

  • Big payloads + warmups: Downloads, model loads, and KV caches make cold starts painful. Retaining session affinity can save seconds.

  • Cost and resiliency: Smarter balancing lets you use fewer GPUs without violating SLOs and fails over seamlessly when a host dies.

What we’ll build

A minimal HAProxy front end routing to two model servers with:

  • Health checks, timeouts, and stats.

  • Least-connections balancing for better tail latency.

  • Optional sticky-by-session for cache/warmth.

  • A Bash script that reads local GPU/CPU load and adjusts server weights at runtime.

You can adapt this to bare metal, VMs, or containers.

1) Install prerequisites

We’ll use HAProxy (LB), socat (admin socket helper), jq and curl (testing and JSON), plus Python 3 to run a tiny mock model server for local testing.

APT (Debian/Ubuntu):

sudo apt update
sudo apt install -y haproxy socat jq curl python3

DNF (Fedora/RHEL/CentOS Stream):

sudo dnf install -y haproxy socat jq curl python3

Zypper (openSUSE/SLE):

sudo zypper refresh
sudo zypper install -y haproxy socat jq curl python3

Optional: NGINX (alternative LB)

  • APT: sudo apt install -y nginx

  • DNF: sudo dnf install -y nginx

  • Zypper: sudo zypper install -y nginx

If you have NVIDIA GPUs and want GPU-aware weighting, ensure nvidia-smi is available (via your NVIDIA driver install). The script below will degrade gracefully if it’s not.

2) Configure HAProxy for AI inference

Back up your HAProxy config and replace it with a sane default that fits AI traffic:

sudo cp /etc/haproxy/haproxy.cfg /etc/haproxy/haproxy.cfg.bak.$(date +%s)

sudo tee /etc/haproxy/haproxy.cfg >/dev/null <<'HAPROXY'
global
    log /dev/log local0
    log /dev/log local1 notice
    stats socket /run/haproxy/admin.sock mode 660 level admin
    stats timeout 30s
    user haproxy
    group haproxy
    daemon
    # Larger buffers can help with big JSON payloads
    tune.bufsize 131072
    tune.maxaccept 200

defaults
    log     global
    mode    http
    option  httplog
    option  dontlognull
    timeout connect 5s
    timeout client  120s
    timeout server  120s
    option  http-keep-alive
    option  redispatch
    retries 2

frontend fe_ai
    bind *:8080
    default_backend be_ai
    # Optional: reject overly large payloads to protect backends
    http-request set-header X-Forwarded-For %[src]
    http-request set-header X-Request-Start t=%Ts

backend be_ai
    balance leastconn
    # Health check your model servers' /health endpoint
    option httpchk GET /health
    http-check expect status 200
    default-server inter 2s fall 3 rise 2
    # Replace 127.0.0.1:9001/9002 with your model servers
    server s1 127.0.0.1:9001 check weight 128 maxconn 200
    server s2 127.0.0.1:9002 check weight 128 maxconn 200

# Sticky option for warm caches / session affinity (swap in if needed):
# backend be_ai_sticky
#     balance hdr(X-Session-Id)
#     hash-type consistent
#     option httpchk GET /health
#     http-check expect status 200
#     default-server inter 2s fall 3 rise 2
#     server s1 127.0.0.1:9001 check weight 128 maxconn 200
#     server s2 127.0.0.1:9002 check weight 128 maxconn 200

listen stats
    bind *:8404
    mode http
    stats enable
    stats uri /
    stats refresh 2s
HAPROXY

sudo systemctl enable --now haproxy
sudo systemctl restart haproxy

  • 8080: client traffic.

  • 8404: HAProxy stats UI and metrics (view with curl or a browser).

  • Admin socket: /run/haproxy/admin.sock for dynamic weight tuning.

Test the LB (it’ll 502 until backends are up):

curl -i http://127.0.0.1:8080/health

3) Spin up two tiny mock “AI” servers (local test)

We’ll simulate inference with a minimal Python HTTP server. It supports:

  • GET /health => 200

  • POST /infer => sleeps a bit (to mimic compute), returns JSON with which “device” handled it

cat > model.py <<'PY'
#!/usr/bin/env python3
import os, json, time
from http.server import BaseHTTPRequestHandler, HTTPServer

DEVICE = os.environ.get("DEVICE", "CPU")
PORT = int(os.environ.get("PORT", "9001"))

class H(BaseHTTPRequestHandler):
    def _send(self, code, body, ctype="application/json"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path == "/health":
            self._send(200, b"OK\n", ctype="text/plain")
        else:
            self._send(404, b"not found\n", ctype="text/plain")

    def do_POST(self):
        if self.path == "/infer":
            n = int(self.headers.get("Content-Length", "0"))
            data = self.rfile.read(n) if n > 0 else b""
            # Simulate variable latency; longer inputs are slower
            delay = min(2.0, 0.001 * len(data))
            time.sleep(delay)
            out = json.dumps({
                "device": DEVICE,
                "delay_s": delay,
                "bytes_in": len(data)
            }).encode("utf-8")
            self._send(200, out)
        else:
            self._send(404, b"not found\n", ctype="text/plain")

HTTPServer(("0.0.0.0", PORT), H).serve_forever()
PY

chmod +x model.py

# Terminal 1:
PORT=9001 DEVICE=GPU0 ./model.py &
# Terminal 2 (or next line):
PORT=9002 DEVICE=GPU1 ./model.py &

Now test through the load balancer:

for i in $(seq 1 6); do
  curl -s -X POST http://127.0.0.1:8080/infer -d "$(printf '%*s' $((RANDOM%500+100)) x | tr ' ' x)" | jq .
done

Watch HAProxy stats:

curl -s http://127.0.0.1:8404/ | sed -n '1,120p'

You should see traffic spread across s1 and s2 with live health and connection counts.

4) Make balancing adaptive to GPU/CPU load (Bash)

This script:

  • Reads GPU utilization via nvidia-smi if available; otherwise falls back to CPU load average.

  • Computes a weight (higher weight for less utilized servers).

  • Calls the HAProxy admin socket to update weights on the fly.

Assumptions:

  • s1 is bound to GPU0 (PORT 9001), s2 to GPU1 (PORT 9002) on the same host.

  • If you’re CPU-only, it still works, just uses loadavg.

sudo tee /usr/local/bin/ai-weights.sh >/dev/null <<'BASH'
#!/usr/bin/env bash
set -euo pipefail

SOCK="/run/haproxy/admin.sock"
BACKEND="be_ai"

# Helpers
has_nvidia_smi() { command -v nvidia-smi >/dev/null 2>&1; }
set_weight() {
  local server="$1" weight="$2"
  echo "set server ${BACKEND}/${server} weight ${weight}" | socat stdio "${SOCK}" >/dev/null 2>&1 || true
}

# Map utilization (0..100) -> weight (1..256), inverse
util_to_weight() {
  local util="$1"
  local inv=$((100 - util))
  ((inv < 1)) && inv=1
  # Scale roughly to 1..256
  local w=$(( inv * 256 / 100 ))
  ((w < 1)) && w=1
  ((w > 256)) && w=256
  echo "${w}"
}

if has_nvidia_smi; then
  # Read GPU0 and GPU1 utilization; default to 50 if missing
  mapfile -t utils < <(nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits 2>/dev/null || true)
  u0="${utils[0]:-50}"
  u1="${utils[1]:-50}"
else
  # Use CPU loadavg per core as a proxy
  cores=$(getconf _NPROCESSORS_ONLN)
  la1=$(awk '{print $1}' /proc/loadavg)
  # simple percentage estimate: min(100, la1/cores*100)
  pc=$(awk -v la="$la1" -v c="$cores" 'BEGIN{p=int((la/c)*100); if(p>100)p=100; if(p<0)p=0; print p}')
  u0="$pc"
  u1="$pc"
fi

w0=$(util_to_weight "$u0")
w1=$(util_to_weight "$u1")

set_weight "s1" "$w0"
set_weight "s2" "$w1"

echo "Updated weights: s1=${w0} (util=${u0}%), s2=${w1} (util=${u1}%)"
BASH

sudo chmod +x /usr/local/bin/ai-weights.sh

Run it manually:

/usr/local/bin/ai-weights.sh && echo "OK"

Set it to run every 10 seconds with systemd:

sudo tee /etc/systemd/system/ai-weights.service >/dev/null <<'UNIT'
[Unit]
Description=Adaptive HAProxy weights based on GPU/CPU usage
Wants=ai-weights.timer

[Service]
Type=oneshot
ExecStart=/usr/local/bin/ai-weights.sh
User=root
Group=root
UNIT

sudo tee /etc/systemd/system/ai-weights.timer >/dev/null <<'UNIT'
[Unit]
Description=Run ai-weights every 10s

[Timer]
OnBootSec=5s
OnUnitActiveSec=10s
AccuracySec=1s
Unit=ai-weights.service

[Install]
WantedBy=timers.target
UNIT

sudo systemctl daemon-reload
sudo systemctl enable --now ai-weights.timer
systemctl list-timers | grep ai-weights

As load shifts between GPUs (or CPU), HAProxy will bias new connections automatically.

Tip:

  • For session-heavy generation, prefer sticky routing by X-Session-Id or Authorization hash to keep KV caches warm. Use the be_ai_sticky backend shown above and add the header from your client.

5) Real-world knobs that matter

  • Algorithm: leastconn usually beats round-robin for AI. If your traffic uses long-held HTTP/2 streams or SSE, also tune maxconn per server.

  • Timeouts: raise client/server timeouts to match worst-case generation times, but keep them finite to prevent wedging.

  • Health checks: actively check /health and use rise/fall to avoid flapping. Return 200 only when the model is loaded and memory is sufficient.

  • Weights: cap at 256 in HAProxy. Use admin socket to adjust without reloads.

  • Observability: scrape the stats page (port 8404) or export HAProxy metrics into Prometheus/Grafana. Watch p95/p99 latency and queue depth.

NGINX alternative (quick sketch)

NGINX Open Source lacks native active HTTP health checks, but passive health + least_conn works for many setups:

Install:

  • APT: sudo apt install -y nginx

  • DNF: sudo dnf install -y nginx

  • Zypper: sudo zypper install -y nginx

Config:

sudo tee /etc/nginx/conf.d/ai.conf >/dev/null <<'NGX'
upstream ai_backend {
    least_conn;
    server 127.0.0.1:9001 max_fails=2 fail_timeout=5s;
    server 127.0.0.1:9002 max_fails=2 fail_timeout=5s;
    keepalive 64;
}

server {
    listen 8080;
    location / {
        proxy_pass http://ai_backend;
        proxy_http_version 1.1;
        proxy_set_header Connection "";
        proxy_set_header X-Forwarded-For $remote_addr;
    }
}
NGX

sudo nginx -t && sudo systemctl enable --now nginx && sudo systemctl reload nginx

You can add sticky sessions with NGINX Plus or with community modules; for pure OSS, consider HAProxy for stickiness and rich health checks.

Troubleshooting quick hits

  • 502/503 from HAProxy: check backend health with curl 127.0.0.1:9001/health; check logs journalctl -u haproxy -e.

  • Weights not changing: ensure the admin socket exists and permissions allow root (or adjust mode/group).

  • One server still overloaded: raise maxconn on that server or add more instances; if using long-lived connections, consider HTTP/1.1 keep-alive reuse or HTTP/2 tuning.

Conclusion and next steps

AI traffic is not “just web” traffic. With a few Linux-friendly tools—HAProxy, a Bash script, and sensible health/timeouts—you can dramatically reduce tail latency and improve GPU utilization.

Your next steps: 1. Point the backend servers to your real model hosts. 2. Enable sticky routing if you rely on warm KV/caches. 3. Wire stats into your monitoring and track p99 latency. 4. Experiment with the weighting script for your hardware mix (A100 vs. L4, CPUs, etc.).

If you want a follow-up, I can share:

  • A GPU-to-port pinning pattern for multi-GPU hosts.

  • Prometheus + Grafana dashboard JSON for HAProxy and nvidia-smi exporters.

  • A Kubernetes version using DaemonSets and Runtime Class for GPU pods.

Happy balancing—and may your p99s drop like a rock.

Related posts