[router] add grpc pd and regular router init (#9893)

This commit is contained in:
Chang Su
2025-09-01 20:06:15 -07:00
committed by GitHub
parent b5245064f6
commit 9a0cac1be0
14 changed files with 783 additions and 58 deletions

View File

@@ -349,6 +349,16 @@ impl ConfigValidator {
return Ok(());
}
// Validate gRPC connection mode requires tokenizer configuration
if config.connection_mode == ConnectionMode::Grpc
&& config.tokenizer_path.is_none()
&& config.model_path.is_none()
{
return Err(ConfigError::ValidationFailed {
reason: "gRPC connection mode requires either --tokenizer-path or --model-path to be specified".to_string(),
});
}
// All policies are now supported for both router types thanks to the unified trait design
// No mode/policy restrictions needed anymore
@@ -419,11 +429,14 @@ impl ConfigValidator {
});
}
if !url.starts_with("http://") && !url.starts_with("https://") {
if !url.starts_with("http://")
&& !url.starts_with("https://")
&& !url.starts_with("grpc://")
{
return Err(ConfigError::InvalidValue {
field: "worker_url".to_string(),
value: url.clone(),
reason: "URL must start with http:// or https://".to_string(),
reason: "URL must start with http://, https://, or grpc://".to_string(),
});
}
@@ -684,4 +697,60 @@ mod tests {
assert!(e.to_string().contains("prefill requires at least 2"));
}
}
#[test]
fn test_validate_grpc_requires_tokenizer() {
// Test that gRPC connection mode requires tokenizer configuration
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
// Set connection mode to gRPC without tokenizer config
config.connection_mode = ConnectionMode::Grpc;
config.tokenizer_path = None;
config.model_path = None;
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
if let Err(e) = result {
assert!(e.to_string().contains("gRPC connection mode requires"));
}
}
#[test]
fn test_validate_grpc_with_model_path() {
// Test that gRPC works with model_path
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
config.connection_mode = ConnectionMode::Grpc;
config.model_path = Some("meta-llama/Llama-3-8B".to_string());
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_grpc_with_tokenizer_path() {
// Test that gRPC works with tokenizer_path
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
config.connection_mode = ConnectionMode::Grpc;
config.tokenizer_path = Some("/path/to/tokenizer.json".to_string());
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
}