review classifier: model, effort and review concurrency as swarm-wide settings
The classifier ran every classification on a hard-coded `haiku` with no
`--effort`, and the controller ran at most two requested reviews at once,
also hard-coded, so changing any of them took a code change and a
redeploy.
swarm-controller now keeps one settings document,
`{model, effort, max_concurrent_reviews}`, in its own
`review-classifier-settings` KV bucket under `all`, the way it keeps the
schedule definitions. No document means today's values: haiku, no
`--effort`, two reviews. `GET /api/review-classifier/settings` reads it;
`PUT` replaces it whole, needs the gateway's `Remote-User` like the
try-out POST, and refuses a bad model (the inference route's check), an
effort outside the five levels or a limit outside 1..=8 with
problem+json 400.
Every run request (try-out and requested review) now carries the model
and effort, and the classifier passes them to claude instead of its
constant, which is gone. A try-out reads the settings when it is posted;
a review reads them for the limit before it waits for a place and again
for the model and effort once it has one. The review semaphore is
replaced by a running count under a resizable limit, so a PUT, or the
next review's read, changes how many run without a restart, and a lower
limit lets the runs already going finish. MAX_WAITING_REVIEWS is
unchanged.
Closes #5151, part of #5096.
This commit is contained in:
atlas
mara
parent
78412b3d3b
commit
2c5eef73cd
10 changed files with 868 additions and 69 deletions
391
swarm-controller/src/review_classifier/settings.rs
Normal file
391
swarm-controller/src/review_classifier/settings.rs
Normal file
|
|
@ -0,0 +1,391 @@
|
|||
//! `GET` and `PUT /api/review-classifier/settings`: the model and effort
|
||||
//! every run uses, and how many requested reviews run at once, for the whole
|
||||
//! swarm.
|
||||
//!
|
||||
//! The settings are one document under one key ([`KEY`] in [`BUCKET`]),
|
||||
//! which only this daemon reads or writes, as the schedule definitions are.
|
||||
//! No document is [`Settings::default`]. A try-out reads them when it starts
|
||||
//! and a review when it gets its place, and each sends the model and effort
|
||||
//! in its run request, so a change applies to the next run with no restart.
|
||||
|
||||
use std::future::Future;
|
||||
use std::ops::RangeInclusive;
|
||||
|
||||
use anyhow::Context as _;
|
||||
use async_nats::jetstream::kv::Store;
|
||||
use axum::Json;
|
||||
use axum::extract::State;
|
||||
use axum::extract::rejection::JsonRejection;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use problem_details::ProblemDetails;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use swarm_queue_client::agent_inference::Effort;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
use crate::schedules::{queue, read, refused, write};
|
||||
use crate::{AppState, error_problem, wanted_error_status};
|
||||
|
||||
/// The bucket holding the settings. No other service holds a grant on it.
|
||||
pub const BUCKET: &str = "review-classifier-settings";
|
||||
|
||||
/// The one key in [`BUCKET`].
|
||||
pub const KEY: &str = "all";
|
||||
|
||||
/// How many times a losing writer re-reads and re-applies, as in
|
||||
/// `wanted::MAX_ATTEMPTS`.
|
||||
const MAX_ATTEMPTS: usize = 5;
|
||||
|
||||
/// What `max_concurrent_reviews` may be. Each running review is one `claude`
|
||||
/// at a time in the classifier's container, which has no bound of its own.
|
||||
pub(crate) const CONCURRENT_REVIEWS: RangeInclusive<usize> = 1..=8;
|
||||
|
||||
/// The review classifier's settings: the value under [`KEY`], and the body of
|
||||
/// both routes.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub(crate) struct Settings {
|
||||
/// The model every classification runs on, as claude's `--model` takes
|
||||
/// it: one word, not starting with `-`.
|
||||
pub(crate) model: String,
|
||||
/// One of `low`, `medium`, `high`, `xhigh`, `max`; `null` or absent runs
|
||||
/// at claude's own default.
|
||||
#[serde(default)]
|
||||
#[schema(value_type = Option<String>)]
|
||||
pub(crate) effort: Option<Effort>,
|
||||
/// How many requested reviews run at once, 1 to 8. Try-outs don't count.
|
||||
pub(crate) max_concurrent_reviews: usize,
|
||||
}
|
||||
|
||||
impl Default for Settings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model: "haiku".to_owned(),
|
||||
effort: None,
|
||||
max_concurrent_reviews: 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
/// What is wrong with these settings, if anything.
|
||||
fn check(&self) -> Result<(), String> {
|
||||
crate::inference::check_model(&self.model, "model")?;
|
||||
if !CONCURRENT_REVIEWS.contains(&self.max_concurrent_reviews) {
|
||||
return Err(format!(
|
||||
"max_concurrent_reviews {} is not {} to {}",
|
||||
self.max_concurrent_reviews,
|
||||
CONCURRENT_REVIEWS.start(),
|
||||
CONCURRENT_REVIEWS.end()
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a run reads the settings: the queue in production, a fake in the
|
||||
/// tests.
|
||||
pub(crate) trait SettingsSource {
|
||||
fn settings(&self) -> impl Future<Output = anyhow::Result<Settings>> + Send;
|
||||
}
|
||||
|
||||
impl SettingsSource for async_nats::Client {
|
||||
async fn settings(&self) -> anyhow::Result<Settings> {
|
||||
let (settings, _) = read::<Settings>(&open(self).await?, KEY)
|
||||
.await
|
||||
.context("reading the review classifier's settings")?;
|
||||
Ok(settings)
|
||||
}
|
||||
}
|
||||
|
||||
/// Open [`BUCKET`], creating it if nothing has yet.
|
||||
async fn open(client: &async_nats::Client) -> anyhow::Result<Store> {
|
||||
let js = async_nats::jetstream::new(client.clone());
|
||||
if let Ok(store) = js.get_key_value(BUCKET).await {
|
||||
return Ok(store);
|
||||
}
|
||||
js.create_key_value(async_nats::jetstream::kv::Config {
|
||||
bucket: BUCKET.to_owned(),
|
||||
description: "The review classifier's swarm-wide settings".to_owned(),
|
||||
history: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.with_context(|| format!("creating the {BUCKET} bucket"))
|
||||
}
|
||||
|
||||
fn queue_problem(e: &anyhow::Error) -> ProblemDetails {
|
||||
tracing::warn!(error = %format!("{e:#}"), "review classifier settings: queue failure");
|
||||
error_problem(wanted_error_status(e), &format!("{e:#}"))
|
||||
}
|
||||
|
||||
/// Read the settings.
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/api/review-classifier/settings",
|
||||
responses(
|
||||
(status = 200, description = "the settings in force; with none stored, \
|
||||
`{\"model\":\"haiku\",\"effort\":null,\"max_concurrent_reviews\":2}`",
|
||||
body = Settings),
|
||||
(status = 503, description = "no swarm queue is configured on this host, or it is \
|
||||
not connected (problem+json)", body = String),
|
||||
(status = 500, description = "the settings could not be read (problem+json)",
|
||||
body = String),
|
||||
),
|
||||
tag = "review-classifier"
|
||||
)]
|
||||
pub(crate) async fn get_settings(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<Json<Settings>, ProblemDetails> {
|
||||
let client = queue(&state).map_err(refused)?;
|
||||
client
|
||||
.settings()
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(|e| queue_problem(&e))
|
||||
}
|
||||
|
||||
/// Replace the settings.
|
||||
#[utoipa::path(
|
||||
put,
|
||||
path = "/api/review-classifier/settings",
|
||||
request_body = Settings,
|
||||
responses(
|
||||
(status = 200, description = "the settings as stored; they replace the previous \
|
||||
ones whole and apply from the next run", body = Settings),
|
||||
(status = 400, description = "not JSON or not this shape: a missing `model` or \
|
||||
`max_concurrent_reviews`, an unknown field, an effort other than the five \
|
||||
levels, a negative or non-integer `max_concurrent_reviews`; or a model that is \
|
||||
empty, has whitespace or starts with `-`; or `max_concurrent_reviews` outside \
|
||||
1 to 8 (problem+json)", body = String),
|
||||
(status = 401, description = "the request carries no `Remote-User`, so there is no \
|
||||
operator to name (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 settings could not be written (problem+json)",
|
||||
body = String),
|
||||
),
|
||||
tag = "review-classifier"
|
||||
)]
|
||||
pub(crate) async fn put_settings(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
body: Result<Json<Settings>, JsonRejection>,
|
||||
) -> Result<Json<Settings>, ProblemDetails> {
|
||||
let operator = super::operator(&headers).ok_or_else(|| {
|
||||
error_problem(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"no Remote-User header: the gateway names the operator, and this request did not \
|
||||
come through it",
|
||||
)
|
||||
})?;
|
||||
let Json(value) = body.map_err(|e| error_problem(StatusCode::BAD_REQUEST, &e.body_text()))?;
|
||||
value
|
||||
.check()
|
||||
.map_err(|why| error_problem(StatusCode::BAD_REQUEST, &why))?;
|
||||
let client = queue(&state).map_err(refused)?;
|
||||
let store = open(&client).await.map_err(|e| queue_problem(&e))?;
|
||||
for _ in 0..MAX_ATTEMPTS {
|
||||
let (current, revision) = read::<Settings>(&store, KEY)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?;
|
||||
if current == value
|
||||
|| write(&store, KEY, &value, revision)
|
||||
.await
|
||||
.map_err(|e| queue_problem(&e))?
|
||||
{
|
||||
state.reviews.resize(value.max_concurrent_reviews);
|
||||
tracing::info!(%operator, ?value, "set the review classifier's settings");
|
||||
return Ok(Json(value));
|
||||
}
|
||||
}
|
||||
Err(queue_problem(&anyhow::anyhow!(
|
||||
"gave up writing the review classifier's settings after {MAX_ATTEMPTS} conflicting \
|
||||
writes"
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::*;
|
||||
use crate::test_nats::TestNats;
|
||||
|
||||
fn settings(json: &str) -> serde_json::Result<Settings> {
|
||||
serde_json::from_str(json)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_document_is_todays_values() {
|
||||
assert_eq!(
|
||||
serde_json::to_value(Settings::default()).unwrap(),
|
||||
serde_json::json!({"model": "haiku", "effort": null, "max_concurrent_reviews": 2})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_settings_are_refused() {
|
||||
for json in [
|
||||
r#"{"model":"","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"opus 4","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"--print","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":0}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":9}"#,
|
||||
] {
|
||||
assert!(settings(json).unwrap().check().is_err(), "{json}");
|
||||
}
|
||||
// Control: the bounds themselves, and a full model id.
|
||||
for json in [
|
||||
r#"{"model":"claude-haiku-4-5","max_concurrent_reviews":1}"#,
|
||||
r#"{"model":"haiku","effort":"max","max_concurrent_reviews":8}"#,
|
||||
] {
|
||||
assert!(settings(json).unwrap().check().is_ok(), "{json}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bad_shape_does_not_decode() {
|
||||
for json in [
|
||||
r#"{"max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"haiku"}"#,
|
||||
r#"{"model":"haiku","effort":"extreme","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"haiku","effort":"HIGH","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":-1}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":2.5}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":2,"api_key":"x"}"#,
|
||||
] {
|
||||
assert!(settings(json).is_err(), "{json}");
|
||||
}
|
||||
// Control: effort may be left out or null.
|
||||
assert_eq!(
|
||||
settings(r#"{"model":"haiku","max_concurrent_reviews":2}"#).unwrap(),
|
||||
Settings::default()
|
||||
);
|
||||
assert_eq!(
|
||||
settings(r#"{"model":"haiku","effort":null,"max_concurrent_reviews":2}"#).unwrap(),
|
||||
Settings::default()
|
||||
);
|
||||
}
|
||||
|
||||
/// An app whose queue is `nats`.
|
||||
async fn app(nats: &TestNats) -> (axum::Router, AppState) {
|
||||
let (mut state, _) = crate::tests::state_with_roster();
|
||||
state.status = Some(std::sync::Arc::new(crate::status::StatusReader::new(
|
||||
nats.client().await,
|
||||
std::time::Duration::from_mins(1),
|
||||
)));
|
||||
(crate::build_app(state.clone()), state)
|
||||
}
|
||||
|
||||
async fn send(
|
||||
app: &axum::Router,
|
||||
req: Request<Body>,
|
||||
) -> (StatusCode, serde_json::Value, Option<String>) {
|
||||
let res = app.clone().oneshot(req).await.expect("infallible");
|
||||
let status = res.status();
|
||||
let kind = res
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
let bytes = axum::body::to_bytes(res.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("a body");
|
||||
(
|
||||
status,
|
||||
serde_json::from_slice(&bytes).unwrap_or_default(),
|
||||
kind,
|
||||
)
|
||||
}
|
||||
|
||||
fn get() -> Request<Body> {
|
||||
Request::get("/api/review-classifier/settings")
|
||||
.body(Body::empty())
|
||||
.expect("a request")
|
||||
}
|
||||
|
||||
fn put(body: &str) -> Request<Body> {
|
||||
Request::put("/api/review-classifier/settings")
|
||||
.header("content-type", "application/json")
|
||||
.header("remote-user", "alice")
|
||||
.body(Body::from(body.to_owned()))
|
||||
.expect("a request")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_put_is_what_the_get_and_the_next_run_read() {
|
||||
let nats = TestNats::start();
|
||||
let (app, state) = app(&nats).await;
|
||||
let (status, before, _) = send(&app, get()).await;
|
||||
assert_eq!(status, StatusCode::OK, "{before}");
|
||||
assert_eq!(before, serde_json::to_value(Settings::default()).unwrap());
|
||||
|
||||
let wanted = serde_json::json!({
|
||||
"model": "sonnet", "effort": "high", "max_concurrent_reviews": 5,
|
||||
});
|
||||
let (status, stored, _) = send(&app, put(&wanted.to_string())).await;
|
||||
assert_eq!(status, StatusCode::OK, "{stored}");
|
||||
assert_eq!(stored, wanted);
|
||||
let (_, after, _) = send(&app, get()).await;
|
||||
assert_eq!(after, wanted);
|
||||
assert_eq!(state.reviews.limit(), 5);
|
||||
|
||||
let client = nats.client().await;
|
||||
let read = client.settings().await.expect("readable");
|
||||
assert_eq!(
|
||||
(read.model.as_str(), read.effort),
|
||||
("sonnet", Some(Effort::High))
|
||||
);
|
||||
|
||||
// The same value again is no conflict.
|
||||
let (status, _, _) = send(&app, put(&wanted.to_string())).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_refused_put_is_a_problem_and_changes_nothing() {
|
||||
let nats = TestNats::start();
|
||||
let (app, state) = app(&nats).await;
|
||||
for bad in [
|
||||
"not json",
|
||||
r#"{"model":"haiku"}"#,
|
||||
r#"{"model":"haiku","effort":"extreme","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"-x","max_concurrent_reviews":2}"#,
|
||||
r#"{"model":"haiku","max_concurrent_reviews":9}"#,
|
||||
] {
|
||||
let (status, problem, kind) = send(&app, put(bad)).await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "{bad}: {problem}");
|
||||
assert!(problem["detail"].is_string(), "{bad}: {problem}");
|
||||
assert_eq!(kind.as_deref(), Some("application/problem+json"), "{bad}");
|
||||
}
|
||||
let (_, after, _) = send(&app, get()).await;
|
||||
assert_eq!(after, serde_json::to_value(Settings::default()).unwrap());
|
||||
assert_eq!(state.reviews.limit(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_put_without_an_operator_is_401() {
|
||||
let (state, _) = crate::tests::state_with_roster();
|
||||
let req = Request::put("/api/review-classifier/settings")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
serde_json::to_string(&Settings::default()).unwrap(),
|
||||
))
|
||||
.expect("a request");
|
||||
let (status, _, _) = send(&crate::build_app(state), req).await;
|
||||
assert_eq!(status, StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn without_a_queue_both_verbs_are_503() {
|
||||
let (state, _) = crate::tests::state_with_roster();
|
||||
let app = crate::build_app(state);
|
||||
let (status, _, _) = send(&app, get()).await;
|
||||
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
let body = serde_json::to_string(&Settings::default()).unwrap();
|
||||
let (status, _, _) = send(&app, put(&body)).await;
|
||||
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue