Load Balancing AI/ML API with Apache Apisix

This blog provides a step-by-step guide to configure Apache APISIX for AI traffic splitting and load balancing between API versions, covering security setup, canary testing, and deployment monitoring.

AI/ML API is a one-stop, OpenAI-compatible endpoint that is trusted by 150,000+ developers to 300+ state-of-the-art models—chat, vision, image/video/music generation, embeddings, OCR, and more—from Google, Meta, OpenAI, Anthropic, Mistral, and others.

Apache APISIX is a dynamic, real-time, high-performance API Gateway. APISIX API Gateway provides rich traffic management features and can serve as an AI Gateway through its flexible plugin system.

Modern AI workloads often require smooth version migrations, A/B testing, and rolling updates. This guide shows you how to:

1. Install Apache APISIX with Docker quickstart.
2. Secure the Admin API with keys and IP whitelisting.
3. Define separate routes for API versions v1 and v2.
4. Implement weighted traffic splitting (50/50) via the traffic-split plugin.
5. Verify the newly created split endpoint functionality.
6. Load test and monitor distribution accuracy.

To perform authenticated requests, you'll need an AI/ML API key. You can get one at https://aimlapi.com/app/keys/ and use it as a Bearer token in your Authorization headers.

Quickstart Installation​

Tip: If you encounter port conflicts, adjust Docker host networking or map to different ports in the quickstart script.

By default, quickstart bypasses Admin API authentication. For any non-development environment, enforce security:

1. Set an Admin Key​

Edit conf/config.yaml inside the APISIX container or local install directory, replacing the example key with your own API key obtained from the link above:

Security Best Practice: Use at least 32 characters, mix letters/numbers/symbols, and rotate keys quarterly.

2. Whitelist Management IPs (allow_admin)​

Add your management or local networks under the admin: section:

Warning: 0.0.0.0/0 opens Admin API to the world! Lock this down to specific subnets in production.

3. Restart APISIX​

Check Logs: docker logs apisix-quickstart --tail 50 to ensure no errors about admin authentication.

Define Basic Routes for v1 and v2​

Before splitting traffic, ensure each version works individually.

1. Route for v1​

Tip: Use id fields if you want to manage or delete routes easily later.

2. Route for v2​

Use the traffic-split plugin for controlled distribution between v1 and v2. In the admin request below, replace YOUR_ADMIN_KEY_HERE with your actual key.

Tip: Adjust the weight values to shift traffic ratios (e.g., 80/20 for canary).

Note: rewrite must match the internal API path exactly.

Verify Split Endpoint Functionality​

Test the /chat/completions endpoint you just created. Replace <AIML_API_KEY> with the key obtained earlier and use it as a Bearer token:

Expected Output:

Tip: Use -v for verbose output to troubleshoot headers or TLS issues.

After configuring the split route, use the following commands to validate distribution. Replace <AIML_API_KEY> with your Bearer token.

Expected: Approximately 50 requests to each upstream.

Tip: Use Prometheus or OpenTelemetry plugins for real‑time metrics instead of manual log parsing.

Best Practices & Next Steps​

• Rate Limiting & Quotas: Add limit-count plugin to protect your upstream from spikes.
• Authentication: Layer on the key-auth plugin for consumer management.
• Circuit Breaker: Prevent cascading failures with the api-breaker plugin.
• Observability: Integrate Prometheus, Skywalking, or Loki for dashboards and alerts.
• Infrastructure as Code: Consider managing APISIX config via Kubernetes CRDs or ADC for reproducibility.

• APISIX Load Balancing Documentation
• AI/ML API Documentation