runpod: let any capability rent, scale and reap its own pod

The pod pool never cared that it was serving one LLM — it provisions,
health-checks, scales and reaps whatever image it is given. Two things
were hardcoded: the image and the /v1/models readiness probe. Both are
now per-pool, via pod_plan():

  auto      vLLM, or llama.cpp for a GGUF   -> /v1/models
  coderai   your image (required)           -> /healthz
  custom    your image + docker_args        -> health_path

So a capability can now be a pool instead of a fixed URL. Configure
`remotes.pods: {"images": {...}}` (the same fields a model's runpod block
takes) and the gateway borrows a pod per request and hands it back;
`"endpoints": {"images": "runpod"}` is the shorthand. Everything that
already guards LLM pods now guards these: $/hr ceilings, rolling spend
limits, the global budget, min/max autoscaling, idle teardown, the stale
-pod reaper, and the stats page (as `capability:<name>`).

That closes the gap this work left open: an idle image or video pod is
exactly the kind you forget about, and until now nothing would reap it.

A pod that cannot be provisioned — RunPod disabled, budget hit, no GPU
under the ceiling — returns 502 rather than falling back to running the
model locally. Silent local fallback is a failure mode you only notice by
watching your own VRAM vanish.

22 tests, including pod acquire/release (the handle is returned even on
the happy path, so the reaper can see it), the runpod shorthand, the
provisioning-failure 502, and every pod_plan branch.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY
parent 696ebc25
......@@ -3532,6 +3532,7 @@ def build_settings_dict(c, gpu_cards):
"api_key": c.remotes.api_key,
"max_body_mb": c.remotes.max_body_mb,
"endpoints": dict(c.remotes.endpoints or {}),
"pods": dict(c.remotes.pods or {}),
},
"ds4": {
"enabled": c.ds4.enabled,
......@@ -3986,6 +3987,11 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
c.remotes.endpoints = ({str(k).strip().lower(): str(v).strip()
for k, v in eps.items() if str(v).strip()}
if isinstance(eps, dict) else {})
if "pods" in d:
pods = d.get("pods")
c.remotes.pods = ({str(k).strip().lower(): v
for k, v in pods.items() if isinstance(v, dict)}
if isinstance(pods, dict) else {})
if "ds4" in data:
d = data["ds4"]
......
......@@ -124,6 +124,24 @@ def capability_endpoints() -> dict:
return out
def capability_pods() -> dict:
"""The configured capability -> RunPod pod block map (``remotes.pods``).
A capability listed here is served by a coderai-managed pod: provisioned on
the first request, health-checked, scaled to ``max_pods``, and reaped after
``idle_timeout_s`` — the same budgets and lifecycle a RunPod-served LLM gets,
which is what stops a forgotten image/video pod billing all night.
"""
cfg = _remotes_config()
if cfg is None or not getattr(cfg, "enabled", True):
return {}
pods = getattr(cfg, "pods", None)
if not isinstance(pods, dict):
return {}
return {str(k).strip().lower(): (v if isinstance(v, dict) else {})
for k, v in pods.items()}
def _model_remote(model: str) -> str:
if not model:
return ""
......@@ -153,7 +171,7 @@ def any_remote_configured() -> bool:
def _any_remote_configured_uncached() -> bool:
if capability_endpoints():
if capability_endpoints() or capability_pods():
return True
try:
from codai.admin.routes import config_manager
......@@ -191,9 +209,25 @@ def _model_from_body(body: bytes, content_type: str) -> str:
return ""
class Target:
"""Where a request is going: a fixed URL, or a pod pool to borrow one from."""
def __init__(self, url: str = "", pool=None, reason: str = ""):
self.url = (url or "").rstrip("/")
self.pool = pool
self.reason = reason
def acquire(self):
"""Return (base_url, release). A pool provisions on demand and bills."""
if self.pool is None:
return self.url, (lambda: None)
handle, url = self.pool.acquire()
return url.rstrip("/"), (lambda: self.pool.release(handle))
def resolve_target(path: str, method: str, query: str, body: bytes,
content_type: str) -> Optional[Tuple[str, str]]:
"""Decide where this request should go. Returns (base_url, reason) or None."""
content_type: str) -> Optional["Target"]:
"""Decide where this request should go, or None to serve it locally."""
if not path.startswith("/v1/") or path in _SKIP_PATHS:
return None
cap = capability_for(path)
......@@ -203,11 +237,19 @@ def resolve_target(path: str, method: str, query: str, body: bytes,
model = (parse_qs(query).get("model") or [""])[0]
url = _model_remote(model)
if url:
return url, f"model {model!r}"
if cap:
url = capability_endpoints().get(cap, "")
if url:
return url, f"capability {cap!r}"
return Target(url=url, reason=f"model {model!r}")
if not cap:
return None
url = capability_endpoints().get(cap, "")
# "runpod" as the endpoint is shorthand for "use the pod block for this
# capability" — so the common case needs no second config block.
if url and url.strip().lower() != "runpod":
return Target(url=url, reason=f"capability {cap!r}")
pods = capability_pods()
if url or cap in pods:
from codai.api.runpod_worker import get_capability_pool
return Target(pool=get_capability_pool(cap, pods.get(cap) or {}),
reason=f"capability {cap!r} on a RunPod pod")
return None
......@@ -256,16 +298,31 @@ class RemoteGatewayMiddleware:
if body is None: # over the cap mid-stream
return await self._app(scope, replay, send)
target = resolve_target(path, scope.get("method", "GET"),
(scope.get("query_string") or b"").decode("latin-1"),
body, headers.get("content-type", ""))
try:
target = resolve_target(path, scope.get("method", "GET"),
(scope.get("query_string") or b"").decode("latin-1"),
body, headers.get("content-type", ""))
except Exception as exc:
# A misconfigured remote must not silently fall back to running the
# model here — that is how you discover it by watching local VRAM.
return await _error(send, 502, f"remote routing failed: {exc}")
if not target:
return await self._app(scope, replay, send)
await self._forward(scope, headers, body, target, send)
url, reason = target
await self._forward(scope, headers, body, url, reason, send)
async def _forward(self, scope, headers, body, target, send):
import asyncio
try:
base, release = await asyncio.to_thread(target.acquire)
except Exception as exc:
return await _error(send, 502,
f"could not reach a remote for this request: {exc}")
try:
await self._send_upstream(scope, headers, body, base, target.reason, send)
finally:
release()
async def _forward(self, scope, headers, body, base, reason, send):
async def _send_upstream(self, scope, headers, body, base, reason, send):
import asyncio
path = scope.get("path")
qs = (scope.get("query_string") or b"").decode("latin-1")
......@@ -288,12 +345,7 @@ class RemoteGatewayMiddleware:
try:
resp = await asyncio.to_thread(_call)
except Exception as exc:
msg = json.dumps({"error": {"message": f"remote endpoint {base} unreachable: {exc}",
"type": "remote_gateway_error"}}).encode()
await send({"type": "http.response.start", "status": 502,
"headers": [(b"content-type", b"application/json"),
(b"content-length", str(len(msg)).encode())]})
return await send({"type": "http.response.body", "body": msg})
return await _error(send, 502, f"remote endpoint {base} unreachable: {exc}")
out = [(b"content-type", (resp.headers.get("Content-Type") or "application/json")
.encode("latin-1"))]
......@@ -312,6 +364,15 @@ class RemoteGatewayMiddleware:
resp.close()
async def _error(send, status: int, message: str):
msg = json.dumps({"error": {"message": message,
"type": "remote_gateway_error"}}).encode()
await send({"type": "http.response.start", "status": status,
"headers": [(b"content-type", b"application/json"),
(b"content-length", str(len(msg)).encode())]})
await send({"type": "http.response.body", "body": msg})
async def _drain(receive, cap: int):
"""Buffer the request body so it can be inspected and then replayed.
......
......@@ -49,7 +49,13 @@ class RunpodModelConfig:
# a GGUF goes to llama.cpp (vLLM's GGUF support is experimental — single-file
# only, a limited architecture list, and it still wants the original repo's
# tokenizer — so it is not a drop-in for the llama.cpp catalogue).
engine: str = "auto" # auto | vllm | llamacpp
engine: str = "auto" # auto | vllm | llamacpp | coderai | custom
# Readiness probe path. Blank = the engine's default (/v1/models for the
# OpenAI servers, /healthz for a coderai pod).
health_path: str = ""
# Verbatim docker args, replacing whatever the engine would generate. The
# escape hatch for `engine: custom` and for images with their own CLI.
docker_args: str = ""
served_model: str = "" # HF id the pod/endpoint serves
# GGUF to pull on a llama.cpp pod, in llama.cpp's own `-hf` form:
# "user/repo:Q4_K_M" or "user/repo:file.gguf". Needed because a local
......@@ -106,7 +112,7 @@ def resolve_pod_engine(mcfg: "RunpodModelConfig", model_key: str = "",
repo id — goes to vLLM.
"""
want = (mcfg.engine or "auto").strip().lower()
if want in ("vllm", "llamacpp"):
if want in ("vllm", "llamacpp", "coderai", "custom"):
return want
if mcfg.hf_gguf:
return "llamacpp"
......@@ -116,6 +122,38 @@ def resolve_pod_engine(mcfg: "RunpodModelConfig", model_key: str = "",
return "vllm"
#: Readiness probe per engine. A coderai pod answers /healthz long before any
#: model is loaded, which is exactly what we want: it is ready to be asked.
_HEALTH_PATHS = {"vllm": "/v1/models", "llamacpp": "/v1/models",
"coderai": "/healthz", "custom": "/healthz"}
def pod_plan(mcfg: "RunpodModelConfig", served: str, model_key: str = "",
model_path: str = "") -> dict:
"""Everything needed to launch one pod: image, docker args, health path.
This is the seam that lets a pool serve something other than an LLM — a whole
coderai (so images/video/TTS/… can run on a rented GPU with the same budgets,
autoscaling and idle reaping), or any image at all with `engine: custom`.
"""
engine = resolve_pod_engine(mcfg, model_key, model_path)
if engine == "llamacpp":
image = mcfg.image or LLAMACPP_POD_IMAGE
args = mcfg.docker_args or _llamacpp_docker_args(mcfg, served)
elif engine in ("coderai", "custom"):
if not mcfg.image:
raise RuntimeError(
f"RunPod {engine} pod: set `image` on the runpod block to the image "
"to run (coderai does not publish one to a registry for you).")
image = mcfg.image
args = mcfg.docker_args # usually blank: the image's entrypoint serves
else:
image = mcfg.image or DEFAULT_POD_IMAGE
args = mcfg.docker_args or _vllm_docker_args(mcfg, served)
return {"engine": engine, "image": image, "args": args,
"health_path": mcfg.health_path or _HEALTH_PATHS.get(engine, "/v1/models")}
def _llamacpp_docker_args(mcfg: "RunpodModelConfig", served: str) -> str:
"""Args for the llama.cpp server image so it downloads and serves the GGUF.
......@@ -186,6 +224,8 @@ def parse_model_runpod(block: Optional[dict]) -> RunpodModelConfig:
cfg.allow_spot = _as_bool(b.get("allow_spot"), cfg.allow_spot)
cfg.image = (b.get("image") or "").strip()
cfg.engine = (b.get("engine") or cfg.engine).strip().lower() or "auto"
cfg.health_path = (b.get("health_path") or "").strip()
cfg.docker_args = (b.get("docker_args") or "").strip()
cfg.served_model = (b.get("served_model") or "").strip()
cfg.hf_gguf = (b.get("hf_gguf") or "").strip()
cfg.container_disk_gb = _as_int(b.get("container_disk_gb"), cfg.container_disk_gb)
......@@ -260,11 +300,11 @@ import uuid
from dataclasses import dataclass as _dc_pod
def _pod_health_ok(url: str, timeout: float = 4.0) -> bool:
"""True when the pod's OpenAI server answers /v1/models (vLLM is up)."""
def _pod_health_ok(url: str, timeout: float = 4.0, path: str = "/v1/models") -> bool:
"""True when the server inside the pod answers its readiness path."""
import requests
try:
r = requests.get(url.rstrip("/") + "/v1/models", timeout=timeout)
r = requests.get(url.rstrip("/") + "/" + path.lstrip("/"), timeout=timeout)
return r.status_code == 200
except Exception:
return False
......@@ -380,6 +420,8 @@ class RunpodPodPool:
self.mcfg = mcfg
self.served = served_name
self.pods: list = []
#: Set from the pod plan at provision time; the default suits an LLM pod.
self.health_path = mcfg.health_path or "/v1/models"
self._cv = threading.Condition(threading.RLock())
self._provisioning = False
self._closed = False
......@@ -426,14 +468,10 @@ class RunpodPodPool:
(pod_id, sel, next_index). Raises if all fail."""
from codai.api.runpod_client import RunpodError
port = self.mcfg.port or 8000
engine = resolve_pod_engine(self.mcfg, str(self.model_key),
_model_path_for(self.model_key))
if engine == "llamacpp":
image = self.mcfg.image or LLAMACPP_POD_IMAGE
args = _llamacpp_docker_args(self.mcfg, self.served)
else:
image = self.mcfg.image or DEFAULT_POD_IMAGE
args = _vllm_docker_args(self.mcfg, self.served)
plan = pod_plan(self.mcfg, self.served, str(self.model_key),
_model_path_for(self.model_key))
image, args = plan["image"], plan["args"]
self.health_path = plan["health_path"]
env = dict(self.mcfg.env or {})
dc = getattr(self.account, "data_center", "") or ""
last = None
......@@ -522,7 +560,7 @@ class RunpodPodPool:
# Wait for the OpenAI server inside the pod (image pull + model load).
deadline = time.time() + load_to
while time.time() < deadline:
if _pod_health_ok(url):
if _pod_health_ok(url, path=self.health_path):
break
time.sleep(5)
else:
......@@ -685,6 +723,31 @@ def get_pod_pool(model_key, account_cfg, mcfg: "RunpodModelConfig", served_name)
return pool
def get_capability_pool(capability: str, block: dict):
"""A managed pod pool for a whole capability (images, video, tts, …).
The pool machinery never cared that it was serving one model: it provisions,
health-checks, scales and reaps whatever image it is given. Pointing it at a
capability is what brings budgets, autoscaling and idle teardown to the
subsystems that have no local engine to pin — the remote gateway asks for a
URL here exactly the way the chat backend does.
Defaults to `engine: coderai`, since a capability is served by a whole
coderai on the far side rather than by an LLM server.
"""
from codai.models.manager import get_active_runpod_config
acct = get_active_runpod_config()
if acct is None or not getattr(acct, "enabled", False):
raise RuntimeError(
"RunPod is not enabled — cannot provision a pod for capability "
f"{capability!r}. Enable it (and set an API key) or point the "
"capability at a URL instead.")
b = dict(block or {})
b.setdefault("engine", "coderai")
mcfg = parse_model_runpod(b)
return get_pod_pool(f"capability:{capability}", acct, mcfg, capability)
def _all_pools_hourly_rate() -> float:
with _pools_lock:
return sum(p.hourly_rate() for p in _pools.values())
......
......@@ -380,6 +380,12 @@ class RemotesConfig:
api_key: str = "" # bearer token for the remote(s)
max_body_mb: int = 512 # bigger requests are served locally
endpoints: dict = field(default_factory=dict) # capability -> base URL
# capability -> RunPod pod block (same fields as a model's `runpod` block).
# coderai then provisions, health-checks, scales and reaps that pod itself,
# so an image or video capability gets the budgets and idle teardown a
# RunPod-served LLM has. `"endpoints": {"images": "runpod"}` is shorthand
# for "use the pod block for images" (defaults apply when there is none).
pods: dict = field(default_factory=dict)
@dataclass
......@@ -1101,6 +1107,7 @@ class ConfigManager:
"api_key": self.config.remotes.api_key,
"max_body_mb": self.config.remotes.max_body_mb,
"endpoints": dict(self.config.remotes.endpoints or {}),
"pods": dict(self.config.remotes.pods or {}),
},
"ds4": {
"enabled": self.config.ds4.enabled,
......
......@@ -142,6 +142,56 @@ Capability keys: `images`, `video`, `embeddings`, `rerank`, `ocr`, `tts`, `stt`,
Precedence: a model's own `service_url` wins over the capability map.
### Letting coderai rent the pod
A URL means you started the machine and you own its bill. To have coderai
provision, scale and reap the pod instead, give the capability a pod block —
the same fields a model's `runpod` block takes:
```json
{
"remotes": {
"endpoints": {"images": "runpod"},
"pods": {
"images": {
"image": "registry.example.com/coderai:base",
"min_vram_gb": 24, "max_hourly_usd": 0.6,
"max_pods": 2, "idle_timeout_s": 300,
"cost_limit_usd": 20, "cost_period": "day"
}
}
}
}
```
`"endpoints": {"images": "runpod"}` is shorthand for "use the pod block for
images"; a capability listed under `pods` is routed there even without it.
The first matching request provisions a pod, the gateway borrows it for the
duration of the request and hands it back, and the scaler tears it down after
`idle_timeout_s` — so an idle capability costs nothing. Rate caps, rolling spend
limits, the global budget, the stale-pod reaper and the stats page all apply
exactly as they do for a RunPod-served LLM; capability pools show up there as
`capability:<name>`.
Pod engines (`engine`):
| value | image | ready when |
|---|---|---|
| `auto` | vLLM, or llama.cpp for a GGUF | `/v1/models` |
| `vllm` / `llamacpp` | forced | `/v1/models` |
| `coderai` | your `image` (required) | `/healthz` |
| `custom` | your `image`, `docker_args` verbatim | `health_path` |
`coderai` is the default for a capability pool: the far side is a whole coderai,
reached through the same API. You must supply `image` — a registry the pod can
pull from — because coderai does not publish one for you.
If a pod cannot be provisioned (RunPod disabled, budget cap hit, no GPU under
the ceiling) the request fails with 502. It does **not** quietly fall back to
running the model locally — that failure mode is only discovered by noticing
your own VRAM disappear.
Notes:
* `/v1/chat/completions` and `/v1/completions` are deliberately **not** gatewayed
......@@ -159,6 +209,9 @@ Notes:
* **One small coderai, many remote GPUs** — the local instance holds the API,
auth, catalogue and UI; each capability points at a box or pod that owns one
model family.
* **Rent per capability, pay only while used** — capabilities with a pod block
scale from zero on demand and are reaped when idle, under the same budgets as
a RunPod-served LLM.
* **Per-engine pod images** — vLLM's and llama.cpp's own images serve LLMs with
no coderai on the far side; the coderai image is only needed where a coderai
subsystem runs (the mux service, or a gateway target).
......
......@@ -373,6 +373,99 @@ def test_chat_is_never_gatewayed(gateway_app):
json={"model": "anything"}).json()["served_by"] == "local"
def test_capability_served_by_a_managed_pod(gateway_app, monkeypatch):
"""A capability can point at a RunPod pool instead of a fixed URL: the
gateway borrows a pod for the request and gives it back afterwards."""
client, gw, url = gateway_app
released = []
class _Handle:
pod_id = "pod-1"
class _Pool:
def acquire(self):
return _Handle(), url
def release(self, handle):
released.append(handle.pod_id)
import codai.api.runpod_worker as rw
monkeypatch.setattr(rw, "get_capability_pool", lambda cap, block: _Pool())
gw.capability_pods = lambda: {"images": {"max_pods": 2}}
gw.capability_endpoints = lambda: {}
r = client.post("/v1/images/generations", json={"model": "sdxl", "prompt": "cat"})
assert r.json()["served_by"] == "remote"
# The pod is handed back even on the happy path, so the idle reaper can see it.
assert released == ["pod-1"]
def test_runpod_shorthand_endpoint_uses_the_pool(gateway_app, monkeypatch):
client, gw, url = gateway_app
import codai.api.runpod_worker as rw
asked = []
class _Pool:
def acquire(self):
return object(), url
def release(self, handle):
pass
monkeypatch.setattr(rw, "get_capability_pool",
lambda cap, block: asked.append(cap) or _Pool())
gw.capability_endpoints = lambda: {"video": "runpod"}
gw.capability_pods = lambda: {}
assert client.post("/v1/video/generations",
json={"model": "wan"}).json()["served_by"] == "remote"
assert asked == ["video"]
def test_pod_that_cannot_be_provisioned_is_a_502_not_a_local_run(gateway_app, monkeypatch):
"""Falling back to local would quietly load a 30 GB model on the wrong box."""
client, gw, _ = gateway_app
import codai.api.runpod_worker as rw
def _boom(cap, block):
raise RuntimeError("RunPod is not enabled")
monkeypatch.setattr(rw, "get_capability_pool", _boom)
gw.capability_pods = lambda: {"images": {}}
gw.capability_endpoints = lambda: {}
r = client.post("/v1/images/generations", json={"model": "sdxl"})
assert r.status_code == 502 and "RunPod is not enabled" in r.json()["error"]["message"]
def test_pod_plan_per_engine():
from codai.api.runpod_worker import (parse_model_runpod, pod_plan,
LLAMACPP_POD_IMAGE, DEFAULT_POD_IMAGE)
vllm = pod_plan(parse_model_runpod({"served_model": "Qwen/Qwen3.5-9B"}),
"Qwen/Qwen3.5-9B")
assert vllm["image"] == DEFAULT_POD_IMAGE and vllm["health_path"] == "/v1/models"
gguf = pod_plan(parse_model_runpod({"hf_gguf": "user/repo:Q4_K_M"}), "m",
model_path="/AI/m.gguf")
assert gguf["image"] == LLAMACPP_POD_IMAGE and "-hf user/repo:Q4_K_M" in gguf["args"]
# A coderai pod serves whole capabilities; it is ready as soon as /healthz answers.
cod = pod_plan(parse_model_runpod({"engine": "coderai", "image": "reg/coderai:base"}),
"images")
assert cod["image"] == "reg/coderai:base" and cod["health_path"] == "/healthz"
# …but it needs an image: we cannot invent a registry the pod can pull from.
with pytest.raises(RuntimeError, match="set `image`"):
pod_plan(parse_model_runpod({"engine": "coderai"}), "images")
# `custom` takes verbatim args and an explicit probe.
cus = pod_plan(parse_model_runpod({"engine": "custom", "image": "me/thing:1",
"docker_args": "--serve", "health_path": "/ping"}),
"x")
assert cus["args"] == "--serve" and cus["health_path"] == "/ping"
def test_unreachable_remote_is_a_clean_502(gateway_app):
client, gw, _ = gateway_app
gw.capability_endpoints = lambda: {"images": "http://127.0.0.1:1"}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment