View : 122

03/08/2026 05:00am

Building an AI Load Balancer in Go with Reverse Proxy and Background Health Check

Golang The Series EP.165: Load Balancing AI Servers - Scaling AI Workloads for High Concurrency

#Golang Load Balancer

#AI Load Balancing

#AI Infrastructure

#Reverse Proxy Go

#Ollama Scaling

#Health Check Go

#Local LLM Infrastructure

Welcome to EP.165! In our previous episode, we implemented Rate Limiting to prevent spam and control API request traffic. However, as your enterprise AI system gains popularity and legitimate user requests flood in, a single instance can no longer handle the heavy workload.

Especially if your organization runs its own Local LLM Infrastructure (such as hosting models with Ollama, vLLM, or TGI across multiple GPU nodes), each request involving long prompts or document ingestion consumes massive amounts of GPU and VRAM resources. Routing all requests to a single machine will quickly lead to soaring latency, freezing, or system crashes.

The solution is Load Balancing—using Go to build a proxy intermediary that distributes requests to AI server workers efficiently and reliably!

Load Balancing Algorithms for AI

When distributing workloads across AI inference nodes, there are three popular strategies:

  • Round Robin: Distributes requests sequentially (1, 2, 3, then loops back to 1). Ideal for nodes with identical hardware specs and similar task sizes.

  • Least Connections: Directs new requests to the node with the fewest active processing connections. This is best suited for AI workloads since prompt token processing times vary.

  • Weighted Round Robin: Assigns higher request proportions to nodes with superior GPU specifications (e.g., higher weights for more powerful hardware).

Architecture of an AI Load Balancer

We will build a lightweight Reverse Proxy in Go that performs periodic Health Checks on each AI node, routing traffic using a Round Robin with Health Check strategy to ensure requests never hit dead or down nodes.

Plaintext

[Client Requests]
                                      │
                                      ▼
                        ┌──────────────────────────┐
                        │    Go Load Balancer      │
                        │  (Round Robin + Proxy)   │
                        └─────────────┬────────────┘
                                      │
           ┌──────────────────────────┼──────────────────────────┐
           │ (Active)                 │ (Active)                 │ (Down ❌)
           ▼                          ▼                          ▼
   ┌──────────────┐           ┌──────────────┐           ┌──────────────┐
   │  AI Node 1   │           │  AI Node 2   │           │  AI Node 3   │
   │ (GPU Node A) │           │ (GPU Node B) │           │ (GPU Node C) │
   └──────────────┘           └──────────────┘           └──────────────┘

Component & Proxy Logic

To maintain Clean Architecture principles, we separate our node data structures and reverse proxy logic from the core execution flow.

Part 1: Node Structures and Health Checker

Go

package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

// AINode stores details and status for each AI server
type AINode struct {
	URL          *url.URL
	Alive        bool
	ReverseProxy *httputil.ReverseProxy
	mu           sync.RWMutex
}

// SetAlive updates the availability status of the node
func (node *AINode) SetAlive(alive bool) {
	node.mu.Lock()
	node.Alive = alive
	node.mu.Unlock()
}

// IsAlive checks if the node is operational
func (node *AINode) IsAlive() bool {
	node.mu.RLock()
	defer node.mu.RUnlock()
	return node.Alive
}

// AILoadBalancer manages a group of AI nodes
type AILoadBalancer struct {
	nodes   []*AINode
	current uint64
}

// GetNextNode selects the next available node using Round Robin
func (lb *AILoadBalancer) GetNextNode() *AINode {
	nodeCount := len(lb.nodes)
	if nodeCount == 0 {
		return nil
	}

	next := atomic.AddUint64(&lb.current, 1)
	
	for i := 0; i < nodeCount; i++ {
		idx := int((next + uint64(i)) % uint64(nodeCount))
		if lb.nodes[idx].IsAlive() {
			return lb.nodes[idx]
		}
	}
	return nil
}

// HealthCheck pings the /health endpoint of each node at set intervals
func (lb *AILoadBalancer) HealthCheck() {
	client := http.Client{
		Timeout: 2 * time.Second,
	}

	for _, node := range lb.nodes {
		go func(n *AINode) {
			resp, err := client.Get(n.URL.String() + "/health")
			if err != nil || resp.StatusCode != http.StatusOK {
				if n.IsAlive() {
					log.Printf("⚠️ [Health Check] AI Node %s is unresponsive -> Marked as DOWN", n.URL.String())
					n.SetAlive(false)
				}
				return
			}
			_ = resp.Body.Close()

			if !n.IsAlive() {
				log.Printf("✅ [Health Check] AI Node %s is back online -> Marked as UP", n.URL.String())
				n.SetAlive(true)
			}
		}(node)
	}
}

Part 2: Core Server Logic and Proxy Routing: main.go

Go

func main() {
	rawURLs := []string{
		"http://localhost:11434", // AI Node 1 (Ollama Instance A)
		"http://localhost:11435", // AI Node 2 (Ollama Instance B)
		"http://localhost:11436", // AI Node 3 (Ollama Instance C)
	}

	var nodes []*AINode
	for _, rawURL := range rawURLs {
		targetURL, err := url.Parse(rawURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v", err)
		}

		proxy := httputil.NewSingleHostReverseProxy(targetURL)
		nodes = append(nodes, &AINode{
			URL:          targetURL,
			Alive:        true,
			ReverseProxy: proxy,
		})
	}

	lb := &AILoadBalancer{nodes: nodes}

	// Background Health Checker running every 10 seconds
	go func() {
		ticker := time.NewTicker(10 * time.Second)
		for range ticker.C {
			lb.HealthCheck()
		}
	}()

	// HTTP Handler acting as a Reverse Proxy
	http.HandleFunc("/api/ai/generate", func(w http.ResponseWriter, r *http.Request) {
		targetNode := lb.GetNextNode()
		if targetNode == nil {
			http.Error(w, "❌ No available AI servers at the moment", http.StatusServiceUnavailable)
			return
		}

		fmt.Printf("🔀 [Load Balancer] Proxying request to -> %s\n", targetNode.URL.String())
		targetNode.ReverseProxy.ServeHTTP(w, r)
	})

	log.Println("🚀 AI Load Balancer running on port :8080...")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

Important Considerations for AI Load Balancing

  1. Server-Sent Events (SSE) / Streaming Responses: Modern AI models often stream tokens back incrementally via SSE. Go's reverse proxy flush intervals must be carefully handled so chunks stream to the client in real-time without buffering in memory.

  2. Stateless Node Design: Design AI workers to be as stateless as possible. For chat history or conversation memory, offload context storage to a centralized database (such as Redis or PostgreSQL) so subsequent requests from the same user can hit any AI node without losing context.

🎯 Daily Mission

Try running the code above and test real-world production scenarios:

Challenge: If your system includes one ultra-fast GPU node (e.g., NVIDIA H100) and two standard GPU nodes (e.g., RTX 4090), traditional Round Robin may overload the slower nodes while the powerful node sits idle.

As a systems engineer, how would you modify the GetNextNode() function and AINode fields to implement Weighted Round Robin based on hardware capabilities? Give it a try!

❓ Frequently Asked Questions (FAQ)

Why build a load balancer in Go instead of using Nginx or HAProxy?

For standard AI infrastructure, Nginx and HAProxy are exceptional reverse proxies and load balancers. However, writing a load balancer in Go lets you tailor custom logic specifically for AI workloads—such as handling LLM token streaming, implementing specialized health checks, or integrating authentication and rate limiting directly without complex modules.

If an AI node hangs rather than completely crashing, will health checks catch it?

In our basic example, we use a short timeout (2 seconds) to ping the /health endpoint. If a node responds too slowly due to processing a long queue, it is temporarily marked as down. In production environments, you may also want to monitor GPU VRAM utilization or active connection counts.

What are the disadvantages of Round Robin load balancing for AI?

The main drawback is processing time inequality. Since prompts vary in length and token generation time, long requests piling up on Node A while Node B sits idle can cause resource imbalances. This is why Least Connections or Weighted Round Robin strategies are recommended for advanced setups.

Does maintaining chat history affect the load balancer?

Yes. If AI nodes store conversation memory internally, a follow-up request routed to a different node by the load balancer will lose context. The solution is to design stateless AI workers and persist chat history in a centralized cache like Redis.

Summary

In this article, we explored how to resolve bottlenecks when enterprise AI systems face high concurrent traffic by building a Go-based load balancer with a reverse proxy. We also implemented a background health check system to monitor node availability in real-time, preventing requests from hitting dead servers and improving overall enterprise AI infrastructure stability.

In the next episode (EP.166): Even with a great load balancer, if your primary AI API or backend destination completely goes down, slow timeouts can trigger Cascading Failures across your entire backend architecture. Next time, we will look at installing automated circuit breakers with "Circuit Breaker Patterns: Handling Failures When AI APIs Collapse". See you there, Gophers!

Follow Superdev Academy on all platforms: