Set resource caps and inference when creating an agent
POST /api/agents takes optional limits and inference objects, checked by the same rules as the two PUT routes before anything is queued and stored ahead of the job graph. The create form gets collapsed sections for both, reusing the fields of the detail-panel editors.
This commit is contained in:
parent
aaac754356
commit
e36a15de91
8 changed files with 389 additions and 100 deletions
|
|
@ -51,7 +51,7 @@ pub struct Caps {
|
|||
|
||||
impl Caps {
|
||||
/// The value to store, or why it is refused.
|
||||
fn into_stored(self) -> Result<AgentCaps, String> {
|
||||
pub(crate) 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()?,
|
||||
|
|
@ -253,18 +253,28 @@ pub(crate) async fn put_agent_limits(
|
|||
&format!("no hive's wanted state declares an agent {agent:?} that is not destroyed"),
|
||||
));
|
||||
}
|
||||
let store = open_or_create(&client).await?;
|
||||
store(&client, &agent, &value).await?;
|
||||
Ok(Json(value.into()))
|
||||
}
|
||||
|
||||
/// Replace `agent`'s stored caps with `value`.
|
||||
pub(crate) async fn store(
|
||||
client: &async_nats::Client,
|
||||
agent: &str,
|
||||
value: &AgentCaps,
|
||||
) -> Result<(), ProblemDetails> {
|
||||
let store = open_or_create(client).await?;
|
||||
for _ in 0..MAX_ATTEMPTS {
|
||||
let (current, revision, undecodable) = read_or_default::<AgentCaps>(&store, &agent)
|
||||
let (current, revision, undecodable) = read_or_default::<AgentCaps>(&store, agent)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?;
|
||||
if (current == value && !undecodable)
|
||||
|| write(&store, &agent, &value, revision)
|
||||
if (current == *value && !undecodable)
|
||||
|| write(&store, agent, value, revision)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?
|
||||
{
|
||||
tracing::info!(%agent, ?value, "set agent limits");
|
||||
return Ok(Json(value.into()));
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Err(queue_problem(&anyhow::anyhow!(
|
||||
|
|
|
|||
|
|
@ -87,7 +87,7 @@ pub(crate) fn check_model(name: &str, what: &str) -> Result<(), String> {
|
|||
|
||||
impl InferenceSettings {
|
||||
/// The value to store, or why it is refused.
|
||||
fn into_stored(self) -> Result<AgentInference, String> {
|
||||
pub(crate) fn into_stored(self) -> Result<AgentInference, String> {
|
||||
if let Some(model) = &self.model {
|
||||
check_model(model, "model")?;
|
||||
}
|
||||
|
|
@ -313,18 +313,28 @@ pub(crate) async fn put_agent_inference(
|
|||
&format!("no hive's wanted state declares an agent {agent:?} that is not destroyed"),
|
||||
));
|
||||
}
|
||||
let store = open(&client).await?;
|
||||
store(&client, &agent, &value).await?;
|
||||
Ok(Json(value.into()))
|
||||
}
|
||||
|
||||
/// Replace `agent`'s stored inference settings with `value`.
|
||||
pub(crate) async fn store(
|
||||
client: &async_nats::Client,
|
||||
agent: &str,
|
||||
value: &AgentInference,
|
||||
) -> Result<(), ProblemDetails> {
|
||||
let store = open(client).await?;
|
||||
for _ in 0..MAX_ATTEMPTS {
|
||||
let (current, revision, undecodable) = read_or_default::<AgentInference>(&store, &agent)
|
||||
let (current, revision, undecodable) = read_or_default::<AgentInference>(&store, agent)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?;
|
||||
if (current == value && !undecodable)
|
||||
|| write(&store, &agent, &value, revision)
|
||||
if (current == *value && !undecodable)
|
||||
|| write(&store, agent, value, revision)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?
|
||||
{
|
||||
tracing::info!(%agent, ?value, "set inference settings");
|
||||
return Ok(Json(value.into()));
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Err(queue_problem(&anyhow::anyhow!(
|
||||
|
|
|
|||
|
|
@ -1727,6 +1727,15 @@ async fn get_agents_status(
|
|||
struct CreateAgentRequest {
|
||||
name: String,
|
||||
hive: String,
|
||||
/// CPU and memory caps to set for the agent, as `PUT
|
||||
/// /api/agents/{name}/limits` takes them. Absent leaves any caps already
|
||||
/// stored for the name alone.
|
||||
#[serde(default)]
|
||||
limits: Option<agent_limits::Caps>,
|
||||
/// Inference settings to set for the agent, as `PUT
|
||||
/// /api/agents/{name}/inference` takes them. Absent as `limits`.
|
||||
#[serde(default)]
|
||||
inference: Option<inference::InferenceSettings>,
|
||||
}
|
||||
|
||||
/// Where the queued job landed — a caller polls `/api/jobq/graph` (or
|
||||
|
|
@ -1780,9 +1789,9 @@ struct CreateAgentResponse {
|
|||
request_body = CreateAgentRequest,
|
||||
responses(
|
||||
(status = 200, description = "job chain queued", body = CreateAgentResponse),
|
||||
(status = 400, description = "`name` or `hive` is not a valid identifier, `name` is new and reserved or refused by the forge, or `hive` is not in this swarm (problem+json)", body = String),
|
||||
(status = 400, description = "`name` or `hive` is not a valid identifier, `name` is new and reserved or refused by the forge, `hive` is not in this swarm, or `limits` or `inference` holds a value the matching `PUT` refuses (problem+json)", body = String),
|
||||
(status = 409, description = "the swarm has already placed `name` on a different hive (problem+json)", body = String),
|
||||
(status = 503, description = "the swarm queue is not connected, so whether `name` is placed on another hive is unknown (problem+json)", body = String),
|
||||
(status = 503, description = "the swarm queue is not connected, so whether `name` is placed on another hive is unknown, or `limits` or `inference` is set and cannot be stored (problem+json)", body = String),
|
||||
(status = 500, description = "the job chain could not be queued, or another hive's wanted state could not be read (problem+json)", body = String),
|
||||
),
|
||||
tag = "agents"
|
||||
|
|
@ -1795,6 +1804,24 @@ async fn create_agent(
|
|||
.map_err(|reason| error_problem(axum::http::StatusCode::BAD_REQUEST, reason))?
|
||||
.into_string();
|
||||
|
||||
// Settings are checked before anything is created, so a bad value stops
|
||||
// the request rather than leaving an agent without them.
|
||||
let limits = req
|
||||
.limits
|
||||
.map(agent_limits::Caps::into_stored)
|
||||
.transpose()
|
||||
.map_err(|why| error_problem(axum::http::StatusCode::BAD_REQUEST, &why))?;
|
||||
let inference = req
|
||||
.inference
|
||||
.map(inference::InferenceSettings::into_stored)
|
||||
.transpose()
|
||||
.map_err(|why| error_problem(axum::http::StatusCode::BAD_REQUEST, &why))?;
|
||||
let settings_queue = if limits.is_some() || inference.is_some() {
|
||||
Some(schedules::queue(&state).map_err(schedules::refused)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Two checks, and the second is the one that makes the field worth
|
||||
// having: `Ident::parse` says the string is *shaped* like a hive name,
|
||||
// and the roster says it *is* one. Without the roster check a typo is
|
||||
|
|
@ -1874,6 +1901,16 @@ async fn create_agent(
|
|||
);
|
||||
return Err(error_problem(axum::http::StatusCode::CONFLICT, &detail));
|
||||
}
|
||||
// Stored ahead of the job graph: the hive reads these by agent name when it
|
||||
// runs the agent, and a failed write here leaves nothing created.
|
||||
if let Some(client) = &settings_queue {
|
||||
if let Some(limits) = &limits {
|
||||
agent_limits::store(client, &agent, limits).await?;
|
||||
}
|
||||
if let Some(inference) = &inference {
|
||||
inference::store(client, &agent, inference).await?;
|
||||
}
|
||||
}
|
||||
let mut sched = state
|
||||
.jobq
|
||||
.lock()
|
||||
|
|
@ -3658,6 +3695,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1maa".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -3696,6 +3735,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -3710,6 +3751,72 @@ mod tests {
|
|||
assert!(queued > 0, "an accepted creation must queue work");
|
||||
}
|
||||
|
||||
/// A request with no settings is the one callers sent before the fields
|
||||
/// existed, and each field is checked by the same rules as its `PUT`.
|
||||
#[test]
|
||||
fn a_creation_request_takes_optional_settings() {
|
||||
let plain: super::CreateAgentRequest =
|
||||
serde_json::from_str(r#"{"name":"atlas","hive":"pr1ma"}"#).expect("plain body");
|
||||
assert!(plain.limits.is_none() && plain.inference.is_none());
|
||||
|
||||
let full: super::CreateAgentRequest = serde_json::from_str(
|
||||
r#"{"name":"atlas","hive":"pr1ma","limits":{"cpu_quota":"400%"},
|
||||
"inference":{"model":"opus","effort":"high"}}"#,
|
||||
)
|
||||
.expect("body with settings");
|
||||
assert!(full.limits.is_some() && full.inference.is_some());
|
||||
|
||||
let typo = serde_json::from_str::<super::CreateAgentRequest>(
|
||||
r#"{"name":"atlas","hive":"pr1ma","limits":{"cpu":"400%"}}"#,
|
||||
);
|
||||
assert!(typo.is_err(), "an unknown cap field must be refused");
|
||||
}
|
||||
|
||||
/// A bad cap or model must stop the creation before the graph is queued:
|
||||
/// the point of taking settings in the creating request is that a refusal
|
||||
/// leaves no agent without them.
|
||||
#[tokio::test]
|
||||
async fn bad_settings_are_refused_before_anything_is_queued() {
|
||||
use axum::http::StatusCode;
|
||||
let bad: [(&str, &str); 2] = [
|
||||
(r#"{"cpu_quota":"lots"}"#, "limits"),
|
||||
(r#"{"model":"-opus"}"#, "inference"),
|
||||
];
|
||||
for (value, field) in bad {
|
||||
let (state, sched) = state_with_roster();
|
||||
let body = format!(r#"{{"name":"atlas","hive":"pr1ma","{field}":{value}}}"#);
|
||||
let req: super::CreateAgentRequest = serde_json::from_str(&body).expect("body parses");
|
||||
let err = super::create_agent(axum::extract::State(state), axum::Json(req))
|
||||
.await
|
||||
.expect_err("a bad setting must be refused");
|
||||
assert_eq!(
|
||||
axum::response::IntoResponse::into_response(err).status(),
|
||||
StatusCode::BAD_REQUEST,
|
||||
"{field}"
|
||||
);
|
||||
assert_eq!(node_count(&sched), 0, "{field}: nothing may be queued");
|
||||
}
|
||||
}
|
||||
|
||||
/// Settings that cannot be stored, here because there is no queue, refuse
|
||||
/// the creation rather than leaving an agent without them.
|
||||
#[tokio::test]
|
||||
async fn settings_that_cannot_be_stored_refuse_the_creation() {
|
||||
let (state, sched) = state_with_roster();
|
||||
let req: super::CreateAgentRequest = serde_json::from_str(
|
||||
r#"{"name":"atlas","hive":"pr1ma","limits":{"cpu_quota":"400%"}}"#,
|
||||
)
|
||||
.expect("body parses");
|
||||
let err = super::create_agent(axum::extract::State(state), axum::Json(req))
|
||||
.await
|
||||
.expect_err("no queue means the caps cannot be stored");
|
||||
assert_eq!(
|
||||
axum::response::IntoResponse::into_response(err).status(),
|
||||
axum::http::StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(node_count(&sched), 0, "nothing may be queued");
|
||||
}
|
||||
|
||||
/// `state_with_roster` with a second hive, `sec0nd`.
|
||||
fn state_with_two_hives() -> (super::AppState, SharedSched) {
|
||||
let (state, sched) = state_with_roster();
|
||||
|
|
@ -3735,6 +3842,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: name.to_owned(),
|
||||
hive: hive.to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -4855,6 +4964,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -4906,6 +5017,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -4953,6 +5066,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -5008,6 +5123,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
@ -5210,6 +5327,8 @@ mod tests {
|
|||
axum::Json(super::CreateAgentRequest {
|
||||
name: "atlas".to_owned(),
|
||||
hive: "pr1ma".to_owned(),
|
||||
limits: None,
|
||||
inference: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
|
|
|
|||
Loading…
Reference in a new issue