Skip to content

refactor(worker): land Pending start semantics and state machine #1102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slin1237 merged 8 commits into from
Apr 12, 2026
Merged

refactor(worker): land Pending start semantics and state machine #1102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
53e5233
refactor(worker): land Pending start semantics and state machine
slin1237 Apr 12, 2026
e2c8b68
fix(review): address PR feedback on Pending state machine
slin1237 Apr 12, 2026
cec40bf
refactor(worker): move status preservation from registry to builder
slin1237 Apr 12, 2026
cb307ad
fix(review): log transport errors before treating as failed probe
slin1237 Apr 12, 2026
60cca3f
fix(review): early-return Failed in check_health_async + replace test…
slin1237 Apr 12, 2026
0c48755
fix(review): address remaining PR feedback (4 real bugs)
slin1237 Apr 12, 2026
042ab0e
fix(worker): activate workers in workflow path, simplify update step
slin1237 Apr 12, 2026
6eaae8e
fix(worker): activate all workflow-created workers, not just disable_…
slin1237 Apr 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions model_gateway/src/policies/bucket.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -562,9 +562,18 @@ impl Bucket {

#[cfg(test)]
mod tests {
use openai_protocol::worker::HealthCheckConfig;

use super::*;
use crate::worker::{BasicWorkerBuilder, WorkerType};

fn no_health_check() -> HealthCheckConfig {
HealthCheckConfig {
disable_health_check: true,
..Default::default()
}
}
Comment thread
slin1237 marked this conversation as resolved.

#[tokio::test]
async fn test_load_balancing_conditions() {
// Test 1: Basic load balancing trigger
Expand All @@ -579,18 +588,21 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -793,18 +805,21 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -996,18 +1011,21 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1104,18 +1122,21 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
];
Expand Down
40 changes: 40 additions & 0 deletions model_gateway/src/policies/cache_aware.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -918,10 +918,18 @@ impl Default for CacheAwarePolicy {
#[cfg(test)]
mod tests {
use kv_index::{compute_content_hash, SequenceHash, StoredBlock, WorkerBlockMap};
use openai_protocol::worker::HealthCheckConfig;

use super::*;
use crate::worker::{BasicWorkerBuilder, WorkerType};

fn no_health_check() -> HealthCheckConfig {
HealthCheckConfig {
disable_health_check: true,
..Default::default()
}
}

#[test]
fn test_cache_aware_with_balanced_load() {
// Create policy without eviction thread for testing
Expand All @@ -935,12 +943,14 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -997,9 +1007,11 @@ mod tests {

let worker1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
let worker2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();

// Create significant load imbalance
Expand Down Expand Up @@ -1033,11 +1045,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1097,6 +1111,7 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
)];

Expand Down Expand Up @@ -1268,6 +1283,7 @@ mod tests {
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.health_config(no_health_check())
.build(),
)];

Expand Down Expand Up @@ -1331,11 +1347,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1370,6 +1388,7 @@ mod tests {
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
)];
policy.init_workers(&workers);
Expand All @@ -1394,9 +1413,11 @@ mod tests {

let w1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
let w2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();

// Give w1 higher load
Expand Down Expand Up @@ -1440,11 +1461,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1489,6 +1512,7 @@ mod tests {
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
)];

Expand All @@ -1506,11 +1530,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1574,11 +1600,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1610,9 +1638,11 @@ mod tests {

let w1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
let w2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
// Give w1 higher load so min-load picks w2
for _ in 0..3 {
Expand Down Expand Up @@ -1647,9 +1677,11 @@ mod tests {

let w1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
let w2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
for _ in 0..3 {
w1.increment_load();
Expand Down Expand Up @@ -1683,11 +1715,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1743,11 +1777,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down Expand Up @@ -1799,9 +1835,11 @@ mod tests {

let w1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();
let w2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build();

// Create heavy imbalance: w1 has 20 load, w2 has 0
Expand Down Expand Up @@ -1839,11 +1877,13 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
),
];
Expand Down
10 changes: 10 additions & 0 deletions model_gateway/src/policies/consistent_hashing.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,9 +198,18 @@ impl LoadBalancingPolicy for ConsistentHashingPolicy {
mod tests {
use std::collections::HashMap;

use openai_protocol::worker::HealthCheckConfig;

use super::*;
use crate::worker::{BasicWorkerBuilder, HashRing, WorkerType};

fn no_health_check() -> HealthCheckConfig {
HealthCheckConfig {
disable_health_check: true,
..Default::default()
}
}

fn headers_with_routing_key(key: &str) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
Expand All @@ -219,6 +228,7 @@ mod tests {
Arc::new(
BasicWorkerBuilder::new(*url)
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
) as Arc<dyn Worker>
})
Expand Down
10 changes: 10 additions & 0 deletions model_gateway/src/policies/manual.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,15 +301,25 @@ fn min_group_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> u
mod tests {
use std::collections::HashMap;

use openai_protocol::worker::HealthCheckConfig;

use super::*;
use crate::worker::{BasicWorkerBuilder, WorkerType};

fn no_health_check() -> HealthCheckConfig {
HealthCheckConfig {
disable_health_check: true,
..Default::default()
}
}

fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
urls.iter()
.map(|url| {
Arc::new(
BasicWorkerBuilder::new(*url)
.worker_type(WorkerType::Regular)
.health_config(no_health_check())
.build(),
) as Arc<dyn Worker>
})
Expand Down
Loading
Loading