Skip to content

Commit

Permalink
[ENH] Support delayed workloads in chroma-load-start.
Browse files Browse the repository at this point in the history
  • Loading branch information
rescrv committed Dec 17, 2024
1 parent 6558172 commit 71ed33b
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion rust/load/src/bin/chroma-load-start.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ struct Args {
#[arg(long)]
expires: String,
#[arg(long)]
delay: Option<String>,
#[arg(long)]
data_set: String,
#[arg(long)]
workload: String,
Expand Down Expand Up @@ -84,11 +86,20 @@ async fn main() {
let args = Args::parse();
let client = reqwest::Client::new();
let throughput = args.throughput();
let mut workload = Workload::ByName(args.workload);
if let Some(delay) = args.delay {
let delay = humanize_expires(&delay).expect("delay must be humanizable");
let delay = delay.parse().expect("delay must be a date time");
workload = Workload::Delay {
after: delay,
wrap: Box::new(workload),
};
}
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),
workload,
throughput,
};
match client
Expand Down

0 comments on commit 71ed33b

Please sign in to comment.