Skip to content

Darrow8/foundry-load-balancer

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Azure Foundry Load-Balancing Proxy

Open-source OpenAI-compatible HTTP proxy for Azure AI Foundry (and similar Azure OpenAI-style endpoints). It spreads traffic across multiple projects or resources so you use the combined token-per-minute (TPM) and requests-per-minute (RPM) limits instead of a single quota.

Stack: Python 3.10+, FastAPI, httpx. Optional AWS CDK deployment to ECS Fargate (see below).

Example: Two endpoints at 100k TPM / 100 RPM each → 200k TPM / 200 RPM to your app as one API.

Client ──► localhost:8787 ──► Proxy ──┬──► Azure Project 1 (100k TPM, 100 RPM)
                                      └──► Azure Project 2 (100k TPM, 100 RPM)

Features

  • Dual-limit routing — balances on both TPM and RPM so one limit cannot silently saturate an endpoint.
  • In-flight reservations — concurrent clients do not all pile onto the same endpoint.
  • Graceful backoff — waits for rolling-window capacity instead of failing immediately when limits are tight.
  • Failover — on upstream errors or 429s, tries the next configured endpoint.
  • Optional proxy authAuthorization: Bearer or api-key for the proxy itself; Azure keys stay server-side.
  • AWS-ready — Dockerfile + CDK stack for ALB + Fargate + Secrets Manager.

Repository layout

Path Purpose
main.py Uvicorn entrypoint
proxy/ Config, TokenTracker, FastAPI app
tests/ Unit tests and load test script
infra/ AWS CDK app (Fargate + ALB)
Dockerfile Container image for deployment

How It Works

Dual-limit routing

The proxy tracks both tokens per minute (TPM) and requests per minute (RPM) in a rolling 60-second window per endpoint. When a request arrives, pick_best() scores each endpoint by min(remaining_tpm%, remaining_rpm%) — whichever limit is closer to exhaustion drives the routing decision. This prevents overloading an endpoint on either dimension.

In-flight reservations

Every request is assigned an estimated token cost (input chars / 4 + max_tokens) before it is forwarded. That estimate is immediately reserved against the endpoint's capacity. This prevents concurrent requests from all seeing the same "available" headroom and piling onto one endpoint. When a response completes, the reservation is replaced with the actual usage.total_tokens reported by Azure.

Graceful capacity handling

When all endpoints are at capacity (TPM or RPM), the proxy does not immediately return an error. Instead it:

  1. Calculates exactly when the oldest entry in the rolling window will expire, freeing a slot
  2. Sleeps until capacity opens up (up to 60 seconds max)
  3. Retries the request on the now-available endpoint

This means clients never see 429 errors unless sustained traffic exceeds the combined capacity ceiling for over a minute.

Transparent failover

If an upstream endpoint returns a 429 despite the proxy's tracking (estimates are approximate), the proxy waits for capacity and tries the next endpoint. All of this happens before any response bytes are sent to the client — from the client's perspective there is no error, just a slightly slower response.

Request lifecycle

Request arrives
  ├─ estimate tokens (input chars/4 + max_tokens)
  ├─ wait for capacity if all endpoints are full
  ├─ pick best endpoint (highest min of TPM% and RPM% remaining)
  ├─ reserve estimate on chosen endpoint
  └─ forward to Azure
       ├─ success → record actual usage, settle reservation
       ├─ 429 → cancel reservation, wait for capacity, try next endpoint
       └─ other error → cancel reservation, try next endpoint

Prerequisites

  • Python 3.10+
  • Two (or more) Azure Foundry projects with your model deployed in separate projects (deployments within the same project share quota)
  • Each project's endpoint URL and API key

Setup

1. Install dependencies

pip install -r requirements.txt

2. Configure endpoints

cp .env.example .env

Edit .env:

# Comma-separated base URLs (one per Azure Foundry project)
AZURE_ENDPOINTS=https://your-resource-1.openai.azure.com/openai/v1,https://your-resource-2.openai.azure.com/openai/v1

# Comma-separated API keys (same order as endpoints)
AZURE_API_KEYS=key-for-project-1,key-for-project-2

# Per-endpoint limits
TPM_LIMIT=100000
RPM_LIMIT=100

# Proxy authentication — clients must send this as a Bearer token
# Generate one: python3 -c "import secrets; print(secrets.token_urlsafe(32))"
# Leave empty to disable auth (only safe when bound to localhost)
PROXY_API_KEY=your-proxy-key-here

# Max request body size (default 10 MB)
MAX_BODY_BYTES=10485760

# Local server
HOST=127.0.0.1
PORT=8787

3. Run the proxy

python main.py

You should see:

INFO:     Uvicorn running on http://127.0.0.1:8787

4. Configure your client

Point any OpenAI-compatible client at the proxy (IDEs such as Zed or Cursor, custom apps, or curl). Use your PROXY_API_KEY as the API key when auth is enabled. Example for Zed (Cmd+, on macOS):

{
  "language_models": {
    "openai": {
      "api_url": "http://127.0.0.1:8787/v1",
      "api_key": "your-proxy-key-here",
      "available_models": [
        {
          "name": "Kimi-K2.5",
          "display_name": "Kimi K2.5 (Load Balanced)",
          "max_tokens": 131072,
          "max_output_tokens": 8192
        }
      ]
    }
  }
}

API Endpoints

Method Path Auth Description
POST /v1/chat/completions Yes Proxied chat completions (streaming and non-streaming)
GET /v1/models Yes Proxied model list from the first available endpoint
GET /status Yes Real-time usage dashboard
GET /health No Unauthenticated health check

Security

The proxy protects your Azure API keys and controls access:

  • API key authentication — set PROXY_API_KEY in .env and clients must send it as Authorization: Bearer <key> or api-key: <key>. Both OpenAI and Azure header conventions are supported.
  • URL masking — the /status endpoint only exposes hostnames, never full endpoint URLs or API keys.
  • Request size limitMAX_BODY_BYTES (default 10 MB) rejects oversized payloads before they reach Azure.
  • Localhost binding — defaults to 127.0.0.1, not reachable from the network. Change HOST to 0.0.0.0 only if you need remote access, and always set a PROXY_API_KEY when doing so.
  • Constant-time key comparison — uses hmac.compare_digest to prevent timing attacks on the API key.

If PROXY_API_KEY is empty or unset, auth is disabled entirely (suitable for local-only use).

Reporting a security vulnerability

Please do not open a public issue for undisclosed security bugs. Use GitHub Security Advisories (private report) if enabled on the repository, or contact the maintainers through a private channel. Never paste Azure keys, PROXY_API_KEY, or other secrets in issues or discussions.

Monitoring

GET http://127.0.0.1:8787/status

Returns per-endpoint usage for both tokens and requests:

{
  "endpoints": [
    {
      "endpoint_index": 0,
      "endpoint_url": "project-1.openai.azure.com",
      "used_tokens": 42000,
      "completed_tokens": 40000,
      "in_flight_tokens": 2000,
      "remaining_tokens": 58000,
      "tpm_limit": 100000,
      "used_requests": 12,
      "completed_requests": 10,
      "in_flight_requests": 2,
      "remaining_requests": 88,
      "rpm_limit": 100
    }
  ]
}
Field What it tells you
completed_tokens Tokens from settled requests (actual counts from Azure)
in_flight_tokens Estimated tokens for requests still in progress
remaining_tokens TPM headroom in the current 60s window
used_requests / remaining_requests RPM usage and headroom in the current 60s window

Load Testing

A load test script is included to verify the proxy distributes traffic across both endpoints:

# Start the proxy
python main.py

# In another terminal — send 4 large concurrent requests
python3 tests/load_test.py

# Or customize
python3 tests/load_test.py --total 6 --max-tokens 60000 --model Kimi-K2.5

The test sends a few large requests concurrently. The proxy should route them across endpoints. The test passes when total tokens exceed a single endpoint's 100k limit and both endpoints received traffic.

Scaling

Add more URLs and keys to .env — the proxy handles N endpoints automatically:

Endpoints Effective TPM Effective RPM
2 200k 200
3 300k 300
4 400k 400

Running as a Background Service

macOS (launchd)

Create ~/Library/LaunchAgents/com.azure-lb-proxy.plist pointing to python /path/to/main.py, or run in a tmux/screen session.

Linux (systemd)

# /etc/systemd/user/azure-lb-proxy.service
[Unit]
Description=Azure LB Proxy

[Service]
ExecStart=/usr/bin/python3 /path/to/main.py
Restart=on-failure

[Install]
WantedBy=default.target
systemctl --user enable --now azure-lb-proxy

Deploying to AWS (ECS Fargate)

The infra/ directory contains an AWS CDK app that deploys the proxy to ECS Fargate behind an Application Load Balancer.

Prerequisites

  • AWS CLI configured with credentials
  • AWS CDK CLI (npm install -g aws-cdk)
  • Docker (CDK builds the container image locally)

1. Store secrets in Secrets Manager

aws secretsmanager create-secret \
  --name ai-load-balancer/secrets \
  --secret-string '{"AZURE_API_KEYS":"your-key-1,your-key-2","PROXY_API_KEY":"your-proxy-key"}'

If the secret already exists, update it:

aws secretsmanager put-secret-value \
  --secret-id ai-load-balancer/secrets \
  --secret-string '{"AZURE_API_KEYS":"your-key-1,your-key-2","PROXY_API_KEY":"your-proxy-key"}'

2. Install CDK dependencies

cd infra
pip install -r requirements.txt

3. Bootstrap CDK (first time only)

cdk bootstrap

4. Deploy

Create the Secrets Manager secret before deploying — the stack references ai-load-balancer/secrets and does not create it (so it will not conflict with a secret you already created).

Pass your Azure endpoint URLs as context:

cdk deploy \
  -c azure_endpoints="https://resource-1.openai.azure.com/openai/v1,https://resource-2.openai.azure.com/openai/v1" \
  -c tpm_limit=100000 \
  -c rpm_limit=100

Use -c secret_name=my/custom/name if your secret uses a different name.

CDK will build the Docker image, push it to ECR, and deploy the Fargate service. The ALB DNS name is printed as an output — use it as your API base URL:

http://<alb-dns-name>/v1/chat/completions

5. Tear down

cdk destroy

Contributing

Contributions are welcome: bug reports, documentation improvements, and pull requests.

  • For larger changes (new routing strategy, major AWS changes), open an issue first so design can be discussed.
  • Keep PRs focused; match existing style in proxy/.
  • Do not commit .env or real API keys — use .env.example only in the repo.

Development

pip install -r requirements.txt
pip install pytest pytest-asyncio   # for unit tests
pytest tests/test_routing.py -v

The load test (tests/load_test.py) calls a running proxy and real endpoints — use only with your own credentials.

License

Add a LICENSE file at the repository root before publishing (for example MIT or Apache-2.0). Until then, clarify terms in your repository settings if you need a default license for contributors.

Troubleshooting

Symptom Fix
Client says "connection refused" Make sure the proxy is running on the configured host/port
Getting 401 from proxy Check PROXY_API_KEY in .env matches what your client is sending
Getting 401 from Azure Check your AZURE_API_KEYS in .env
Still hitting rate limits Verify you have separate projects — deployments in the same project share quota
Streaming responses cut off Increase the httpx timeout in server.py (default 600s)
remaining_tokens looks wrong Check /status — if in_flight_tokens is large, estimates may be over-counting
Requests queuing for a long time You're exceeding the combined TPM/RPM ceiling; add more endpoints or reduce traffic
cdk deploy failed with secret ... already exists Fixed in current CDK: the stack imports the secret by name instead of creating it. Delete any failed CloudFormation stack in ROLLBACK_COMPLETE, then run cdk deploy again

About

No description, website, or topics provided.

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors