i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 18:59:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improve gateway scheduling and runtime admission

Browse Source
This commit is contained in:
elky
2026-06-24 01:53:45 +08:00
parent cf0af8fa1e
commit d336d1a7fa
87 changed files with 9671 additions and 804 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-testkit/src/bin/gateway_pressure_probe.rs
+49 -2
View File
@@ -176,25 +176,56 @@ fn spawn_metrics_sampler(
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut target_url: Option<String> = None;
let mut warmup_url: Option<String> = None;
let mut metrics_url: Option<String> = None;
let mut total_requests: Option<usize> = None;
let mut concurrency: Option<usize> = None;
let mut warmup_connections: usize = 0;
let mut timeout_ms: Option<u64> = None;
let mut connect_timeout_ms: Option<u64> = None;
let mut client_shards: Option<usize> = None;
let mut pool_max_idle_per_host: Option<usize> = None;
let mut start_ramp_ms: u64 = 0;
let mut first_body_hold_ms: u64 = 0;
let mut sample_interval_ms: u64 = 500;
let mut method = Method::GET;
let mut headers = BTreeMap::new();
let mut body: Option<Vec<u8>> = None;
let mut response_mode = HttpLoadProbeResponseMode::HeadersOnly;
let mut http1_only = false;
let mut http2_prior_knowledge = false;
let mut output_path = None;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--url" => target_url = Some(next_value(&mut iter, "--url")?),
"--warmup-url" => warmup_url = Some(next_value(&mut iter, "--warmup-url")?),
"--metrics-url" => metrics_url = Some(next_value(&mut iter, "--metrics-url")?),
"--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?),
"--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?),
"--warmup-connections" => {
warmup_connections = next_value(&mut iter, "--warmup-connections")?.parse()?
}
"--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?),
"--connect-timeout-ms" => {
connect_timeout_ms = Some(next_value(&mut iter, "--connect-timeout-ms")?.parse()?)
}
"--client-shards" => {
client_shards = Some(next_value(&mut iter, "--client-shards")?.parse()?)
}
"--pool-max-idle-per-host" => {
pool_max_idle_per_host =
Some(next_value(&mut iter, "--pool-max-idle-per-host")?.parse()?)
}
"--start-ramp-ms" => {
start_ramp_ms = next_value(&mut iter, "--start-ramp-ms")?.parse()?
}
"--first-body-hold-ms" => {
first_body_hold_ms = next_value(&mut iter, "--first-body-hold-ms")?.parse()?
}
"--http1-only" => http1_only = true,
"--http2-prior-knowledge" => http2_prior_knowledge = true,
"--sample-interval-ms" => {
sample_interval_ms = next_value(&mut iter, "--sample-interval-ms")?.parse()?
}
@@ -247,9 +278,20 @@ fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
response_mode,
..HttpLoadProbeConfig::default()
};
load.warmup_url = warmup_url;
load.warmup_connections = warmup_connections;
if let Some(timeout_ms) = timeout_ms {
load.timeout = Duration::from_millis(timeout_ms);
}
load.connect_timeout = connect_timeout_ms.map(Duration::from_millis);
if let Some(client_shards) = client_shards {
load.client_shards = client_shards;
}
load.pool_max_idle_per_host = pool_max_idle_per_host;
load.start_ramp = Duration::from_millis(start_ramp_ms);
load.first_body_hold = Duration::from_millis(first_body_hold_ms);
load.http1_only = http1_only;
load.http2_prior_knowledge = http2_prior_knowledge;
load.validate()
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
if sample_interval_ms == 0 {
@@ -298,10 +340,15 @@ fn parse_response_mode(
) -> Result<HttpLoadProbeResponseMode, Box<dyn std::error::Error>> {
match value.trim().to_ascii_lowercase().as_str() {
"headers" | "headers-only" | "header" => Ok(HttpLoadProbeResponseMode::HeadersOnly),
"first-body-byte" | "first-body" | "first-byte" | "first-chunk" => {
Ok(HttpLoadProbeResponseMode::FirstBodyByte)
}
"full" | "full-body" | "body" => Ok(HttpLoadProbeResponseMode::FullBody),
other => Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unsupported --response-mode {other}; expected headers or full"),
format!(
"unsupported --response-mode {other}; expected headers, first-body-byte, or full"
),
)
.into()),
}
@@ -348,6 +395,6 @@ fn metric_name_matches(actual: &str, expected: &str) -> bool {
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin gateway_pressure_probe -- --url <URL> --metrics-url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000] [--sample-interval-ms 500] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full] [--output /tmp/gateway_pressure.json]"
"usage: cargo run -p aether-testkit --bin gateway_pressure_probe -- --url <URL> --metrics-url <URL> --requests <N> --concurrency <N> [--warmup-url <URL>] [--warmup-connections N] [--method GET] [--timeout-ms 30000] [--connect-timeout-ms 10000] [--client-shards 1] [--pool-max-idle-per-host N] [--start-ramp-ms 0] [--first-body-hold-ms 0] [--http1-only | --http2-prior-knowledge] [--sample-interval-ms 500] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|first-body-byte|full] [--output /tmp/gateway_pressure.json]"
);
}
crates/aether-testkit/src/bin/http_load_probe.rs
+49 -2
View File
@@ -15,21 +15,52 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::error::Error>> {
let mut url: Option<String> = None;
let mut warmup_url: Option<String> = None;
let mut total_requests: Option<usize> = None;
let mut concurrency: Option<usize> = None;
let mut warmup_connections: usize = 0;
let mut timeout_ms: Option<u64> = None;
let mut connect_timeout_ms: Option<u64> = None;
let mut client_shards: Option<usize> = None;
let mut pool_max_idle_per_host: Option<usize> = None;
let mut start_ramp_ms: u64 = 0;
let mut first_body_hold_ms: u64 = 0;
let mut method = Method::GET;
let mut headers = std::collections::BTreeMap::new();
let mut body: Option<Vec<u8>> = None;
let mut response_mode = aether_testkit::HttpLoadProbeResponseMode::HeadersOnly;
let mut http1_only = false;
let mut http2_prior_knowledge = false;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--url" => url = Some(next_value(&mut iter, "--url")?),
"--warmup-url" => warmup_url = Some(next_value(&mut iter, "--warmup-url")?),
"--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?),
"--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?),
"--warmup-connections" => {
warmup_connections = next_value(&mut iter, "--warmup-connections")?.parse()?
}
"--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?),
"--connect-timeout-ms" => {
connect_timeout_ms = Some(next_value(&mut iter, "--connect-timeout-ms")?.parse()?)
}
"--client-shards" => {
client_shards = Some(next_value(&mut iter, "--client-shards")?.parse()?)
}
"--pool-max-idle-per-host" => {
pool_max_idle_per_host =
Some(next_value(&mut iter, "--pool-max-idle-per-host")?.parse()?)
}
"--start-ramp-ms" => {
start_ramp_ms = next_value(&mut iter, "--start-ramp-ms")?.parse()?
}
"--first-body-hold-ms" => {
first_body_hold_ms = next_value(&mut iter, "--first-body-hold-ms")?.parse()?
}
"--http1-only" => http1_only = true,
"--http2-prior-knowledge" => http2_prior_knowledge = true,
"--method" => {
method = Method::from_bytes(next_value(&mut iter, "--method")?.as_bytes())?
}
@@ -78,9 +109,20 @@ fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::err
response_mode,
..HttpLoadProbeConfig::default()
};
config.warmup_url = warmup_url;
config.warmup_connections = warmup_connections;
if let Some(timeout_ms) = timeout_ms {
config.timeout = Duration::from_millis(timeout_ms);
}
config.connect_timeout = connect_timeout_ms.map(Duration::from_millis);
if let Some(client_shards) = client_shards {
config.client_shards = client_shards;
}
config.pool_max_idle_per_host = pool_max_idle_per_host;
config.start_ramp = Duration::from_millis(start_ramp_ms);
config.first_body_hold = Duration::from_millis(first_body_hold_ms);
config.http1_only = http1_only;
config.http2_prior_knowledge = http2_prior_knowledge;
config
.validate()
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
@@ -115,10 +157,15 @@ fn parse_response_mode(
"headers" | "headers-only" | "header" => {
Ok(aether_testkit::HttpLoadProbeResponseMode::HeadersOnly)
}
"first-body-byte" | "first-body" | "first-byte" | "first-chunk" => {
Ok(aether_testkit::HttpLoadProbeResponseMode::FirstBodyByte)
}
"full" | "full-body" | "body" => Ok(aether_testkit::HttpLoadProbeResponseMode::FullBody),
other => Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unsupported --response-mode {other}; expected headers or full"),
format!(
"unsupported --response-mode {other}; expected headers, first-body-byte, or full"
),
)
.into()),
}
@@ -139,6 +186,6 @@ fn next_value(
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin http_load_probe -- --url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full]"
"usage: cargo run -p aether-testkit --bin http_load_probe -- --url <URL> --requests <N> --concurrency <N> [--warmup-url <URL>] [--warmup-connections N] [--method GET] [--timeout-ms 30000] [--connect-timeout-ms 10000] [--client-shards 1] [--pool-max-idle-per-host N] [--start-ramp-ms 0] [--first-body-hold-ms 0] [--http1-only | --http2-prior-knowledge] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|first-body-byte|full]"
);
}
crates/aether-testkit/src/bin/mock_openai_upstream.rs
+53 -8
View File
@@ -14,7 +14,7 @@ use serde_json::json;
#[derive(Debug, Clone)]
struct Config {
bind: SocketAddr,
binds: Vec<SocketAddr>,
chunks: u64,
first_byte_delay: Duration,
chunk_delay: Duration,
@@ -25,9 +25,9 @@ struct Config {
impl Default for Config {
fn default() -> Self {
Self {
bind: "127.0.0.1:18181"
binds: vec!["127.0.0.1:18181"
.parse()
.expect("default bind address should parse"),
.expect("default bind address should parse")],
chunks: 8,
first_byte_delay: Duration::from_millis(0),
chunk_delay: Duration::from_millis(20),
@@ -68,9 +68,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.route("/responses", post(responses))
.with_state(app_state);
let listener = tokio::net::TcpListener::bind(config.bind).await?;
eprintln!("mock OpenAI upstream listening on http://{}", config.bind);
axum::serve(listener, app).await?;
serve_listeners(&config.binds, app).await?;
Ok(())
}
async fn serve_listeners(
binds: &[SocketAddr],
app: Router,
) -> Result<(), Box<dyn std::error::Error>> {
let mut listeners = Vec::with_capacity(binds.len());
for bind in binds {
listeners.push((*bind, tokio::net::TcpListener::bind(bind).await?));
}
if listeners.len() == 1 {
let (bind, listener) = listeners
.into_iter()
.next()
.ok_or_else(|| std::io::Error::other("mock upstream listener set is empty"))?;
eprintln!("mock OpenAI upstream listening on http://{bind}");
axum::serve(listener, app).await?;
return Ok(());
}
let mut servers = tokio::task::JoinSet::new();
for (bind, listener) in listeners {
let app = app.clone();
eprintln!("mock OpenAI upstream listening on http://{bind}");
servers.spawn(async move { axum::serve(listener, app).await });
}
if let Some(result) = servers.join_next().await {
servers.abort_all();
let serve_result = result
.map_err(|err| std::io::Error::other(format!("mock listener task failed: {err}")))?;
serve_result?;
}
Ok(())
}
@@ -292,10 +323,17 @@ fn current_unix_secs() -> u64 {
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut config = Config::default();
let mut binds_overridden = false;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--bind" => config.bind = next_value(&mut iter, "--bind")?.parse()?,
"--bind" => {
if !binds_overridden {
config.binds.clear();
binds_overridden = true;
}
config.binds.push(next_value(&mut iter, "--bind")?.parse()?);
}
"--chunks" => config.chunks = next_value(&mut iter, "--chunks")?.parse()?,
"--first-byte-delay-ms" => {
config.first_byte_delay =
@@ -325,6 +363,13 @@ fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
}
}
}
if config.binds.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"at least one --bind is required",
)
.into());
}
Ok(config)
}
@@ -343,6 +388,6 @@ fn next_value(
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181] [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]"
"usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181]... [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]"
);
}