Skip to content

Commit

Permalink
[ENH] Make the error messages out of chroma-load decipherable. (#3316)
Browse files Browse the repository at this point in the history
Staring was failing with a 400 that was really a 404.  This PR changes
that case to a 404 instead of 400 and prints the body in the
chroma-load-start tool.  Also, a tool to stop all workloads.
  • Loading branch information
rescrv authored Dec 17, 2024
1 parent 7d0f95d commit 6558172
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 3 deletions.
5 changes: 4 additions & 1 deletion rust/load/src/bin/chroma-load-start.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,10 +116,13 @@ async fn main() {
);
} else {
eprintln!(
"Failed to start workload on {}: {}",
"Categorically failed to start workload on {}: {}",
args.host,
resp.status()
);
if let Ok(text) = resp.text().await {
eprintln!("{}", text.trim());
}
}
}
Err(e) => eprintln!("Failed to start workload on {}: {}", args.host, e),
Expand Down
60 changes: 60 additions & 0 deletions rust/load/src/bin/chroma-load-stop-all.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
//! Stop all workloads on the chroma-load server.
//!
//! If you are looking to stop traffic for a SEV, see chroma-load-inhibit.
use clap::Parser;

use chroma_load::rest::StopRequest;

#[derive(Parser, Debug)]
struct Args {
#[arg(long)]
host: String,
}

#[tokio::main]
async fn main() {
let args = Args::parse();
let client = reqwest::Client::new();
match client
.get(format!("{}/", args.host))
.header(reqwest::header::ACCEPT, "application/json")
.send()
.await
{
Ok(resp) => {
let resp = resp.error_for_status().expect("Failed to get status");
let resp = resp
.json::<chroma_load::rest::Status>()
.await
.expect("Failed to parse status");
for workload in resp.running {
let req = StopRequest {
uuid: workload.uuid,
};
match client
.post(format!("{}/stop", args.host))
.json(&req)
.send()
.await
{
Ok(resp) => {
if resp.status().is_success() {
println!("Stopped workload on {}", args.host);
} else {
eprintln!(
"Failed to stop workload on {}: {}",
args.host,
resp.status()
);
}
}
Err(e) => eprintln!("Failed to stop workload on {}: {}", args.host, e),
}
}
}
Err(e) => {
eprintln!("Failed to get status: {}", e);
}
}
}
3 changes: 1 addition & 2 deletions rust/load/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -478,7 +478,7 @@ impl Workload {
if let Some(workload) = workloads.get(name) {
*self = workload.clone();
} else {
return Err(Error::InvalidRequest(format!("workload not found: {name}")));
return Err(Error::NotFound(format!("workload not found: {name}")));
}
}
Workload::Get(_) => {}
Expand Down Expand Up @@ -1412,7 +1412,6 @@ pub async fn entrypoint() {
"http_request",
method = ?request.method(),
matched_path,
some_other_field = tracing::field::Empty,
)
}),
)
Expand Down

0 comments on commit 6558172

Please sign in to comment.