vinzenz/hyperhive
vinzenz/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Activity

swarm-controller: GET/PUT/DELETE /api/agents/{name}/limits

Browse source 
An agent's CPU and memory caps, set from the swarm. Keyed by the agent
alone, in a new `agent-limits` KV bucket that swarm-controller alone
writes, the same shape as `agent-inference` (#5034): mara on #4268,
"agents are not tied to a hive. dont tie the limits to the hive."

  PUT    {"cpu_quota": "400%", "memory_max": "8G"}   every field optional,
                                                     replaces the stored value
  GET    the stored caps, {} when there are none
  DELETE 204

PUT refuses an agent no hive's wanted state declares (as the inference
route does), unknown fields (422), zero and anything systemd does not take,
before the queue is asked, and stores memory in the spelling systemd takes
(8gb -> 8G). The hive hands these values to systemd verbatim.

The routes sit under /api/, behind the swarm UI vhost's authelia
subrequest like every sibling; the core-toggle module-eval case pins that
gate.

The bucket's reader on the hive comes in the next commit.

Contributes to #3865.
This commit is contained in:
atlas 2026-10-12 11:26:18 +02:00 • committed by mara
commit ebeef4b84d
7 changed files with 728 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

5
swarm-controller/README.md
View file

@ -52,6 +52,11 @@ The swarm's control plane. The swarm UI and `swarmctl` are its clients.
agent applies them itself. `GET` also answers the model and effort the
agent last reported in its turn-state header, kept in memory →
[`docs/swarm/inference.md`](../docs/swarm/inference.md).
- **Resource caps** — each agent's CPU and memory cap, one key per agent in
the `agent-limits` KV bucket, which this daemon alone writes
(`GET`/`PUT`/`DELETE /api/agents/{name}/limits`); the hive running the
agent reads it, and each cap replaces the hive's own answer for that field →
[`docs/scheduler/coordinator.md`](../docs/scheduler/coordinator.md#container-resource-limits).
- **Relays** — each agent's terminal and turn-state header as SSE, agent
icons, the cross-repo issue report, and the UI's quick links.

527
swarm-controller/src/agent_limits.rs Normal file
View file

@ -0,0 +1,527 @@
//! `/api/agents/{name}/limits` — the CPU and memory caps the swarm sets for
//! one agent.
//!
//! Each agent's caps are one value under its own key of `agent-limits`
//! (`swarm_queue_client::agent_limits`), which this daemon alone writes. The
//! hive running the agent reads the key; on the hive, a cap replaces the
//! hive's own answer for that field, and a field left unset stays with the
//! hive.
//!
//! The hive hands these values to systemd verbatim, and a value systemd
//! rejects stops the agent's container from starting, so they are validated
//! here as well as on the hive.
use axum::Json;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use problem_details::ProblemDetails;
use serde::{Deserialize, Serialize};
use swarm_queue_client::agent_limits::{self, AgentCaps};
use utoipa::ToSchema;
use crate::schedules::{live, queue, read, refused, write};
use crate::{AppState, error_problem, wanted_error_status};
/// How many times a losing writer re-reads and re-applies, as in
/// `wanted::MAX_ATTEMPTS`.
const MAX_ATTEMPTS: usize = 5;
/// One agent's caps. A field left out is not capped by the swarm.
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
#[serde(deny_unknown_fields)]
pub struct Caps {
/// systemd `CPUQuota=`: a percentage above zero, `"100%"` being one core.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schema(example = "400%")]
cpu_quota: Option<String>,
/// systemd `MemoryMax=`: a size above zero with an optional `K`/`M`/`G`/`T`
/// suffix, a percentage of physical memory up to `100%`, or `infinity`.
/// Stored in the one spelling systemd accepts (`8gb` is stored as `8G`).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schema(example = "8G")]
memory_max: Option<String>,
}
impl Caps {
/// The value to store, or why it is refused.
fn into_stored(self) -> Result<AgentCaps, String> {
Ok(AgentCaps {
cpu_quota: self.cpu_quota.as_deref().map(cpu_quota).transpose()?,
memory_max: self.memory_max.as_deref().map(memory_max).transpose()?,
})
}
}
impl From<AgentCaps> for Caps {
fn from(v: AgentCaps) -> Self {
Self {
cpu_quota: v.cpu_quota,
memory_max: v.memory_max,
}
}
}
// The accepted shapes are hive-c0re's `resource_limits::validate_cpu_quota` /
// `validate_memory_max`, the operator's own input path for the same drop-in,
// with zero and an over-100% memory share refused as well.
fn cpu_quota(value: &str) -> Result<String, String> {
match value.strip_suffix('%') {
Some(n) if is_positive_number(n) => Ok(value.to_owned()),
_ => Err(format!(
"invalid cpu_quota {value:?}: expected a percentage above zero such as \"200%\" \
(100% = one full core)"
)),
}
}
fn memory_max(value: &str) -> Result<String, String> {
if value == "infinity" {
return Ok(value.to_owned());
}
let normalized = match value.strip_suffix('%') {
Some(n) => (is_positive_number(n) && n.parse::<f64>().is_ok_and(|p| p <= 100.0))
.then(|| value.to_owned()),
None => byte_size(value),
};
normalized.ok_or_else(|| {
format!(
"invalid memory_max {value:?}: expected a size above zero such as \"8G\", a \
percentage up to \"100%\", or \"infinity\""
)
})
}
/// A byte count with an optional `K`/`M`/`G`/`T` multiplier, in any case and
/// with an optional trailing `B`, in the form systemd takes: `8gb` → `8G`, a
/// bare `1048576B` kept as it is.
fn byte_size(value: &str) -> Option<String> {
const UNITS: [char; 8] = ['K', 'M', 'G', 'T', 'k', 'm', 'g', 't'];
let (body, b) = match value.strip_suffix(['B', 'b']) {
Some(body) => (body, true),
None => (value, false),
};
let (mantissa, unit) = match body.strip_suffix(UNITS) {
Some(mantissa) => (mantissa, body[mantissa.len()..].to_ascii_uppercase()),
None if b => (body, "B".to_owned()),
None => (body, String::new()),
};
is_positive_number(mantissa).then(|| format!("{mantissa}{unit}"))
}
/// Digits with at most one `.` between digits, and not zero. No sign or
/// exponent, which systemd does not take either.
fn is_positive_number(value: &str) -> bool {
let (int, frac) = value.split_once('.').unwrap_or((value, "1"));
let digits = |s: &str| !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit());
digits(int) && digits(frac) && value.bytes().any(|b| (b'1'..=b'9').contains(&b))
}
fn agent_name(name: &str) -> Result<String, (StatusCode, String)> {
hive_types::Ident::parse(name)
.map(hive_types::Ident::into_string)
.map_err(|reason| (StatusCode::BAD_REQUEST, reason.to_owned()))
}
fn queue_problem(e: &anyhow::Error) -> ProblemDetails {
tracing::warn!(error = %format!("{e:#}"), "agent limits: queue failure");
error_problem(wanted_error_status(e), &format!("{e:#}"))
}
async fn open(
client: &async_nats::Client,
) -> Result<async_nats::jetstream::kv::Store, ProblemDetails> {
agent_limits::open_or_create(client)
.await
.map_err(|e| queue_problem(&anyhow::Error::new(e)))
}
#[utoipa::path(
get,
path = "/api/agents/{name}/limits",
params(("name" = String, Path, description = "agent name")),
responses(
(status = 200, description = "the caps the swarm sets for the agent; `{}` when it \
sets none", body = Caps),
(status = 400, description = "not an identifier (problem+json)", body = String),
(status = 503, description = "no swarm queue is configured on this host, or it is \
not connected (problem+json)", body = String),
(status = 500, description = "the caps could not be read (problem+json)",
body = String),
),
tag = "agents"
)]
pub(crate) async fn get_agent_limits(
State(state): State<AppState>,
Path(name): Path<String>,
) -> Result<Json<Caps>, ProblemDetails> {
let agent = agent_name(&name).map_err(refused)?;
let client = queue(&state).map_err(refused)?;
let store = open(&client).await?;
let (current, _) = read::<AgentCaps>(&store, &agent)
.await
.map_err(|e| queue_problem(&e))?;
Ok(Json(current.into()))
}
/// Set the CPU and memory caps the swarm declares for one agent.
///
/// The value replaces the stored one whole. Whichever hive runs the agent
/// applies it; a running agent picks the caps up only at its next spawn or
/// rebuild.
#[utoipa::path(
put,
path = "/api/agents/{name}/limits",
params(("name" = String, Path, description = "agent name")),
request_body = Caps,
responses(
(status = 200, description = "the caps as stored", body = Caps),
(status = 400, description = "not an identifier, not declared on any hive, or a cap \
is not a value systemd takes or is zero (problem+json)", body = String),
(status = 422, description = "an unknown field"),
(status = 503, description = "no swarm queue is configured on this host, or it is \
not connected (problem+json)", body = String),
(status = 500, description = "the caps could not be written (problem+json)",
body = String),
),
tag = "agents"
)]
pub(crate) async fn put_agent_limits(
State(state): State<AppState>,
Path(name): Path<String>,
Json(req): Json<Caps>,
) -> Result<Json<Caps>, ProblemDetails> {
let agent = agent_name(&name).map_err(refused)?;
let value = req
.into_stored()
.map_err(|why| error_problem(StatusCode::BAD_REQUEST, &why))?;
let client = queue(&state).map_err(refused)?;
if !live(&state).await?.contains(&agent) {
return Err(error_problem(
StatusCode::BAD_REQUEST,
&format!("no hive's wanted state declares an agent {agent:?} that is not destroyed"),
));
}
let store = open(&client).await?;
for _ in 0..MAX_ATTEMPTS {
let (current, revision) = read::<AgentCaps>(&store, &agent)
.await
.map_err(|e| queue_problem(&e))?;
if current == value
|| write(&store, &agent, &value, revision)
.await
.map_err(|e| queue_problem(&e))?
{
tracing::info!(%agent, ?value, "set agent limits");
return Ok(Json(value.into()));
}
}
Err(queue_problem(&anyhow::anyhow!(
"gave up writing the limits of {agent} after {MAX_ATTEMPTS} conflicting writes"
)))
}
#[utoipa::path(
delete,
path = "/api/agents/{name}/limits",
params(("name" = String, Path, description = "agent name")),
responses(
(status = 204, description = "cleared; the hive's own answer applies again"),
(status = 400, description = "not an identifier (problem+json)", body = String),
(status = 503, description = "no swarm queue is configured on this host, or it is \
not connected (problem+json)", body = String),
(status = 500, description = "the caps could not be cleared (problem+json)",
body = String),
),
tag = "agents"
)]
pub(crate) async fn delete_agent_limits(
State(state): State<AppState>,
Path(name): Path<String>,
) -> Result<StatusCode, ProblemDetails> {
let agent = agent_name(&name).map_err(refused)?;
let client = queue(&state).map_err(refused)?;
let store = open(&client).await?;
store
.delete(&agent)
.await
.map_err(|e| queue_problem(&anyhow::Error::new(e)))?;
tracing::info!(%agent, "cleared agent limits");
Ok(StatusCode::NO_CONTENT)
}
#[cfg(test)]
mod tests {
use super::{
Caps, cpu_quota, delete_agent_limits, get_agent_limits, memory_max, put_agent_limits,
};
use crate::test_nats::TestNats;
use crate::wanted::WantedWriter;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use std::sync::Arc;
use swarm_queue_client::wanted::AgentState;
#[test]
fn cpu_quota_takes_a_percentage_above_zero() {
for ok in ["400%", "100%", "0.5%", "1500%"] {
assert_eq!(cpu_quota(ok).as_deref(), Ok(ok));
}
for bad in [
"0%", "0.0%", "-50%", "400", "4 00%", "%", "lots", "", "1e3%", "+5%",
] {
assert!(cpu_quota(bad).is_err(), "{bad:?} was accepted");
}
}
#[test]
fn memory_max_takes_a_size_a_share_or_infinity_and_stores_the_systemd_spelling() {
for (given, stored) in [
("8G", "8G"),
("8gb", "8G"),
("512m", "512M"),
("1048576", "1048576"),
("1048576B", "1048576B"),
("1.5T", "1.5T"),
("50%", "50%"),
("100%", "100%"),
("infinity", "infinity"),
] {
assert_eq!(memory_max(given).as_deref(), Ok(stored), "{given:?}");
}
for bad in [
"0", "0G", "0.0G", "0%", "-8G", "150%", "8X", "G", "8 G", "", "eight", "8GG",
] {
assert!(memory_max(bad).is_err(), "{bad:?} was accepted");
}
}
/// One hive, `pr1ma`, on `client`.
fn state(client: &async_nats::Client) -> crate::AppState {
crate::AppState {
hives: Arc::new(vec![crate::HiveEntry {
name: "pr1ma".to_owned(),
domain: "pr1ma.example".to_owned(),
}]),
links: Arc::new(Vec::new()),
status: Some(Arc::new(crate::status::StatusReader::new(
client.clone(),
std::time::Duration::from_mins(1),
))),
wanted: Some(Arc::new(WantedWriter::new(client.clone()))),
agent_status: None,
agent_icons: None,
reported_inference: None,
jobq: Arc::new(std::sync::Mutex::new(hive_jobq::scheduler::Scheduler::new(
hive_jobq::Graph::new(),
hive_jobq::resources::ResourceTable::new(),
))),
webhook_secret: None,
config_prs: Arc::default(),
swarm_name: None,
auth: None,
forge: Arc::default(),
create_gate: Arc::default(),
accounts: crate::linked_accounts::Accounts::Fake(Arc::default()),
try_outs: crate::review_classifier::TryOuts::default(),
}
}
/// A state on a real server whose `pr1ma` declares `atlas`.
async fn declared(client: &async_nats::Client) -> crate::AppState {
WantedWriter::new(client.clone())
.set("pr1ma", "atlas", AgentState::Up)
.await
.expect("declare atlas");
state(client)
}
fn caps(cpu: Option<&str>, mem: Option<&str>) -> Caps {
Caps {
cpu_quota: cpu.map(str::to_owned),
memory_max: mem.map(str::to_owned),
}
}
async fn put(
state: &crate::AppState,
agent: &str,
value: Caps,
) -> Result<axum::Json<Caps>, problem_details::ProblemDetails> {
put_agent_limits(
State(state.clone()),
Path(agent.to_owned()),
axum::Json(value),
)
.await
}
async fn get(state: &crate::AppState, agent: &str) -> Caps {
get_agent_limits(State(state.clone()), Path(agent.to_owned()))
.await
.expect("caps are read")
.0
}
/// The raw value a hive reads for `agent`, or `None` for no key or a
/// deleted one.
async fn published(client: &async_nats::Client, agent: &str) -> Option<String> {
let store = swarm_queue_client::agent_limits::open_or_create(client)
.await
.expect("bucket");
store
.get(agent)
.await
.expect("read")
.map(|v| String::from_utf8(v.to_vec()).expect("utf-8"))
}
/// Through the real handlers and a real `nats-server`: a cap lands under
/// the agent's own key as the hive reads it, `GET` answers it, and
/// `DELETE` takes it out again.
///
/// The published bytes are pinned literally because hive-c0re's
/// `workers::swarm_limits` tests decode this same string; a change to
/// either side's reading of it fails one of the two.
#[tokio::test]
async fn the_route_writes_the_caps_a_hive_reads_and_delete_clears_them() {
let nats = TestNats::start();
let client = nats.client().await;
let state = declared(&client).await;
let answer = put(&state, "atlas", caps(Some("400%"), Some("8gb")))
.await
.expect("caps are set");
assert_eq!(answer.0, caps(Some("400%"), Some("8G")));
assert_eq!(
published(&client, "atlas").await.as_deref(),
Some(r#"{"cpu_quota":"400%","memory_max":"8G"}"#)
);
assert_eq!(get(&state, "atlas").await, caps(Some("400%"), Some("8G")));
assert_eq!(
delete_agent_limits(State(state.clone()), Path("atlas".to_owned())).await,
Ok(StatusCode::NO_CONTENT)
);
assert_eq!(published(&client, "atlas").await, None);
assert_eq!(get(&state, "atlas").await, Caps::default());
}
/// A name no hive declares gets no key, so a typo cannot cap an agent
/// that does not exist yet.
#[tokio::test]
async fn an_agent_no_hive_declares_is_refused_and_nothing_is_written() {
let nats = TestNats::start();
let client = nats.client().await;
let state = declared(&client).await;
let problem = put(&state, "iris", caps(Some("400%"), None))
.await
.expect_err("no hive declares iris");
assert_eq!(problem.status, Some(StatusCode::BAD_REQUEST), "{problem:?}");
assert_eq!(published(&client, "iris").await, None);
}
/// A refused cap is answered before the queue is asked: the client here
/// was never connected, and would answer 503.
#[tokio::test]
async fn a_refused_cap_is_a_400_before_the_queue() {
let client = async_nats::ConnectOptions::new()
.retry_on_initial_connect()
.connect("127.0.0.1:1")
.await
.expect("returns without a connection");
let state = state(&client);
let problem = put(&state, "atlas", caps(Some("0%"), None))
.await
.expect_err("a zero cap is refused");
assert_eq!(problem.status, Some(StatusCode::BAD_REQUEST), "{problem:?}");
// The control: a valid cap on the same state does reach the queue.
let problem = put(&state, "atlas", caps(Some("400%"), None))
.await
.expect_err("the queue is not connected");
assert_eq!(
problem.status,
Some(StatusCode::SERVICE_UNAVAILABLE),
"{problem:?}"
);
}
const CAPPED: &str = r#"{"cpu_quota":"400%","memory_max":"8G"}"#;
/// `atlas` capped, behind the real router.
async fn capped_router(client: &async_nats::Client) -> axum::Router {
let state = declared(client).await;
let _ = put(&state, "atlas", caps(Some("400%"), Some("8G")))
.await
.expect("cap atlas");
crate::build_app(state)
}
async fn put_raw(app: &axum::Router, body: &'static str) -> StatusCode {
use tower::ServiceExt as _;
app.clone()
.oneshot(
axum::http::Request::put("/api/agents/atlas/limits")
.header(axum::http::header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(body))
.expect("request"),
)
.await
.expect("router answers")
.status()
}
/// A body that is not JSON is refused before the handler runs, and the
/// caps already set stay. The well-formed `{}` after it is the control:
/// the same router and request do reach the bucket.
#[tokio::test]
async fn a_malformed_body_is_a_400_and_the_caps_stay() {
let nats = TestNats::start();
let client = nats.client().await;
let app = capped_router(&client).await;
for body in [r#"{"cpu_quota":"#, "cpu_quota=400%", ""] {
assert_eq!(
put_raw(&app, body).await,
StatusCode::BAD_REQUEST,
"{body:?}"
);
assert_eq!(
published(&client, "atlas").await.as_deref(),
Some(CAPPED),
"{body:?}"
);
}
assert_eq!(put_raw(&app, "{}").await, StatusCode::OK);
assert_eq!(published(&client, "atlas").await.as_deref(), Some("{}"));
}
/// A field the body does not define is a 422 and writes nothing: a
/// misspelt `memory_max` must not read as "no memory cap".
#[tokio::test]
async fn an_unknown_field_is_a_422_and_the_caps_stay() {
let nats = TestNats::start();
let client = nats.client().await;
let app = capped_router(&client).await;
for body in [
r#"{"cpu_quota":"100%","memory":"2G"}"#,
r#"{"limits":{"cpu_quota":"100%"}}"#,
] {
assert_eq!(
put_raw(&app, body).await,
StatusCode::UNPROCESSABLE_ENTITY,
"{body}"
);
assert_eq!(
published(&client, "atlas").await.as_deref(),
Some(CAPPED),
"{body}"
);
}
}
}

6
swarm-controller/src/main.rs
View file

@ -44,6 +44,7 @@ use utoipa_axum::{router::OpenApiRouter, routes};
mod agent_icon;
mod agent_identity;
mod agent_input;
mod agent_limits;
mod agent_renewal;
mod agent_state_stream;
mod agent_status;
@ -3300,6 +3301,11 @@ fn build_app(state: AppState) -> axum::Router {
inference::put_agent_inference,
inference::delete_agent_inference
))
.routes(routes!(
agent_limits::get_agent_limits,
agent_limits::put_agent_limits,
agent_limits::delete_agent_limits
))
.routes(routes!(get_config_prs))
.routes(routes!(create_agent))
.routes(routes!(mint_agent_identity))