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

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.

nix/module-eval/core-toggle.nix

@@ -286,6 +286,22 @@ let
         && !(lib.hasInfix "mint swarmctl" s)
         && lib.hasInfix "mint swarm-authelia-bridge" s;
     }
+    {
+      # swarm-controller checks no caller itself: this location's authelia
+      # subrequest is the operator gate for every route under `/api/`. The
+      # webhook location is the control, the one path to the same socket that
+      # skips the gate, so the first clause reads a string that can lack it.
+      name = "the swarm UI's API location asks authelia before reaching the controller";
+      ok =
+        let
+          l = allLocal.services.nginx.virtualHosts."t.local".locations;
+          socket = allLocal.services.hyperhive.deploy.swarm-controller.socketPath;
+        in
+        lib.hasInfix "auth_request /__hive_authelia;" l."/api/".extraConfig
+        && lib.hasInfix socket l."/api/".proxyPass
+        && lib.hasInfix socket l."/webhook/forge/".proxyPass
+        && !(lib.hasInfix "auth_request" l."/webhook/forge/".extraConfig);
+    }
     {
       name = "a hive with matrix off serves no matrix discovery endpoint";
       ok =

swarm-controller/README.md

@@ -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.
 

swarm-controller/src/agent_limits.rs

@@ -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}"
+            );
+        }
+    }
+}

swarm-controller/src/main.rs

@@ -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))

swarm-queue-client/Cargo.toml

@@ -47,6 +47,9 @@ agent-inference = ["async-nats/jetstream"]
 # `ToSchema` on the review classifier's shared types, for swarm-controller,
 # whose HTTP API serves the same shapes the queue carries.
 schema = ["dep:utoipa"]
+# `agent_limits::read`, for hive-c0re, which reads the key of each agent it
+# runs the same way. Explicit for the same reason as `notices`.
+agent-limits = ["async-nats/jetstream"]
 
 [dependencies]
 # Bare (no `kv`/`jetstream`) unless a consumer opts into the `kv` feature

swarm-queue-client/src/agent_limits.rs

@@ -0,0 +1,166 @@
+//! The per-agent resource-cap KV bucket: one key per agent, holding the CPU
+//! and memory caps the swarm sets for it.
+//!
+//! **swarm-controller is the only writer**, and creates the bucket on its
+//! first write. Every field is optional: an absent field leaves the agent on
+//! its hive's own answer, so an agent with no key runs exactly as it would
+//! with no swarm.
+//!
+//! **The hive running the agent reads the key**, not the agent: the caps go
+//! into the container's systemd drop-in, which the host writes. A hive reads
+//! with one `DIRECT.GET` per agent
+//! ([`crate::agent_limits::direct_get_subject`]), and its grant covers every
+//! key of this bucket, read-only. Scoping it to the hive's own agents would
+//! need a roster in the auth responder.
+//!
+//! Keyed by the agent alone, like [`crate::agent_inference`]: an agent can
+//! move between hives, and its caps go with it.
+
+/// The KV bucket caps are published into, keyed by agent name.
+pub const BUCKET: &str = "agent-limits";
+
+/// The stream backing [`BUCKET`]. `create_key_value` prefixes `KV_`.
+pub const STREAM: &str = "KV_agent-limits";
+
+/// The subject `agent`'s entry is stored under.
+#[must_use]
+pub fn key_subject(agent: &str) -> String {
+    format!("$KV.{BUCKET}.{agent}")
+}
+
+/// The request subject that reads `agent`'s current entry.
+#[must_use]
+pub fn direct_get_subject(agent: &str) -> String {
+    format!("$JS.API.DIRECT.GET.{STREAM}.{}", key_subject(agent))
+}
+
+/// One agent's caps: the value under `BUCKET`/`<agent>`, as systemd values
+/// carried verbatim. `None` is "not capped by the swarm", never "uncapped".
+#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct AgentCaps {
+    /// systemd `CPUQuota=`, e.g. `"400%"`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cpu_quota: Option<String>,
+    /// systemd `MemoryMax=`, e.g. `"8G"`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub memory_max: Option<String>,
+}
+
+impl AgentCaps {
+    /// True when no field is set, which a hive reads the same as no key.
+    #[must_use]
+    pub fn is_empty(&self) -> bool {
+        self.cpu_quota.is_none() && self.memory_max.is_none()
+    }
+}
+
+/// Open the caps bucket for writing, creating it if nothing has yet.
+///
+/// **Controller-side only.** `history: 1`: a hive wants the current value.
+#[cfg(feature = "kv")]
+pub async fn open_or_create(
+    client: &async_nats::Client,
+) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
+    let js = async_nats::jetstream::new(client.clone());
+    crate::kv::open_or_create(
+        &js,
+        BUCKET,
+        "agent-limits",
+        "Resource caps the swarm controller sets for each agent".to_owned(),
+    )
+    .await
+}
+
+/// What reading one agent's key found.
+#[cfg(feature = "agent-limits")]
+#[derive(Debug)]
+pub enum Read {
+    /// The bucket answered; an absent or deleted key has every field unset.
+    Current(AgentCaps),
+    /// Nothing answers for the bucket's stream: no cap has been written yet,
+    /// or `JetStream` is not serving it right now. Not an unset value, since
+    /// the second case says nothing about the agent's caps.
+    NoBucket,
+}
+
+/// Read `agent`'s current caps with one direct get.
+///
+/// # Errors
+///
+/// A failed request or an undecodable value. A request the queue refused for
+/// permissions reaches the client as a timeout, not as a refusal.
+#[cfg(feature = "agent-limits")]
+pub async fn read(
+    client: &async_nats::Client,
+    agent: &str,
+) -> Result<Read, Box<dyn std::error::Error + Send + Sync>> {
+    use async_nats::jetstream::stream::DirectGetErrorKind;
+
+    let js = async_nats::jetstream::new(client.clone());
+    let stream = js.get_stream_no_info(STREAM).await?;
+    let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
+        Ok(message) => message,
+        Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => {
+            return Ok(Read::Current(AgentCaps::default()));
+        }
+        Err(e) if is_no_responders(&e) => return Ok(Read::NoBucket),
+        Err(e) => return Err(e.into()),
+    };
+    // A KV delete or purge is stored as an empty marker carrying this header.
+    let removed = message
+        .headers
+        .get("KV-Operation")
+        .is_some_and(|op| matches!(op.as_str(), "DEL" | "PURGE"));
+    if removed {
+        return Ok(Read::Current(AgentCaps::default()));
+    }
+    Ok(Read::Current(serde_json::from_slice(&message.payload)?))
+}
+
+#[cfg(feature = "agent-limits")]
+fn is_no_responders(e: &async_nats::jetstream::stream::DirectGetError) -> bool {
+    matches!(
+        e.kind(),
+        async_nats::jetstream::stream::DirectGetErrorKind::ErrorResponse(status, _)
+            if status == async_nats::StatusCode::NO_RESPONDERS
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The names the hive's grant in `swarm-nats-auth` is built from.
+    #[test]
+    fn names_match_the_queue_grant() {
+        assert_eq!(STREAM, format!("KV_{BUCKET}"));
+        assert_eq!(key_subject("iris"), "$KV.agent-limits.iris");
+        assert_eq!(
+            direct_get_subject("iris"),
+            "$JS.API.DIRECT.GET.KV_agent-limits.$KV.agent-limits.iris"
+        );
+    }
+
+    /// An unset field is left out on the wire, and `{}` is every field unset.
+    #[test]
+    fn unset_fields_are_absent() {
+        let value = AgentCaps {
+            memory_max: Some("8G".to_owned()),
+            ..AgentCaps::default()
+        };
+        assert_eq!(
+            serde_json::to_string(&value).unwrap(),
+            r#"{"memory_max":"8G"}"#
+        );
+        let empty: AgentCaps = serde_json::from_str("{}").unwrap();
+        assert!(empty.is_empty());
+    }
+
+    /// Fields a newer controller adds must not cost an older hive its caps.
+    #[test]
+    fn decoding_tolerates_unknown_fields() {
+        let decoded: AgentCaps =
+            serde_json::from_str(r#"{"cpu_quota":"400%","io_weight":"10"}"#).unwrap();
+        assert_eq!(decoded.cpu_quota.as_deref(), Some("400%"));
+    }
+}

swarm-queue-client/src/lib.rs

@@ -216,6 +216,11 @@ pub mod agent_schedules;
 /// unconditional; writing needs `kv` and reading needs `agent-inference`.
 pub mod agent_inference;
 
+/// The bucket swarm-controller sets each agent's CPU and memory caps in — one
+/// key per agent, read by the hive that runs it. The names are unconditional;
+/// writing needs `kv` and reading needs `agent-limits`.
+pub mod agent_limits;
+
 #[cfg(feature = "kv")]
 mod kv;