Skip to content

Commit

Permalink
[ENH] Command-line tools for every chroma-load endpoint.
Browse files Browse the repository at this point in the history
chroma-load-start       start a workload
chroma-load-stop        stop a workload
chroma-load-inhibit     stop all workloads
chroma-load-uninhibit   reverse the effects of chroma-load-inhibit
  • Loading branch information
rescrv committed Dec 2, 2024
1 parent 2c1101e commit 1eb8240
Show file tree
Hide file tree
Showing 7 changed files with 166 additions and 24 deletions.
37 changes: 13 additions & 24 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions rust/load/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ async-trait = "0.1.83"
axum = "0.7"
chromadb = { git = "https://github.com/rescrv/chromadb-rs", rev = "e364e35c34c660d4e8e862436ea600ddc2f46a1e" }
chrono = { version = "0.4.38", features = ["serde"] }
clap = { version = "4", features = ["derive"] }
figment = { version = "0.10.12", features = ["env", "yaml", "test"] }
guacamole = { version = "0.9", default-features = false }
serde.workspace = true
Expand All @@ -27,3 +28,4 @@ opentelemetry-otlp = "0.27"
opentelemetry_sdk = { version = "0.27", features = ["rt-tokio"] }
tracing.workspace = true
tower-http = { version = "0.6.2", features = ["trace"] }
reqwest = { version = "0.12", features = ["json"] }
18 changes: 18 additions & 0 deletions rust/load/src/bin/chroma-load-inhibit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
//! Inhibit chroma-load on every host provided on the command line.
#[tokio::main]
async fn main() {
for host in std::env::args().skip(1) {
let client = reqwest::Client::new();
match client.post(format!("{}/inhibit", host)).send().await {
Ok(resp) => {
if resp.status().is_success() {
println!("Inhibited load on {}", host);
} else {
eprintln!("Failed to inhibit load on {}: {}", host, resp.status());
}
}
Err(e) => eprintln!("Failed to inhibit load on {}: {}", host, e),
}
}
}
54 changes: 54 additions & 0 deletions rust/load/src/bin/chroma-load-start.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
//! Start a workload on the chroma-load server.
use clap::Parser;

use chroma_load::rest::StartRequest;
use chroma_load::{humanize_expires, Workload};

#[derive(Parser, Debug)]
struct Args {
#[arg(long)]
host: String,
#[arg(long)]
name: String,
#[arg(long)]
expires: String,
#[arg(long)]
data_set: String,
#[arg(long)]
workload: String,
#[arg(long)]
throughput: f64,
}

#[tokio::main]
async fn main() {
let args = Args::parse();
let client = reqwest::Client::new();
let req = StartRequest {
name: args.name,
expires: humanize_expires(&args.expires).unwrap_or(args.expires),
data_set: args.data_set,
workload: Workload::ByName(args.workload),
throughput: args.throughput,
};
match client
.post(format!("{}/start", args.host))
.json(&req)
.send()
.await
{
Ok(resp) => {
if resp.status().is_success() {
println!("Started workload on {}", args.host);
} else {
eprintln!(
"Failed to start workload on {}: {}",
args.host,
resp.status()
);
}
}
Err(e) => eprintln!("Failed to start workload on {}: {}", args.host, e),
}
}
44 changes: 44 additions & 0 deletions rust/load/src/bin/chroma-load-stop.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
//! Stop a single workload on the chroma-load server.
//!
//! If you are looking to stop traffic for a SEV, see chroma-load-inhibit.
use clap::Parser;
use uuid::Uuid;

use chroma_load::rest::StopRequest;

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

#[tokio::main]
async fn main() {
let args = Args::parse();
let client = reqwest::Client::new();
let req = StopRequest {
uuid: Uuid::parse_str(&args.uuid).unwrap(),
};
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),
}
}
18 changes: 18 additions & 0 deletions rust/load/src/bin/chroma-load-uninhibit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
//! Uninhibit chroma-load on every host provided on the command line.
#[tokio::main]
async fn main() {
for host in std::env::args().skip(1) {
let client = reqwest::Client::new();
match client.post(format!("{}/uninhibit", host)).send().await {
Ok(resp) => {
if resp.status().is_success() {
println!("Resumed load on {}", host);
} else {
eprintln!("Failed to uninhibit load on {}: {}", host, resp.status());
}
}
Err(e) => eprintln!("Failed to uninhibit load on {}: {}", host, e),
}
}
}
17 changes: 17 additions & 0 deletions rust/load/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,23 @@ pub async fn entrypoint() {
runner.abort();
}

///////////////////////////////////////// humanize_expires /////////////////////////////////////////

pub fn humanize_expires(expires: &str) -> Option<String> {
if let Ok(expires) = chrono::DateTime::parse_from_rfc3339(expires) {
Some(expires.to_rfc3339())
} else if let Some(duration) = expires.strip_suffix("s") {
let expires = chrono::Utc::now() + chrono::Duration::seconds(duration.trim().parse().ok()?);
Some(expires.to_rfc3339())
} else if let Some(duration) = expires.strip_suffix("min") {
let expires = chrono::Utc::now()
+ chrono::Duration::seconds(duration.trim().parse::<i64>().ok()? * 60i64);
Some(expires.to_rfc3339())
} else {
Some(expires.to_string())
}
}

/////////////////////////////////////////////// tests //////////////////////////////////////////////

#[cfg(test)]
Expand Down

0 comments on commit 1eb8240

Please sign in to comment.