Metadata-Version: 2.4
Name: atlas-infra
Version: 0.1.0
Summary: Engine-agnostic control plane and router for large-scale LLM inference.
Author: Dhruv Reddy
License: Apache-2.0
Project-URL: Homepage, https://github.com/dhruvr101/atlas_infra
Keywords: llm,inference,kubernetes,vllm,sglang,router,gateway
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: pydantic>=2.9
Provides-Extra: control-plane
Requires-Dist: kubernetes>=31; extra == "control-plane"
Requires-Dist: fastapi>=0.115; extra == "control-plane"
Requires-Dist: uvicorn[standard]>=0.30; extra == "control-plane"
Provides-Extra: router
Requires-Dist: kubernetes>=31; extra == "router"
Requires-Dist: fastapi>=0.115; extra == "router"
Requires-Dist: uvicorn[standard]>=0.30; extra == "router"
Requires-Dist: httpx>=0.27; extra == "router"
Requires-Dist: prometheus-client>=0.21; extra == "router"
Requires-Dist: pyyaml>=6; extra == "router"
Provides-Extra: dev
Requires-Dist: atlas-infra[control-plane,router]; extra == "dev"
Requires-Dist: pytest>=8; extra == "dev"
Dynamic: license-file

# Atlas

**An open-source, engine-agnostic control plane and router for large-scale LLM inference.**

Atlas deploys inference engines onto your own Kubernetes cluster and puts a single
OpenAI-compatible endpoint in front of the whole fleet — routing every request to the
best worker by a pluggable policy, and streaming the answer back. It supports **vLLM,
SGLang, and any other engine** through a small adapter, without changing a line of
routing code.

Atlas never owns your compute, models, or data. It runs inside your infrastructure and
automates the operational layer on top of it. **It is not an inference engine** — it's
the platform that deploys, routes, and operates the engines you already use. Everything
is built in **Python**.

---

## Table of contents

- [The two halves](#the-two-halves)
- [How it works (the router)](#how-it-works-the-router)
- [The core idea: engine-agnostic by design](#the-core-idea-engine-agnostic-by-design)
- [Repository structure](#repository-structure)
- [Design choices](#design-choices)
- [The control plane (deploying workers)](#the-control-plane-deploying-workers)
- [Running it](#running-it)
- [Extending it](#extending-it)
- [Status & roadmap](#status--roadmap)

---

## The two halves

Atlas splits cleanly into two cooperating halves. You can use either on its own.

| Half | What it does | Where it lives |
|---|---|---|
| **Control plane** | *Deploys* engines onto Kubernetes: takes a validated spec, checks GPU capacity, generates the Kubernetes resources, and waits for the workers to come up. | `engine/`, `api/` |
| **Data plane (router)** | *Serves* traffic: finds the live workers, watches how loaded each one is, routes each request to the best one, and streams the response back. | `router/` |

They share a single word — the **engine name** (`"vllm"`, `"sglang"`). The control plane
uses it to pick a container image and launch flags; the router uses it to pick a
telemetry adapter. One vocabulary, both ends.

```
        DESCRIBE                    DEPLOY                         SERVE
   "run Qwen, 3 replicas"   ──►   control plane   ──►   worker pods   ──►   router
                                  (engine/, api/)       (vLLM/SGLang)       (router/)
                                                                              │
                                                                     one OpenAI endpoint
                                                                              │
                                                                            clients
```

---

## How it works (the router)

The router is the heart of Atlas. Think of it as four stages — **discover, observe,
decide, act** — wrapped in a FastAPI server:

```
   ┌──────────┐   POST /v1/chat/completions                 ┌──────────────────────────┐
   │  client  │ ──────────────────────────────────────────► │       Atlas Router       │
   │  (any    │   { "model": "...", "messages": [ ... ] }    │        (FastAPI)         │
   │  OpenAI  │                                              │                          │
   │  client) │ ◄────────────────────────────────────────── │                          │
   └──────────┘            streamed tokens (SSE)             └────────────┬─────────────┘
                                                                          │  for each request
        ┌─────────────────────────────────────────────────────────────────┘
        ▼
   1. DISCOVER              2. OBSERVE                3. DECIDE             4. ACT
   which workers exist      how loaded each is        pick the best        forward + stream back
   ┌─────────────┐          ┌──────────────┐          ┌──────────────┐     ┌──────────────────┐
   │  Discovery  │          │   Scraper    │          │   Routing    │     │      Proxy       │
   │ static / k8s│ ───────► │ GET /metrics │ ───────► │ policy.select│ ──► │ stream passthrough│
   │ → Endpoints │          │ → EngineStats│          │ → Endpoint   │     │ + failover        │
   └─────────────┘          └──────┬───────┘          └──────────────┘     └────────┬─────────┘
                                   │ normalized per-engine                          │
                                   │ (adapter)                                      ▼
                                   ▼                                       ┌──────────────────┐
                           one common EngineStats                          │   worker pods    │
                           for every engine                                │  vLLM  │  SGLang │
                                                                           └──────────────────┘
```

1. **Discover** — a discovery backend keeps a live list of `Endpoint`s (a worker's URL,
   the model it serves, and which engine it runs). Two modes ship in the box: a **static**
   list from config (no cluster needed) and a **Kubernetes** watcher that tracks worker
   pods as they appear, become ready, and die.
2. **Observe** — a background scraper periodically reads each worker's Prometheus
   `/metrics` and normalizes it into a common `EngineStats` (requests running, requests
   queued, KV-cache usage). It **fails soft**: an unreachable worker keeps its last-known
   load rather than breaking a request.
3. **Decide** — a routing policy picks one endpoint from the healthy workers serving the
   requested model, using that live load. Policies are pluggable (see below).
4. **Act** — the proxy forwards the request to the chosen worker's OpenAI API and streams
   the response back verbatim, token by token. If the chosen worker refuses the
   connection, it fails over to the next candidate.

---

## The core idea: engine-agnostic by design

Every inference engine reports its load under its own metric names — vLLM calls a value
`vllm:num_requests_running`, SGLang calls the same concept `sglang:num_running_reqs`. If
the router read those names directly, it would be locked to one engine forever.

Atlas puts a tiny **engine adapter** in front of each engine whose only job is to
translate that engine's raw metrics into one common `EngineStats`. Everything above the
adapter — discovery, routing, proxying — only ever sees the normalized type and never
knows which engine produced it.

```
   vLLM worker   ──/metrics──►  vllm:num_requests_running 3    ─┐
                                                                 ├─► EngineAdapter ─► EngineStats(running=3, …)
   SGLang worker ──/metrics──►  sglang:num_running_reqs 3      ─┘   (per engine)      (one common type)
```

That single seam is what makes the whole system multi-engine. **Adding a new engine is
one adapter file and one line to register it** — no routing code changes.

---

## Repository structure

Atlas is a `pip`-installable package under a single `atlas` namespace (`src/` layout):

```
atlas/
├── pyproject.toml          package manifest: install extras + plugin entry-point groups
├── Dockerfile              container image for the router
├── src/atlas/
│   ├── engine/             CONTROL PLANE — deploy engines onto Kubernetes
│   │   ├── spec.py         DeploymentSpec (validated Pydantic "order form") + scheduling preferences
│   │   ├── manifests.py    turns a spec into Kubernetes Deployment + headless Service objects
│   │   ├── scheduling.py   turns scheduling preferences into affinity / node-affinity fields
│   │   └── cluster.py      talks to the cluster: capacity check, apply, wait-until-ready
│   │
│   ├── api/                CONTROL PLANE — HTTP layer over engine/
│   │   └── main.py         FastAPI: POST /deployments, /deployments/validate, GET /cluster/nodes
│   │
│   └── router/             DATA PLANE — the router
│       ├── endpoint.py     Endpoint: one worker (url + model + engine)
│       ├── stats.py        EngineStats (scraped load) + RequestStats (measured performance)
│       │
│       ├── engines/        ── PLUGGABLE: inference engines ──
│       │   ├── base.py     EngineAdapter contract: parse_metrics(raw) -> EngineStats
│       │   ├── vllm.py     VLLMEngine
│       │   ├── sglang.py   SGLangEngine
│       │   └── __init__.py registry: get_engine(name) + entry-point plugin loading
│       │
│       ├── discovery/      ── PLUGGABLE: where workers come from ──
│       │   ├── base.py     ServiceDiscovery contract: get_endpoints() -> [Endpoint]
│       │   ├── static.py   StaticDiscovery (from config)
│       │   ├── k8s.py      K8sDiscovery (watches worker pods)
│       │   └── __init__.py registry: get_discovery(mode, …) + entry-point plugin loading
│       │
│       ├── routing/        ── PLUGGABLE: how to choose a worker ──
│       │   ├── base.py     RoutingPolicy contract: select(endpoints, stats, body) -> Endpoint
│       │   ├── roundrobin.py    fair rotation
│       │   ├── least_loaded.py  fewest in-flight/queued (uses EngineStats)
│       │   ├── session.py       sticky per session id (consistent-hash ring)
│       │   ├── prefix_aware.py  cache affinity by shared prompt prefix
│       │   └── __init__.py registry: get_policy(name, …) + entry-point plugin loading
│       │
│       ├── scraper.py      background thread: /metrics -> adapter -> EngineStats
│       ├── proxy.py        streaming passthrough + connection failover
│       ├── config.py       RouterConfig (from a YAML file)
│       ├── app.py          the FastAPI gateway that wires it all together
│       ├── _plugins.py     loads third-party plugins from entry points
│       └── config.example.yaml
│
├── scripts/check_cluster.py    control-plane connectivity check
└── tests/                      unit + end-to-end tests (control plane + router)
```

---

## Design choices

**Three pluggable seams, one pattern.** Engines, discovery backends, and routing policies
are each defined by a small abstract contract plus a name→implementation registry. Adding
a new one is: write a subclass, register it, done — nothing else changes. Third parties
can register their own at runtime without forking Atlas.

**Normalize at the edge, stay generic in the core.** The *only* engine-specific code is
inside `engines/`. Discovery, routing, and the proxy operate purely on the normalized
`Endpoint` and `EngineStats` types, so the same routing logic works across every engine.

**Endpoints and load are stored separately, both keyed by URL.** A metrics refresh
updates load without churning the worker list, and a discovery change updates the worker
list without wiping load history.

**Fail soft on telemetry.** A missing or late metrics scrape is treated as "idle," never
an error — a slow worker can't break routing, and a brand-new worker starts receiving
traffic immediately instead of being starved.

**Stream everything.** Responses pass through the proxy token-by-token (SSE), so clients
see output live and the router never buffers a whole completion. Failover happens only
before the first byte is sent — once tokens are flowing, the router is committed.

**The router is standalone.** It imports nothing from the control plane. You can run it in
front of engines you started any way you like (Docker, bare metal, another orchestrator) —
Kubernetes is one discovery mode, not a requirement.

### Routing policies that ship

| Policy | Chooses… | Good for |
|---|---|---|
| `roundrobin` | the next worker in a fair rotation | even, predictable load |
| `least_loaded` | the worker with the fewest running/queued requests | latency under uneven load |
| `session` | the same worker for a given session id (consistent-hash ring) | reusing a user's KV cache across turns |
| `prefix_aware` | the same worker for requests sharing a prompt prefix | reusing a shared system prompt's cache |

---

## The control plane (deploying workers)

Instead of writing Kubernetes YAML by hand, you describe a deployment as a validated
`DeploymentSpec` (engine, model, replica count, GPUs per replica, and scheduling
preferences). The control plane then:

1. **Validates capacity** against the live cluster — it counts free GPUs per node and
   confirms the replicas can actually be packed onto them, so a request for more than the
   cluster can hold fails fast with a clear message instead of leaving pods stuck
   `Pending`.
2. **Generates Kubernetes resources** — a `Deployment` (N identical GPU pods running the
   engine) and a headless `Service` (a stable handle for the pods; headless so the router,
   not a load-balancer VIP, owns routing).
3. **Applies and waits** — submits them to the Kubernetes API and watches until every
   replica reports `Ready`.

Atlas only ever talks to the Kubernetes API, never to the machines directly, so the same
code runs unchanged against a one-node dev cluster or a production cluster with hundreds
of GPU nodes.

**Scheduling preferences.** A spec can optionally express soft/hard placement hints that
Atlas compiles into standard Kubernetes scheduling fields (Kubernetes still makes the
final decision):

- `spread_replicas` — prefer placing each replica on a different node (soft pod
  anti-affinity), so losing one node takes out at most one replica.
- `gpu_type` — require a specific GPU model (hard node affinity), so a replica never lands
  on the wrong accelerator.

An HTTP layer (`api/`) exposes this over REST: `POST /deployments`,
`POST /deployments/validate`, `GET /cluster/nodes`.

---

## Installing

Atlas ships as a package with opt-in extras, so you install only the half you need:

```bash
pip install -e ".[router]"          # the data-plane router
pip install -e ".[control-plane]"   # the Kubernetes deploy engine + REST API
pip install -e ".[dev]"             # everything, plus test tooling
```

Until it's published to PyPI, install from a checkout as above — or straight from git:
`pip install "atlas-infra[router] @ git+https://github.com/dhruvr101/atlas_infra"`.

## Running it

### Locally (no cluster — static backends)

```bash
pip install -e ".[router]"

# Point at your own vLLM / SGLang servers by editing the example config.
ATLAS_ROUTER_CONFIG=src/atlas/router/config.example.yaml \
  uvicorn atlas.router.app:app --host 0.0.0.0 --port 8080
```

Send it OpenAI-compatible traffic:

```bash
curl http://localhost:8080/v1/chat/completions \
  -H 'content-type: application/json' \
  -d '{"model": "Qwen/Qwen2.5-3B-Instruct",
       "messages": [{"role": "user", "content": "hello"}]}'

curl http://localhost:8080/status     # every worker, its engine, and live load
curl http://localhost:8080/v1/models  # models available across the fleet
```

### With Docker

```bash
docker build -t atlas-router .
docker run -p 8080:8080 \
  -v $(pwd)/my-config.yaml:/app/config.yaml \
  -e ATLAS_ROUTER_CONFIG=/app/config.yaml \
  atlas-router
```

### Against Kubernetes

Set `discovery_mode: k8s` in the config and give it a `k8s_label_selector` matching your
worker pods. The control plane stamps each pod with `atlas/model` and `atlas/engine`
annotations, which discovery reads so every pod is routed through the right engine
adapter automatically. Run the router **in-cluster** (as a pod) so it can reach pod IPs
directly.

A minimal config:

```yaml
discovery_mode: static          # or: k8s
static_backends:
  - { url: http://localhost:8000, model: Qwen/Qwen2.5-3B-Instruct, engine: vllm }
  - { url: http://localhost:8001, model: Qwen/Qwen2.5-3B-Instruct, engine: sglang }
routing_policy: least_loaded    # roundrobin | least_loaded | session | prefix_aware
scrape_interval_seconds: 10
```

---

## Extending it

Atlas has three plug seams — **engines**, **routing policies**, and **discovery
backends** — and you can extend any of them two ways.

**In-tree** (a subclass + one register line):

- *Engine* — add `src/atlas/router/engines/<name>.py` with a `parse_metrics()` that maps
  the engine's metric names into `EngineStats`, then `register_engine(...)`.
- *Routing policy* — add `src/atlas/router/routing/<name>.py` subclassing `RoutingPolicy`
  with a `select(...)`, then `register_policy(...)`.
- *Discovery* — subclass `ServiceDiscovery` with a `get_endpoints()`, then
  `register_discovery(name, cls)`.

**As a separate installable plugin — without forking Atlas.** Publish your own package
that declares an entry point, and Atlas discovers and registers it on import. For example,
a `atlas-tensorrt` package would put this in its `pyproject.toml`:

```toml
[project.entry-points."atlas.engines"]
tensorrt = "atlas_tensorrt:TensorRTEngine"
```

Then any user just runs `pip install atlas-tensorrt` alongside Atlas, and the router can
route to TensorRT-LLM workers — no changes to Atlas itself. The same mechanism works for
`atlas.routing_policies` and `atlas.discovery`.

---

## Status & roadmap

**Built and working:** the full router (discovery, per-engine telemetry adapters, four
routing policies, streaming proxy with failover, config, Docker image) — proven end-to-end
across vLLM and SGLang through one endpoint — plus the Kubernetes control plane (spec →
capacity check → manifests → apply → wait) with GPU scheduling preferences and its REST
layer. Packaged as a `pip`-installable library with an entry-point plugin system, and
covered by a test suite (unit + end-to-end).

**Planned:** fleet metrics/observability (Prometheus + Grafana), a natural-language
deployment assistant that translates a plain-English request into a validated
`DeploymentSpec`, and an operator CLI and dashboard over the same REST API.
