diff --git a/src/libraries/rust/stargate/benches/backend-degradation.yaml b/src/libraries/rust/stargate/benches/backend-degradation.yaml index 26e970f8d..c17ff479a 100644 --- a/src/libraries/rust/stargate/benches/backend-degradation.yaml +++ b/src/libraries/rust/stargate/benches/backend-degradation.yaml @@ -99,15 +99,15 @@ traffic_pattern: target_rps: 6 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: wait-and-widen config: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml index 00db155c1..21d0ce1e9 100644 --- a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml @@ -105,9 +105,9 @@ traffic_pattern: burst_period_requests: 16 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: random config: default: random diff --git a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml index 463b79ebd..482895b8f 100644 --- a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml @@ -64,15 +64,15 @@ traffic_pattern: target_rps: 8 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: wait-and-widen config: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -81,7 +81,7 @@ algorithms: require_input_tokens: true - name: pulsar-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -90,7 +90,7 @@ algorithms: consider_kv_free_tokens: true - name: pulsar-wait-and-widen config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen @@ -99,7 +99,7 @@ algorithms: require_input_tokens: true - name: pulsar-wait-and-widen-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml index b0107d2f8..dc588cfb8 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml @@ -107,9 +107,9 @@ traffic_pattern: target_rps: 5 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -121,7 +121,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml index 00ce62793..3c1d09041 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml @@ -107,9 +107,9 @@ traffic_pattern: target_rps: 40 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -121,7 +121,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml index 3bff125d0..a7cbc5df5 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml @@ -107,9 +107,9 @@ traffic_pattern: burst_period_requests: 40 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -121,7 +121,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml index d84fcafba..7c28f4a43 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml @@ -109,9 +109,9 @@ traffic_pattern: target_rps: 140 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -123,7 +123,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml index 62bc9c0e4..b4fdb7d1a 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml @@ -109,9 +109,9 @@ traffic_pattern: target_rps: 45 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -123,7 +123,7 @@ algorithms: default: wait-and-widen - name: wait-and-widen-affinity config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: wait-and-widen @@ -133,7 +133,7 @@ algorithms: cache_affinity_backend_selection_count: 1 - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -142,7 +142,7 @@ algorithms: require_input_tokens: true - name: pulsar-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -151,7 +151,7 @@ algorithms: consider_kv_free_tokens: true - name: pulsar-wait-and-widen config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen @@ -160,7 +160,7 @@ algorithms: require_input_tokens: true - name: pulsar-wait-and-widen-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml index 7053cd3c7..447e4db95 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml @@ -115,7 +115,7 @@ algorithms: pylon_queue_admission: enabled: false config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: wait-and-widen @@ -127,7 +127,7 @@ algorithms: pylon_queue_admission: enabled: false config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -138,7 +138,7 @@ algorithms: pylon_queue_admission: enabled: false config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml index 49b29a640..2019da9fa 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml @@ -63,9 +63,9 @@ traffic_pattern: target_rps: 20 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -77,7 +77,7 @@ algorithms: default: wait-and-widen - name: wait-and-widen-affinity config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: wait-and-widen @@ -87,7 +87,7 @@ algorithms: cache_affinity_backend_selection_count: 1 - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -96,7 +96,7 @@ algorithms: require_input_tokens: true - name: pulsar-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar @@ -105,7 +105,7 @@ algorithms: consider_kv_free_tokens: true - name: pulsar-wait-and-widen config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen @@ -114,7 +114,7 @@ algorithms: require_input_tokens: true - name: pulsar-wait-and-widen-consider-kv-free-tokens config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar-wait-and-widen diff --git a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml index 0dfaf9754..0e31c2394 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml @@ -62,9 +62,9 @@ traffic_pattern: interval_ms: 40 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -76,7 +76,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml index 10d6ab5a0..06c36fbbc 100644 --- a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml +++ b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml @@ -118,15 +118,15 @@ traffic_pattern: p99_cap: 2048 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: wait-and-widen config: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/overload-6-backends.yaml b/src/libraries/rust/stargate/benches/overload-6-backends.yaml index 63c6838da..6bc9f8206 100644 --- a/src/libraries/rust/stargate/benches/overload-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/overload-6-backends.yaml @@ -91,9 +91,9 @@ traffic_pattern: burst_period_requests: 20 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin @@ -105,7 +105,7 @@ algorithms: default: wait-and-widen - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml index 99161141d..c280d752b 100644 --- a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml +++ b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml @@ -62,15 +62,15 @@ traffic_pattern: step_requests: 20 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml index fbb666546..7b3072764 100644 --- a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml +++ b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml @@ -64,15 +64,15 @@ traffic_pattern: target_rps: 6 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin - name: pulsar config: - default: power-of-two + default: power-of-n models: dummy-model: algorithm: pulsar diff --git a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml index 328005d70..419fb13dc 100644 --- a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml +++ b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml @@ -60,9 +60,9 @@ traffic_pattern: interval_ms: 100 algorithms: - - name: power-of-two + - name: power-of-n config: - default: power-of-two + default: power-of-n - name: round-robin config: default: round-robin diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs index 5e2ce80b8..e48495859 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs @@ -73,8 +73,8 @@ fn config() -> BenchmarkConfig { }), degradation: DegradationConfig::default(), algorithms: vec![AlgorithmConfig { - name: "power-of-two".to_string(), - config: serde_json::json!({"default": "power-of-two"}), + name: "power-of-n".to_string(), + config: serde_json::json!({"default": "power-of-n"}), pylon_queue_admission: None, }], } @@ -111,7 +111,7 @@ fn render_default_test_manifest(config: &BenchmarkConfig) -> RenderedManifests { &config.algorithms[0], "sgbench-sg-power", "sgbench-be-power", - r#"{"default":"power-of-two"}"#, + r#"{"default":"power-of-n"}"#, ) } @@ -275,11 +275,11 @@ fn benchmark_run(run_dir: &Path) -> BenchmarkK8sRun { fs::write(run_dir.join(manifest), "kind: List\n").expect("manifest should write"); } BenchmarkK8sRun { - algorithm_name: "power-of-two".to_string(), + algorithm_name: "power-of-n".to_string(), manifest_path: run_dir.join("manifest.json"), run_dir: run_dir.to_path_buf(), - stargate_ns: "sgbench-sg-power-of-two".to_string(), - backends_ns: "sgbench-be-power-of-two".to_string(), + stargate_ns: "sgbench-sg-power-of-n".to_string(), + backends_ns: "sgbench-be-power-of-n".to_string(), stargate_count: 2, nodeport_host: "node.test".to_string(), stargate_http_endpoint: "http://node.test:30080".to_string(), @@ -316,8 +316,8 @@ fn prepare_benchmark_k8s_run_writes_split_manifests_and_run_info() { ) .expect("benchmark k8s run should prepare"); - assert_eq!(run.stargate_ns, "sgbench-sg-power-of-two"); - assert_eq!(run.backends_ns, "sgbench-be-power-of-two"); + assert_eq!(run.stargate_ns, "sgbench-sg-power-of-n"); + assert_eq!(run.backends_ns, "sgbench-be-power-of-n"); assert_eq!(run.stargate_http_endpoint, "http://node.example:30082"); assert_eq!( run.stargate_metrics_endpoint, @@ -352,7 +352,7 @@ fn prepare_benchmark_k8s_run_writes_split_manifests_and_run_info() { "run info should read", )) .expect("run info should parse"); - assert_eq!(run_info["algorithm_name"], "power-of-two"); + assert_eq!(run_info["algorithm_name"], "power-of-n"); assert_eq!(run_info["http_node_port"], 30082); assert_eq!( run_info["stargate_http_endpoint"], @@ -369,7 +369,7 @@ fn kubectl_runner_executes_readiness_and_maintenance_commands() { let fake = FakeKubectl::new(); let kubectl = fake.runner(); let tempdir = tempfile::tempdir().expect("tempdir should create"); - let run = benchmark_run(&tempdir.path().join("run-power-of-two")); + let run = benchmark_run(&tempdir.path().join("run-power-of-n")); kubectl.apply(&run).expect("stargate manifest should apply"); kubectl @@ -412,7 +412,7 @@ fn kubectl_runner_executes_readiness_and_maintenance_commands() { &stargate_log, &[ "status: exit status: 0", - "logs for -n sgbench-sg-power-of-two logs", + "logs for -n sgbench-sg-power-of-n logs", ], ); let backend_log = read_utf8( @@ -455,7 +455,7 @@ fi ); let kubectl = fake.runner(); let tempdir = tempfile::tempdir().expect("tempdir should create"); - let run = benchmark_run(&tempdir.path().join("run-power-of-two")); + let run = benchmark_run(&tempdir.path().join("run-power-of-n")); assert_error_contains(kubectl.apply(&run), "kubectl apply failed"); assert_error_contains(kubectl.delete(&run), "kubectl delete failed"); diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs index d34ce8e8c..a54340e91 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs @@ -672,7 +672,7 @@ mod tests { let actions = RefCell::new(Vec::new()); finalize_k8s_run_with_actions( - "power-of-two", + "power-of-n", true, false, || { @@ -694,7 +694,7 @@ mod tests { ) -> Vec<&'static str> { let actions = RefCell::new(Vec::new()); finalize_k8s_run_with_actions( - "power-of-two", + "power-of-n", run_failed, keep_resources_on_failure, || { @@ -806,7 +806,7 @@ mod tests { #[test] fn k8s_replay_artifacts_use_canonical_run_filenames() { let tempdir = tempfile::tempdir().expect("tempdir should create"); - let run_dir = tempdir.path().join("run-power-of-two"); + let run_dir = tempdir.path().join("run-power-of-n"); let artifacts = K8sRunArtifacts(&run_dir); @@ -1010,8 +1010,8 @@ traffic_pattern: output_tokens: { distribution: constant, value: 20 } arrival: { distribution: constant, interval_ms: 10 } algorithms: - - name: power-of-two - config: { default: power-of-two } + - name: power-of-n + config: { default: power-of-n } "#, ) .expect("benchmark config fixture should parse") @@ -1069,11 +1069,11 @@ algorithms: ) -> BenchmarkK8sRun { std::fs::create_dir_all(run_dir).expect("run dir should create"); BenchmarkK8sRun { - algorithm_name: "power-of-two".to_string(), + algorithm_name: "power-of-n".to_string(), manifest_path: run_dir.join("manifest.json"), run_dir: run_dir.to_path_buf(), - stargate_ns: "sgbench-sg-power-of-two".to_string(), - backends_ns: "sgbench-be-power-of-two".to_string(), + stargate_ns: "sgbench-sg-power-of-n".to_string(), + backends_ns: "sgbench-be-power-of-n".to_string(), stargate_count: 1, nodeport_host: "node.test".to_string(), stargate_http_endpoint: "http://node.test:30080".to_string(), diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/main.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/main.rs index 106cb6db6..11546d490 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/main.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/main.rs @@ -92,7 +92,7 @@ enum Command { Run(RunArgs), /// Compare Raw QUIC, HTTP/3, and WebTransport tunnel transports on loopback TransportBench(TransportBenchArgs), - /// Measure in-process wait-and-widen/pulsar load-balancer choose-path overhead + /// Measure in-process load-balancer choose-path overhead LbMicrobench { #[arg(long, default_value_t = 100_000, value_name = "N")] iterations: usize, @@ -566,7 +566,7 @@ mod tests { use clap::error::ErrorKind; const TRANSPORT_DISABLED: &str = "stargate-bench transport-bench --disable-quic-send-fairness --disable-http3-grease"; - const MATERIALIZE_ARGS: &str = "stargate-bench materialize --scenario uniform-4-backends --seed 7 --algorithm power-of-two --output-dir out"; + const MATERIALIZE_ARGS: &str = "stargate-bench materialize --scenario uniform-4-backends --seed 7 --algorithm power-of-n --output-dir out"; const PREPARE_ARGS: &str = "stargate-bench prepare-run --config config.yaml"; const RUN_ARGS: &str = "stargate-bench run --scenario uniform-4-backends --keep-resources-on-failure --reliability-mode controlled"; const STARGATE_REQUEST_METRIC: &str = concat!( @@ -612,7 +612,7 @@ mod tests { assert_eq!(conflict.kind(), ErrorKind::ArgumentConflict); assert_eq!( command_json!(MATERIALIZE_ARGS, Materialize), - r#"{"source":{"config":null,"scenario":"uniform-4-backends"},"seed":7,"algorithms":["power-of-two"],"output_dir":"out"}"# + r#"{"source":{"config":null,"scenario":"uniform-4-backends"},"seed":7,"algorithms":["power-of-n"],"output_dir":"out"}"# ); assert_eq!( command_json!(PREPARE_ARGS, PrepareRun), @@ -703,7 +703,7 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 concurrency: 1, candidates: 2, cache_key_count: 1, - scenarios: vec![LbMicrobenchScenario::PowerOfTwo], + scenarios: vec![LbMicrobenchScenario::PowerOfN], } .execute() .expect("real command should run lb microbench"); @@ -729,7 +729,7 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 #[test] fn run_info_reader_ignores_extra_run_metadata_fields() { let tempdir = tempfile::tempdir().expect("tempdir should create"); - let run_dir = tempdir.path().join("run-power-of-two"); + let run_dir = tempdir.path().join("run-power-of-n"); std::fs::create_dir(&run_dir).expect("run dir should create"); std::fs::write( run_dir.join("run-info.json"), @@ -830,7 +830,7 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 BenchmarkInputArgs { source: source(None, Some("uniform-4-backends")), seed: Some(123), - algorithms: ["random", "power-of-two"].map(str::to_string).to_vec(), + algorithms: ["random", "power-of-n"].map(str::to_string).to_vec(), output_dir, } } @@ -971,7 +971,7 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 .expect("benchmark input should load"); assert_eq!(manifest.seed, 123); - assert_eq!(algorithm_names(&config), ["power-of-two", "random"]); + assert_eq!(algorithm_names(&config), ["power-of-n", "random"]); assert_eq!(output_dir, Path::new(".bench-out").join(&config.name)); } @@ -991,11 +991,11 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 .expect("summary should load"); assert_eq!(manifest.seed, 123); - assert_eq!(algorithm_names(&config_copy), ["power-of-two", "random"]); + assert_eq!(algorithm_names(&config_copy), ["power-of-n", "random"]); assert_eq!(summary["seed"], 123); assert_eq!( summary["algorithm_names"], - serde_json::json!(["power-of-two", "random"]) + serde_json::json!(["power-of-n", "random"]) ); } @@ -1030,7 +1030,7 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 filter_algorithms(&mut config, &benchmark_args(None).algorithms) .expect("algorithm filter should succeed"); - assert_eq!(algorithm_names(&config), ["power-of-two", "random"]); + assert_eq!(algorithm_names(&config), ["power-of-n", "random"]); } #[test] @@ -1046,14 +1046,14 @@ pylon_requests_total_total{model="dummy-model",status="complete"} 3 #[test] fn load_balance_sweep_scenarios_cover_grouped_topologies_and_all_algorithms() { const STANDARD_ALGORITHMS: &[&str] = &[ - "power-of-two", + "power-of-n", "round-robin", "random", "wait-and-widen", "pulsar", ]; const PREFIX_ALGORITHMS: &[&str] = &[ - "power-of-two", + "power-of-n", "round-robin", "random", "wait-and-widen", diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs index 01ed8e926..258746d18 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs @@ -374,8 +374,8 @@ mod tests { }), degradation: crate::config::DegradationConfig::default(), algorithms: vec![AlgorithmConfig { - name: "power-of-two".to_string(), - config: serde_json::json!({ "default": "power-of-two" }), + name: "power-of-n".to_string(), + config: serde_json::json!({ "default": "power-of-n" }), pylon_queue_admission: None, }], } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs index 44f53e7b8..78a83abdc 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs @@ -26,6 +26,7 @@ use clap::ValueEnum; use stargate::load_balancer::{ LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerConfig, LoadBalancerModelConfig, LoadBalancerRequest, LoadBalancerRouter, LoadBalancerTargetState, + MAX_POWER_OF_N_SAMPLE_COUNT, }; use stargate::routing::{RoutedClusterSnapshot, RoutingTargetKey}; use stargate_proto::pb::{InferenceServerStatus, ModelStats}; @@ -53,7 +54,7 @@ macro_rules! scenarios { #[value(rename_all = "kebab-case")] pub enum LbMicrobenchScenario { $($scenario,)+ } - const LB_MICROBENCH_SCENARIOS: [LbMicrobenchScenarioMetadata; 19] = [$( + const LB_MICROBENCH_SCENARIOS: [LbMicrobenchScenarioMetadata; 23] = [$( LbMicrobenchScenarioMetadata { model_id: $model_id, algorithm: LoadBalancerAlgorithm::$algorithm, @@ -67,8 +68,12 @@ macro_rules! scenarios { // Keep the scenario matrix row-oriented so differences remain directly comparable. #[rustfmt::skip] scenarios! { - PowerOfTwo, "lb-bench-power-of-two", PowerOfTwo, None, 0; - PowerOfTwoOneExcluded, "lb-bench-power-of-two-one-excluded", PowerOfTwo, None, 1; + PowerOfN, "lb-bench-power-of-n", PowerOfN, None, 0; + PowerOfNSample1, "lb-bench-power-of-n-sample-1", PowerOfN, None, 0; + PowerOfNSample4, "lb-bench-power-of-n-sample-4", PowerOfN, None, 0; + PowerOfNSample8, "lb-bench-power-of-n-sample-8", PowerOfN, None, 0; + PowerOfNFullPool, "lb-bench-power-of-n-full-pool", PowerOfN, None, 0; + PowerOfNOneExcluded, "lb-bench-power-of-n-one-excluded", PowerOfN, None, 1; WaitAndWiden, "lb-bench-wait-and-widen", WaitAndWiden, None, 0; WaitAndWidenOneExcluded, "lb-bench-wait-and-widen-one-excluded", WaitAndWiden, None, 1; WaitAndWidenIgnoreQueue, "lb-bench-wait-and-widen-ignore-queue", WaitAndWiden, Some(IgnoreQueue), 0; @@ -92,6 +97,17 @@ impl LbMicrobenchScenario { fn metadata(self) -> &'static LbMicrobenchScenarioMetadata { &LB_MICROBENCH_SCENARIOS[self as usize] } + + fn configured_sample_count(self, candidate_count: usize) -> Option { + match self { + Self::PowerOfN | Self::PowerOfNOneExcluded => Some(2), + Self::PowerOfNSample1 => Some(1), + Self::PowerOfNSample4 => Some(4), + Self::PowerOfNSample8 => Some(8), + Self::PowerOfNFullPool => Some(candidate_count), + _ => None, + } + } } impl fmt::Display for LbMicrobenchScenario { @@ -156,13 +172,16 @@ pub fn write_lb_microbench_csv( ) -> std::io::Result<()> { writeln!( writer, - "methodology,scenario,candidates,iterations,warmup_iterations,concurrency,total_ns,ns_per_choose,choices,avg_rank_depth,selected_backend_count,top_backend,top_backend_choices,backend_counts,checksum" + "methodology,scenario,configured_sample_count,candidates,iterations,warmup_iterations,concurrency,total_ns,ns_per_choose,choices,avg_rank_depth,selected_backend_count,top_backend,top_backend_choices,backend_counts,checksum" )?; for row in rows { writeln!( writer, - "choose-only-v2,{},{},{},{},{},{},{:.1},{},{:.3},{},{},{},{},{}", + "choose-only-v3,{},{},{},{},{},{},{},{:.1},{},{:.3},{},{},{},{},{}", row.scenario, + row.scenario + .configured_sample_count(row.candidates) + .map_or_else(String::new, |count| count.to_string()), row.candidates, row.iterations, row.warmup_iterations, @@ -192,6 +211,16 @@ fn validate_config(config: &LbMicrobenchConfig) -> anyhow::Result<()> { anyhow::bail!("{flag} must be greater than 0"); } } + if config.candidates > MAX_POWER_OF_N_SAMPLE_COUNT + && (config.scenarios.is_empty() + || config + .scenarios + .contains(&LbMicrobenchScenario::PowerOfNFullPool)) + { + anyhow::bail!( + "power-of-n-full-pool requires --candidates to be at most {MAX_POWER_OF_N_SAMPLE_COUNT}" + ); + } Ok(()) } @@ -201,12 +230,13 @@ fn run_scenario( candidates: &[RoutedClusterSnapshot], cache_keys: &[String], ) -> anyhow::Result { + let algorithm_config = config_for_scenario(scenario, config.candidates); let router = LoadBalancerRouter::from_config(&LoadBalancerConfig { - default: LoadBalancerAlgorithm::PowerOfTwo, + default: LoadBalancerAlgorithm::PowerOfN, request_algorithms: HashMap::new(), models: HashMap::from([( scenario.metadata().model_id.to_string(), - LoadBalancerModelConfig::Detailed(Box::new(config_for_scenario(scenario))), + LoadBalancerModelConfig::Detailed(Box::new(algorithm_config)), )]), }) .with_context(|| format!("failed to build {scenario} load balancer"))?; @@ -456,12 +486,21 @@ fn format_backend_counts(backend_counts: &[(String, usize)]) -> String { .join(";") } -fn config_for_scenario(scenario: LbMicrobenchScenario) -> LoadBalancerAlgorithmConfig { +fn config_for_scenario( + scenario: LbMicrobenchScenario, + candidate_count: usize, +) -> LoadBalancerAlgorithmConfig { let metadata = scenario.metadata(); let mut config = LoadBalancerAlgorithmConfig::from(metadata.algorithm); let is_pulsar = metadata.algorithm == LoadBalancerAlgorithm::Pulsar; config.request_policy_mut().require_cache_affinity_key = is_pulsar; config.request_policy_mut().require_input_tokens = is_pulsar; + if let Some(sample_count) = scenario.configured_sample_count(candidate_count) { + config + .power_of_n_settings_mut() + .expect("power-of-n scenario should expose power-of-n settings") + .sample_count = sample_count; + } if is_pulsar { config .set_seed(Some("lb-microbench-seed".to_string())) @@ -558,15 +597,15 @@ mod tests { #[test] fn lb_microbench_runs_default_scenarios() { - let mut config = config(LbMicrobenchScenario::PowerOfTwo); + let mut config = config(LbMicrobenchScenario::PowerOfN); config.iterations = 8; config.warmup_iterations = 2; config.candidates = 4; config.scenarios.clear(); let rows = run_lb_microbench(&config).expect("microbench should run"); - assert_eq!(rows.len(), 19); - assert_eq!(rows[0].scenario, LbMicrobenchScenario::PowerOfTwo); + assert_eq!(rows.len(), 23); + assert_eq!(rows[0].scenario, LbMicrobenchScenario::PowerOfN); for row in rows { assert_eq!(row.choices, row.iterations); assert_eq!(row.concurrency, 2); @@ -609,6 +648,32 @@ mod tests { } } + #[test] + fn power_of_n_microbench_scenarios_cover_requested_sample_counts() { + for (scenario, expected_sample_count) in [ + (LbMicrobenchScenario::PowerOfNSample1, 1), + (LbMicrobenchScenario::PowerOfN, 2), + (LbMicrobenchScenario::PowerOfNSample4, 4), + (LbMicrobenchScenario::PowerOfNSample8, 8), + (LbMicrobenchScenario::PowerOfNFullPool, 8), + ] { + let rows = run_lb_microbench(&config(scenario)).expect("microbench should run"); + assert_eq!( + rows[0].scenario.configured_sample_count(rows[0].candidates), + Some(expected_sample_count) + ); + } + } + + #[test] + fn full_pool_microbench_rejects_candidate_count_above_sample_limit() { + let mut config = config(LbMicrobenchScenario::PowerOfNFullPool); + config.candidates = MAX_POWER_OF_N_SAMPLE_COUNT + 1; + + let error = run_lb_microbench(&config).expect_err("oversized full pool should fail"); + assert!(error.to_string().contains("power-of-n-full-pool")); + } + #[test] fn lb_microbench_rejects_zero_iterations() { let mut config = config(LbMicrobenchScenario::Pulsar); @@ -639,7 +704,7 @@ mod tests { #[rustfmt::skip] exclusion_tests! { random_one_excluded_scenario_never_selects_excluded_backend: RandomOneExcluded => ["cluster-0000"]; - power_of_two_one_excluded_scenario_never_selects_excluded_backend: PowerOfTwoOneExcluded => ["cluster-0000"]; + power_of_n_one_excluded_scenario_never_selects_excluded_backend: PowerOfNOneExcluded => ["cluster-0000"]; wait_and_widen_one_excluded_scenario_never_selects_excluded_backend: WaitAndWidenOneExcluded => ["cluster-0000"]; wait_and_widen_affinity_one_excluded_scenario_never_selects_excluded_backend: WaitAndWidenAffinityOneExcluded => ["cluster-0000"]; wait_and_widen_affinity_multi_excluded_scenario_never_selects_excluded_backend: WaitAndWidenAffinityMultiExcluded => ["cluster-0000", "cluster-0001"]; @@ -675,10 +740,13 @@ mod tests { write_lb_microbench_csv(&mut output, &rows).expect("csv should render"); let rendered = String::from_utf8(output).expect("csv should be utf8"); - assert!(rendered.starts_with("methodology,scenario,candidates,iterations")); + assert!( + rendered + .starts_with("methodology,scenario,configured_sample_count,candidates,iterations") + ); assert!(rendered.contains("selected_backend_count,top_backend,top_backend_choices")); assert!(rendered.contains( - "choose-only-v2,pulsar,8,10,2,4,1234,123.4,10,1.200,2,cluster-0001,7,cluster-0000:3;cluster-0001:7,42" + "choose-only-v3,pulsar,,8,10,2,4,1234,123.4,10,1.200,2,cluster-0001,7,cluster-0000:3;cluster-0001:7,42" )); } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs index f59462bc3..185127b72 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs @@ -361,7 +361,7 @@ traffic_pattern: output_tokens: { distribution: constant, value: 20 } arrival: { distribution: constant, interval_ms: 10 } algorithms: - - { name: power-of-two, config: { default: power-of-two } } + - { name: power-of-n, config: { default: power-of-n } } - { name: random, config: { default: random } } "#, ) @@ -429,7 +429,7 @@ algorithms: let run = prepared .algorithm_runs .iter() - .find(|run| run.algorithm_name == "power-of-two") + .find(|run| run.algorithm_name == "power-of-n") .expect("configured run should exist"); let run_info: serde_json::Value = serde_json::from_slice( &std::fs::read(&run.run_info_path).expect("run info should read"), @@ -448,7 +448,7 @@ algorithms: let compose = build_compose_spec( &config, &config.algorithms[0], - Path::new(".bench-out/prepare/run-power-of-two/lb-config.json"), + Path::new(".bench-out/prepare/run-power-of-n/lb-config.json"), STARGATE_GRPC_PORT, STARGATE_HTTP_PORT, STARGATE_METRICS_PORT, diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs index 5f94d5c86..0f2c9dd96 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs @@ -678,7 +678,7 @@ mod tests { ..CacheSummary::default() }; - let report = render(&config(), "power-of-two", None, summary); + let report = render(&config(), "power-of-n", None, summary); assert!(report.contains("| Algorithm | Admission Mode | Success |")); assert!(report.contains("Successful RPS")); @@ -753,13 +753,13 @@ mod tests { }]; let successful_report = render(&config(), "round-robin", None, successful_summary); - let failed_report = render(&config(), "power-of-two", None, failed_summary); + let failed_report = render(&config(), "power-of-n", None, failed_summary); assert!(!successful_report.contains("## Failures")); assert!(failed_report.contains("## Failures")); assert!(failed_report.contains("| Algorithm | Status | Backend | Count | Error |")); assert!( - failed_report.contains("| power-of-two | 503 | backend-a | 2 | upstream unavailable |") + failed_report.contains("| power-of-n | 503 | backend-a | 2 | upstream unavailable |") ); } diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs index eddbea1fd..fe456939f 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs @@ -279,7 +279,7 @@ mod tests { #[test] fn proxy_valid_configured_routing_method_uses_request_algorithm() { let lb_router = LoadBalancerRouter::from_config(&LoadBalancerConfig { - default: LoadBalancerAlgorithm::PowerOfTwo, + default: LoadBalancerAlgorithm::PowerOfN, request_algorithms: HashMap::from([( LoadBalancerAlgorithm::RoundRobin, LoadBalancerModelConfig::Name(LoadBalancerAlgorithm::RoundRobin), diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs index cd4b49c86..e5c4ad3ef 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs @@ -286,7 +286,7 @@ mod tests { let mut candidate = cluster_candidate("cluster-a"); candidate.stats.queued_input_size = 300; candidate.stats.last_mean_input_tps = 100.0; - let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfTwo); + let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); let target = routing_target(); let request = input_work_admission_request(&target, 50); @@ -302,7 +302,7 @@ mod tests { candidate.stats.total_query_input_size = 300; candidate.stats.queued_input_size = 0; candidate.stats.last_mean_input_tps = 100.0; - let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfTwo); + let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); let target = routing_target(); let request = input_work_admission_request(&target, 50); @@ -316,7 +316,7 @@ mod tests { fn input_work_admission_rejects_pool_without_valid_capacity() { let mut candidate = cluster_candidate("cluster-a"); candidate.stats.last_mean_input_tps = 0.0; - let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfTwo); + let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); let target = routing_target(); let request = input_work_admission_request(&target, 50); diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs index 9cb040cf6..ee02dcecd 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs @@ -465,7 +465,7 @@ mod tests { rank_depth: 1, selected_after_kv_free_tokens_skip: false, }, - effective_algorithm: LoadBalancerAlgorithm::PowerOfTwo, + effective_algorithm: LoadBalancerAlgorithm::PowerOfN, requested_algorithm: None, }; let selected_cluster = SelectedClusterRun::new( @@ -479,7 +479,7 @@ mod tests { let body = metrics_text(&app.metrics); assert!( body.contains( - r#"stargate_routing_selections_total{algorithm="power-of-two",model="model-a",routing_key="tenant-a",selection="primary"} 1"# + r#"stargate_routing_selections_total{algorithm="power-of-n",model="model-a",routing_key="tenant-a",selection="primary"} 1"# ), "selected cluster should preserve routing selection metric labels, got:\n{body}" ); diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs index 83699f263..987d3be8c 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs @@ -57,6 +57,8 @@ pub(super) fn proxy_openai_request_span(headers: &HeaderMap) -> Span { selected_inst.snapshot_age_ms = field::Empty, routing.algorithm = field::Empty, routing.num_candidates = field::Empty, + routing.sample_count_configured = field::Empty, + routing.sample_count_effective = field::Empty, routing.rank_depth = field::Empty, routing.selected_after_kv_free_tokens_skip = field::Empty, routing.retry_attempts = field::Empty, @@ -137,6 +139,7 @@ pub(super) fn record_routing_to_span(span: &Span, routing: RoutingTraceFields<'_ #[cfg(test)] mod tests { + use super::proxy_openai_request_span; use crate::telemetry::parent_context_from_headers; use axum::http::{HeaderMap, HeaderName, HeaderValue}; use opentelemetry::global; @@ -164,4 +167,16 @@ mod tests { ); assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); } + + #[test] + fn proxy_span_declares_load_balancer_sample_count_fields() { + let span = proxy_openai_request_span(&HeaderMap::new()); + let fields = span + .metadata() + .expect("proxy span should have static metadata") + .fields(); + + assert!(fields.field("routing.sample_count_configured").is_some()); + assert!(fields.field("routing.sample_count_effective").is_some()); + } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs index 9770c2178..c9603a5a7 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs @@ -49,7 +49,14 @@ impl<'de> Deserialize<'de> for LoadBalancerModelConfig { #[serde(rename_all = "kebab-case")] pub enum LoadBalancerAlgorithm { #[default] - PowerOfTwo, + #[serde( + alias = "power-of-two", + alias = "powerOf2", + alias = "powerOfN", + alias = "powerof2", + alias = "powerofn" + )] + PowerOfN, #[serde(alias = "groq-multiregion")] WaitAndWiden, RoundRobin, @@ -61,7 +68,7 @@ pub enum LoadBalancerAlgorithm { impl LoadBalancerAlgorithm { pub const ALL: [Self; 6] = [ - Self::PowerOfTwo, + Self::PowerOfN, Self::WaitAndWiden, Self::RoundRobin, Self::Random, @@ -73,7 +80,7 @@ impl LoadBalancerAlgorithm { impl fmt::Display for LoadBalancerAlgorithm { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let name = match self { - Self::PowerOfTwo => "power-of-two", + Self::PowerOfN => "power-of-n", Self::WaitAndWiden => "wait-and-widen", Self::RoundRobin => "round-robin", Self::Random => "random", @@ -194,10 +201,38 @@ pub struct WaitAndWidenAlgorithmConfig { pub ignore_input_processing_time: Option, } -#[derive(Debug, Clone, Default, PartialEq)] +const DEFAULT_POWER_OF_N_SAMPLE_COUNT: usize = 2; +pub const MAX_POWER_OF_N_SAMPLE_COUNT: usize = 64; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +#[serde(default)] +pub struct PowerOfNAlgorithmConfig { + pub sample_count: usize, +} + +impl Default for PowerOfNAlgorithmConfig { + fn default() -> Self { + Self { + sample_count: DEFAULT_POWER_OF_N_SAMPLE_COUNT, + } + } +} + +impl PowerOfNAlgorithmConfig { + pub(crate) fn validated_sample_count(&self) -> Result { + let sample_count = self.sample_count; + if !(1..=MAX_POWER_OF_N_SAMPLE_COUNT).contains(&sample_count) { + return Err(format!( + "power-of-n sample_count must be between 1 and {MAX_POWER_OF_N_SAMPLE_COUNT}, got {sample_count}" + )); + } + Ok(sample_count) + } +} + +#[derive(Debug, Clone, PartialEq)] pub enum LoadBalancerAlgorithmSettings { - #[default] - PowerOfTwo, + PowerOfN(PowerOfNAlgorithmConfig), WaitAndWiden(WaitAndWidenAlgorithmConfig), RoundRobin, Random, @@ -205,10 +240,16 @@ pub enum LoadBalancerAlgorithmSettings { PulsarWaitAndWiden(WaitAndWidenAlgorithmConfig), } +impl Default for LoadBalancerAlgorithmSettings { + fn default() -> Self { + Self::PowerOfN(PowerOfNAlgorithmConfig::default()) + } +} + impl LoadBalancerAlgorithmSettings { fn algorithm(&self) -> LoadBalancerAlgorithm { match self { - Self::PowerOfTwo => LoadBalancerAlgorithm::PowerOfTwo, + Self::PowerOfN(_) => LoadBalancerAlgorithm::PowerOfN, Self::WaitAndWiden(_) => LoadBalancerAlgorithm::WaitAndWiden, Self::RoundRobin => LoadBalancerAlgorithm::RoundRobin, Self::Random => LoadBalancerAlgorithm::Random, @@ -252,7 +293,7 @@ impl LoadBalancerAlgorithmConfig { LoadBalancerAlgorithmSettings::Pulsar(seed) => seed.as_deref(), LoadBalancerAlgorithmSettings::WaitAndWiden(config) | LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(config) => config.seed.as_deref(), - LoadBalancerAlgorithmSettings::PowerOfTwo + LoadBalancerAlgorithmSettings::PowerOfN(_) | LoadBalancerAlgorithmSettings::RoundRobin | LoadBalancerAlgorithmSettings::Random => None, } @@ -274,7 +315,7 @@ impl LoadBalancerAlgorithmConfig { config.seed = seed; Ok(()) } - LoadBalancerAlgorithmSettings::PowerOfTwo + LoadBalancerAlgorithmSettings::PowerOfN(_) | LoadBalancerAlgorithmSettings::RoundRobin | LoadBalancerAlgorithmSettings::Random => { Err(LoadBalancerSeedError::Unsupported { algorithm }) @@ -286,7 +327,7 @@ impl LoadBalancerAlgorithmConfig { match &self.settings { LoadBalancerAlgorithmSettings::WaitAndWiden(config) | LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(config) => Some(config), - LoadBalancerAlgorithmSettings::PowerOfTwo + LoadBalancerAlgorithmSettings::PowerOfN(_) | LoadBalancerAlgorithmSettings::RoundRobin | LoadBalancerAlgorithmSettings::Random | LoadBalancerAlgorithmSettings::Pulsar(_) => None, @@ -297,12 +338,34 @@ impl LoadBalancerAlgorithmConfig { match &mut self.settings { LoadBalancerAlgorithmSettings::WaitAndWiden(config) | LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(config) => Some(config), - LoadBalancerAlgorithmSettings::PowerOfTwo + LoadBalancerAlgorithmSettings::PowerOfN(_) | LoadBalancerAlgorithmSettings::RoundRobin | LoadBalancerAlgorithmSettings::Random | LoadBalancerAlgorithmSettings::Pulsar(_) => None, } } + + pub fn power_of_n_settings(&self) -> Option<&PowerOfNAlgorithmConfig> { + match &self.settings { + LoadBalancerAlgorithmSettings::PowerOfN(config) => Some(config), + LoadBalancerAlgorithmSettings::WaitAndWiden(_) + | LoadBalancerAlgorithmSettings::RoundRobin + | LoadBalancerAlgorithmSettings::Random + | LoadBalancerAlgorithmSettings::Pulsar(_) + | LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(_) => None, + } + } + + pub fn power_of_n_settings_mut(&mut self) -> Option<&mut PowerOfNAlgorithmConfig> { + match &mut self.settings { + LoadBalancerAlgorithmSettings::PowerOfN(config) => Some(config), + LoadBalancerAlgorithmSettings::WaitAndWiden(_) + | LoadBalancerAlgorithmSettings::RoundRobin + | LoadBalancerAlgorithmSettings::Random + | LoadBalancerAlgorithmSettings::Pulsar(_) + | LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(_) => None, + } + } } impl From for LoadBalancerAlgorithmConfig { @@ -317,7 +380,7 @@ impl From for LoadBalancerAlgorithmConfig { impl From for LoadBalancerAlgorithmSettings { fn from(algorithm: LoadBalancerAlgorithm) -> Self { match algorithm { - LoadBalancerAlgorithm::PowerOfTwo => Self::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN => Self::PowerOfN(Default::default()), LoadBalancerAlgorithm::WaitAndWiden => Self::WaitAndWiden(Default::default()), LoadBalancerAlgorithm::RoundRobin => Self::RoundRobin, LoadBalancerAlgorithm::Random => Self::Random, @@ -374,7 +437,22 @@ impl RawCommonAlgorithmConfig { #[derive(Debug, Deserialize)] #[serde(tag = "algorithm", rename_all = "kebab-case")] enum RawLoadBalancerAlgorithmConfig { - PowerOfTwo(RawCommonAlgorithmConfig), + // Technical debt: these aliases duplicate LoadBalancerAlgorithm because serde parses this + // tagged detailed-config enum independently. Keep both lists in sync until NVIDIA/nvcf#831 + // centralizes algorithm-name parsing. + #[serde( + alias = "power-of-two", + alias = "powerOf2", + alias = "powerOfN", + alias = "powerof2", + alias = "powerofn" + )] + PowerOfN { + #[serde(flatten)] + settings: PowerOfNAlgorithmConfig, + #[serde(flatten)] + common: RawCommonAlgorithmConfig, + }, #[serde(alias = "groq-multiregion")] WaitAndWiden { #[serde(flatten)] @@ -409,7 +487,11 @@ impl RawLoadBalancerAlgorithmConfig { Option, ) { match self { - Self::PowerOfTwo(common) => (common, LoadBalancerAlgorithmSettings::PowerOfTwo, None), + Self::PowerOfN { settings, common } => ( + common, + LoadBalancerAlgorithmSettings::PowerOfN(settings), + None, + ), Self::WaitAndWiden { settings, common } => ( common, LoadBalancerAlgorithmSettings::WaitAndWiden(settings), @@ -440,6 +522,9 @@ impl RawLoadBalancerAlgorithmConfig { fn into_config(self) -> Result { let (common, settings, consider_kv_free_tokens) = self.normalized(); + if let LoadBalancerAlgorithmSettings::PowerOfN(config) = &settings { + config.validated_sample_count()?; + } common.into_config(settings, consider_kv_free_tokens) } } @@ -476,7 +561,7 @@ pub struct LoadBalancerConfig { } impl LoadBalancerConfig { - /// Config used when no config file is given: `power-of-two` by default, + /// Config used when no config file is given: `power-of-n` by default, /// and any built-in algorithm can be selected per request. pub fn permissive_default() -> Self { Self { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/factory.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/factory.rs index 8dbd9fa46..197fe4764 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/factory.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/factory.rs @@ -15,7 +15,7 @@ use std::sync::Arc; -use super::power_of_two::PowerOfTwoLoadBalancer; +use super::power_of_n::PowerOfNLoadBalancer; use super::pulsar::PulsarLoadBalancer; use super::pulsar_wait_and_widen::PulsarWaitAndWidenLoadBalancer; use super::random::RandomLoadBalancer; @@ -38,7 +38,9 @@ pub fn create_load_balancer_with_config( } match config.algorithm() { - LoadBalancerAlgorithm::PowerOfTwo => Ok(Arc::new(PowerOfTwoLoadBalancer)), + LoadBalancerAlgorithm::PowerOfN => Ok(Arc::new( + PowerOfNLoadBalancer::from_algorithm_config(config)?, + )), LoadBalancerAlgorithm::WaitAndWiden => Ok(Arc::new(WaitAndWidenLoadBalancer::new( WaitAndWidenConfig::from_algorithm_config(config), ))), diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs index 02bd8bec6..0259ab147 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs @@ -26,7 +26,7 @@ macro_rules! impl_display { mod algorithm; mod config; mod factory; -mod power_of_two; +mod power_of_n; mod pulsar; mod pulsar_wait_and_widen; mod random; @@ -45,7 +45,7 @@ pub use config::{ LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerAlgorithmOverride, LoadBalancerAlgorithmSettings, LoadBalancerConfig, LoadBalancerModelConfig, LoadBalancerRequestPolicy, LoadBalancerRoutingAlgorithmError, LoadBalancerSeedError, - WaitAndWidenAlgorithmConfig, + MAX_POWER_OF_N_SAMPLE_COUNT, PowerOfNAlgorithmConfig, WaitAndWidenAlgorithmConfig, }; pub use factory::create_load_balancer_with_config; pub use request::{LoadBalancerCandidateChoice, LoadBalancerRequest}; diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs new file mode 100644 index 000000000..4ee8cc6b0 --- /dev/null +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs @@ -0,0 +1,468 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use rand::Rng; +use rand::seq::IteratorRandom; +use stargate_protocol::common::valid_last_mean_input_tps; +use tracing::{Span, debug}; + +#[cfg(test)] +use super::tests::LoadBalancerTestChoiceExt; +use super::{ + LoadBalancer, LoadBalancerAlgorithmConfig, LoadBalancerCandidateChoice, LoadBalancerRequest, + MAX_POWER_OF_N_SAMPLE_COUNT, +}; +use crate::routing_state::RoutedClusterSnapshot; + +pub(super) struct PowerOfNLoadBalancer { + sample_count: usize, +} + +impl PowerOfNLoadBalancer { + pub(super) fn from_algorithm_config( + config: &LoadBalancerAlgorithmConfig, + ) -> anyhow::Result { + let settings = config + .power_of_n_settings() + .expect("power-of-n settings should match load-balancer algorithm"); + let sample_count = settings + .validated_sample_count() + .map_err(anyhow::Error::msg)?; + Ok(Self { sample_count }) + } + + fn choose_candidate_with_rng( + &self, + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], + rng: &mut R, + ) -> Option { + let sampled = sample_candidates(request, candidates, self.sample_count, rng); + let span = Span::current(); + span.record("routing.sample_count_configured", self.sample_count); + span.record("routing.sample_count_effective", sampled.len()); + + choose_least_loaded(candidates, sampled.as_slice(), request.input_tokens, rng) + } +} + +impl_display!(PowerOfNLoadBalancer, "power-of-n"); + +impl LoadBalancer for PowerOfNLoadBalancer { + fn choose_candidate( + &self, + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], + ) -> Option { + let mut rng = rand::rng(); + self.choose_candidate_with_rng(request, candidates, &mut rng) + } +} + +struct CandidateSample { + indices: [usize; MAX_POWER_OF_N_SAMPLE_COUNT], + len: usize, +} + +impl CandidateSample { + fn new() -> Self { + Self { + indices: [0; MAX_POWER_OF_N_SAMPLE_COUNT], + len: 0, + } + } + + fn as_slice(&self) -> &[usize] { + &self.indices[..self.len] + } + + fn len(&self) -> usize { + self.len + } +} + +fn sample_candidates( + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], + sample_count: usize, + rng: &mut R, +) -> CandidateSample { + let mut sampled = CandidateSample::new(); + if candidates.is_empty() { + return sampled; + } + + if !request.has_excluded_clusters() { + sampled.len = sample_count.min(candidates.len()); + match sampled.len { + 0 => {} + 1 => sampled.indices[0] = rng.random_range(0..candidates.len()), + 2 => { + let (first, second) = sample_distinct_pair(candidates.len(), rng); + sampled.indices[..2].copy_from_slice(&[first, second]); + } + len if len == candidates.len() => { + for (index, slot) in sampled.indices[..len].iter_mut().enumerate() { + *slot = index; + } + } + len => { + let filled = + (0..candidates.len()).choose_multiple_fill(rng, &mut sampled.indices[..len]); + debug_assert_eq!(filled, len); + } + } + return sampled; + } + + let eligible_indices = candidates + .iter() + .enumerate() + .filter(|(_, candidate)| !request.excludes_cluster(&candidate.cluster_id)) + .map(|(index, _)| index); + sampled.len = eligible_indices.choose_multiple_fill(rng, &mut sampled.indices[..sample_count]); + sampled +} + +fn sample_distinct_pair(len: usize, rng: &mut R) -> (usize, usize) { + let a_index = rng.random_range(0..len); + let mut b_index = rng.random_range(0..len - 1); + if b_index >= a_index { + b_index += 1; + } + (a_index, b_index) +} + +fn choose_least_loaded( + candidates: &[RoutedClusterSnapshot], + sampled_indices: &[usize], + input_tokens: Option, + rng: &mut R, +) -> Option { + let (&first_index, remaining_indices) = sampled_indices.split_first()?; + let mut selected_index = first_index; + let mut selected_score = load_score(&candidates[first_index], input_tokens); + let mut tied_best_count = 1u32; + + for &candidate_index in remaining_indices { + let score = load_score(&candidates[candidate_index], input_tokens); + if score < selected_score { + selected_index = candidate_index; + selected_score = score; + tied_best_count = 1; + } else if score == selected_score { + tied_best_count += 1; + if rng.random_ratio(1, tied_best_count) { + selected_index = candidate_index; + } + } + } + + debug!( + effective_sample_count = sampled_indices.len(), + selected_candidate_index = selected_index, + selected_load_score = selected_score, + "sampled clusters" + ); + Some(LoadBalancerCandidateChoice::with_rank_depth_1( + selected_index, + )) +} + +fn load_score(candidate: &RoutedClusterSnapshot, input_tokens: Option) -> f64 { + let last_mean_input_tps = candidate.stats.last_mean_input_tps; + if valid_last_mean_input_tps(last_mean_input_tps) { + (super::input_work_units(candidate) + input_tokens.unwrap_or_default() as f64) + / last_mean_input_tps + } else { + f64::INFINITY + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::time::{Duration, Instant}; + + use rand::SeedableRng; + use rand::rngs::StdRng; + use stargate_proto::pb::{InferenceServerStatus, ModelStats}; + + use super::*; + + fn candidate( + id: &str, + last_mean_input_tps: f64, + queued_input_size: u64, + ) -> RoutedClusterSnapshot { + RoutedClusterSnapshot { + cluster_id: id.to_string(), + stats: ModelStats { + last_mean_input_tps, + queued_input_size, + ..ModelStats::default() + }, + rtt: Duration::from_millis(1), + snapshot_updated_at: Instant::now(), + status: InferenceServerStatus::Active, + active_backend_count: 1, + } + } + + fn selected_cluster_id( + sample_count: usize, + candidates: &[RoutedClusterSnapshot], + excluded_cluster_ids: &HashSet, + ) -> Option { + let target = crate::routing_state::RoutingTargetKey::new(None, "model-a"); + let request = LoadBalancerRequest { + routing_target: &target, + cache_affinity_key: None, + input_tokens: Some(1000), + priority: 0, + received_at: Instant::now(), + request_slo: None, + excluded_cluster_ids: Some(excluded_cluster_ids), + }; + PowerOfNLoadBalancer { sample_count } + .choose_for_test(&request, candidates) + .map(|choice| choice.candidate.cluster_id) + } + + fn sampled_indices( + sample_count: usize, + candidates: &[RoutedClusterSnapshot], + excluded_cluster_ids: Option<&HashSet>, + rng: &mut StdRng, + ) -> Vec { + let target = crate::routing_state::RoutingTargetKey::new(None, "model-a"); + let request = LoadBalancerRequest { + routing_target: &target, + cache_affinity_key: None, + input_tokens: Some(1000), + priority: 0, + received_at: Instant::now(), + request_slo: None, + excluded_cluster_ids, + }; + sample_candidates(&request, candidates, sample_count, rng) + .as_slice() + .to_vec() + } + + #[test] + fn load_score_prefers_faster_empty_backend_for_incoming_prefill() { + let fast = candidate("fast", 200.0, 0); + let slow = candidate("slow", 100.0, 0); + + assert!(load_score(&fast, Some(1000)) < load_score(&slow, Some(1000))); + } + + #[test] + fn load_score_accounts_for_queued_prefill_work() { + let busy_fast = candidate("busy-fast", 200.0, 10_000); + let empty_slow = candidate("empty-slow", 100.0, 0); + + assert!(load_score(&empty_slow, Some(1000)) < load_score(&busy_fast, Some(1000))); + } + + #[test] + fn load_score_rejects_invalid_input_throughput() { + for input_tps in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!(load_score(&candidate("invalid", input_tps, 0), Some(100)).is_infinite()); + } + } + + #[test] + fn power_of_n_never_selects_excluded_clusters() { + let candidates = vec![ + candidate("excluded-a", 1_000.0, 0), + candidate("eligible", 1.0, 0), + candidate("excluded-b", 1_000.0, 0), + ]; + let excluded = HashSet::from(["excluded-a".to_string(), "excluded-b".to_string()]); + + for _ in 0..64 { + assert_eq!( + selected_cluster_id(2, &candidates, &excluded).as_deref(), + Some("eligible") + ); + } + } + + #[test] + fn power_of_n_skips_single_excluded_cluster_in_retry_set() { + let candidates = (0..64) + .map(|index| candidate(&format!("cluster-{index:04}"), 1_000.0, 0)) + .collect::>(); + let excluded = HashSet::from(["cluster-0000".to_string()]); + + for _ in 0..512 { + let selected = selected_cluster_id(2, &candidates, &excluded) + .expect("an eligible cluster should be selected"); + assert_ne!(selected, "cluster-0000"); + } + } + + #[test] + fn power_of_n_returns_none_when_all_candidates_are_excluded() { + let candidates = vec![ + candidate("excluded-a", 1_000.0, 0), + candidate("excluded-b", 1_000.0, 0), + ]; + let excluded = HashSet::from(["excluded-a".to_string(), "excluded-b".to_string()]); + + assert!(selected_cluster_id(2, &candidates, &excluded).is_none()); + } + + #[test] + fn sampling_handles_empty_singleton_and_oversized_pools() { + let one = vec![candidate("only", 1.0, 0)]; + let three = vec![ + candidate("a", 1.0, 0), + candidate("b", 1.0, 0), + candidate("c", 1.0, 0), + ]; + let mut rng = StdRng::seed_from_u64(7); + + assert!(sampled_indices(4, &[], None, &mut rng).is_empty()); + assert_eq!(sampled_indices(4, &one, None, &mut rng), [0]); + assert_eq!(sampled_indices(3, &three, None, &mut rng), [0, 1, 2]); + assert_eq!(sampled_indices(8, &three, None, &mut rng), [0, 1, 2]); + } + + #[test] + fn sampling_is_distinct_and_reproducible_for_counts_one_two_and_four() { + let candidates = (0..8) + .map(|index| candidate(&format!("cluster-{index}"), 1.0, 0)) + .collect::>(); + + for sample_count in [1, 2, 4] { + let mut first_rng = StdRng::seed_from_u64(42); + let mut second_rng = StdRng::seed_from_u64(42); + let first = sampled_indices(sample_count, &candidates, None, &mut first_rng); + let second = sampled_indices(sample_count, &candidates, None, &mut second_rng); + let distinct = first.iter().copied().collect::>(); + + assert_eq!(first, second); + assert_eq!(first.len(), sample_count); + assert_eq!(distinct.len(), sample_count); + } + } + + #[test] + fn sampling_applies_sparse_and_dense_exclusions_before_sampling() { + let candidates = (0..8) + .map(|index| candidate(&format!("cluster-{index}"), 1.0, 0)) + .collect::>(); + + for excluded in [ + HashSet::from(["cluster-0".to_string()]), + HashSet::from([ + "cluster-0".to_string(), + "cluster-1".to_string(), + "cluster-2".to_string(), + "cluster-3".to_string(), + "cluster-4".to_string(), + "cluster-5".to_string(), + ]), + ] { + let mut rng = StdRng::seed_from_u64(91); + let sampled = sampled_indices(4, &candidates, Some(&excluded), &mut rng); + let expected_count = 4.min(candidates.len() - excluded.len()); + + assert_eq!(sampled.len(), expected_count); + assert!( + sampled + .iter() + .all(|index| !excluded.contains(&candidates[*index].cluster_id)) + ); + assert_eq!( + sampled.iter().copied().collect::>().len(), + expected_count + ); + } + } + + #[test] + fn sampling_is_uniform_without_replacement() { + let candidates = (0..4) + .map(|index| candidate(&format!("cluster-{index}"), 1.0, 0)) + .collect::>(); + let mut rng = StdRng::seed_from_u64(73); + let mut inclusion_counts = [0usize; 4]; + const ITERATIONS: usize = 20_000; + + for _ in 0..ITERATIONS { + for index in sampled_indices(2, &candidates, None, &mut rng) { + inclusion_counts[index] += 1; + } + } + + let expected = ITERATIONS / 2; + let tolerance = expected / 25; + for count in inclusion_counts { + assert!(count.abs_diff(expected) <= tolerance, "count={count}"); + } + } + + #[test] + fn full_pool_selection_chooses_the_lowest_score() { + let candidates = vec![ + candidate("slow", 100.0, 1000), + candidate("best", 1000.0, 0), + candidate("busy", 1000.0, 10_000), + ]; + let excluded = HashSet::new(); + + assert_eq!( + selected_cluster_id(3, &candidates, &excluded).as_deref(), + Some("best") + ); + } + + #[test] + fn equal_scores_do_not_bias_ties_to_one_candidate() { + let candidates = vec![ + candidate("a", 100.0, 0), + candidate("b", 100.0, 0), + candidate("c", 100.0, 0), + ]; + let target = crate::routing_state::RoutingTargetKey::new(None, "model-a"); + let request = LoadBalancerRequest { + routing_target: &target, + cache_affinity_key: None, + input_tokens: Some(1000), + priority: 0, + received_at: Instant::now(), + request_slo: None, + excluded_cluster_ids: None, + }; + let load_balancer = PowerOfNLoadBalancer { sample_count: 3 }; + let mut rng = StdRng::seed_from_u64(11); + let mut selected = HashSet::new(); + + for _ in 0..64 { + selected.insert( + load_balancer + .choose_candidate_with_rng(&request, &candidates, &mut rng) + .expect("equal-score pool should produce a candidate") + .candidate_index, + ); + } + + assert_eq!(selected, HashSet::from([0, 1, 2])); + } +} diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_two.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_two.rs deleted file mode 100644 index 070aa7fad..000000000 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_two.rs +++ /dev/null @@ -1,313 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use rand::Rng; -use stargate_protocol::common::valid_last_mean_input_tps; -use tracing::debug; - -#[cfg(test)] -use super::tests::LoadBalancerTestChoiceExt; -use super::{LoadBalancer, LoadBalancerCandidateChoice, LoadBalancerRequest}; -use crate::routing_state::RoutedClusterSnapshot; - -const EXCLUSION_REJECTION_ATTEMPTS: usize = 8; -const REJECTION_SAMPLE_EXCLUSION_RATIO_DIVISOR: usize = 4; - -pub(super) struct PowerOfTwoLoadBalancer; - -impl_display!(PowerOfTwoLoadBalancer, "power-of-two"); - -impl LoadBalancer for PowerOfTwoLoadBalancer { - fn choose_candidate( - &self, - request: &LoadBalancerRequest<'_>, - candidates: &[RoutedClusterSnapshot], - ) -> Option { - let mut rng = rand::rng(); - let sampled = if request.has_excluded_clusters() { - sample_with_exclusions(request, candidates, &mut rng) - } else { - sample_without_exclusions(candidates, &mut rng) - }; - - match sampled { - CandidateSample::None => None, - CandidateSample::One(candidate_index) => Some( - LoadBalancerCandidateChoice::with_rank_depth_1(candidate_index), - ), - CandidateSample::Two(a_index, b_index) => Some(choose_less_loaded( - candidates, - a_index, - b_index, - request.input_tokens, - &mut rng, - )), - } - } -} - -enum CandidateSample { - None, - One(usize), - Two(usize, usize), -} - -fn sample_without_exclusions( - candidates: &[RoutedClusterSnapshot], - rng: &mut R, -) -> CandidateSample { - match candidates.len() { - 0 => CandidateSample::None, - 1 => CandidateSample::One(0), - len => { - // First-attempt proxy routing normally has no failed clusters. Pick two - // distinct indices directly so the default production algorithm is O(1) - // instead of scanning every candidate on each request. - let (a_index, b_index) = sample_distinct_pair(len, rng); - CandidateSample::Two(a_index, b_index) - } - } -} - -fn sample_distinct_pair(len: usize, rng: &mut R) -> (usize, usize) { - let a_index = rng.random_range(0..len); - let mut b_index = rng.random_range(0..len - 1); - if b_index >= a_index { - b_index += 1; - } - (a_index, b_index) -} - -fn sample_with_exclusions( - request: &LoadBalancerRequest<'_>, - candidates: &[RoutedClusterSnapshot], - rng: &mut R, -) -> CandidateSample { - sample_sparse_pair(request, candidates, rng) - .unwrap_or_else(|| sample_with_reservoir(request, candidates, rng)) -} - -fn sample_sparse_pair( - request: &LoadBalancerRequest<'_>, - candidates: &[RoutedClusterSnapshot], - rng: &mut R, -) -> Option { - let excluded = request.excluded_cluster_ids?; - if candidates.len() < 2 - || excluded.len() > candidates.len() / REJECTION_SAMPLE_EXCLUSION_RATIO_DIVISOR - { - return None; - } - // Sparse retries use bounded O(1) pair sampling before uniform reservoir fallback. - for _ in 0..EXCLUSION_REJECTION_ATTEMPTS { - let (a_index, b_index) = sample_distinct_pair(candidates.len(), rng); - if !request.excludes_cluster(&candidates[a_index].cluster_id) - && !request.excludes_cluster(&candidates[b_index].cluster_id) - { - return Some(CandidateSample::Two(a_index, b_index)); - } - } - None -} - -fn sample_with_reservoir( - request: &LoadBalancerRequest<'_>, - candidates: &[RoutedClusterSnapshot], - rng: &mut R, -) -> CandidateSample { - let mut sampled: [Option; 2] = [None, None]; - let mut eligible_seen = 0usize; - - for (candidate_index, candidate) in candidates.iter().enumerate() { - if request.excludes_cluster(&candidate.cluster_id) { - continue; - } - - eligible_seen += 1; - match eligible_seen { - 1 => sampled[0] = Some(candidate_index), - 2 => sampled[1] = Some(candidate_index), - _ => { - // Retry/failover attempts may exclude clusters. Reservoir sampling - // preserves the old uniform sample-without-replacement behavior for - // the remaining eligible set without allocating a filtered Vec. - let slot = rng.random_range(0..eligible_seen); - if slot < sampled.len() { - sampled[slot] = Some(candidate_index); - } - } - } - } - - match (sampled[0], sampled[1]) { - (None, _) => CandidateSample::None, - (Some(candidate), None) => CandidateSample::One(candidate), - (Some(a), Some(b)) => CandidateSample::Two(a, b), - } -} - -fn choose_less_loaded( - candidates: &[RoutedClusterSnapshot], - a_index: usize, - b_index: usize, - input_tokens: Option, - rng: &mut R, -) -> LoadBalancerCandidateChoice { - let a = &candidates[a_index]; - let b = &candidates[b_index]; - let a_score = load_score(a, input_tokens); - let b_score = load_score(b, input_tokens); - debug!( - inst_a = %a.cluster_id, - inst_b = %b.cluster_id, - load_score_a = a_score, - load_score_b = b_score, - "sampled two clusters" - ); - - let selected_index = if a_score < b_score { - a_index - } else if b_score < a_score { - b_index - } else if rng.random_bool(0.5) { - a_index - } else { - b_index - }; - LoadBalancerCandidateChoice::with_rank_depth_1(selected_index) -} - -fn load_score(candidate: &RoutedClusterSnapshot, input_tokens: Option) -> f64 { - let last_mean_input_tps = candidate.stats.last_mean_input_tps; - if valid_last_mean_input_tps(last_mean_input_tps) { - (super::input_work_units(candidate) + input_tokens.unwrap_or_default() as f64) - / last_mean_input_tps - } else { - f64::INFINITY - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashSet; - use std::time::{Duration, Instant}; - - use stargate_proto::pb::{InferenceServerStatus, ModelStats}; - - use super::*; - - fn candidate( - id: &str, - last_mean_input_tps: f64, - queued_input_size: u64, - ) -> RoutedClusterSnapshot { - RoutedClusterSnapshot { - cluster_id: id.to_string(), - stats: ModelStats { - last_mean_input_tps, - queued_input_size, - ..ModelStats::default() - }, - rtt: Duration::from_millis(1), - snapshot_updated_at: Instant::now(), - status: InferenceServerStatus::Active, - active_backend_count: 1, - } - } - - fn selected_cluster_id( - candidates: &[RoutedClusterSnapshot], - excluded_cluster_ids: &HashSet, - ) -> Option { - let target = crate::routing_state::RoutingTargetKey::new(None, "model-a"); - let request = LoadBalancerRequest { - routing_target: &target, - cache_affinity_key: None, - input_tokens: Some(1000), - priority: 0, - received_at: Instant::now(), - request_slo: None, - excluded_cluster_ids: Some(excluded_cluster_ids), - }; - PowerOfTwoLoadBalancer - .choose_for_test(&request, candidates) - .map(|choice| choice.candidate.cluster_id) - } - - #[test] - fn load_score_prefers_faster_empty_backend_for_incoming_prefill() { - let fast = candidate("fast", 200.0, 0); - let slow = candidate("slow", 100.0, 0); - - assert!(load_score(&fast, Some(1000)) < load_score(&slow, Some(1000))); - } - - #[test] - fn load_score_accounts_for_queued_prefill_work() { - let busy_fast = candidate("busy-fast", 200.0, 10_000); - let empty_slow = candidate("empty-slow", 100.0, 0); - - assert!(load_score(&empty_slow, Some(1000)) < load_score(&busy_fast, Some(1000))); - } - - #[test] - fn load_score_rejects_invalid_input_throughput() { - for input_tps in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { - assert!(load_score(&candidate("invalid", input_tps, 0), Some(100)).is_infinite()); - } - } - - #[test] - fn power_of_two_never_selects_excluded_clusters() { - let candidates = vec![ - candidate("excluded-a", 1_000.0, 0), - candidate("eligible", 1.0, 0), - candidate("excluded-b", 1_000.0, 0), - ]; - let excluded = HashSet::from(["excluded-a".to_string(), "excluded-b".to_string()]); - - for _ in 0..64 { - assert_eq!( - selected_cluster_id(&candidates, &excluded).as_deref(), - Some("eligible") - ); - } - } - - #[test] - fn power_of_two_skips_single_excluded_cluster_in_retry_set() { - let candidates = (0..64) - .map(|index| candidate(&format!("cluster-{index:04}"), 1_000.0, 0)) - .collect::>(); - let excluded = HashSet::from(["cluster-0000".to_string()]); - - for _ in 0..512 { - let selected = selected_cluster_id(&candidates, &excluded) - .expect("an eligible cluster should be selected"); - assert_ne!(selected, "cluster-0000"); - } - } - - #[test] - fn power_of_two_returns_none_when_all_candidates_are_excluded() { - let candidates = vec![ - candidate("excluded-a", 1_000.0, 0), - candidate("excluded-b", 1_000.0, 0), - ]; - let excluded = HashSet::from(["excluded-a".to_string(), "excluded-b".to_string()]); - - assert!(selected_cluster_id(&candidates, &excluded).is_none()); - } -} diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs index 3ba3c16e3..ee2916ddf 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs @@ -144,7 +144,7 @@ fn seeded_pulsar_algorithm_config(seed: &str) -> LoadBalancerAlgorithmConfig { #[test] fn set_seed_reports_unsupported_algorithms_without_panicking() { for algorithm in [ - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, LoadBalancerAlgorithm::RoundRobin, LoadBalancerAlgorithm::Random, ] { @@ -571,7 +571,7 @@ where fn assert_algorithm_overrides(raw: impl Fn(LoadBalancerAlgorithm) -> String) { for algorithm in [ LoadBalancerAlgorithm::WaitAndWiden, - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, LoadBalancerAlgorithm::Pulsar, LoadBalancerAlgorithm::PulsarWaitAndWiden, LoadBalancerAlgorithm::Random, @@ -604,7 +604,7 @@ fn simple_model_config_parses_to_algorithm_enum() { #[test] fn detailed_model_config_parses_input_work_admission_limit() { let config: LoadBalancerConfig = parse_json( - r#"{"models":{"model-a":{"algorithm":"power-of-two","max_input_work_seconds":2.5}}}"#, + r#"{"models":{"model-a":{"algorithm":"power-of-n","max_input_work_seconds":2.5}}}"#, ); let detailed = config @@ -714,17 +714,115 @@ fn algorithm_specific_load_balancer_fields_are_rejected_for_other_algorithms() { r#"{"algorithm":"wait-and-widen","consider_kv_free_tokens":true}"#, "consider_kv_free_tokens", ), + (r#"{"algorithm":"random","sample_count":4}"#, "sample_count"), ] { assert_json_rejected::(raw, expected_field); } } +#[test] +fn power_of_n_sample_count_defaults_to_two() { + let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); + let settings = config + .power_of_n_settings() + .expect("power-of-n config should expose settings"); + + assert_eq!(settings.sample_count, 2); +} + +#[test] +fn detailed_power_of_n_sample_count_parses_in_every_supported_context() { + let direct: LoadBalancerAlgorithmConfig = + parse_json(r#"{"algorithm":"power-of-n","sample_count":1}"#); + assert_eq!( + direct + .power_of_n_settings() + .expect("direct config should expose settings") + .sample_count, + 1 + ); + + let router = router_from_json( + r#"{"default":"random","request_algorithms":{"power-of-n":{"algorithm":"power-of-n","sample_count":4}},"models":{"model-a":{"algorithm":"power-of-n","sample_count":8,"request_algorithms":{"power-of-n":{"algorithm":"power-of-n","sample_count":64}}}}}"#, + ); + assert_eq!( + router + .algorithm_config("model-a") + .power_of_n_settings() + .expect("model config should expose settings") + .sample_count, + 8 + ); + + let override_header = LoadBalancerAlgorithmOverride::parse("power-of-n") + .expect("power-of-n override should parse"); + let model_override = router + .resolve_algorithm_override("model-a", Some(&override_header)) + .expect("model override should resolve"); + let default_override = router + .resolve_algorithm_override("model-b", Some(&override_header)) + .expect("top-level override should resolve"); + assert_eq!( + model_override + .config() + .power_of_n_settings() + .expect("model override should expose settings") + .sample_count, + 8, + "the configured model algorithm takes precedence over its same-algorithm override" + ); + assert_eq!( + default_override + .config() + .power_of_n_settings() + .expect("top-level override should expose settings") + .sample_count, + 4 + ); + + let nested_router = router_from_json( + r#"{"default":"random","models":{"model-a":{"algorithm":"random","request_algorithms":{"power-of-n":{"algorithm":"power-of-n","sample_count":64}}}}}"#, + ); + let nested_override = nested_router + .resolve_algorithm_override("model-a", Some(&override_header)) + .expect("nested override should resolve"); + assert_eq!( + nested_override + .config() + .power_of_n_settings() + .expect("nested override should expose settings") + .sample_count, + 64 + ); +} + +#[test] +fn invalid_power_of_n_sample_counts_are_rejected_with_field_context() { + for sample_count in [0, MAX_POWER_OF_N_SAMPLE_COUNT + 1] { + assert_json_rejected::( + &format!(r#"{{"algorithm":"power-of-n","sample_count":{sample_count}}}"#), + "power-of-n sample_count must be between 1 and 64", + ); + } + + let mut config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); + config + .power_of_n_settings_mut() + .expect("power-of-n config should expose mutable settings") + .sample_count = 0; + let error = match create_load_balancer_with_config(&config) { + Ok(_) => panic!("programmatic invalid sample count should fail"), + Err(error) => error, + }; + assert!(error.to_string().contains("power-of-n sample_count")); +} + #[test] fn detailed_algorithm_configs_preserve_all_variant_identities() { use LoadBalancerAlgorithm::*; for (raw, expected, expected_seed, considers_kv_free_tokens) in [ - (r#"{"algorithm":"power-of-two"}"#, PowerOfTwo, None, false), + (r#"{"algorithm":"power-of-n"}"#, PowerOfN, None, false), ( r#"{"algorithm":"wait-and-widen","seed":"wait-and-widen-seed"}"#, WaitAndWiden, @@ -756,7 +854,7 @@ fn detailed_algorithm_configs_preserve_all_variant_identities() { #[test] fn unknown_load_balancer_config_fields_are_rejected() { assert_json_rejected::( - r#"{"default":"power-of-two","unused_top_level_field":true,"models":{"model-a":{"algorithm":"pulsar","unused_model_field":123}}}"#, + r#"{"default":"power-of-n","unused_top_level_field":true,"models":{"model-a":{"algorithm":"pulsar","unused_model_field":123}}}"#, "unused_top_level_field", ); } @@ -828,7 +926,7 @@ fn published_load_balancer_configuration_examples_parse() { #[test] fn detailed_model_config_parses_for_pulsar() { let router = router_from_json( - r#"{"default":"power-of-two","models":{"model-a":{"algorithm":"pulsar","seed":"seed-1","require_cache_affinity_key":true,"consider_kv_free_tokens":true}}}"#, + r#"{"default":"power-of-n","models":{"model-a":{"algorithm":"pulsar","seed":"seed-1","require_cache_affinity_key":true,"consider_kv_free_tokens":true}}}"#, ); let model_config = router.algorithm_config("model-a"); assert_eq!(model_config.algorithm(), LoadBalancerAlgorithm::Pulsar); @@ -849,7 +947,7 @@ fn kv_free_token_consideration_is_rejected_for_non_pulsar_algorithms() { #[test] fn detailed_model_config_parses_for_pulsar_wait_and_widen() { let router = router_from_json( - r#"{"default":"power-of-two","models":{"model-a":{"algorithm":"pulsar-wait-and-widen","seed":"seed-1","require_cache_affinity_key":true,"require_input_tokens":true,"max_queue_time_floor_ms":100,"max_queue_time_ceil_ms":100,"ttft_bucket_size_ms":50,"n":2}}}"#, + r#"{"default":"power-of-n","models":{"model-a":{"algorithm":"pulsar-wait-and-widen","seed":"seed-1","require_cache_affinity_key":true,"require_input_tokens":true,"max_queue_time_floor_ms":100,"max_queue_time_ceil_ms":100,"ttft_bucket_size_ms":50,"n":2}}}"#, ); let model_config = router.algorithm_config("model-a"); assert_eq!( @@ -907,7 +1005,7 @@ fn legacy_algorithm_names_remain_compatible_in_short_configs() { #[test] fn legacy_algorithm_names_remain_compatible_in_detailed_configs() { let router = router_from_json( - r#"{"default":"power-of-two","models":{"wait-model":{"algorithm":"groq-multiregion","seed":"seed-1"},"pulsar-model":{"algorithm":"pulsar-multiregion","seed":"seed-2","max_queue_time_floor_ms":100,"max_queue_time_ceil_ms":200}}}"#, + r#"{"default":"power-of-n","models":{"wait-model":{"algorithm":"groq-multiregion","seed":"seed-1"},"pulsar-model":{"algorithm":"pulsar-multiregion","seed":"seed-2","max_queue_time_floor_ms":100,"max_queue_time_ceil_ms":200}}}"#, ); assert_eq!( @@ -927,7 +1025,7 @@ fn legacy_algorithm_names_remain_compatible_in_detailed_configs() { #[test] fn detailed_model_config_parses_wait_and_widen_cache_affinity() { let router = router_from_json( - r#"{"default":"power-of-two","models":{"model-a":{"algorithm":"wait-and-widen","seed":"seed-1","require_cache_affinity_key":true,"cache_affinity_virtual_nodes":64,"cache_affinity_backend_selection_count":2}}}"#, + r#"{"default":"power-of-n","models":{"model-a":{"algorithm":"wait-and-widen","seed":"seed-1","require_cache_affinity_key":true,"cache_affinity_virtual_nodes":64,"cache_affinity_backend_selection_count":2}}}"#, ); let model_config = router.algorithm_config("model-a"); assert_eq!( @@ -947,7 +1045,7 @@ fn detailed_model_config_parses_wait_and_widen_cache_affinity() { #[test] fn request_algorithms_parse_and_override_default_selection() { let router = router_from_json( - r#"{"default":"power-of-two","request_algorithms":{"round-robin":"round-robin"}}"#, + r#"{"default":"power-of-n","request_algorithms":{"round-robin":"round-robin"}}"#, ); let target = target_with_model("model-a"); let request = request(&target, None, None); @@ -987,7 +1085,7 @@ fn choose_candidate_returns_slice_index_for_selected_cluster() { #[test] fn choose_candidate_with_resolution_preserves_algorithm_metadata() { let router = router_from_json( - r#"{"default":"power-of-two","request_algorithms":{"round-robin":"round-robin"}}"#, + r#"{"default":"power-of-n","request_algorithms":{"round-robin":"round-robin"}}"#, ); let target = target_with_model("model-a"); let request = request(&target, None, None); @@ -1022,7 +1120,7 @@ fn choose_candidate_with_resolution_preserves_algorithm_metadata() { #[test] fn model_request_algorithms_override_top_level_request_algorithms() { let router = router_from_json( - r#"{"default":"power-of-two","request_algorithms":{"round-robin":"round-robin"},"models":{"model-a":{"algorithm":"power-of-two","request_algorithms":{"round-robin":{"algorithm":"round-robin","require_input_tokens":true}}}}}"#, + r#"{"default":"power-of-n","request_algorithms":{"round-robin":"round-robin"},"models":{"model-a":{"algorithm":"power-of-n","request_algorithms":{"round-robin":{"algorithm":"round-robin","require_input_tokens":true}}}}}"#, ); let algorithm_override = LoadBalancerAlgorithmOverride::parse("round-robin") .expect("routing algorithm override should parse"); @@ -1041,7 +1139,7 @@ fn model_request_algorithms_override_top_level_request_algorithms() { #[test] fn request_algorithm_key_must_match_configured_algorithm() { let config: LoadBalancerConfig = - parse_json(r#"{"default":"power-of-two","request_algorithms":{"random":"round-robin"}}"#); + parse_json(r#"{"default":"power-of-n","request_algorithms":{"random":"round-robin"}}"#); let err = match LoadBalancerRouter::from_config(&config) { Ok(_) => panic!("mismatched request algorithm should fail"), @@ -1091,7 +1189,7 @@ fn wait_and_widen_config_resolves_internal_defaults() { #[test] fn router_reports_wait_and_widen_algorithm_name() { let router = router_with_model( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, "model-a", LoadBalancerModelConfig::Name(LoadBalancerAlgorithm::WaitAndWiden), ); @@ -1170,7 +1268,7 @@ fn target_state_distinguishes_independent_router_definitions() { #[test] fn configured_round_robin_uses_independent_sequences_per_routing_target() { let router = router_with_model( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, "shared-model", LoadBalancerModelConfig::Name(LoadBalancerAlgorithm::RoundRobin), ); @@ -1203,7 +1301,7 @@ fn choose_with_no_candidates_does_not_cache_default_lb_for_target() { #[test] fn request_round_robin_override_uses_stable_per_target_sequence() { let router = router_with_options( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, &[LoadBalancerAlgorithm::RoundRobin], None, ); @@ -1232,24 +1330,19 @@ fn request_round_robin_override_uses_stable_per_target_sequence() { fn configured_request_override_creates_target_local_balancer() { let router = router_with_options( LoadBalancerAlgorithm::RoundRobin, - &[LoadBalancerAlgorithm::PowerOfTwo], + &[LoadBalancerAlgorithm::PowerOfN], None, ); let target = target_with_model("model-a"); let request = request(&target, None, None); let candidates = candidates(&["cluster-0", "cluster-1"]); let target_state = LoadBalancerTargetState::default(); - let selection = choose_with_override( - &router, - &target_state, - &request, - &candidates, - "power-of-two", - ); + let selection = + choose_with_override(&router, &target_state, &request, &candidates, "power-of-n"); assert_eq!( selection.effective_algorithm, - LoadBalancerAlgorithm::PowerOfTwo + LoadBalancerAlgorithm::PowerOfN ); assert_eq!(target_state.instance_count(), 1); } @@ -1277,7 +1370,7 @@ fn matching_round_robin_override_reuses_configured_target_sequence() { #[test] fn request_round_robin_override_keeps_routing_targets_isolated() { let router = router_with_options( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, &[LoadBalancerAlgorithm::RoundRobin], None, ); @@ -1302,8 +1395,8 @@ fn request_round_robin_override_keeps_routing_targets_isolated() { #[test] fn request_override_beats_configured_model_algorithm() { let router = router_with_options( - LoadBalancerAlgorithm::PowerOfTwo, - &[LoadBalancerAlgorithm::PowerOfTwo], + LoadBalancerAlgorithm::PowerOfN, + &[LoadBalancerAlgorithm::PowerOfN], Some(( "shared-model", LoadBalancerModelConfig::Name(LoadBalancerAlgorithm::RoundRobin), @@ -1313,17 +1406,12 @@ fn request_override_beats_configured_model_algorithm() { let request = request(&target, None, None); let candidates = candidates(&["cluster-0", "cluster-1"]); let target_state = LoadBalancerTargetState::default(); - let selection = choose_with_override( - &router, - &target_state, - &request, - &candidates, - "power_of_two", - ); + let selection = + choose_with_override(&router, &target_state, &request, &candidates, "power_of_n"); assert_eq!( selection.effective_algorithm, - LoadBalancerAlgorithm::PowerOfTwo + LoadBalancerAlgorithm::PowerOfN ); } @@ -1333,7 +1421,7 @@ fn matching_request_override_reuses_configured_algorithm_config() { LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::RoundRobin); round_robin_config.request_policy_mut().require_input_tokens = true; let router = router_with_model( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, "shared-model", LoadBalancerModelConfig::Detailed(Box::new(round_robin_config)), ); @@ -1356,7 +1444,7 @@ fn matching_model_algorithm_beats_top_level_request_config() { let mut pulsar_config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::Pulsar); pulsar_config.request_policy_mut().require_input_tokens = true; let router = router_with_options( - LoadBalancerAlgorithm::PowerOfTwo, + LoadBalancerAlgorithm::PowerOfN, &[LoadBalancerAlgorithm::Pulsar], Some(( "shared-model", @@ -1376,7 +1464,7 @@ fn matching_model_algorithm_beats_top_level_request_config() { #[test] fn known_unavailable_request_override_returns_error() { - let router = router_with_default(LoadBalancerAlgorithm::PowerOfTwo); + let router = router_with_default(LoadBalancerAlgorithm::PowerOfN); let target = target_with_model("shared-model"); let request = request(&target, None, None); let candidates = candidates(&["cluster-0", "cluster-1"]); @@ -1437,7 +1525,8 @@ fn permissive_default_resolves_alias_and_underscore_spellings() { let router = LoadBalancerRouter::from_config(&LoadBalancerConfig::permissive_default()) .expect("permissive default config should build"); let spellings = [ - ("power_of_two", LoadBalancerAlgorithm::PowerOfTwo), + ("power_of_n", LoadBalancerAlgorithm::PowerOfN), + ("power_of_two", LoadBalancerAlgorithm::PowerOfN), ("round_robin", LoadBalancerAlgorithm::RoundRobin), ("groq-multiregion", LoadBalancerAlgorithm::WaitAndWiden), ("groq_multiregion", LoadBalancerAlgorithm::WaitAndWiden), @@ -1462,23 +1551,20 @@ fn permissive_default_resolves_alias_and_underscore_spellings() { } #[test] -fn permissive_default_keeps_power_of_two_without_override() { +fn permissive_default_keeps_power_of_n_without_override() { let router = LoadBalancerRouter::from_config(&LoadBalancerConfig::permissive_default()) .expect("permissive default config should build"); let config = router .resolve_algorithm_override("any-model", None) .expect("default algorithm should resolve"); - assert_eq!( - config.config().algorithm(), - LoadBalancerAlgorithm::PowerOfTwo - ); + assert_eq!(config.config().algorithm(), LoadBalancerAlgorithm::PowerOfN); } #[test] fn explicit_config_stays_restrictive() { let router = router_from_json( - r#"{"default":"power-of-two","request_algorithms":{"round-robin":"round-robin"}}"#, + r#"{"default":"power-of-n","request_algorithms":{"round-robin":"round-robin"}}"#, ); let algorithm_override = LoadBalancerAlgorithmOverride::parse("pulsar") .expect("routing algorithm override should parse"); diff --git a/src/libraries/rust/stargate/crates/stargate/src/main.rs b/src/libraries/rust/stargate/crates/stargate/src/main.rs index 9c6efee32..28d4a8594 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/main.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/main.rs @@ -192,7 +192,7 @@ struct Args { /// Skip QUIC TLS certificate verification for outbound connections and relays. #[arg(long, default_value_t = false, env = "STARGATE_QUIC_INSECURE")] quic_insecure: bool, - /// Path to load balancer config JSON file. If omitted, uses power-of-two + /// Path to load balancer config JSON file. If omitted, uses power-of-n /// with every built-in algorithm selectable per request. #[arg(long, value_name = "PATH")] lb_config_path: Option, diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs index a758bc89d..522f417a1 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs @@ -2255,7 +2255,7 @@ async fn registered_backend_rtt_means_drive_cluster_load_balancer_selection() { ..LoadBalancerAlgorithmConfig::default() }; let router = LoadBalancerRouter::from_config(&LoadBalancerConfig { - default: LoadBalancerAlgorithm::PowerOfTwo, + default: LoadBalancerAlgorithm::PowerOfN, request_algorithms: HashMap::new(), models: HashMap::from([( model_id.to_string(), diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/integration.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/integration.rs index 5d15d9efb..f0469446a 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/integration.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/integration.rs @@ -269,7 +269,7 @@ async fn round_robin_load_balancing() { let mut tmp_file = tempfile::NamedTempFile::new().expect("failed to create temp file"); write!( tmp_file, - r#"{{"default": "power-of-two", "models": {{"rr-model": "round-robin"}}}}"# + r#"{{"default": "power-of-n", "models": {{"rr-model": "round-robin"}}}}"# ) .expect("failed to write config"); let config_path = tmp_file.path().to_str().unwrap().to_string(); diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs index c78b1edde..f52a4fb56 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs @@ -403,7 +403,7 @@ fn float_eq(actual: f64, expected: f64) -> bool { } #[tokio::test] -async fn power_of_two_prefers_less_input_work() { +async fn power_of_n_prefers_less_input_work() { let stargate = RunningStargate::start("test-sg-p2c", None).await; let mut low = RegisteredBackend::active(stargate.grpc_addr, "p2c-model", "inst-low-headroom").await; @@ -476,7 +476,7 @@ async fn input_work_admission_rejects_overloaded_pool_and_registered_unavailable let stargate = RunningStargate::start( "test-sg-input-work-admission", Some( - r#"{"models": {"admission-model": {"algorithm": "power-of-two", "max_input_work_seconds": 0.5}}}"#, + r#"{"models": {"admission-model": {"algorithm": "power-of-n", "max_input_work_seconds": 0.5}}}"#, ), ) .await; @@ -655,7 +655,7 @@ async fn random_load_balancing_uses_all_instances() { } #[tokio::test] -async fn power_of_two_uses_cluster_aggregated_metrics_and_backend_round_robin() { +async fn power_of_n_uses_cluster_aggregated_metrics_and_backend_round_robin() { let stargate = RunningStargate::start("test-sg-p2c-clusters", None).await; let state = stargate.handle.state(); let mut shared_a = RegisteredBackend::active_in_cluster( @@ -850,9 +850,7 @@ async fn power_of_two_uses_cluster_aggregated_metrics_and_backend_round_robin() async fn wait_and_widen_load_balancing_prefers_lower_estimated_ttft() { let stargate = RunningStargate::start( "test-sg-wait-and-widen", - Some( - r#"{"default": "power-of-two", "models": {"wait-and-widen-model": "wait-and-widen"}}"#, - ), + Some(r#"{"default": "power-of-n", "models": {"wait-and-widen-model": "wait-and-widen"}}"#), ) .await; @@ -923,7 +921,7 @@ async fn wait_and_widen_waits_for_later_bucket_when_fastest_is_full() { let stargate = RunningStargate::start( "test-sg-wait-and-widen-wait", Some( - r#"{"default": "power-of-two", "models": {"wait-and-widen-wait-model": "wait-and-widen"}}"#, + r#"{"default": "power-of-n", "models": {"wait-and-widen-wait-model": "wait-and-widen"}}"#, ), ) .await; @@ -991,7 +989,7 @@ async fn wait_and_widen_cache_affinity_prefers_stable_subset_then_falls_back() { "test-sg-wait-and-widen-affinity", Some( r#"{ - "default": "power-of-two", + "default": "power-of-n", "models": { "wait-and-widen-affinity-model": { "algorithm": "wait-and-widen", @@ -1123,7 +1121,7 @@ async fn wait_and_widen_requires_cache_affinity_key_when_configured() { "test-sg-wait-and-widen-affinity-required", Some( r#"{ - "default": "power-of-two", + "default": "power-of-n", "models": { "wait-and-widen-affinity-required-model": { "algorithm": "wait-and-widen", @@ -1204,7 +1202,7 @@ async fn wait_and_widen_priority_header_uses_matching_queue_estimate() { "test-sg-wait-and-widen-priority", Some( r#"{ - "default": "power-of-two", + "default": "power-of-n", "models": { "wait-and-widen-priority-model": { "algorithm": "wait-and-widen", @@ -1311,7 +1309,7 @@ async fn pulsar_routes_same_affinity_key_consistently() { "test-sg-pulsar", Some( r#"{ - "default": "power-of-two", + "default": "power-of-n", "models": { "pulsar-model": { "algorithm": "pulsar", diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs index 54096d918..407851c39 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs @@ -1009,7 +1009,7 @@ impl PulsarHeaderFixture { let mut config_file = tempfile::NamedTempFile::new().expect("failed to create temp file"); std::io::Write::write_all( &mut config_file, - br#"{"default":"power-of-two","models":{"pulsar-model":{"algorithm":"pulsar","seed":"test-seed","require_cache_affinity_key":true,"require_input_tokens":true}}}"#, + br#"{"default":"power-of-n","models":{"pulsar-model":{"algorithm":"pulsar","seed":"test-seed","require_cache_affinity_key":true,"require_input_tokens":true}}}"#, ) .expect("failed to write config"); let config_path = config_file.path().to_str().unwrap().to_string(); @@ -2027,7 +2027,7 @@ async fn transport_local_shared_cluster_failover_stays_within_selected_cluster() init_crypto(); let mut tmp_file = tempfile::NamedTempFile::new().expect("failed to create temp file"); - std::io::Write::write_all(&mut tmp_file, br#"{"default":"power-of-two"}"#) + std::io::Write::write_all(&mut tmp_file, br#"{"default":"power-of-n"}"#) .expect("failed to write config"); let config_path = tmp_file.path().to_str().unwrap().to_string(); @@ -2597,7 +2597,7 @@ async fn queue_estimate_mismatch_retries_sibling_in_selected_shared_cluster() { ); assert_delta!( &before_metrics, &metrics, "stargate_routing_selections_total", 1.0; - r#"algorithm="power-of-two""#, + r#"algorithm="power-of-n""#, r#"model="queue-mismatch-shared-model""#, r#"routing_key="""#, r#"selection="primary""# diff --git a/src/libraries/rust/stargate/docs/diagrams/chat-completions-e2e.puml b/src/libraries/rust/stargate/docs/diagrams/chat-completions-e2e.puml index 9682b0430..f1aec1ff0 100644 --- a/src/libraries/rust/stargate/docs/diagrams/chat-completions-e2e.puml +++ b/src/libraries/rust/stargate/docs/diagrams/chat-completions-e2e.puml @@ -35,7 +35,7 @@ Proxy -> LB: cluster_candidates_for_target(\n RoutingTargetKey{my-rk, my-model} LB --> Proxy: [cluster-A (rtt=5ms), cluster-B (rtt=12ms)] note over LB - Default: power-of-two + Default: power-of-n (sample 2, pick lower prompt-work time). Only clusters with a backend that has an open QUIC connection and a healthy RTT diff --git a/src/libraries/rust/stargate/docs/load-balancer-configuration.md b/src/libraries/rust/stargate/docs/load-balancer-configuration.md index e2b3f8979..5aeca7fe1 100644 --- a/src/libraries/rust/stargate/docs/load-balancer-configuration.md +++ b/src/libraries/rust/stargate/docs/load-balancer-configuration.md @@ -4,8 +4,8 @@ Stargate selects one load-balancing algorithm for each model. A request can select another preconfigured algorithm through a trusted header. This page defines the `lb-config.json` schema and the behavior of -`wait-and-widen`, `pulsar`, and `pulsar-wait-and-widen`. Deployment systems own -the file mount and the `--lb-config-path` argument. +`power-of-n`, `wait-and-widen`, `pulsar`, and `pulsar-wait-and-widen`. +Deployment systems own the file mount and the `--lb-config-path` argument. ## Load the configuration @@ -15,7 +15,7 @@ Start Stargate with an optional JSON file: --lb-config-path=/config/lb-config.json ``` -When the argument is absent, Stargate uses `power-of-two` for every model and +When the argument is absent, Stargate uses `power-of-n` for every model and accepts a routing-method override when it is in the allowlist of built-in algorithms, each with its default settings. When the argument is present, the file defines the allowlist. @@ -30,24 +30,25 @@ The top-level object has three fields: | Field | Type | Default | Meaning | | --- | --- | --- | --- | -| `default` | algorithm name | `power-of-two` | Algorithm for models without an entry in `models`. | +| `default` | algorithm name | `power-of-n` | Algorithm for models without an entry in `models`. | | `request_algorithms` | object | `{}` | Algorithms that `x-routing-method` may select for every model. | | `models` | object | `{}` | Exact model ID to algorithm configuration. | -Valid algorithm names are `power-of-two`, `wait-and-widen`, `round-robin`, +Valid algorithm names are `power-of-n`, `wait-and-widen`, `round-robin`, `random`, `pulsar`, and `pulsar-wait-and-widen`. -For backward compatibility, Stargate also accepts `groq-multiregion` as an -alias for `wait-and-widen` and `pulsar-multiregion` as an alias for -`pulsar-wait-and-widen` in `default`, `models`, `request_algorithms`, detailed -algorithm objects, and routing-method overrides. Use the canonical names for -new configurations. +For backward compatibility, Stargate also accepts `powerOfN`, `powerOf2`, and +`power-of-two` as aliases for `power-of-n`, `groq-multiregion` as an alias for +`wait-and-widen`, and `pulsar-multiregion` as an alias for +`pulsar-wait-and-widen`. These aliases work in `default`, `models`, +`request_algorithms`, detailed algorithm objects, and routing-method overrides. +Use the canonical names for new configurations. An entry in `models` or `request_algorithms` can be an algorithm name: ```json { - "default": "power-of-two", + "default": "power-of-n", "models": { "model-a": "wait-and-widen" } @@ -58,7 +59,7 @@ Use a detailed object to set algorithm fields: ```json { - "default": "power-of-two", + "default": "power-of-n", "models": { "model-a": { "algorithm": "wait-and-widen", @@ -111,10 +112,34 @@ Choose based on the routing goal and available backend statistics: | Keep the same prefix on a stable, capacity-weighted cluster. | `pulsar` | Positive finite `last_mean_input_tps` for every participating cluster. | | Keep Pulsar affinity when possible, but escape to lower-latency capacity when the primary cannot meet queue policy. | `pulsar-wait-and-widen` | Pulsar capacity plus the RTT and queue statistics used by `wait-and-widen`. | -Use `power-of-two` when these statistics or affinity requirements are not +Use `power-of-n` when these statistics or affinity requirements are not available. Use `round-robin` for deterministic cycling and `random` for uniform random selection. +## `power-of-n` + +`power-of-n` uniformly samples distinct eligible clusters and selects the +cluster with the lowest sum of queued and request input tokens divided by its +last mean input TPS. It breaks equal scores randomly. Retried clusters are +excluded before sampling. + +The default sample count is `2`. A larger sample can improve routing decisions +in a heterogeneous pool, but it compares more clusters on every request. Valid +values are `1` through `64`. If fewer eligible clusters exist, the algorithm +compares every eligible cluster once. + +```json +{ + "default": "power-of-n", + "models": { + "model-a": { + "algorithm": "power-of-n", + "sample_count": 4 + } + } +} +``` + ## `wait-and-widen` `wait-and-widen` estimates time to first token (TTFT) as: @@ -143,7 +168,7 @@ Minimal configuration: ```json { - "default": "power-of-two", + "default": "power-of-n", "models": { "model-a": { "algorithm": "wait-and-widen", @@ -170,7 +195,7 @@ Minimal configuration: ```json { - "default": "power-of-two", + "default": "power-of-n", "models": { "model-a": { "algorithm": "pulsar", @@ -194,7 +219,7 @@ Minimal configuration: ```json { - "default": "power-of-two", + "default": "power-of-n", "models": { "model-a": { "algorithm": "pulsar-wait-and-widen", @@ -209,6 +234,12 @@ Minimal configuration: ## Algorithm fields +`power-of-n` supports this field: + +| Field | Type | Default | Constraint and effect | +| --- | --- | --- | --- | +| `sample_count` | unsigned integer | `2` | Number of distinct eligible clusters sampled. Must be from `1` through `64`. Values above the eligible cluster count compare the complete eligible pool. | + `wait-and-widen` supports these cache-affinity fields: | Field | Type | Default | Constraint and effect | @@ -258,7 +289,7 @@ Preconfigure every algorithm that a request may select: ```json { - "default": "power-of-two", + "default": "power-of-n", "request_algorithms": { "wait-and-widen": "wait-and-widen", "pulsar": { diff --git a/src/libraries/rust/stargate/docs/multi-backend-clusters.md b/src/libraries/rust/stargate/docs/multi-backend-clusters.md index 4e6051ccf..c75df11db 100644 --- a/src/libraries/rust/stargate/docs/multi-backend-clusters.md +++ b/src/libraries/rust/stargate/docs/multi-backend-clusters.md @@ -104,7 +104,7 @@ source backend disappears, recompute from remaining active snapshots. All algorithms choose from cluster snapshots. -- `power-of-two`: compares aggregated cluster load. +- `power-of-n`: compares aggregated cluster load. - `wait-and-widen`: uses aggregated stats and representative RTT. - `pulsar-wait-and-widen`: combines Pulsar ranking with WaitAndWiden fallback. - `round-robin`: rounds across clusters.