Compare commits
14 Commits
fix/neuron
...
feat/73-ca
| Author | SHA1 | Date | |
|---|---|---|---|
|
7984d27553
|
|||
|
43ffffdccb
|
|||
|
5fd7736abd
|
|||
|
03fd4960c3
|
|||
|
5ed6bc3390
|
|||
|
cabec1d08a
|
|||
|
881fc85a4c
|
|||
|
b2ed20b55a
|
|||
|
bee27e9b9c
|
|||
|
87d9c291ce
|
|||
|
d4742467e0
|
|||
|
e7f7e376fc
|
|||
|
3b9a6e37f6
|
|||
|
526b662c5e
|
26
CLAUDE.md
26
CLAUDE.md
@@ -185,6 +185,32 @@ Run these locally before pushing. `cargo fmt --all` fixes formatting
|
||||
automatically. Clippy warnings must be resolved, not suppressed with
|
||||
`#[allow(...)]` unless there is a clear rationale.
|
||||
|
||||
## Development workflow
|
||||
|
||||
Work each change on its own branch; `main` stays releasable.
|
||||
|
||||
1. Implement on a feature branch (`fix/<issue>-…`, `feat/<issue>-…`).
|
||||
2. Run the CI triad locally (`cargo fmt --check --all`,
|
||||
`cargo clippy --workspace -- -D warnings`, `cargo test --workspace`).
|
||||
Local builds are **CPU-only** — the `#[cfg(feature = "cuda")]` neuron/TP
|
||||
paths do NOT compile locally. The branch CI's **CUDA type-check** job is
|
||||
the only thing that validates them, so for any neuron change the push to
|
||||
Gitea is the real gate, not a rubber stamp.
|
||||
3. Push the branch on local-green (no need to ask first), and background-watch
|
||||
its CI run via the gitea-mcp `actions_run_read` tools. Start the next piece
|
||||
of work meanwhile.
|
||||
4. Merge to `main` when the four **validation** jobs are green — Format,
|
||||
Clippy, Test, CUDA type-check. The SRPM / COPR / version-bump jobs are the
|
||||
deploy pipeline (they run on `main`), not validation — don't wait on them.
|
||||
5. Merging/pushing to `main` triggers the auto-deploy pipeline.
|
||||
|
||||
Docs-only changes (no `#[cfg(feature = "cuda")]` impact) can go straight to
|
||||
`main` — there's nothing for the CUDA type-check to prove.
|
||||
|
||||
SSH note: the gitea remote host offers multiple agent keys and cuts the
|
||||
connection before reaching the right one. This repo pins the working key via
|
||||
`git config core.sshCommand "ssh -i ~/.ssh/id_grenade -o IdentitiesOnly=yes"`.
|
||||
|
||||
## Environment
|
||||
|
||||
- Targets Fedora 43 (systemd, SELinux enforcing)
|
||||
|
||||
34
Cargo.lock
generated
34
Cargo.lock
generated
@@ -800,6 +800,7 @@ dependencies = [
|
||||
"cortex-core",
|
||||
"eventsource-stream",
|
||||
"futures",
|
||||
"helexa-stream",
|
||||
"metrics",
|
||||
"metrics-exporter-prometheus",
|
||||
"reqwest",
|
||||
@@ -1922,6 +1923,39 @@ dependencies = [
|
||||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "helexa-router"
|
||||
version = "0.1.16"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"axum",
|
||||
"chrono",
|
||||
"clap",
|
||||
"cortex-core",
|
||||
"figment",
|
||||
"helexa-stream",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tower-http",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "helexa-stream"
|
||||
version = "0.1.16"
|
||||
dependencies = [
|
||||
"async-stream",
|
||||
"axum",
|
||||
"futures",
|
||||
"reqwest",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hermit-abi"
|
||||
version = "0.5.2"
|
||||
|
||||
@@ -7,6 +7,8 @@ members = [
|
||||
"crates/neuron",
|
||||
"crates/helexa-acp",
|
||||
"crates/helexa-bench",
|
||||
"crates/helexa-router",
|
||||
"crates/helexa-stream",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
|
||||
@@ -54,10 +54,26 @@ pub struct ModelLimit {
|
||||
pub output: usize,
|
||||
}
|
||||
|
||||
/// Operator-set pricing in USD per 1M tokens.
|
||||
/// Operator-set pricing, **USD per 1,000,000 tokens, as JSON numbers**
|
||||
/// (`float`) — the models.dev/opencode `cost` convention, which is what
|
||||
/// helexa's primary client reads. NOT per-token, NOT decimal strings (that
|
||||
/// is OpenRouter's `pricing` shape, which helexa deliberately does not emit
|
||||
/// — see #68). A client must not rescale by 10⁶.
|
||||
///
|
||||
/// Self-hosted deployments typically leave both at `0.0`. Cache fields are
|
||||
/// optional — set when the backend supports a prefix-cache discount tier.
|
||||
/// `cost` is sourced from the operator's `models.toml` catalogue profile and
|
||||
/// surfaced verbatim on `/v1/models`. The *absent* vs *zero* distinction is
|
||||
/// intentional and load-bearing (#68):
|
||||
/// - **`cost` absent** (the whole object omitted) — the model is **not
|
||||
/// priced**: the operator has not declared a rate. Clients should treat
|
||||
/// spend as unknown, not free.
|
||||
/// - **`cost` present with `input`/`output` = `0.0`** — the model is
|
||||
/// **intentionally free** (self-hosted, no charge). opencode renders `$0`.
|
||||
///
|
||||
/// Cache fields are optional — set them only when the backend supports a
|
||||
/// prefix-cache discount tier (relevant once cache-token reporting, #64,
|
||||
/// lands). The advertised rate here must equal the rate metering (#51) and
|
||||
/// reconciliation (#58/#59) bill against; today both read this catalogue
|
||||
/// value.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelCost {
|
||||
/// USD per 1M input (prompt) tokens.
|
||||
@@ -98,7 +114,8 @@ pub struct ModelInfo {
|
||||
/// `None` when neither the catalogue nor the loaded model can provide it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub limit: Option<ModelLimit>,
|
||||
/// Operator-set pricing in USD per 1M tokens (0.0 = free/self-hosted).
|
||||
/// Operator-set pricing — see [`ModelCost`] for units and the
|
||||
/// absent (not priced) vs `0.0` (intentionally free) distinction.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cost: Option<ModelCost>,
|
||||
/// `true` when the model's tokenizer contains recognised tool-call
|
||||
|
||||
@@ -32,6 +32,12 @@ pub struct NodeState {
|
||||
/// least-busy replica when a model is loaded on more than one neuron.
|
||||
/// Empty until the first /health poll reports load.
|
||||
pub model_load: HashMap<String, ModelLoad>,
|
||||
/// Consecutive failed `/models` polls. The poller marks a node
|
||||
/// unhealthy only once this crosses a threshold, so a single transient
|
||||
/// miss (e.g. a neuron momentarily slow to answer while busy) doesn't
|
||||
/// yank the node — and all its models — out of routing. Reset to 0 on
|
||||
/// any successful poll.
|
||||
pub consecutive_poll_failures: u32,
|
||||
}
|
||||
|
||||
/// A model registered on a node, with its runtime status.
|
||||
@@ -130,7 +136,9 @@ pub struct CortexModelEntry {
|
||||
/// at load time. `None` when neither source provides it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub limit: Option<ModelLimit>,
|
||||
/// Operator-set pricing in USD per 1M tokens (0.0 = free/self-hosted).
|
||||
/// Operator-set pricing from the catalogue profile — see
|
||||
/// [`cortex_core::harness::ModelCost`] for units (USD per 1M tokens) and
|
||||
/// the absent (not priced) vs `0.0` (intentionally free) distinction.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cost: Option<ModelCost>,
|
||||
/// `true` when any neuron reports this model supports tool calls.
|
||||
|
||||
@@ -6,6 +6,7 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
cortex-core.workspace = true
|
||||
helexa-stream = { path = "../helexa-stream" }
|
||||
async-trait.workspace = true
|
||||
tokio.workspace = true
|
||||
axum.workspace = true
|
||||
|
||||
@@ -5,12 +5,29 @@ use crate::state::CortexState;
|
||||
use chrono::Utc;
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::ModelInfo;
|
||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||
use cortex_core::node::{ModelEntry, ModelStatus, NodeState};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
const POLL_INTERVAL: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Consecutive failed `/models` polls before a node is marked unhealthy.
|
||||
/// Debounces transient misses (a busy neuron briefly slow to answer) so a
|
||||
/// single blip can't yank a node — and its models — out of routing. At the
|
||||
/// 10s poll interval this tolerates ~20s of flapping before evicting.
|
||||
const POLL_FAILURE_THRESHOLD: u32 = 3;
|
||||
|
||||
/// Record a failed poll for `node`, marking it unhealthy only once failures
|
||||
/// reach [`POLL_FAILURE_THRESHOLD`]. Below the threshold the node keeps its
|
||||
/// last-known health, riding over transient misses. A successful poll resets
|
||||
/// the counter (see the success arm in `poll_once`).
|
||||
fn record_poll_failure(node: &mut NodeState) {
|
||||
node.consecutive_poll_failures = node.consecutive_poll_failures.saturating_add(1);
|
||||
if node.consecutive_poll_failures >= POLL_FAILURE_THRESHOLD {
|
||||
node.healthy = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs forever, polling all neurons on a fixed interval.
|
||||
pub async fn poll_loop(fleet: Arc<CortexState>) {
|
||||
loop {
|
||||
@@ -138,13 +155,14 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
|
||||
// Remove models no longer reported by the neuron.
|
||||
node.models.retain(|id, _| seen.contains(id));
|
||||
|
||||
node.consecutive_poll_failures = 0;
|
||||
node.healthy = true;
|
||||
node.last_poll = Some(Utc::now());
|
||||
tracing::debug!(node = name, models = models.len(), "poll ok");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(node = name, error = %e, "failed to parse /models response");
|
||||
node.healthy = false;
|
||||
record_poll_failure(node);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -154,11 +172,11 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
|
||||
status = %resp.status(),
|
||||
"neuron returned non-success status"
|
||||
);
|
||||
node.healthy = false;
|
||||
record_poll_failure(node);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(node = name, error = %e, "failed to reach neuron");
|
||||
node.healthy = false;
|
||||
record_poll_failure(node);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,21 +1,27 @@
|
||||
//! Streaming HTTP reverse proxy to neuron backends.
|
||||
//!
|
||||
//! For streaming requests, SSE chunks are forwarded as they arrive.
|
||||
//! The proxy captures timing information for metrics but does not
|
||||
//! buffer the full response.
|
||||
//! The streaming *mechanism* — forward an SSE body chunk-for-chunk without
|
||||
//! buffering, observing the bytes for metrics — lives in the shared
|
||||
//! [`helexa_stream`] crate (#71), so cortex and helexa-router use one
|
||||
//! implementation. This module supplies cortex's *policy*: the
|
||||
//! [`CortexMetrics`] observer (per-request token metrics + per-principal
|
||||
//! reservation settle), cortex's logging contract, and the cortex error
|
||||
//! envelope. The usage-extraction helper is re-exported from the shared
|
||||
//! crate so existing call sites keep working.
|
||||
|
||||
use crate::router::RouteDecision;
|
||||
use anyhow::Result;
|
||||
use axum::body::Body;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use futures::Stream;
|
||||
use futures::stream::BoxStream;
|
||||
use helexa_stream::{BodyTail, ChunkObserver, StreamError};
|
||||
use reqwest::Client;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Re-export the shared usage-extraction helper. Several cortex modules
|
||||
/// (`handlers`, `anthropic_sse`) pull token counts out of a buffered body
|
||||
/// tail via this function; it lives in `helexa-stream` now.
|
||||
pub use helexa_stream::last_count_for;
|
||||
|
||||
/// Proxy a request body to the resolved backend node and stream the response.
|
||||
///
|
||||
/// Logging contract: every call emits exactly one structured event at
|
||||
@@ -42,66 +48,41 @@ pub async fn forward_request(
|
||||
"proxying request"
|
||||
);
|
||||
|
||||
let mut req_builder = client.post(&url).body(body);
|
||||
let observer = CortexMetrics::new(model_id, &route.node_name, request_start, usage_sink);
|
||||
|
||||
// Forward relevant headers.
|
||||
for (key, value) in headers.iter() {
|
||||
if key == "host" || key == "content-length" {
|
||||
continue; // reqwest sets these
|
||||
}
|
||||
req_builder = req_builder.header(key, value);
|
||||
}
|
||||
let response = helexa_stream::forward_streaming(client, &url, headers, body, observer)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
match &e {
|
||||
StreamError::Upstream(err) => tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %err,
|
||||
"proxy: upstream request failed (network)"
|
||||
),
|
||||
StreamError::ResponseBuild(err) => tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %err,
|
||||
"proxy: failed to build response"
|
||||
),
|
||||
}
|
||||
ProxyError::from(e)
|
||||
})?;
|
||||
|
||||
let upstream_resp = match req_builder.send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %e,
|
||||
"proxy: upstream request failed (network)"
|
||||
);
|
||||
return Err(ProxyError::Upstream(e));
|
||||
}
|
||||
};
|
||||
|
||||
let upstream_status = upstream_resp.status();
|
||||
if !upstream_status.is_success() {
|
||||
if !response.status().is_success() {
|
||||
// Streaming body — can't snippet without breaking the stream
|
||||
// pass-through. Log status + URL; the client still gets the
|
||||
// upstream status, just without the leaked body.
|
||||
tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
status = upstream_status.as_u16(),
|
||||
status = response.status().as_u16(),
|
||||
"proxy: upstream returned non-2xx"
|
||||
);
|
||||
}
|
||||
|
||||
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||
|
||||
let resp_headers = upstream_resp.headers().clone();
|
||||
let stream = TokenMetricsStream::new(
|
||||
Box::pin(upstream_resp.bytes_stream()),
|
||||
TokenMetrics::new(model_id, &route.node_name, request_start, usage_sink),
|
||||
);
|
||||
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
let mut response = Response::builder().status(status);
|
||||
for (key, value) in resp_headers.iter() {
|
||||
response = response.header(key, value);
|
||||
}
|
||||
|
||||
response.body(body).map_err(|e| {
|
||||
tracing::warn!(
|
||||
node = %route.node_name,
|
||||
url = %url,
|
||||
error = %e,
|
||||
"proxy: failed to build response"
|
||||
);
|
||||
ProxyError::ResponseBuild(e.to_string())
|
||||
})
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -112,6 +93,15 @@ pub enum ProxyError {
|
||||
ResponseBuild(String),
|
||||
}
|
||||
|
||||
impl From<StreamError> for ProxyError {
|
||||
fn from(e: StreamError) -> Self {
|
||||
match e {
|
||||
StreamError::Upstream(err) => ProxyError::Upstream(err),
|
||||
StreamError::ResponseBuild(msg) => ProxyError::ResponseBuild(msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for ProxyError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, code, message) = match &self {
|
||||
@@ -139,9 +129,10 @@ impl IntoResponse for ProxyError {
|
||||
//
|
||||
// The proxy never buffers or re-serialises the upstream body — chunks
|
||||
// are forwarded verbatim. For metrics it observes each chunk's arrival
|
||||
// time and keeps a bounded tail of the body text, from which the final
|
||||
// OpenAI `usage` object (present on the last SSE chunk and on
|
||||
// non-streaming JSON bodies alike) yields engine-truth token counts.
|
||||
// time and keeps a bounded tail of the body text (via the shared
|
||||
// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
|
||||
// (present on the last SSE chunk and on non-streaming JSON bodies alike)
|
||||
// yields engine-truth token counts.
|
||||
//
|
||||
// Emitted per request, labelled {model, node}:
|
||||
// cortex_time_to_first_token_seconds (histogram) — first body chunk
|
||||
@@ -155,37 +146,15 @@ impl IntoResponse for ProxyError {
|
||||
/// non-streaming bodies.
|
||||
const TAIL_CAP_BYTES: usize = 64 * 1024;
|
||||
|
||||
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
||||
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
||||
/// The quoted-needle form means `completion_tokens` never matches
|
||||
/// `completion_tokens_details`.
|
||||
pub(crate) fn last_count_for(tail: &str, key: &str) -> Option<u64> {
|
||||
let needle = format!("\"{key}\"");
|
||||
let mut result = None;
|
||||
for (idx, _) in tail.match_indices(&needle) {
|
||||
let rest = tail[idx + needle.len()..].trim_start();
|
||||
let Some(rest) = rest.strip_prefix(':') else {
|
||||
continue;
|
||||
};
|
||||
let rest = rest.trim_start();
|
||||
let digits: &str = &rest[..rest
|
||||
.char_indices()
|
||||
.find(|(_, c)| !c.is_ascii_digit())
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(rest.len())];
|
||||
if let Ok(v) = digits.parse::<u64>() {
|
||||
result = Some(v);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
struct TokenMetrics {
|
||||
/// cortex's [`ChunkObserver`]: per-request token metrics plus the
|
||||
/// per-principal reservation settle. Drives cortex policy over the shared
|
||||
/// streaming mechanism.
|
||||
struct CortexMetrics {
|
||||
labels: [(&'static str, String); 2],
|
||||
request_start: Instant,
|
||||
first_chunk: Option<Instant>,
|
||||
last_chunk: Option<Instant>,
|
||||
tail: String,
|
||||
tail: BodyTail,
|
||||
finished: bool,
|
||||
/// Per-principal metering hook (#51). Invoked exactly once in `finish`
|
||||
/// with the observed `(prompt, completion)` so the reservation can be
|
||||
@@ -193,7 +162,7 @@ struct TokenMetrics {
|
||||
usage_sink: Option<crate::metering::UsageSink>,
|
||||
}
|
||||
|
||||
impl TokenMetrics {
|
||||
impl CortexMetrics {
|
||||
fn new(
|
||||
model_id: &str,
|
||||
node_name: &str,
|
||||
@@ -208,26 +177,19 @@ impl TokenMetrics {
|
||||
request_start,
|
||||
first_chunk: None,
|
||||
last_chunk: None,
|
||||
tail: String::new(),
|
||||
tail: BodyTail::new(TAIL_CAP_BYTES),
|
||||
finished: false,
|
||||
usage_sink,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ChunkObserver for CortexMetrics {
|
||||
fn observe(&mut self, chunk: &[u8]) {
|
||||
let now = Instant::now();
|
||||
self.first_chunk.get_or_insert(now);
|
||||
self.last_chunk = Some(now);
|
||||
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||
if self.tail.len() > TAIL_CAP_BYTES {
|
||||
// Keep the newest half; the usage object is always at the
|
||||
// very end of the body. Split at a char boundary.
|
||||
let mut cut = self.tail.len() - TAIL_CAP_BYTES / 2;
|
||||
while !self.tail.is_char_boundary(cut) {
|
||||
cut += 1;
|
||||
}
|
||||
self.tail.drain(..cut);
|
||||
}
|
||||
self.tail.push(chunk);
|
||||
}
|
||||
|
||||
/// Emit the metrics exactly once — called on clean stream end and
|
||||
@@ -239,8 +201,8 @@ impl TokenMetrics {
|
||||
}
|
||||
self.finished = true;
|
||||
|
||||
let prompt = last_count_for(&self.tail, "prompt_tokens");
|
||||
let completion = last_count_for(&self.tail, "completion_tokens");
|
||||
let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
|
||||
let completion = last_count_for(self.tail.as_str(), "completion_tokens");
|
||||
|
||||
// Per-model metrics — only when body chunks actually arrived.
|
||||
if let Some(first) = self.first_chunk {
|
||||
@@ -280,97 +242,3 @@ impl TokenMetrics {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
|
||||
/// clean end-of-stream; the Drop impl covers client disconnects.
|
||||
struct TokenMetricsStream {
|
||||
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
|
||||
metrics: TokenMetrics,
|
||||
}
|
||||
|
||||
impl TokenMetricsStream {
|
||||
fn new(
|
||||
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
|
||||
metrics: TokenMetrics,
|
||||
) -> Self {
|
||||
Self { inner, metrics }
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for TokenMetricsStream {
|
||||
type Item = Result<bytes::Bytes, reqwest::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let this = self.get_mut();
|
||||
match this.inner.as_mut().poll_next(cx) {
|
||||
Poll::Ready(Some(Ok(chunk))) => {
|
||||
this.metrics.observe(&chunk);
|
||||
Poll::Ready(Some(Ok(chunk)))
|
||||
}
|
||||
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
||||
Poll::Ready(None) => {
|
||||
this.metrics.finish();
|
||||
Poll::Ready(None)
|
||||
}
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TokenMetricsStream {
|
||||
fn drop(&mut self) {
|
||||
self.metrics.finish();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::last_count_for;
|
||||
|
||||
#[test]
|
||||
fn extracts_counts_from_final_sse_usage_chunk() {
|
||||
let tail = concat!(
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
||||
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_counts_from_non_streaming_body() {
|
||||
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
||||
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_details_variants_and_takes_last_occurrence() {
|
||||
// completion_tokens_details must not shadow completion_tokens,
|
||||
// and the LAST usage object wins (matters when content echoes
|
||||
// a usage-shaped string earlier in the stream).
|
||||
let tail = concat!(
|
||||
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
||||
"data: {\"usage\":{\"completion_tokens\":99,",
|
||||
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
||||
);
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_keys_yield_none() {
|
||||
assert_eq!(
|
||||
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
||||
None
|
||||
);
|
||||
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
||||
// key present but non-numeric value
|
||||
assert_eq!(
|
||||
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,6 +50,10 @@ pub enum RouteError {
|
||||
"model '{model_id}' is in the catalogue but no healthy neuron's topology satisfies its constraints"
|
||||
)]
|
||||
NoFeasibleNeuron { model_id: String },
|
||||
#[error(
|
||||
"model '{model_id}' is feasible on a neuron that is currently unhealthy — retry shortly"
|
||||
)]
|
||||
FeasibleNodeUnhealthy { model_id: String },
|
||||
#[error("cold-load of '{model_id}' on '{node}' failed: {message}")]
|
||||
ColdLoadFailed {
|
||||
model_id: String,
|
||||
@@ -68,7 +72,9 @@ impl RouteError {
|
||||
/// safe to retry the same request); everything else is 404.
|
||||
pub fn http_status(&self) -> u16 {
|
||||
match self {
|
||||
RouteError::NoHealthyNodes | RouteError::ModelRecovering { .. } => 503,
|
||||
RouteError::NoHealthyNodes
|
||||
| RouteError::ModelRecovering { .. }
|
||||
| RouteError::FeasibleNodeUnhealthy { .. } => 503,
|
||||
_ => 404,
|
||||
}
|
||||
}
|
||||
@@ -81,7 +87,8 @@ impl RouteError {
|
||||
| RouteError::EndpointResolveFailed(_, _)
|
||||
| RouteError::NoFeasibleNeuron { .. }
|
||||
| RouteError::ColdLoadFailed { .. }
|
||||
| RouteError::ModelRecovering { .. } => "api_error",
|
||||
| RouteError::ModelRecovering { .. }
|
||||
| RouteError::FeasibleNodeUnhealthy { .. } => "api_error",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -94,6 +101,7 @@ impl RouteError {
|
||||
RouteError::NoFeasibleNeuron { .. } => "service_unavailable",
|
||||
RouteError::ColdLoadFailed { .. } => "service_unavailable",
|
||||
RouteError::ModelRecovering { .. } => "service_unavailable",
|
||||
RouteError::FeasibleNodeUnhealthy { .. } => "service_unavailable",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,6 +113,7 @@ impl RouteError {
|
||||
pub fn retry_after_secs(&self) -> Option<u64> {
|
||||
match self {
|
||||
RouteError::ModelRecovering { .. } => Some(2),
|
||||
RouteError::FeasibleNodeUnhealthy { .. } => Some(3),
|
||||
RouteError::NoHealthyNodes => Some(5),
|
||||
_ => None,
|
||||
}
|
||||
@@ -252,11 +261,32 @@ async fn pick_feasible_neuron(
|
||||
b.2.cmp(&a.2) // pinned first (true > false)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
let pick = candidates.into_iter().next();
|
||||
pick.map(|(n, e, _)| (n, e))
|
||||
.ok_or_else(|| RouteError::NoFeasibleNeuron {
|
||||
if let Some((n, e, _)) = candidates.into_iter().next() {
|
||||
return Ok((n, e));
|
||||
}
|
||||
|
||||
// No *healthy* feasible neuron. Distinguish a transient outage from a
|
||||
// permanent misconfiguration: if some neuron is topologically feasible
|
||||
// but currently unhealthy (e.g. it briefly missed polls while busy),
|
||||
// this is retryable — return 503 + Retry-After so the client backs off
|
||||
// and retries instead of treating a 404 as a hard failure. Only when no
|
||||
// neuron could *ever* satisfy the topology is it a permanent 404.
|
||||
let feasible_but_unhealthy = nodes.values().any(|node| {
|
||||
!node.healthy
|
||||
&& node
|
||||
.discovery
|
||||
.as_ref()
|
||||
.is_some_and(|disc| profile.is_feasible_on(&node.name, &disc.devices))
|
||||
});
|
||||
if feasible_but_unhealthy {
|
||||
Err(RouteError::FeasibleNodeUnhealthy {
|
||||
model_id: profile.id.clone(),
|
||||
})
|
||||
} else {
|
||||
Err(RouteError::NoFeasibleNeuron {
|
||||
model_id: profile.id.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Issue `POST {endpoint}/models/load` for this profile on this neuron,
|
||||
|
||||
@@ -38,6 +38,7 @@ impl CortexState {
|
||||
discovery: None,
|
||||
activation: None,
|
||||
model_load: HashMap::new(),
|
||||
consecutive_poll_failures: 0,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
124
crates/cortex-gateway/tests/feasibility_routing.rs
Normal file
124
crates/cortex-gateway/tests/feasibility_routing.rs
Normal file
@@ -0,0 +1,124 @@
|
||||
//! Router: a catalogued model whose only topologically-feasible neuron is
|
||||
//! currently unhealthy is a *transient* condition (retryable 503), not a
|
||||
//! permanent 404. This is the exact shape of the beast incident: benjy/
|
||||
//! quadbrat (1 GPU, healthy) can't host the 27B, and beast (2 GPU) — the
|
||||
//! sole feasible node — briefly drops out → clients must back off and retry,
|
||||
//! not hard-fail.
|
||||
|
||||
use cortex_core::config::{
|
||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_core::discovery::{DeviceInfo, DiscoveryResponse};
|
||||
use cortex_gateway::router::{self, RouteError};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use std::sync::Arc;
|
||||
|
||||
fn devices(n: usize) -> Vec<DeviceInfo> {
|
||||
(0..n)
|
||||
.map(|i| DeviceInfo {
|
||||
index: i as u32,
|
||||
name: "RTX 5090".into(),
|
||||
vram_total_mb: 32_768,
|
||||
compute_capability: "9.0".into(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn discovery(host: &str, n_devices: usize) -> DiscoveryResponse {
|
||||
DiscoveryResponse {
|
||||
hostname: host.into(),
|
||||
os: "Linux".into(),
|
||||
kernel: "7.0".into(),
|
||||
cuda_version: Some("13.0".into()),
|
||||
driver_version: Some("999".into()),
|
||||
devices: devices(n_devices),
|
||||
harnesses: vec!["candle".into()],
|
||||
cuda_unavailable_reason: None,
|
||||
max_prompt_tokens: 49_152,
|
||||
}
|
||||
}
|
||||
|
||||
/// Catalogue with one model needing 2 devices. Returns a temp path.
|
||||
fn write_catalogue() -> std::path::PathBuf {
|
||||
let toml = r#"
|
||||
[[models]]
|
||||
id = "big-model"
|
||||
harness = "candle"
|
||||
min_devices = 2
|
||||
"#;
|
||||
let path = std::env::temp_dir().join("cortex_test_feasibility_models.toml");
|
||||
std::fs::write(&path, toml).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
async fn fleet_with(big_healthy: bool, big_devices: usize) -> Arc<CortexState> {
|
||||
let cat = write_catalogue();
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![
|
||||
NeuronEndpoint {
|
||||
name: "small".into(),
|
||||
endpoint: "http://127.0.0.1:1".into(),
|
||||
},
|
||||
NeuronEndpoint {
|
||||
name: "big".into(),
|
||||
endpoint: "http://127.0.0.1:2".into(),
|
||||
},
|
||||
],
|
||||
models_config: cat.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
// "small" is healthy but only has 1 GPU → not feasible for the model.
|
||||
let small = nodes.get_mut("small").unwrap();
|
||||
small.healthy = true;
|
||||
small.discovery = Some(discovery("small", 1));
|
||||
// "big" has enough GPUs but its health is the variable under test.
|
||||
let big = nodes.get_mut("big").unwrap();
|
||||
big.healthy = big_healthy;
|
||||
big.discovery = Some(discovery("big", big_devices));
|
||||
}
|
||||
fleet
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn feasible_node_unhealthy_is_transient_503() {
|
||||
// big (2 GPU, the only feasible node) is unhealthy; small (1 GPU) is
|
||||
// healthy but can't host the model → retryable, not a permanent 404.
|
||||
let fleet = fleet_with(false, 2).await;
|
||||
let err = router::resolve(&fleet, "big-model")
|
||||
.await
|
||||
.expect_err("model can't be served right now");
|
||||
assert!(
|
||||
matches!(err, RouteError::FeasibleNodeUnhealthy { .. }),
|
||||
"expected FeasibleNodeUnhealthy, got {err:?}"
|
||||
);
|
||||
assert_eq!(err.http_status(), 503);
|
||||
assert_eq!(err.retry_after_secs(), Some(3));
|
||||
assert_eq!(err.code(), "service_unavailable");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_node_can_ever_satisfy_is_permanent_404() {
|
||||
// big is healthy but only has 1 GPU now (e.g. topology genuinely can't
|
||||
// satisfy min_devices=2 anywhere) → permanent, non-retryable 404.
|
||||
let fleet = fleet_with(true, 1).await;
|
||||
let err = router::resolve(&fleet, "big-model")
|
||||
.await
|
||||
.expect_err("no feasible topology");
|
||||
assert!(
|
||||
matches!(err, RouteError::NoFeasibleNeuron { .. }),
|
||||
"expected NoFeasibleNeuron, got {err:?}"
|
||||
);
|
||||
assert_eq!(err.http_status(), 404);
|
||||
assert_eq!(err.retry_after_secs(), None);
|
||||
}
|
||||
131
crates/cortex-gateway/tests/model_cost.rs
Normal file
131
crates/cortex-gateway/tests/model_cost.rs
Normal file
@@ -0,0 +1,131 @@
|
||||
//! Issue #68: the `cost` wire contract on `GET /v1/models`.
|
||||
//!
|
||||
//! `cost` is operator-set pricing sourced from the `models.toml` catalogue
|
||||
//! profile (the source of truth today; the marketplace clearing house #59
|
||||
//! later — both must read the same value metering/#51 bills against). The
|
||||
//! shape is the models.dev/opencode convention: **USD per 1,000,000 tokens,
|
||||
//! as JSON numbers**, with optional `cache_read`/`cache_write` tiers. This
|
||||
//! test pins:
|
||||
//! - the units/shape (per-million floats, not per-token, not strings);
|
||||
//! - that cache fields flow through when present and are omitted otherwise;
|
||||
//! - the load-bearing **absent vs `0.0`** distinction (#68): a model with
|
||||
//! no catalogue `cost` omits the key entirely (price unknown), distinct
|
||||
//! from an explicit `0.0` (intentionally free).
|
||||
//!
|
||||
//! Catalogue-only models surface via Pass 1 of `list_models` even with no
|
||||
//! feasible neuron, so this is hermetic — no nodes or poller needed.
|
||||
|
||||
use cortex_core::config::{
|
||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
#[tokio::test]
|
||||
async fn v1_models_cost_units_shape_and_absent_vs_zero() {
|
||||
// Three catalogue models exercise the whole contract: a priced model
|
||||
// with cache tiers, an intentionally-free model (explicit 0.0), and an
|
||||
// unpriced model (no `cost` block at all).
|
||||
let models_toml = r#"
|
||||
[[models]]
|
||||
id = "priced-model"
|
||||
harness = "candle"
|
||||
cost.input = 0.5
|
||||
cost.output = 1.5
|
||||
cost.cache_read = 0.05
|
||||
cost.cache_write = 0.6
|
||||
|
||||
[[models]]
|
||||
id = "free-model"
|
||||
harness = "candle"
|
||||
cost.input = 0.0
|
||||
cost.output = 0.0
|
||||
|
||||
[[models]]
|
||||
id = "unpriced-model"
|
||||
harness = "candle"
|
||||
"#;
|
||||
let cat_path = std::env::temp_dir().join("cortex_test_issue68_models.toml");
|
||||
std::fs::write(&cat_path, models_toml).unwrap();
|
||||
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
// Never contacted: build_app does not spawn the poller, so the
|
||||
// catalogue alone drives /v1/models.
|
||||
neurons: vec![NeuronEndpoint {
|
||||
name: "mock-node".into(),
|
||||
endpoint: "http://127.0.0.1:1".into(),
|
||||
}],
|
||||
models_config: cat_path.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
let body: serde_json::Value = reqwest::Client::new()
|
||||
.get(format!("http://{addr}/v1/models"))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let data = body["data"].as_array().expect("data is an array");
|
||||
let entry = |id: &str| {
|
||||
data.iter()
|
||||
.find(|m| m["id"] == id)
|
||||
.unwrap_or_else(|| panic!("{id} present in /v1/models"))
|
||||
.clone()
|
||||
};
|
||||
|
||||
// Priced model: exact values flow through as JSON numbers (USD per 1M
|
||||
// tokens). If anything rescaled by 10⁶ or stringified, these fail.
|
||||
let priced = entry("priced-model");
|
||||
assert_eq!(priced["cost"]["input"], 0.5);
|
||||
assert_eq!(priced["cost"]["output"], 1.5);
|
||||
assert_eq!(priced["cost"]["cache_read"], 0.05);
|
||||
assert_eq!(priced["cost"]["cache_write"], 0.6);
|
||||
assert!(
|
||||
priced["cost"]["input"].is_number(),
|
||||
"cost.input must be a JSON number, not a string"
|
||||
);
|
||||
|
||||
// Intentionally free: cost present, rates explicitly 0.0. Unset cache
|
||||
// tiers are omitted (skip_serializing_if), not emitted as null/0.
|
||||
let free = entry("free-model");
|
||||
assert_eq!(free["cost"]["input"], 0.0);
|
||||
assert_eq!(free["cost"]["output"], 0.0);
|
||||
assert!(
|
||||
free["cost"].get("cache_read").is_none(),
|
||||
"absent cache tiers must be omitted, not null"
|
||||
);
|
||||
assert!(free["cost"].get("cache_write").is_none());
|
||||
|
||||
// Unpriced: the whole `cost` object is omitted — "price unknown",
|
||||
// distinct from the free model's explicit 0.0. This is the #68
|
||||
// distinction opencode needs to avoid showing $0 for a model whose
|
||||
// price simply hasn't been declared.
|
||||
let unpriced = entry("unpriced-model");
|
||||
assert!(
|
||||
unpriced.get("cost").is_none(),
|
||||
"a model with no catalogue cost must omit `cost` entirely, got {:?}",
|
||||
unpriced.get("cost")
|
||||
);
|
||||
|
||||
let _ = std::fs::remove_file(&cat_path);
|
||||
}
|
||||
@@ -228,10 +228,26 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
|
||||
nodes.get_mut("dead-node").unwrap().healthy = true;
|
||||
}
|
||||
|
||||
// Debounce (#53 follow-up): a single missed poll must NOT evict a
|
||||
// previously-healthy node — a busy neuron briefly slow to answer
|
||||
// shouldn't yank its models out of routing.
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
assert!(
|
||||
fleet.nodes.read().await.get("dead-node").unwrap().healthy,
|
||||
"one failed poll should not mark a healthy node unhealthy"
|
||||
);
|
||||
|
||||
let nodes = fleet.nodes.read().await;
|
||||
assert!(!nodes.get("dead-node").unwrap().healthy);
|
||||
// It flips unhealthy only after POLL_FAILURE_THRESHOLD (3) consecutive
|
||||
// failures.
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
assert!(
|
||||
!fleet.nodes.read().await.get("dead-node").unwrap().healthy,
|
||||
"three consecutive failed polls should mark the node unhealthy"
|
||||
);
|
||||
|
||||
// A subsequent successful poll would reset the counter and restore
|
||||
// health; covered implicitly by the discovery tests above.
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
35
crates/helexa-router/Cargo.toml
Normal file
35
crates/helexa-router/Cargo.toml
Normal file
@@ -0,0 +1,35 @@
|
||||
[package]
|
||||
name = "helexa-router"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "helexa-router"
|
||||
path = "src/main.rs"
|
||||
|
||||
[lib]
|
||||
name = "helexa_router"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
cortex-core = { workspace = true }
|
||||
helexa-stream = { path = "../helexa-stream" }
|
||||
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
tower-http = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
figment = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
clap = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
# Jail (isolated cwd + env) for config tests.
|
||||
figment = { workspace = true, features = ["test"] }
|
||||
85
crates/helexa-router/src/config.rs
Normal file
85
crates/helexa-router/src/config.rs
Normal file
@@ -0,0 +1,85 @@
|
||||
use figment::{
|
||||
Figment,
|
||||
providers::{Env, Format, Toml},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
/// Top-level `helexa-router` configuration.
|
||||
///
|
||||
/// Loaded from TOML with `HELEXA_ROUTER_`-prefixed env overrides (using
|
||||
/// `__` as the nesting separator), matching the cortex/neuron convention.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RouterConfig {
|
||||
pub router: RouterSettings,
|
||||
/// Downstream cortex endpoints the router can dispatch to. The skeleton
|
||||
/// (#70) only loads these; capacity/catalogue polling (#72) and
|
||||
/// capacity-aware dispatch (#73) consume them later.
|
||||
#[serde(default)]
|
||||
pub cortexes: Vec<CortexEndpoint>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RouterSettings {
|
||||
/// Address to listen on for the inbound API (e.g. "0.0.0.0:8088").
|
||||
///
|
||||
/// Plaintext only — operator/edge nginx terminates client TLS in front
|
||||
/// of the router (see #69's TLS posture). The router never owns an
|
||||
/// inbound TLS listener.
|
||||
pub listen: String,
|
||||
/// How often (seconds) the background poller refreshes each cortex's
|
||||
/// health + `/v1/models` topology (#72). Defaults to 10s, matching the
|
||||
/// cortex↔neuron poll cadence one tier down.
|
||||
#[serde(default = "default_poll_interval_secs")]
|
||||
pub poll_interval_secs: u64,
|
||||
/// This router instance's region (e.g. "eu-west"). When set, dispatch
|
||||
/// (#73) prefers cortexes whose `region` matches, before falling back to
|
||||
/// any feasible cortex. `None` → no geo affinity.
|
||||
#[serde(default)]
|
||||
pub region: Option<String>,
|
||||
}
|
||||
|
||||
fn default_poll_interval_secs() -> u64 {
|
||||
10
|
||||
}
|
||||
|
||||
/// One downstream cortex the router may proxy to. The router verifies the
|
||||
/// cortex's outbound TLS cert (#74) and routes on capacity (#73); it holds
|
||||
/// no entitlement logic of its own and forwards the client bearer verbatim.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct CortexEndpoint {
|
||||
/// Human-readable label (e.g. "lair-cafe").
|
||||
pub name: String,
|
||||
/// Base URL of the cortex gateway (e.g. "https://cortex.example.com").
|
||||
pub endpoint: String,
|
||||
/// Optional region tag (e.g. "eu-west") for geo affinity in dispatch
|
||||
/// (#73). `None` → no region preference applies to this cortex.
|
||||
#[serde(default)]
|
||||
pub region: Option<String>,
|
||||
}
|
||||
|
||||
impl RouterConfig {
|
||||
/// Load configuration from a TOML file, with environment variable
|
||||
/// overrides prefixed with `HELEXA_ROUTER_` and `__` as the separator
|
||||
/// (e.g. `HELEXA_ROUTER_ROUTER__LISTEN=0.0.0.0:8088`).
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self, Box<figment::Error>> {
|
||||
Figment::new()
|
||||
.merge(Toml::file(path))
|
||||
.merge(Env::prefixed("HELEXA_ROUTER_").split("__"))
|
||||
.extract()
|
||||
.map_err(Box::new)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RouterConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
router: RouterSettings {
|
||||
listen: "0.0.0.0:8088".into(),
|
||||
poll_interval_secs: default_poll_interval_secs(),
|
||||
region: None,
|
||||
},
|
||||
cortexes: vec![],
|
||||
}
|
||||
}
|
||||
}
|
||||
215
crates/helexa-router/src/dispatch.rs
Normal file
215
crates/helexa-router/src/dispatch.rs
Normal file
@@ -0,0 +1,215 @@
|
||||
//! Capacity-aware dispatch (#73) — the router's data path.
|
||||
//!
|
||||
//! Given an inbound request's `model`, pick a reachable cortex that can
|
||||
//! serve it (preferring warm/loaded, region-affine, higher-headroom),
|
||||
//! forward the client's bearer **unchanged** (auth stays at cortex), and
|
||||
//! stream the response back verbatim via the shared [`helexa_stream`]
|
||||
//! module. Cortex's #63-shaped rejections (`429 rate_limit_exceeded`,
|
||||
//! `400 context_length_exceeded`, …) pass through untouched. Transport
|
||||
//! failures fail over to the next feasible cortex; a genuine HTTP response —
|
||||
//! any status — is returned as-is and never retried away.
|
||||
//!
|
||||
//! The router holds **no entitlement logic**: it routes on capacity, not
|
||||
//! budget.
|
||||
|
||||
use crate::config::CortexEndpoint;
|
||||
use crate::error::envelope_response;
|
||||
use crate::state::RouterState;
|
||||
use axum::body::Bytes;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::Response;
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use helexa_stream::{ChunkObserver, StreamError};
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Retry-After hint (seconds) on the router's own transient rejections.
|
||||
const RETRY_AFTER_SECS: u64 = 5;
|
||||
|
||||
/// Outcome of choosing where to send a request.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum Selection {
|
||||
/// Feasible reachable cortexes, best-first (failover order).
|
||||
Candidates(Vec<CortexEndpoint>),
|
||||
/// Some cortex knows the model but none are reachable right now → 503.
|
||||
NoReachableCapacity,
|
||||
/// No configured cortex serves the model at all → 404.
|
||||
UnknownModel,
|
||||
}
|
||||
|
||||
/// Rank the reachable cortexes that can serve `model`, best-first.
|
||||
///
|
||||
/// Ordering (each a tie-break for the next): loaded/warm before cold-loadable
|
||||
/// · region match before not · more healthy nodes before fewer · name for
|
||||
/// determinism.
|
||||
pub async fn select_cortexes(state: &RouterState, model: &str) -> Selection {
|
||||
let topo = state.topology.read().await;
|
||||
let by_name: HashMap<&str, &CortexEndpoint> = state
|
||||
.cortexes
|
||||
.iter()
|
||||
.map(|c| (c.name.as_str(), c))
|
||||
.collect();
|
||||
|
||||
let mut ranked: Vec<Ranked> = Vec::new();
|
||||
let mut known_anywhere = false;
|
||||
|
||||
for (name, t) in topo.iter() {
|
||||
let Some(status) = t.models.get(model) else {
|
||||
continue;
|
||||
};
|
||||
if !status.feasible {
|
||||
continue;
|
||||
}
|
||||
// Known even via an unreachable cortex's last-good poll — lets us
|
||||
// tell "temporarily down" (503) from "nobody serves it" (404).
|
||||
known_anywhere = true;
|
||||
if !t.reachable {
|
||||
continue;
|
||||
}
|
||||
let Some(ep) = by_name.get(name.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
let region_match = match (&state.region, &ep.region) {
|
||||
(Some(r), Some(cr)) => r == cr,
|
||||
_ => false,
|
||||
};
|
||||
ranked.push(Ranked {
|
||||
loaded: status.loaded,
|
||||
region_match,
|
||||
healthy_nodes: t.healthy_nodes,
|
||||
endpoint: (*ep).clone(),
|
||||
});
|
||||
}
|
||||
|
||||
if ranked.is_empty() {
|
||||
return if known_anywhere {
|
||||
Selection::NoReachableCapacity
|
||||
} else {
|
||||
Selection::UnknownModel
|
||||
};
|
||||
}
|
||||
|
||||
ranked.sort_by(|a, b| {
|
||||
// false < true, so negate the "good" booleans to sort good first.
|
||||
(
|
||||
!a.loaded,
|
||||
!a.region_match,
|
||||
Reverse(a.healthy_nodes),
|
||||
&a.endpoint.name,
|
||||
)
|
||||
.cmp(&(
|
||||
!b.loaded,
|
||||
!b.region_match,
|
||||
Reverse(b.healthy_nodes),
|
||||
&b.endpoint.name,
|
||||
))
|
||||
});
|
||||
|
||||
Selection::Candidates(ranked.into_iter().map(|r| r.endpoint).collect())
|
||||
}
|
||||
|
||||
struct Ranked {
|
||||
loaded: bool,
|
||||
region_match: bool,
|
||||
healthy_nodes: u32,
|
||||
endpoint: CortexEndpoint,
|
||||
}
|
||||
|
||||
/// Proxy an inbound inference request to a capacity-bearing cortex.
|
||||
///
|
||||
/// `path` is the inference path to forward to (same on the cortex, e.g.
|
||||
/// `/v1/chat/completions`). The body is parsed only to extract `model`.
|
||||
pub async fn dispatch(
|
||||
state: &RouterState,
|
||||
path: &str,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let Some(model) = extract_model(&body) else {
|
||||
return envelope_response(OpenAiError::new(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"missing_model_field",
|
||||
"missing 'model' field in request body",
|
||||
));
|
||||
};
|
||||
|
||||
let candidates = match select_cortexes(state, &model).await {
|
||||
Selection::Candidates(c) => c,
|
||||
Selection::UnknownModel => {
|
||||
return envelope_response(
|
||||
OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"model_not_found",
|
||||
format!("no operator serves model '{model}'"),
|
||||
)
|
||||
.with_param("model"),
|
||||
);
|
||||
}
|
||||
Selection::NoReachableCapacity => {
|
||||
return envelope_response(OpenAiError::service_unavailable(
|
||||
format!("model '{model}' is temporarily unavailable on all operators"),
|
||||
Some(RETRY_AFTER_SECS),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Try candidates in order, failing over only on transport errors. A
|
||||
// genuine HTTP response (any status — including cortex's #63 429/400)
|
||||
// is returned verbatim and never retried away.
|
||||
for ep in &candidates {
|
||||
let url = format!("{}{}", ep.endpoint, path);
|
||||
tracing::info!(cortex = %ep.name, url = %url, model = %model, "dispatching");
|
||||
match helexa_stream::forward_streaming(
|
||||
&state.http_client,
|
||||
&url,
|
||||
headers.clone(),
|
||||
body.clone(),
|
||||
NoopObserver,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => return resp,
|
||||
Err(StreamError::Upstream(e)) => {
|
||||
tracing::warn!(
|
||||
cortex = %ep.name,
|
||||
url = %url,
|
||||
error = %e,
|
||||
"cortex unreachable; failing over"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
Err(StreamError::ResponseBuild(msg)) => {
|
||||
tracing::error!(cortex = %ep.name, error = %msg, "failed to build proxied response");
|
||||
return envelope_response(OpenAiError::without_code(
|
||||
500,
|
||||
"api_error",
|
||||
"failed to build proxied response",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Every feasible cortex failed to connect.
|
||||
tracing::warn!(model = %model, tried = candidates.len(), "all feasible operators unreachable");
|
||||
envelope_response(OpenAiError::service_unavailable(
|
||||
format!("all operators able to serve '{model}' are unreachable"),
|
||||
Some(RETRY_AFTER_SECS),
|
||||
))
|
||||
}
|
||||
|
||||
/// Pull the `model` field out of a request body without re-serialising it.
|
||||
fn extract_model(body: &Bytes) -> Option<String> {
|
||||
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||
v.get("model")?.as_str().map(str::to_string)
|
||||
}
|
||||
|
||||
/// The router proxies bytes verbatim and keeps no per-request policy, so it
|
||||
/// needs no observation hooks. (Token metrics/metering stay at cortex.)
|
||||
struct NoopObserver;
|
||||
|
||||
impl ChunkObserver for NoopObserver {
|
||||
fn observe(&mut self, _chunk: &[u8]) {}
|
||||
fn finish(&mut self) {}
|
||||
}
|
||||
27
crates/helexa-router/src/error.rs
Normal file
27
crates/helexa-router/src/error.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
//! Router adapter from the shared, axum-agnostic
|
||||
//! [`cortex_core::error_envelope::OpenAiError`] (#60/#63) to an axum
|
||||
//! [`Response`], setting `Retry-After` when the envelope carries one.
|
||||
//!
|
||||
//! cortex-core owns the envelope shape; this is the only place the router
|
||||
//! crosses from that data into axum. Mirrors cortex-gateway's adapter so
|
||||
//! the router's own rejections (no feasible operator, all unreachable) are
|
||||
//! the same #63-shaped envelopes clients already understand — distinct from
|
||||
//! cortex's rejections, which the router proxies through verbatim.
|
||||
|
||||
use axum::http::{HeaderValue, StatusCode, header};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
|
||||
/// Render an [`OpenAiError`] as an axum response (status + JSON envelope +
|
||||
/// optional `Retry-After`).
|
||||
pub fn envelope_response(err: OpenAiError) -> Response {
|
||||
let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
let retry_after = err.retry_after_secs;
|
||||
let mut response = (status, Json(err.body())).into_response();
|
||||
if let Some(secs) = retry_after
|
||||
&& let Ok(value) = HeaderValue::from_str(&secs.to_string())
|
||||
{
|
||||
response.headers_mut().insert(header::RETRY_AFTER, value);
|
||||
}
|
||||
response
|
||||
}
|
||||
87
crates/helexa-router/src/handlers.rs
Normal file
87
crates/helexa-router/src/handlers.rs
Normal file
@@ -0,0 +1,87 @@
|
||||
use crate::dispatch;
|
||||
use crate::state::RouterState;
|
||||
use axum::body::Bytes;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::Response;
|
||||
use axum::{Json, Router, extract::State, routing::get, routing::post};
|
||||
use cortex_core::openai::ModelsResponse;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Routes served by the router. Inference paths are capacity-aware-dispatched
|
||||
/// (#73) to a downstream cortex; `/health` and a stub `/v1/models` are local.
|
||||
pub fn api_routes() -> Router<Arc<RouterState>> {
|
||||
Router::new()
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/messages", post(messages))
|
||||
.route("/v1/models", get(list_models))
|
||||
.route("/health", get(health))
|
||||
.route("/", get(health))
|
||||
}
|
||||
|
||||
// ── Inference paths — forwarded verbatim to a chosen cortex ──────────
|
||||
//
|
||||
// Each handler dispatches to the same path on a capacity-bearing cortex.
|
||||
// The body is parsed only to read `model`; the bearer and bytes are
|
||||
// forwarded unchanged, and the SSE response streams back verbatim.
|
||||
|
||||
async fn chat_completions(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch(&state, "/v1/chat/completions", headers, body).await
|
||||
}
|
||||
|
||||
async fn completions(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch(&state, "/v1/completions", headers, body).await
|
||||
}
|
||||
|
||||
async fn responses(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch(&state, "/v1/responses", headers, body).await
|
||||
}
|
||||
|
||||
async fn messages(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch(&state, "/v1/messages", headers, body).await
|
||||
}
|
||||
|
||||
/// `GET /health` — router liveness plus a summary of downstream cortex
|
||||
/// reachability from the topology poller (#72). `status` reflects the
|
||||
/// router process itself (always `ok` if it answers); downstream health is
|
||||
/// the informational `cortexes` block, so a fully-degraded fleet doesn't
|
||||
/// make the router look dead to its own liveness probe.
|
||||
async fn health(State(state): State<Arc<RouterState>>) -> Json<Value> {
|
||||
let topo = state.topology.read().await;
|
||||
let reachable = topo.values().filter(|t| t.reachable).count();
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"cortexes": {
|
||||
"configured": state.cortexes.len(),
|
||||
"reachable": reachable,
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
/// `GET /v1/models` — empty catalogue stub. The real cross-operator union
|
||||
/// (catalogue × topology feasibility, aggregated from each cortex) is the
|
||||
/// federation-catalogue issue (#75).
|
||||
async fn list_models() -> Json<ModelsResponse> {
|
||||
Json(ModelsResponse {
|
||||
object: "list".into(),
|
||||
data: vec![],
|
||||
})
|
||||
}
|
||||
59
crates/helexa-router/src/lib.rs
Normal file
59
crates/helexa-router/src/lib.rs
Normal file
@@ -0,0 +1,59 @@
|
||||
//! helexa-router — public multi-operator ingress proxy (router.helexa.ai).
|
||||
//!
|
||||
//! The router is the data-plane *ingress* tier: a geo-distributed,
|
||||
//! capacity-aware, OpenAI/Anthropic-compatible reverse proxy in front of
|
||||
//! many operator-run cortexes ("cortex-of-cortexes"). End users configure
|
||||
//! one `baseURL` and the router forwards their request to a cortex with
|
||||
//! capacity, proxying #63-shaped rejections back verbatim.
|
||||
//!
|
||||
//! It holds **zero entitlement logic** — auth/budget stays at cortex
|
||||
//! (epic #47); the router forwards the client bearer unchanged and routes
|
||||
//! on capacity (epic #69). A background [`poller`] keeps a live
|
||||
//! per-cortex topology (#72) that the dispatcher (#73) will route on.
|
||||
|
||||
pub mod config;
|
||||
pub mod dispatch;
|
||||
pub mod error;
|
||||
pub mod handlers;
|
||||
pub mod poller;
|
||||
pub mod state;
|
||||
|
||||
use anyhow::Result;
|
||||
use config::RouterConfig;
|
||||
use std::sync::Arc;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::trace::TraceLayer;
|
||||
|
||||
/// Build the axum application: handlers + CORS + tracing. No auth layer —
|
||||
/// the router asserts no identity of its own and forwards the client bearer
|
||||
/// to the downstream cortex, which authenticates it (#69).
|
||||
pub fn build_app(state: Arc<state::RouterState>) -> axum::Router {
|
||||
axum::Router::new()
|
||||
.merge(handlers::api_routes())
|
||||
.layer(CorsLayer::permissive())
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
/// Start the router: build state from config and bind the plaintext HTTP
|
||||
/// listener. TLS is terminated by edge nginx ahead of this process.
|
||||
pub async fn run(config: RouterConfig) -> Result<()> {
|
||||
let state = Arc::new(state::RouterState::from_config(&config));
|
||||
|
||||
// Background topology poller (#72): refresh each cortex's health +
|
||||
// catalogue so routing decisions see live capacity.
|
||||
let poller_state = Arc::clone(&state);
|
||||
tokio::spawn(async move {
|
||||
poller::poll_loop(poller_state).await;
|
||||
});
|
||||
|
||||
let app = build_app(Arc::clone(&state));
|
||||
|
||||
let listen_addr = config.router.listen.parse::<std::net::SocketAddr>()?;
|
||||
tracing::info!("helexa-router listening on {listen_addr}");
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(listen_addr).await?;
|
||||
axum::serve(listener, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
52
crates/helexa-router/src/main.rs
Normal file
52
crates/helexa-router/src/main.rs
Normal file
@@ -0,0 +1,52 @@
|
||||
use anyhow::Result;
|
||||
use clap::{Parser, Subcommand};
|
||||
use helexa_router::config::RouterConfig;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "helexa-router")]
|
||||
#[command(about = "Public multi-operator ingress proxy for helexa")]
|
||||
#[command(version)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
/// Start the router server.
|
||||
Serve {
|
||||
/// Path to the router config file.
|
||||
#[arg(short, long, default_value = "helexa-router.toml")]
|
||||
config: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("info,helexa_router=debug")),
|
||||
)
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
match cli.command {
|
||||
Commands::Serve { config } => {
|
||||
let cfg = RouterConfig::load(&config)
|
||||
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||
|
||||
tracing::info!(
|
||||
cortexes = cfg.cortexes.len(),
|
||||
listen = %cfg.router.listen,
|
||||
"starting helexa-router"
|
||||
);
|
||||
|
||||
helexa_router::run(cfg).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
157
crates/helexa-router/src/poller.rs
Normal file
157
crates/helexa-router/src/poller.rs
Normal file
@@ -0,0 +1,157 @@
|
||||
//! Background poller that refreshes the multi-operator topology (#72).
|
||||
//!
|
||||
//! The same pattern as cortex↔neuron, one tier up: periodically poll each
|
||||
//! configured cortex's `GET /v1/models` (catalogue × topology feasibility +
|
||||
//! loaded state) and `GET /health` (coarse node-health/load), building the
|
||||
//! live map the dispatcher (#73) routes on. An unreachable or erroring
|
||||
//! cortex is debounced over [`POLL_FAILURE_THRESHOLD`] consecutive misses,
|
||||
//! then flipped unhealthy and excluded from routing; it recovers on the
|
||||
//! next successful poll.
|
||||
|
||||
use crate::state::{RouterModelStatus, RouterState};
|
||||
use chrono::Utc;
|
||||
use cortex_core::node::CortexModelEntry;
|
||||
use serde::Deserialize;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Per-cortex HTTP timeout for each poll request.
|
||||
const POLL_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Consecutive failed polls before a cortex is marked unreachable. Mirrors
|
||||
/// cortex's neuron-poll debounce: a single blip (a busy cortex briefly slow
|
||||
/// to answer) can't yank it — and all its models — out of routing.
|
||||
pub const POLL_FAILURE_THRESHOLD: u32 = 3;
|
||||
|
||||
/// cortex's `/v1/models` envelope — `{ "object": "list", "data": [...] }`.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ModelsEnvelope {
|
||||
#[serde(default)]
|
||||
data: Vec<CortexModelEntry>,
|
||||
}
|
||||
|
||||
/// The subset of cortex's `/health` the router reads.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CortexHealth {
|
||||
nodes: CortexHealthNodes,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CortexHealthNodes {
|
||||
healthy: u32,
|
||||
total: u32,
|
||||
}
|
||||
|
||||
/// Run forever, polling all cortexes on the configured interval.
|
||||
pub async fn poll_loop(state: std::sync::Arc<RouterState>) {
|
||||
loop {
|
||||
poll_once(&state).await;
|
||||
tokio::time::sleep(state.poll_interval).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll every configured cortex once. Public for testing.
|
||||
pub async fn poll_once(state: &RouterState) {
|
||||
for cortex in &state.cortexes {
|
||||
poll_cortex(state, &cortex.name, &cortex.endpoint).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll one cortex: refresh its model map from `/v1/models`, then its node
|
||||
/// health from `/health`. A `/v1/models` failure debounces toward
|
||||
/// unreachable; the `/health` poll is best-effort and never flips
|
||||
/// reachability on its own (a cortex serving `/v1/models` is routable even
|
||||
/// if `/health` momentarily isn't).
|
||||
async fn poll_cortex(state: &RouterState, name: &str, endpoint: &str) {
|
||||
let models = fetch_models(state, endpoint).await;
|
||||
|
||||
let mut topo = state.topology.write().await;
|
||||
let Some(entry) = topo.get_mut(name) else {
|
||||
return; // not a configured cortex (shouldn't happen)
|
||||
};
|
||||
|
||||
match models {
|
||||
Ok(models) => {
|
||||
entry.models = models
|
||||
.into_iter()
|
||||
.map(|m| {
|
||||
let feasible = m.loaded || !m.feasible_on.is_empty();
|
||||
(
|
||||
m.id,
|
||||
RouterModelStatus {
|
||||
loaded: m.loaded,
|
||||
feasible,
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
entry.reachable = true;
|
||||
entry.consecutive_failures = 0;
|
||||
entry.last_poll = Some(Utc::now());
|
||||
tracing::debug!(cortex = name, models = entry.models.len(), "poll ok");
|
||||
}
|
||||
Err(reason) => {
|
||||
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||
if entry.consecutive_failures >= POLL_FAILURE_THRESHOLD {
|
||||
entry.reachable = false;
|
||||
}
|
||||
tracing::warn!(
|
||||
cortex = name,
|
||||
failures = entry.consecutive_failures,
|
||||
reachable = entry.reachable,
|
||||
reason,
|
||||
"cortex poll failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
drop(topo);
|
||||
|
||||
// Best-effort health (node counts). Never flips reachability.
|
||||
if let Some((healthy, total)) = fetch_health(state, endpoint).await {
|
||||
let mut topo = state.topology.write().await;
|
||||
if let Some(entry) = topo.get_mut(name) {
|
||||
entry.healthy_nodes = healthy;
|
||||
entry.total_nodes = total;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET `/v1/models`, returning the parsed entries or a short failure reason.
|
||||
async fn fetch_models(
|
||||
state: &RouterState,
|
||||
endpoint: &str,
|
||||
) -> Result<Vec<CortexModelEntry>, &'static str> {
|
||||
let url = format!("{endpoint}/v1/models");
|
||||
let resp = state
|
||||
.http_client
|
||||
.get(&url)
|
||||
.timeout(POLL_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| "unreachable")?;
|
||||
if !resp.status().is_success() {
|
||||
return Err("non-success status");
|
||||
}
|
||||
let envelope = resp
|
||||
.json::<ModelsEnvelope>()
|
||||
.await
|
||||
.map_err(|_| "bad json")?;
|
||||
Ok(envelope.data)
|
||||
}
|
||||
|
||||
/// GET `/health`, returning `(healthy, total)` node counts. `None` on any
|
||||
/// failure — the caller leaves the previous counts in place.
|
||||
async fn fetch_health(state: &RouterState, endpoint: &str) -> Option<(u32, u32)> {
|
||||
let url = format!("{endpoint}/health");
|
||||
let resp = state
|
||||
.http_client
|
||||
.get(&url)
|
||||
.timeout(POLL_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.ok()?;
|
||||
if !resp.status().is_success() {
|
||||
return None;
|
||||
}
|
||||
let health = resp.json::<CortexHealth>().await.ok()?;
|
||||
Some((health.nodes.healthy, health.nodes.total))
|
||||
}
|
||||
88
crates/helexa-router/src/state.rs
Normal file
88
crates/helexa-router/src/state.rs
Normal file
@@ -0,0 +1,88 @@
|
||||
use crate::config::{CortexEndpoint, RouterConfig};
|
||||
use chrono::{DateTime, Utc};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Shared router state: the configured cortex list plus the live topology
|
||||
/// map the poller (#72) maintains and the dispatcher (#73) will route on.
|
||||
///
|
||||
/// This is the router tier of the fractal neuron ← cortex ← router design:
|
||||
/// just as cortex polls each neuron for capacity/catalogue, the router
|
||||
/// polls each cortex's `/health` + `/v1/models`.
|
||||
#[derive(Debug)]
|
||||
pub struct RouterState {
|
||||
/// Downstream cortex endpoints, as configured.
|
||||
pub cortexes: Vec<CortexEndpoint>,
|
||||
/// Shared client for polling (and proxying to) cortexes.
|
||||
pub http_client: reqwest::Client,
|
||||
/// This router instance's region, for dispatch geo affinity (#73).
|
||||
pub region: Option<String>,
|
||||
/// How often the poller refreshes the topology.
|
||||
pub poll_interval: Duration,
|
||||
/// Live per-cortex topology, keyed by cortex name. Pre-populated from
|
||||
/// config (every configured cortex present, `reachable = false`) so the
|
||||
/// poller and handlers always find an entry; the poller flips
|
||||
/// reachability and fills the model map.
|
||||
pub topology: RwLock<HashMap<String, CortexTopology>>,
|
||||
}
|
||||
|
||||
/// Live view of one downstream cortex, refreshed each poll.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CortexTopology {
|
||||
/// Whether the cortex is currently routable. Flipped `false` only after
|
||||
/// [`crate::poller::POLL_FAILURE_THRESHOLD`] consecutive failed polls
|
||||
/// (debounces transient blips); restored on the next successful poll.
|
||||
pub reachable: bool,
|
||||
/// Consecutive failed polls; reset to 0 on success.
|
||||
pub consecutive_failures: u32,
|
||||
/// Timestamp of the last successful poll.
|
||||
pub last_poll: Option<DateTime<Utc>>,
|
||||
/// Healthy / total neuron counts from the cortex's `/health` (coarse
|
||||
/// load signal; #73 refines headroom). 0/0 until first health poll.
|
||||
pub healthy_nodes: u32,
|
||||
pub total_nodes: u32,
|
||||
/// Per-model serveability, keyed by model id, from `/v1/models`.
|
||||
pub models: HashMap<String, RouterModelStatus>,
|
||||
}
|
||||
|
||||
/// What a cortex can do with one model, distilled from its `/v1/models`
|
||||
/// entry to the two facts the router routes on.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RouterModelStatus {
|
||||
/// The model is loaded on at least one of the cortex's neurons.
|
||||
pub loaded: bool,
|
||||
/// The cortex can serve it — loaded now, or feasible to cold-load
|
||||
/// (catalogue × topology says some neuron can host it).
|
||||
pub feasible: bool,
|
||||
}
|
||||
|
||||
impl RouterState {
|
||||
pub fn from_config(config: &RouterConfig) -> Self {
|
||||
let topology = config
|
||||
.cortexes
|
||||
.iter()
|
||||
.map(|c| (c.name.clone(), CortexTopology::default()))
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
cortexes: config.cortexes.clone(),
|
||||
http_client: reqwest::Client::new(),
|
||||
region: config.router.region.clone(),
|
||||
poll_interval: Duration::from_secs(config.router.poll_interval_secs),
|
||||
topology: RwLock::new(topology),
|
||||
}
|
||||
}
|
||||
|
||||
/// Names of reachable cortexes that can serve `model_id` (loaded or
|
||||
/// feasible to cold-load). Groundwork for capacity-aware dispatch (#73);
|
||||
/// unreachable cortexes are excluded by construction.
|
||||
pub async fn cortexes_serving(&self, model_id: &str) -> Vec<String> {
|
||||
let topo = self.topology.read().await;
|
||||
topo.iter()
|
||||
.filter(|(_, t)| t.reachable)
|
||||
.filter(|(_, t)| t.models.get(model_id).is_some_and(|m| m.feasible))
|
||||
.map(|(name, _)| name.clone())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
277
crates/helexa-router/tests/dispatch.rs
Normal file
277
crates/helexa-router/tests/dispatch.rs
Normal file
@@ -0,0 +1,277 @@
|
||||
//! Capacity-aware dispatch acceptance tests for #73.
|
||||
//!
|
||||
//! Covers: a request routes to a cortex serving the model; the client's
|
||||
//! bearer reaches the cortex; cortex's #63 rejections pass through verbatim
|
||||
//! and are NOT retried away; transport failure fails over to another
|
||||
//! feasible cortex; unknown model → 404, no reachable capacity → 503; and
|
||||
//! the selection ranking (warm/region/headroom).
|
||||
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::State;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||
use helexa_router::dispatch::{Selection, dispatch, select_cortexes};
|
||||
use helexa_router::state::{CortexTopology, RouterModelStatus, RouterState};
|
||||
use serde_json::{Value, json};
|
||||
use std::collections::HashMap;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const MODEL: &str = "Qwen/Qwen3-Coder-30B";
|
||||
|
||||
// ── Mock cortex backend ──────────────────────────────────────────────
|
||||
|
||||
/// Behaviour of a mock cortex, carried in axum State.
|
||||
#[derive(Clone)]
|
||||
struct MockCortex {
|
||||
/// Identifies which cortex answered, echoed in the 200 body.
|
||||
name: &'static str,
|
||||
/// When true, return a genuine #63-shaped `429 rate_limit_exceeded`.
|
||||
rate_limited: bool,
|
||||
}
|
||||
|
||||
async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap) -> Response {
|
||||
if m.rate_limited {
|
||||
return (
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
Json(json!({"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down","param":null}})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let auth = headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
Json(json!({ "served_by": m.name, "auth_seen": auth })).into_response()
|
||||
}
|
||||
|
||||
async fn spawn_cortex(mock: MockCortex) -> String {
|
||||
let app = Router::new()
|
||||
.route("/v1/chat/completions", post(mock_handler))
|
||||
.with_state(mock);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
fn ok_cortex(name: &'static str) -> MockCortex {
|
||||
MockCortex {
|
||||
name,
|
||||
rate_limited: false,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Helpers to build state with a hand-set topology ──────────────────
|
||||
|
||||
fn state_with(cortexes: Vec<CortexEndpoint>, region: Option<String>) -> RouterState {
|
||||
let cfg = RouterConfig {
|
||||
cortexes,
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = RouterState::from_config(&cfg);
|
||||
state.region = region;
|
||||
state
|
||||
}
|
||||
|
||||
/// Overwrite the topology entry for `name` so tests control reachability and
|
||||
/// model serveability directly (no live poll).
|
||||
async fn set_topology(
|
||||
state: &RouterState,
|
||||
name: &str,
|
||||
reachable: bool,
|
||||
loaded: bool,
|
||||
feasible: bool,
|
||||
healthy_nodes: u32,
|
||||
) {
|
||||
let mut topo = state.topology.write().await;
|
||||
let mut models = HashMap::new();
|
||||
models.insert(MODEL.to_string(), RouterModelStatus { loaded, feasible });
|
||||
topo.insert(
|
||||
name.to_string(),
|
||||
CortexTopology {
|
||||
reachable,
|
||||
consecutive_failures: 0,
|
||||
last_poll: None,
|
||||
healthy_nodes,
|
||||
total_nodes: healthy_nodes,
|
||||
models,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn ep(name: &str, endpoint: &str, region: Option<&str>) -> CortexEndpoint {
|
||||
CortexEndpoint {
|
||||
name: name.into(),
|
||||
endpoint: endpoint.into(),
|
||||
region: region.map(str::to_string),
|
||||
}
|
||||
}
|
||||
|
||||
fn chat_body() -> Bytes {
|
||||
Bytes::from(format!("{{\"model\":\"{MODEL}\",\"stream\":false}}"))
|
||||
}
|
||||
|
||||
async fn body_json(resp: Response) -> (StatusCode, Value) {
|
||||
let status = resp.status();
|
||||
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||
.await
|
||||
.unwrap();
|
||||
let v = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
|
||||
(status, v)
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn routes_to_serving_cortex_and_forwards_bearer() {
|
||||
let url = spawn_cortex(ok_cortex("c1")).await;
|
||||
let state = state_with(vec![ep("c1", &url, None)], None);
|
||||
set_topology(&state, "c1", true, true, true, 2).await;
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("authorization", "Bearer sk-test-123".parse().unwrap());
|
||||
|
||||
let resp = dispatch(&state, "/v1/chat/completions", headers, chat_body()).await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["served_by"], "c1");
|
||||
// Bearer reached the cortex unchanged.
|
||||
assert_eq!(body["auth_seen"], "Bearer sk-test-123");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cortex_429_passes_through_and_is_not_retried() {
|
||||
// c1 (ranked first: loaded) returns a genuine 429; c2 would return 200.
|
||||
let c1 = spawn_cortex(MockCortex {
|
||||
name: "c1",
|
||||
rate_limited: true,
|
||||
})
|
||||
.await;
|
||||
let c2 = spawn_cortex(ok_cortex("c2")).await;
|
||||
let state = state_with(vec![ep("c1", &c1, None), ep("c2", &c2, None)], None);
|
||||
// Both reachable + loaded; c1 has more headroom so it ranks first.
|
||||
set_topology(&state, "c1", true, true, true, 5).await;
|
||||
set_topology(&state, "c2", true, true, true, 1).await;
|
||||
|
||||
let resp = dispatch(
|
||||
&state,
|
||||
"/v1/chat/completions",
|
||||
HeaderMap::new(),
|
||||
chat_body(),
|
||||
)
|
||||
.await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
|
||||
// The genuine 4xx is returned verbatim — NOT retried to c2.
|
||||
assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||
assert!(body.get("served_by").is_none(), "must not have hit c2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fails_over_to_next_cortex_on_transport_error() {
|
||||
// c_dead ranks first (more headroom) but its endpoint is a closed port;
|
||||
// c_live is the fallback. The router must fail over and c_live serves.
|
||||
let live = spawn_cortex(ok_cortex("c_live")).await;
|
||||
let state = state_with(
|
||||
vec![
|
||||
ep("c_dead", "http://127.0.0.1:1", None),
|
||||
ep("c_live", &live, None),
|
||||
],
|
||||
None,
|
||||
);
|
||||
set_topology(&state, "c_dead", true, true, true, 9).await;
|
||||
set_topology(&state, "c_live", true, true, true, 1).await;
|
||||
|
||||
let resp = dispatch(
|
||||
&state,
|
||||
"/v1/chat/completions",
|
||||
HeaderMap::new(),
|
||||
chat_body(),
|
||||
)
|
||||
.await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["served_by"], "c_live");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_model_is_404() {
|
||||
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||
// Topology has no entry for MODEL at all.
|
||||
let resp = dispatch(
|
||||
&state,
|
||||
"/v1/chat/completions",
|
||||
HeaderMap::new(),
|
||||
chat_body(),
|
||||
)
|
||||
.await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn known_but_all_unreachable_is_503() {
|
||||
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||
// Cortex knows the model (from a prior good poll) but is now unreachable.
|
||||
set_topology(&state, "c1", false, true, true, 2).await;
|
||||
let resp = dispatch(
|
||||
&state,
|
||||
"/v1/chat/completions",
|
||||
HeaderMap::new(),
|
||||
chat_body(),
|
||||
)
|
||||
.await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(body["error"]["code"], "service_unavailable");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_model_field_is_400() {
|
||||
let state = state_with(vec![ep("c1", "http://127.0.0.1:1", None)], None);
|
||||
let resp = dispatch(
|
||||
&state,
|
||||
"/v1/chat/completions",
|
||||
HeaderMap::new(),
|
||||
Bytes::from_static(b"{\"messages\":[]}"),
|
||||
)
|
||||
.await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(body["error"]["code"], "missing_model_field");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ranking_prefers_loaded_then_region_then_headroom() {
|
||||
// Router is in eu-west. Candidates:
|
||||
// warm-eu : loaded, region match, 1 node → best
|
||||
// warm-us : loaded, no region, 9 nodes
|
||||
// cold-eu : feasible only, region match → worst (cold)
|
||||
let state = state_with(
|
||||
vec![
|
||||
ep("warm-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||
ep("warm-us", "http://127.0.0.1:1", Some("us-east")),
|
||||
ep("cold-eu", "http://127.0.0.1:1", Some("eu-west")),
|
||||
],
|
||||
Some("eu-west".into()),
|
||||
);
|
||||
set_topology(&state, "warm-eu", true, true, true, 1).await;
|
||||
set_topology(&state, "warm-us", true, true, true, 9).await;
|
||||
set_topology(&state, "cold-eu", true, false, true, 5).await;
|
||||
|
||||
let Selection::Candidates(order) = select_cortexes(&state, MODEL).await else {
|
||||
panic!("expected candidates");
|
||||
};
|
||||
let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
|
||||
assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
|
||||
}
|
||||
95
crates/helexa-router/tests/skeleton.rs
Normal file
95
crates/helexa-router/tests/skeleton.rs
Normal file
@@ -0,0 +1,95 @@
|
||||
//! Skeleton acceptance tests for #70: the router builds, serves `/health`
|
||||
//! and `/v1/models` on a plaintext port, and loads its cortex-endpoint list
|
||||
//! from TOML with env overrides.
|
||||
|
||||
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||
use helexa_router::state::RouterState;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Bind the router app on an ephemeral port and return its base URL.
|
||||
async fn spawn_router(cortexes: Vec<CortexEndpoint>) -> String {
|
||||
let cfg = RouterConfig {
|
||||
cortexes,
|
||||
..Default::default()
|
||||
};
|
||||
let state = Arc::new(RouterState::from_config(&cfg));
|
||||
let app = helexa_router::build_app(state);
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_reports_configured_cortex_count() {
|
||||
let base = spawn_router(vec![
|
||||
CortexEndpoint {
|
||||
name: "a".into(),
|
||||
endpoint: "https://a.example.com".into(),
|
||||
region: None,
|
||||
},
|
||||
CortexEndpoint {
|
||||
name: "b".into(),
|
||||
endpoint: "https://b.example.com".into(),
|
||||
region: None,
|
||||
},
|
||||
])
|
||||
.await;
|
||||
|
||||
let body: serde_json::Value = reqwest::get(format!("{base}/health"))
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(body["status"], "ok");
|
||||
assert_eq!(body["cortexes"]["configured"], 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_returns_empty_openai_list() {
|
||||
let base = spawn_router(vec![]).await;
|
||||
|
||||
let resp = reqwest::get(format!("{base}/v1/models")).await.unwrap();
|
||||
assert!(resp.status().is_success());
|
||||
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["object"], "list");
|
||||
assert_eq!(body["data"].as_array().unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[allow(clippy::result_large_err)]
|
||||
fn config_loads_from_toml_with_env_override() {
|
||||
figment::Jail::expect_with(|jail| {
|
||||
jail.create_file(
|
||||
"helexa-router.toml",
|
||||
r#"
|
||||
[router]
|
||||
listen = "127.0.0.1:8088"
|
||||
|
||||
[[cortexes]]
|
||||
name = "lair-cafe"
|
||||
endpoint = "https://cortex.lair.cafe"
|
||||
"#,
|
||||
)?;
|
||||
|
||||
// Env override wins over the TOML value.
|
||||
jail.set_env("HELEXA_ROUTER_ROUTER__LISTEN", "0.0.0.0:9099");
|
||||
|
||||
let cfg = RouterConfig::load("helexa-router.toml").expect("load config");
|
||||
|
||||
assert_eq!(cfg.router.listen, "0.0.0.0:9099");
|
||||
assert_eq!(cfg.cortexes.len(), 1);
|
||||
assert_eq!(cfg.cortexes[0].name, "lair-cafe");
|
||||
assert_eq!(cfg.cortexes[0].endpoint, "https://cortex.lair.cafe");
|
||||
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
170
crates/helexa-router/tests/topology.rs
Normal file
170
crates/helexa-router/tests/topology.rs
Normal file
@@ -0,0 +1,170 @@
|
||||
//! Topology-poller acceptance tests for #72: the router maintains a live
|
||||
//! map of which cortexes serve which models, marks an unreachable/erroring
|
||||
//! cortex unhealthy and excludes it from routing, and recovers it once
|
||||
//! reachable again.
|
||||
|
||||
use axum::extract::State;
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||
use helexa_router::poller::{POLL_FAILURE_THRESHOLD, poll_once};
|
||||
use helexa_router::state::RouterState;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Shared "is this mock cortex up?" flag, toggled by tests to simulate
|
||||
/// outage and recovery.
|
||||
#[derive(Clone)]
|
||||
struct MockState {
|
||||
up: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
async fn mock_models(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||
if !s.up.load(Ordering::SeqCst) {
|
||||
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
Ok(Json(json!({
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "Qwen/Qwen3-Coder-30B",
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "helexa",
|
||||
"loaded": true,
|
||||
"feasible_on": ["beast"],
|
||||
"locations": [{"node": "beast", "status": "loaded", "vram_estimate_mb": 19000}]
|
||||
},
|
||||
{
|
||||
"id": "Qwen/Qwen3-VL-8B",
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "helexa",
|
||||
"loaded": false,
|
||||
"feasible_on": ["beast"],
|
||||
"locations": []
|
||||
}
|
||||
]
|
||||
})))
|
||||
}
|
||||
|
||||
async fn mock_health(State(s): State<MockState>) -> Result<Json<Value>, StatusCode> {
|
||||
if !s.up.load(Ordering::SeqCst) {
|
||||
return Err(StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
Ok(Json(json!({
|
||||
"status": "ok",
|
||||
"nodes": { "healthy": 2, "total": 3 }
|
||||
})))
|
||||
}
|
||||
|
||||
/// Spawn a mock cortex; returns (base_url, up_flag).
|
||||
async fn spawn_mock_cortex() -> (String, Arc<AtomicBool>) {
|
||||
let up = Arc::new(AtomicBool::new(true));
|
||||
let state = MockState { up: up.clone() };
|
||||
let app = Router::new()
|
||||
.route("/v1/models", get(mock_models))
|
||||
.route("/health", get(mock_health))
|
||||
.with_state(state);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{addr}"), up)
|
||||
}
|
||||
|
||||
fn state_for(name: &str, endpoint: &str) -> RouterState {
|
||||
let cfg = RouterConfig {
|
||||
cortexes: vec![CortexEndpoint {
|
||||
name: name.into(),
|
||||
endpoint: endpoint.into(),
|
||||
region: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
RouterState::from_config(&cfg)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_builds_live_topology() {
|
||||
let (base, _up) = spawn_mock_cortex().await;
|
||||
let state = state_for("c1", &base);
|
||||
|
||||
poll_once(&state).await;
|
||||
|
||||
let topo = state.topology.read().await;
|
||||
let c1 = topo.get("c1").expect("cortex present");
|
||||
assert!(c1.reachable, "should be reachable after a good poll");
|
||||
assert_eq!(c1.consecutive_failures, 0);
|
||||
assert!(c1.last_poll.is_some());
|
||||
assert_eq!((c1.healthy_nodes, c1.total_nodes), (2, 3));
|
||||
|
||||
// Loaded model: loaded + feasible. Catalogue-only model: feasible only.
|
||||
let coder = c1.models.get("Qwen/Qwen3-Coder-30B").unwrap();
|
||||
assert!(coder.loaded && coder.feasible);
|
||||
let vl = c1.models.get("Qwen/Qwen3-VL-8B").unwrap();
|
||||
assert!(!vl.loaded && vl.feasible);
|
||||
drop(topo);
|
||||
|
||||
// The routing helper sees both serveable models on the reachable cortex.
|
||||
assert_eq!(
|
||||
state.cortexes_serving("Qwen/Qwen3-VL-8B").await,
|
||||
vec!["c1".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unreachable_cortex_excluded_then_recovers() {
|
||||
let (base, up) = spawn_mock_cortex().await;
|
||||
let state = state_for("c1", &base);
|
||||
|
||||
// Healthy first.
|
||||
poll_once(&state).await;
|
||||
assert!(state.topology.read().await["c1"].reachable);
|
||||
|
||||
// Take it down. The first failures debounce (stay reachable) until the
|
||||
// threshold; only then is it excluded.
|
||||
up.store(false, Ordering::SeqCst);
|
||||
for i in 1..POLL_FAILURE_THRESHOLD {
|
||||
poll_once(&state).await;
|
||||
assert!(
|
||||
state.topology.read().await["c1"].reachable,
|
||||
"still reachable after {i} failure(s) (below threshold)"
|
||||
);
|
||||
}
|
||||
poll_once(&state).await; // crosses the threshold
|
||||
{
|
||||
let topo = state.topology.read().await;
|
||||
assert!(!topo["c1"].reachable, "excluded after threshold failures");
|
||||
assert!(topo["c1"].consecutive_failures >= POLL_FAILURE_THRESHOLD);
|
||||
}
|
||||
// Excluded from routing.
|
||||
assert!(
|
||||
state
|
||||
.cortexes_serving("Qwen/Qwen3-Coder-30B")
|
||||
.await
|
||||
.is_empty()
|
||||
);
|
||||
|
||||
// Bring it back: the next successful poll restores it.
|
||||
up.store(true, Ordering::SeqCst);
|
||||
poll_once(&state).await;
|
||||
let topo = state.topology.read().await;
|
||||
assert!(topo["c1"].reachable, "recovered after a good poll");
|
||||
assert_eq!(topo["c1"].consecutive_failures, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unconfigured_endpoint_is_unreachable() {
|
||||
// Nothing listening on this port → polls fail; below threshold it stays
|
||||
// at its initial unreachable state, and never panics.
|
||||
let state = state_for("dead", "http://127.0.0.1:1");
|
||||
poll_once(&state).await;
|
||||
let topo = state.topology.read().await;
|
||||
assert!(!topo["dead"].reachable);
|
||||
assert_eq!(topo["dead"].consecutive_failures, 1);
|
||||
}
|
||||
21
crates/helexa-stream/Cargo.toml
Normal file
21
crates/helexa-stream/Cargo.toml
Normal file
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "helexa-stream"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "helexa_stream"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
axum = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true }
|
||||
tokio-stream = { workspace = true }
|
||||
async-stream = "0.3"
|
||||
290
crates/helexa-stream/src/lib.rs
Normal file
290
crates/helexa-stream/src/lib.rs
Normal file
@@ -0,0 +1,290 @@
|
||||
//! Shared streaming reverse-proxy mechanism (#71).
|
||||
//!
|
||||
//! cortex and helexa-router both need to proxy an OpenAI/Anthropic SSE
|
||||
//! response from a downstream backend **verbatim** — chunks forwarded as
|
||||
//! they arrive, never buffering the full body — while observing the bytes
|
||||
//! for metrics/metering. This crate owns that mechanism so there is one
|
||||
//! implementation, not one per tier.
|
||||
//!
|
||||
//! The split is mechanism vs policy:
|
||||
//!
|
||||
//! - **Mechanism (here):** [`forward_streaming`] POSTs to a backend and
|
||||
//! streams the response body back through an [`ObservedStream`], which
|
||||
//! feeds every chunk to a caller-supplied [`ChunkObserver`] and calls
|
||||
//! [`ChunkObserver::finish`] exactly once on clean end-of-stream or on
|
||||
//! drop (client disconnect mid-stream). [`BodyTail`] and
|
||||
//! [`last_count_for`] are the reusable pieces an observer uses to pull
|
||||
//! the trailing OpenAI `usage` object out of the streamed bytes.
|
||||
//! - **Policy (caller):** what to *do* with the observed bytes — which
|
||||
//! metric names to emit, which labels, whether to settle a per-principal
|
||||
//! reservation — lives in the consumer's `ChunkObserver` impl, not here.
|
||||
//!
|
||||
//! The proxy is status-agnostic: a non-2xx upstream response (e.g. a
|
||||
//! cortex `429 rate_limit_exceeded`) is streamed back with its status and
|
||||
//! headers intact, so honest backpressure reaches the client unchanged.
|
||||
//! Only a network failure or a malformed response build is an error.
|
||||
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::Response;
|
||||
use futures::Stream;
|
||||
use futures::stream::BoxStream;
|
||||
use reqwest::Client;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// Observes the bytes of a streamed proxy response without altering them.
|
||||
///
|
||||
/// `observe` is called for each forwarded chunk; `finish` is called
|
||||
/// exactly once — on clean end-of-stream or on drop — and implementations
|
||||
/// must be idempotent (the [`ObservedStream`] guards against a double call,
|
||||
/// but a `finish` that runs side effects should still self-guard).
|
||||
pub trait ChunkObserver: Send + Unpin + 'static {
|
||||
/// A body chunk has been forwarded downstream. The slice is the exact
|
||||
/// bytes the client receives.
|
||||
fn observe(&mut self, chunk: &[u8]);
|
||||
|
||||
/// The stream has ended (cleanly or via client disconnect). Called once.
|
||||
fn finish(&mut self);
|
||||
}
|
||||
|
||||
/// A bounded accumulator for the tail of a streamed body.
|
||||
///
|
||||
/// The OpenAI `usage` object rides on the final SSE chunk (and sits at the
|
||||
/// end of a non-streaming JSON body), so retaining a generous tail is
|
||||
/// enough to recover token counts via [`last_count_for`]; the cap bounds
|
||||
/// memory on huge bodies. Appends are char-boundary-safe.
|
||||
#[derive(Debug)]
|
||||
pub struct BodyTail {
|
||||
tail: String,
|
||||
cap: usize,
|
||||
}
|
||||
|
||||
impl BodyTail {
|
||||
/// Create a tail retaining at most `cap` bytes.
|
||||
pub fn new(cap: usize) -> Self {
|
||||
Self {
|
||||
tail: String::new(),
|
||||
cap,
|
||||
}
|
||||
}
|
||||
|
||||
/// Append a chunk, trimming from the front past the cap. When trimming,
|
||||
/// the newest half is kept (the usage object is always at the very end).
|
||||
pub fn push(&mut self, chunk: &[u8]) {
|
||||
self.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||
if self.tail.len() > self.cap {
|
||||
let mut cut = self.tail.len() - self.cap / 2;
|
||||
while !self.tail.is_char_boundary(cut) {
|
||||
cut += 1;
|
||||
}
|
||||
self.tail.drain(..cut);
|
||||
}
|
||||
}
|
||||
|
||||
/// The retained tail text.
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.tail
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
|
||||
///
|
||||
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
|
||||
/// The quoted-needle form means `completion_tokens` never matches
|
||||
/// `completion_tokens_details`, and taking the last occurrence means the
|
||||
/// final `usage` object wins even if content earlier in the stream echoed
|
||||
/// a usage-shaped string.
|
||||
pub fn last_count_for(tail: &str, key: &str) -> Option<u64> {
|
||||
let needle = format!("\"{key}\"");
|
||||
let mut result = None;
|
||||
for (idx, _) in tail.match_indices(&needle) {
|
||||
let rest = tail[idx + needle.len()..].trim_start();
|
||||
let Some(rest) = rest.strip_prefix(':') else {
|
||||
continue;
|
||||
};
|
||||
let rest = rest.trim_start();
|
||||
let digits: &str = &rest[..rest
|
||||
.char_indices()
|
||||
.find(|(_, c)| !c.is_ascii_digit())
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(rest.len())];
|
||||
if let Ok(v) = digits.parse::<u64>() {
|
||||
result = Some(v);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Error from [`forward_streaming`]. Distinguishes a network/transport
|
||||
/// failure reaching the backend from a failure assembling the downstream
|
||||
/// response. A non-2xx upstream *status* is not an error — it is streamed
|
||||
/// through verbatim.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum StreamError {
|
||||
#[error("upstream request failed")]
|
||||
Upstream(reqwest::Error),
|
||||
#[error("failed to build response")]
|
||||
ResponseBuild(String),
|
||||
}
|
||||
|
||||
/// POST `body` to `url` and stream the response back verbatim through
|
||||
/// `observer`.
|
||||
///
|
||||
/// Request headers are forwarded except `host` / `content-length` (reqwest
|
||||
/// sets these). The returned [`Response`] carries the upstream status and
|
||||
/// headers unchanged — including non-2xx — with a body that streams the
|
||||
/// upstream bytes chunk-for-chunk, feeding each chunk to `observer`.
|
||||
pub async fn forward_streaming<O: ChunkObserver>(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
observer: O,
|
||||
) -> Result<Response, StreamError> {
|
||||
let mut req_builder = client.post(url).body(body);
|
||||
for (key, value) in headers.iter() {
|
||||
if key == "host" || key == "content-length" {
|
||||
continue; // reqwest sets these
|
||||
}
|
||||
req_builder = req_builder.header(key, value);
|
||||
}
|
||||
|
||||
let upstream = req_builder.send().await.map_err(StreamError::Upstream)?;
|
||||
|
||||
let status =
|
||||
StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||
let resp_headers = upstream.headers().clone();
|
||||
|
||||
let stream = ObservedStream::new(Box::pin(upstream.bytes_stream()), observer);
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
let mut response = Response::builder().status(status);
|
||||
for (key, value) in resp_headers.iter() {
|
||||
response = response.header(key, value);
|
||||
}
|
||||
response
|
||||
.body(body)
|
||||
.map_err(|e| StreamError::ResponseBuild(e.to_string()))
|
||||
}
|
||||
|
||||
/// Pass-through stream wrapper that feeds a [`ChunkObserver`]. Forwards
|
||||
/// each chunk verbatim, calls `observe` per chunk, and `finish` once on
|
||||
/// clean end-of-stream; the `Drop` impl covers client disconnects.
|
||||
pub struct ObservedStream<O: ChunkObserver> {
|
||||
inner: BoxStream<'static, Result<Bytes, reqwest::Error>>,
|
||||
observer: O,
|
||||
finished: bool,
|
||||
}
|
||||
|
||||
impl<O: ChunkObserver> ObservedStream<O> {
|
||||
/// Wrap a byte stream with an observer.
|
||||
pub fn new(inner: BoxStream<'static, Result<Bytes, reqwest::Error>>, observer: O) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
observer,
|
||||
finished: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(&mut self) {
|
||||
if self.finished {
|
||||
return;
|
||||
}
|
||||
self.finished = true;
|
||||
self.observer.finish();
|
||||
}
|
||||
}
|
||||
|
||||
impl<O: ChunkObserver> Stream for ObservedStream<O> {
|
||||
type Item = Result<Bytes, reqwest::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let this = self.get_mut();
|
||||
match this.inner.as_mut().poll_next(cx) {
|
||||
Poll::Ready(Some(Ok(chunk))) => {
|
||||
this.observer.observe(&chunk);
|
||||
Poll::Ready(Some(Ok(chunk)))
|
||||
}
|
||||
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
||||
Poll::Ready(None) => {
|
||||
this.finish();
|
||||
Poll::Ready(None)
|
||||
}
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<O: ChunkObserver> Drop for ObservedStream<O> {
|
||||
fn drop(&mut self) {
|
||||
self.finish();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn extracts_counts_from_final_sse_usage_chunk() {
|
||||
let tail = concat!(
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
|
||||
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_counts_from_non_streaming_body() {
|
||||
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
|
||||
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
|
||||
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_details_variants_and_takes_last_occurrence() {
|
||||
// completion_tokens_details must not shadow completion_tokens,
|
||||
// and the LAST usage object wins (matters when content echoes
|
||||
// a usage-shaped string earlier in the stream).
|
||||
let tail = concat!(
|
||||
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
|
||||
"data: {\"usage\":{\"completion_tokens\":99,",
|
||||
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
|
||||
);
|
||||
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_keys_yield_none() {
|
||||
assert_eq!(
|
||||
last_count_for("data: [DONE]\n\n", "completion_tokens"),
|
||||
None
|
||||
);
|
||||
assert_eq!(last_count_for("", "prompt_tokens"), None);
|
||||
// key present but non-numeric value
|
||||
assert_eq!(
|
||||
last_count_for("\"completion_tokens\": null", "completion_tokens"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_tail_retains_usage_after_cap_trim() {
|
||||
// Cap small enough that the filler forces several front-trims, but
|
||||
// (as in production, where cap ≫ the usage object) large enough that
|
||||
// the trailing usage object survives the newest-half retention.
|
||||
let mut tail = BodyTail::new(512);
|
||||
for _ in 0..100 {
|
||||
tail.push(b"data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n");
|
||||
}
|
||||
assert!(tail.as_str().len() <= 512, "cap must bound the tail");
|
||||
tail.push(b"data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":9}}\n\n");
|
||||
assert_eq!(last_count_for(tail.as_str(), "prompt_tokens"), Some(5));
|
||||
assert_eq!(last_count_for(tail.as_str(), "completion_tokens"), Some(9));
|
||||
}
|
||||
}
|
||||
162
crates/helexa-stream/tests/streaming.rs
Normal file
162
crates/helexa-stream/tests/streaming.rs
Normal file
@@ -0,0 +1,162 @@
|
||||
//! Integration tests for the shared streaming proxy (#71): proves a backend
|
||||
//! SSE response is forwarded chunk-for-chunk (no buffering), the observer
|
||||
//! sees every byte and finishes once, and non-2xx is streamed through with
|
||||
//! its status intact — the behaviours both cortex and helexa-router rely on.
|
||||
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::Response;
|
||||
use axum::routing::post;
|
||||
use helexa_stream::{BodyTail, ChunkObserver, forward_streaming, last_count_for};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Observer that records what it saw, for assertions.
|
||||
#[derive(Clone, Default)]
|
||||
struct RecordingObserver {
|
||||
inner: Arc<Mutex<Recorded>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Recorded {
|
||||
chunks: usize,
|
||||
finished: usize,
|
||||
tail: String,
|
||||
}
|
||||
|
||||
impl ChunkObserver for RecordingObserver {
|
||||
fn observe(&mut self, chunk: &[u8]) {
|
||||
let mut r = self.inner.lock().unwrap();
|
||||
r.chunks += 1;
|
||||
r.tail.push_str(&String::from_utf8_lossy(chunk));
|
||||
}
|
||||
fn finish(&mut self) {
|
||||
self.inner.lock().unwrap().finished += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Mock backend that streams 5 SSE chunks with 30ms gaps, then a usage
|
||||
/// chunk and `[DONE]`.
|
||||
async fn sse_handler() -> Response {
|
||||
let chunks: Vec<&'static str> = vec![
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"d\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"e\"}}]}\n\n",
|
||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}}\n\n",
|
||||
"data: [DONE]\n\n",
|
||||
];
|
||||
let stream = async_stream::stream! {
|
||||
for c in chunks {
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
yield Ok::<_, std::io::Error>(axum::body::Bytes::from_static(c.as_bytes()));
|
||||
}
|
||||
};
|
||||
Response::new(Body::from_stream(stream))
|
||||
}
|
||||
|
||||
async fn rate_limited_handler() -> Response {
|
||||
Response::builder()
|
||||
.status(StatusCode::TOO_MANY_REQUESTS)
|
||||
.body(Body::from("{\"error\":{\"type\":\"rate_limit_exceeded\"}}"))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn spawn_backend(router: Router) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, router).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streams_chunks_incrementally_and_observes_usage() {
|
||||
let base = spawn_backend(Router::new().route("/v1/chat/completions", post(sse_handler))).await;
|
||||
let observer = RecordingObserver::default();
|
||||
let probe = observer.clone();
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = forward_streaming(
|
||||
&client,
|
||||
&format!("{base}/v1/chat/completions"),
|
||||
HeaderMap::new(),
|
||||
axum::body::Bytes::from_static(b"{\"model\":\"x\",\"stream\":true}"),
|
||||
observer,
|
||||
)
|
||||
.await
|
||||
.expect("forward ok");
|
||||
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
// Read the proxied body as a stream, timestamping arrivals.
|
||||
let mut body = resp.into_body().into_data_stream();
|
||||
let mut arrivals: Vec<Instant> = Vec::new();
|
||||
let mut collected = String::new();
|
||||
use futures::StreamExt;
|
||||
while let Some(item) = body.next().await {
|
||||
let bytes = item.unwrap();
|
||||
arrivals.push(Instant::now());
|
||||
collected.push_str(&String::from_utf8_lossy(&bytes));
|
||||
}
|
||||
|
||||
// Incremental delivery: first and last chunk are meaningfully apart
|
||||
// (5×30ms gaps), proving no full-response buffering.
|
||||
let spread = *arrivals.last().unwrap() - arrivals[0];
|
||||
assert!(
|
||||
spread >= Duration::from_millis(100),
|
||||
"expected incremental delivery, spread was {spread:?}"
|
||||
);
|
||||
|
||||
// The client received the terminator and the usage object verbatim.
|
||||
assert!(collected.contains("data: [DONE]"));
|
||||
|
||||
// The observer saw the bytes and finished exactly once.
|
||||
let r = probe.inner.lock().unwrap();
|
||||
assert!(r.chunks >= 5, "observer saw {} chunks", r.chunks);
|
||||
assert_eq!(r.finished, 1, "finish must run exactly once");
|
||||
assert_eq!(last_count_for(&r.tail, "prompt_tokens"), Some(11));
|
||||
assert_eq!(last_count_for(&r.tail, "completion_tokens"), Some(5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_2xx_is_streamed_through_verbatim() {
|
||||
let base =
|
||||
spawn_backend(Router::new().route("/v1/chat/completions", post(rate_limited_handler)))
|
||||
.await;
|
||||
let observer = RecordingObserver::default();
|
||||
let probe = observer.clone();
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = forward_streaming(
|
||||
&client,
|
||||
&format!("{base}/v1/chat/completions"),
|
||||
HeaderMap::new(),
|
||||
axum::body::Bytes::new(),
|
||||
observer,
|
||||
)
|
||||
.await
|
||||
.expect("forward ok");
|
||||
|
||||
// Backpressure status reaches the client unchanged.
|
||||
assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(String::from_utf8_lossy(&body).contains("rate_limit_exceeded"));
|
||||
|
||||
// finish still runs once even with a tiny non-streaming body.
|
||||
assert_eq!(probe.inner.lock().unwrap().finished, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_tail_smoke() {
|
||||
let mut tail = BodyTail::new(128);
|
||||
tail.push(b"hello ");
|
||||
tail.push(b"world");
|
||||
assert_eq!(tail.as_str(), "hello world");
|
||||
}
|
||||
@@ -1346,25 +1346,67 @@ fn validate_vision_prefill(prompt_len: usize, vram_free_mb: u64) -> Result<(), I
|
||||
/// the caller as `max`), or if free VRAM is below the floor. Enforcing
|
||||
/// the *derived* cap means a VRAM-tight host rejects a prompt that
|
||||
/// wouldn't fit, instead of accepting it and OOMing mid-prefill.
|
||||
///
|
||||
/// The third VRAM check — the length-aware backstop (#65) — closes the
|
||||
/// poll-vs-request snapshot gap #67 leaves open. `max` is
|
||||
/// `effective_prompt_cap()`, the input budget derived at **/models poll
|
||||
/// time** from the tightest card's free VRAM *then*. If free VRAM has
|
||||
/// since dropped (a co-resident model loaded, a concurrent prefill grew
|
||||
/// its KV), a prompt at-or-below that now-stale cap still clears the
|
||||
/// static floor yet no longer fits — and OOMs mid-prefill, poisoning the
|
||||
/// device context (the 2026-05-26 beast incident the #47 work exists to
|
||||
/// eliminate). So we re-run the same length×KV-vs-VRAM physics #67 uses
|
||||
/// for the cap, but against **request-time** free VRAM, reusing the
|
||||
/// model's [`ContextProfile`] rather than re-deriving the KV cost. This
|
||||
/// gives the text path the live-VRAM guard the vision path already has
|
||||
/// (`validate_vision_prefill`). `profile`/`kv_bytes_per_token_per_card`
|
||||
/// are per-card and `vram_free_mb` is the tightest card's free VRAM, so
|
||||
/// the two are commensurable on both single-GPU and TP loads.
|
||||
fn validate_request(
|
||||
prompt_len: usize,
|
||||
vram_free_mb: u64,
|
||||
max: usize,
|
||||
profile: Option<&super::context_limit::ContextProfile>,
|
||||
cfg: &crate::config::ContextLimitConfig,
|
||||
) -> Result<(), InferenceError> {
|
||||
if prompt_len > max {
|
||||
return Err(InferenceError::PromptTooLong { prompt_len, max });
|
||||
}
|
||||
// VRAM check is skipped on CPU loads (vram_free_mb == 0 sentinel)
|
||||
// VRAM checks are skipped on CPU loads (vram_free_mb == 0 sentinel)
|
||||
// because the (0, 0) reply from `query_vram` is also what a missing
|
||||
// worker returns. The CPU path has no per-GPU memory limit anyway —
|
||||
// host RAM is bounded by the OOM killer, not this check.
|
||||
if vram_free_mb == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let min = min_free_vram_mb();
|
||||
if vram_free_mb != 0 && vram_free_mb < min {
|
||||
if vram_free_mb < min {
|
||||
return Err(InferenceError::InsufficientVram {
|
||||
free_mb: vram_free_mb,
|
||||
required_mb: min,
|
||||
});
|
||||
}
|
||||
// Length-aware backstop (#65): KV the whole sequence (prompt +
|
||||
// generation reserve) will occupy, plus the prefill activation
|
||||
// headroom, plus the static floor as an additive cushion — all per
|
||||
// card. A degenerate zero-KV profile (no full-attention layers) or a
|
||||
// model with no captured profile skips this and rides the floor
|
||||
// check above, mirroring `derive_limit`'s VRAM-ceiling fallback.
|
||||
if let Some(profile) = profile
|
||||
&& profile.kv_bytes_per_token_per_card > 0
|
||||
{
|
||||
let tokens = (prompt_len as u64).saturating_add(cfg.output_reserve_tokens as u64);
|
||||
let kv_mb = profile.kv_bytes_per_token_per_card.saturating_mul(tokens) / (1024 * 1024);
|
||||
let required_mb = kv_mb
|
||||
.saturating_add(cfg.activation_headroom_mb)
|
||||
.saturating_add(min);
|
||||
if required_mb > vram_free_mb {
|
||||
return Err(InferenceError::InsufficientVram {
|
||||
free_mb: vram_free_mb,
|
||||
required_mb,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -2250,7 +2292,13 @@ impl CandleHarness {
|
||||
"chat_completion: starting"
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb, loaded.effective_prompt_cap())?;
|
||||
validate_request(
|
||||
prompt_len,
|
||||
vram_free_mb,
|
||||
loaded.effective_prompt_cap(),
|
||||
loaded.context_profile.as_ref(),
|
||||
&self.context_limit_cfg,
|
||||
)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
@@ -2704,7 +2752,13 @@ impl CandleHarness {
|
||||
);
|
||||
}
|
||||
|
||||
validate_request(prompt_len, vram_free_mb, loaded.effective_prompt_cap())?;
|
||||
validate_request(
|
||||
prompt_len,
|
||||
vram_free_mb,
|
||||
loaded.effective_prompt_cap(),
|
||||
loaded.context_profile.as_ref(),
|
||||
&self.context_limit_cfg,
|
||||
)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
@@ -3633,8 +3687,11 @@ impl CandleHarness {
|
||||
}
|
||||
|
||||
let tp_for_marker = Arc::clone(&tp);
|
||||
let handle =
|
||||
tokio::spawn(chat_completion_tp_inner(tp, request, principal).instrument(span.clone()));
|
||||
let context_limit_cfg = self.context_limit_cfg.clone();
|
||||
let handle = tokio::spawn(
|
||||
chat_completion_tp_inner(tp, request, principal, context_limit_cfg)
|
||||
.instrument(span.clone()),
|
||||
);
|
||||
match handle.await {
|
||||
Ok(Ok(resp)) => Ok(resp),
|
||||
Ok(Err(e)) => {
|
||||
@@ -3846,7 +3903,13 @@ impl CandleHarness {
|
||||
"TP chat_completion (stream): starting"
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb, tp.effective_prompt_cap())?;
|
||||
validate_request(
|
||||
prompt_len,
|
||||
vram_free_mb,
|
||||
tp.effective_prompt_cap(),
|
||||
tp.context_profile.as_ref(),
|
||||
&self.context_limit_cfg,
|
||||
)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
@@ -4367,6 +4430,7 @@ async fn chat_completion_tp_inner(
|
||||
tp: Arc<TpLoadedModel>,
|
||||
request: ChatCompletionRequest,
|
||||
principal: Option<String>,
|
||||
context_limit_cfg: crate::config::ContextLimitConfig,
|
||||
) -> Result<ChatCompletionResponse, InferenceError> {
|
||||
let req_start = std::time::Instant::now();
|
||||
let model_id = request.model.clone();
|
||||
@@ -4450,7 +4514,13 @@ async fn chat_completion_tp_inner(
|
||||
"TP chat_completion: starting"
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb, tp.effective_prompt_cap())?;
|
||||
validate_request(
|
||||
prompt_len,
|
||||
vram_free_mb,
|
||||
tp.effective_prompt_cap(),
|
||||
tp.context_profile.as_ref(),
|
||||
&context_limit_cfg,
|
||||
)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
@@ -6767,6 +6837,110 @@ mod tests {
|
||||
assert!(validate_vision_prefill(12_960, 12_445).is_ok());
|
||||
}
|
||||
|
||||
// ── #65: request-time length-aware VRAM backstop (text prefill) ──
|
||||
|
||||
/// A beast-like profile: 16 full-attn layers, 4 kv heads, head_dim
|
||||
/// 256, f16, TP=2 → 32 KiB/token/card (same numbers as the
|
||||
/// `context_limit` unit tests). At defaults this makes the
|
||||
/// length-aware footprint `(prompt_len + 8192)/32 + 2048 + 1500` MiB
|
||||
/// per card.
|
||||
fn backstop_profile() -> super::super::context_limit::ContextProfile {
|
||||
super::super::context_limit::ContextProfile {
|
||||
max_position_embeddings: 262_144,
|
||||
kv_bytes_per_token_per_card: super::super::context_limit::kv_bytes_per_token(
|
||||
16, 4, 256, 2, 2,
|
||||
),
|
||||
world_size: 2,
|
||||
}
|
||||
}
|
||||
|
||||
/// A prompt under the cap with ample free VRAM passes; the same
|
||||
/// prompt over the cap is `PromptTooLong` before any VRAM math.
|
||||
#[test]
|
||||
fn validate_request_cap_and_fit() {
|
||||
let cfg = crate::config::ContextLimitConfig::default();
|
||||
let profile = backstop_profile();
|
||||
// Under cap, 40 GB free → fits.
|
||||
assert!(validate_request(8_000, 40_000, 100_000, Some(&profile), &cfg).is_ok());
|
||||
// Over the cap → PromptTooLong, independent of VRAM.
|
||||
assert!(matches!(
|
||||
validate_request(100_001, 40_000, 100_000, Some(&profile), &cfg),
|
||||
Err(InferenceError::PromptTooLong { .. })
|
||||
));
|
||||
}
|
||||
|
||||
/// The CPU sentinel (`vram_free_mb == 0`) skips every VRAM check,
|
||||
/// including the new length-aware one — host RAM is the OOM killer's
|
||||
/// problem, not this guard's.
|
||||
#[test]
|
||||
fn validate_request_cpu_sentinel_skips_vram() {
|
||||
let cfg = crate::config::ContextLimitConfig::default();
|
||||
let profile = backstop_profile();
|
||||
assert!(validate_request(1_000_000, 0, 2_000_000, Some(&profile), &cfg).is_ok());
|
||||
}
|
||||
|
||||
/// The static floor remains a backstop: free VRAM below
|
||||
/// `min_free_vram_mb()` is rejected before the length-aware estimate
|
||||
/// even runs (so `required_mb` is the floor, not the KV footprint).
|
||||
#[test]
|
||||
fn validate_request_static_floor_still_binds() {
|
||||
let cfg = crate::config::ContextLimitConfig::default();
|
||||
let profile = backstop_profile();
|
||||
assert!(matches!(
|
||||
validate_request(10, 800, 100_000, Some(&profile), &cfg),
|
||||
Err(InferenceError::InsufficientVram {
|
||||
free_mb: 800,
|
||||
required_mb: 1500
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
/// A model with no captured profile (non-qwen3_5 arch) has no
|
||||
/// length-aware physics to apply, so it rides only the static floor —
|
||||
/// a fitting prompt with VRAM above the floor passes.
|
||||
#[test]
|
||||
fn validate_request_no_profile_rides_floor() {
|
||||
let cfg = crate::config::ContextLimitConfig::default();
|
||||
assert!(validate_request(500_000, 5_000, 1_000_000, None, &cfg).is_ok());
|
||||
}
|
||||
|
||||
/// The acceptance test (#65): a cap derived against *ample* free VRAM
|
||||
/// is later applied at request time against *tightened* free VRAM. A
|
||||
/// prompt sized exactly at the now-stale `effective_prompt_cap()`
|
||||
/// clears the cap and the static floor, yet no longer fits — the
|
||||
/// length-aware backstop catches it with a clean `InsufficientVram`
|
||||
/// instead of an OOM-poisoned context. Same prompt with the original
|
||||
/// ample VRAM still passes, proving the guard only bites on staleness.
|
||||
#[test]
|
||||
fn validate_request_catches_poll_vs_request_staleness() {
|
||||
let cfg = crate::config::ContextLimitConfig::default();
|
||||
let profile = backstop_profile();
|
||||
|
||||
// Cap derived at /models poll time with 40 GB free on the tightest
|
||||
// card — throughput binds, giving input = 87040 (the issue's
|
||||
// worked beast figure).
|
||||
let limit = super::super::context_limit::derive_limit(&profile, 40_000, 800.0, None, &cfg);
|
||||
let cap = limit.input.expect("input budget derived");
|
||||
assert_eq!(cap, 87_040);
|
||||
|
||||
// With that same ample VRAM, a prompt at the cap still fits.
|
||||
assert!(validate_request(cap, 40_000, cap, Some(&profile), &cfg).is_ok());
|
||||
|
||||
// Now free VRAM has dropped to 5 GB between the poll and the
|
||||
// request (a co-resident model loaded). The prompt is still ≤ cap
|
||||
// and clears the 1500 MiB floor, but its footprint —
|
||||
// (87040 + 8192)/32 + 2048 + 1500 = 6524 MiB — exceeds 5000 MiB.
|
||||
let err = validate_request(cap, 5_000, cap, Some(&profile), &cfg)
|
||||
.expect_err("stale cap must not let an over-VRAM prompt through");
|
||||
assert!(matches!(
|
||||
err,
|
||||
InferenceError::InsufficientVram {
|
||||
free_mb: 5_000,
|
||||
required_mb: 6_524
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
// ── Tool-call body parsing ───────────────────────────────────────
|
||||
|
||||
fn weather_schemas() -> ToolSchemas {
|
||||
|
||||
@@ -100,9 +100,9 @@ pub const KV_CACHE_DTYPE_BYTES: usize = 2;
|
||||
/// state, not a growing cache). Sharded across the TP world: per-rank
|
||||
/// KV-head count is `n_kv_heads / world_size`.
|
||||
///
|
||||
/// `2 ×` accounts for K and V. Shared by the limit derivation here and
|
||||
/// the per-rank load-time logging in the TP paths (and, in future, by
|
||||
/// #65's length-aware pre-flight guard).
|
||||
/// `2 ×` accounts for K and V. Shared by the limit derivation here, the
|
||||
/// per-rank load-time logging in the TP paths, and #65's request-time
|
||||
/// length-aware pre-flight guard (`candle::validate_request`).
|
||||
pub fn kv_bytes_per_token(
|
||||
n_full_attn_layers: usize,
|
||||
n_kv_heads: usize,
|
||||
|
||||
31
helexa-router.example.toml
Normal file
31
helexa-router.example.toml
Normal file
@@ -0,0 +1,31 @@
|
||||
# helexa-router.example.toml — example configuration
|
||||
#
|
||||
# Copy to helexa-router.toml and adjust for your environment.
|
||||
#
|
||||
# Environment variable overrides use the HELEXA_ROUTER_ prefix with __
|
||||
# separators:
|
||||
# HELEXA_ROUTER_ROUTER__LISTEN=0.0.0.0:8088
|
||||
|
||||
[router]
|
||||
# Plaintext listener. Operator/edge nginx terminates client TLS in front of
|
||||
# the router — the router never owns an inbound TLS listener.
|
||||
listen = "0.0.0.0:8088"
|
||||
|
||||
# How often (seconds) to refresh each cortex's health + /v1/models topology.
|
||||
# poll_interval_secs = 10
|
||||
|
||||
# -- Downstream cortexes -------------------------------------------------
|
||||
# Each [[cortexes]] entry is an operator-run cortex the router may dispatch
|
||||
# to. The router forwards the client's bearer verbatim (auth stays at
|
||||
# cortex) and routes on capacity. Outbound TLS to each cortex is verified.
|
||||
#
|
||||
# The skeleton only loads this list; capacity/catalogue polling and
|
||||
# capacity-aware dispatch arrive in later issues.
|
||||
|
||||
# [[cortexes]]
|
||||
# name = "lair-cafe"
|
||||
# endpoint = "https://cortex.lair.cafe"
|
||||
|
||||
# [[cortexes]]
|
||||
# name = "example-operator"
|
||||
# endpoint = "https://cortex.example.com"
|
||||
@@ -26,6 +26,18 @@
|
||||
# the load to neuron as `scheme:id` so the daemon
|
||||
# fetches from the right registry. Omit to let
|
||||
# neuron substitute its own `default_source`.
|
||||
# cost.* - optional operator-set pricing, surfaced verbatim on
|
||||
# GET /v1/models for clients (opencode) to display
|
||||
# spend. USD per 1,000,000 tokens, as numbers:
|
||||
# cost.input prompt tokens
|
||||
# cost.output completion tokens
|
||||
# cost.cache_read cache-hit tokens (optional tier)
|
||||
# cost.cache_write cache-write tokens (optional tier)
|
||||
# Absent vs zero is intentional (#68): OMIT the whole
|
||||
# cost block to mean "price not declared / unknown";
|
||||
# set cost.input/output = 0.0 to mean "intentionally
|
||||
# free" (self-hosted). The advertised rate must match
|
||||
# what metering bills against.
|
||||
|
||||
# Tensor-parallel target — needs a neuron with at least 2 large GPUs.
|
||||
# The example pins to a specific neuron name; adjust or remove the
|
||||
@@ -41,13 +53,16 @@ pinned_on = ["your-multi-gpu-neuron"]
|
||||
limit.context = 32768
|
||||
limit.input = 28672
|
||||
limit.output = 4096
|
||||
# Pricing in USD per 1M tokens — 0.0 for self-hosted.
|
||||
# Pricing in USD per 1M tokens. Explicit 0.0 = intentionally free
|
||||
# (self-hosted) — distinct from omitting `cost`, which means "not priced".
|
||||
cost.input = 0.0
|
||||
cost.output = 0.0
|
||||
# Static capability hints (unioned with runtime-detected flags).
|
||||
capabilities = ["text", "reasoning"]
|
||||
|
||||
# Mid-size dense model — fits on any single GPU with ≥16 GB VRAM.
|
||||
# No `cost` block here: this model is "not priced" — /v1/models omits the
|
||||
# `cost` key for it, so opencode shows spend as unknown rather than $0.
|
||||
[[models]]
|
||||
id = "Qwen/Qwen3-8B"
|
||||
harness = "candle"
|
||||
|
||||
Reference in New Issue
Block a user