[model-gateway] Add Redis support as a history backend (#16300)
This commit is contained in:
@@ -31,6 +31,7 @@ pub enum HistoryBackendType {
|
||||
None,
|
||||
Oracle,
|
||||
Postgres,
|
||||
Redis,
|
||||
}
|
||||
|
||||
#[pyclass(eq)]
|
||||
@@ -270,6 +271,40 @@ impl PyOracleConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PyRedisConfig {
|
||||
#[pyo3(get, set)]
|
||||
pub url: String,
|
||||
#[pyo3(get, set)]
|
||||
pub pool_max: usize,
|
||||
#[pyo3(get, set)]
|
||||
pub retention_days: Option<u64>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyRedisConfig {
|
||||
#[new]
|
||||
#[pyo3(signature = (url, pool_max = 16, retention_days = Some(30)))]
|
||||
fn new(url: String, pool_max: usize, retention_days: Option<u64>) -> PyResult<Self> {
|
||||
Ok(PyRedisConfig {
|
||||
url,
|
||||
pool_max,
|
||||
retention_days,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl PyRedisConfig {
|
||||
pub fn to_config_redis(&self) -> config::RedisConfig {
|
||||
config::RedisConfig {
|
||||
url: self.url.clone(),
|
||||
pool_max: self.pool_max,
|
||||
retention_days: self.retention_days,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PyPostgresConfig {
|
||||
@@ -374,6 +409,7 @@ struct Router {
|
||||
history_backend: HistoryBackendType,
|
||||
oracle_config: Option<PyOracleConfig>,
|
||||
postgres_config: Option<PyPostgresConfig>,
|
||||
redis_config: Option<PyRedisConfig>,
|
||||
client_cert_path: Option<String>,
|
||||
client_key_path: Option<String>,
|
||||
ca_cert_paths: Vec<String>,
|
||||
@@ -486,6 +522,7 @@ impl Router {
|
||||
HistoryBackendType::None => config::HistoryBackend::None,
|
||||
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
|
||||
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
|
||||
HistoryBackendType::Redis => config::HistoryBackend::Redis,
|
||||
};
|
||||
|
||||
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
|
||||
@@ -504,6 +541,14 @@ impl Router {
|
||||
None
|
||||
};
|
||||
|
||||
let redis_config = if matches!(self.history_backend, HistoryBackendType::Redis) {
|
||||
self.redis_config
|
||||
.as_ref()
|
||||
.map(|cfg| cfg.to_config_redis())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
config::RouterConfig::builder()
|
||||
.mode(mode)
|
||||
.policy(policy)
|
||||
@@ -558,6 +603,7 @@ impl Router {
|
||||
.maybe_chat_template(self.chat_template.as_ref())
|
||||
.maybe_oracle(oracle)
|
||||
.maybe_postgres(postgres_config)
|
||||
.maybe_redis(redis_config)
|
||||
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
||||
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
||||
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
||||
@@ -654,6 +700,7 @@ impl Router {
|
||||
history_backend = HistoryBackendType::Memory,
|
||||
oracle_config = None,
|
||||
postgres_config = None,
|
||||
redis_config = None,
|
||||
client_cert_path = None,
|
||||
client_key_path = None,
|
||||
ca_cert_paths = vec![],
|
||||
@@ -737,6 +784,7 @@ impl Router {
|
||||
history_backend: HistoryBackendType,
|
||||
oracle_config: Option<PyOracleConfig>,
|
||||
postgres_config: Option<PyPostgresConfig>,
|
||||
redis_config: Option<PyRedisConfig>,
|
||||
client_cert_path: Option<String>,
|
||||
client_key_path: Option<String>,
|
||||
ca_cert_paths: Vec<String>,
|
||||
@@ -834,6 +882,7 @@ impl Router {
|
||||
history_backend,
|
||||
oracle_config,
|
||||
postgres_config,
|
||||
redis_config,
|
||||
client_cert_path,
|
||||
client_key_path,
|
||||
ca_cert_paths,
|
||||
@@ -946,6 +995,7 @@ fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyControlPlaneAuthConfig>()?;
|
||||
m.add_class::<PyOracleConfig>()?;
|
||||
m.add_class::<PyPostgresConfig>()?;
|
||||
m.add_class::<PyRedisConfig>()?;
|
||||
m.add_class::<Router>()?;
|
||||
m.add_function(wrap_pyfunction!(get_version_string, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(get_verbose_version_string, m)?)?;
|
||||
|
||||
@@ -10,6 +10,7 @@ from sglang_router.sglang_router_rs import (
|
||||
PyJwtConfig,
|
||||
PyOracleConfig,
|
||||
PyPostgresConfig,
|
||||
PyRedisConfig,
|
||||
PyRole,
|
||||
)
|
||||
from sglang_router.sglang_router_rs import Router as _Router
|
||||
@@ -63,6 +64,8 @@ def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType:
|
||||
return HistoryBackendType.Oracle
|
||||
elif backend_lower == "postgres":
|
||||
return HistoryBackendType.Postgres
|
||||
elif backend_lower == "redis":
|
||||
return HistoryBackendType.Redis
|
||||
else:
|
||||
raise ValueError(f"Unknown history backend: {backend_str}")
|
||||
|
||||
@@ -262,6 +265,20 @@ class Router:
|
||||
)
|
||||
args_dict["postgres_config"] = postgres_config
|
||||
|
||||
# Convert Redis config if needed
|
||||
redis_config = None
|
||||
if history_backend == HistoryBackendType.Redis:
|
||||
retention_days = args_dict.get("redis_retention_days", 30)
|
||||
# If retention_days is negative, it means persistent storage (None in Rust)
|
||||
retention_arg = None if retention_days < 0 else retention_days
|
||||
|
||||
redis_config = PyRedisConfig(
|
||||
url=args_dict.get("redis_url"),
|
||||
pool_max=args_dict.get("redis_pool_max", 16),
|
||||
retention_days=retention_arg,
|
||||
)
|
||||
args_dict["redis_config"] = redis_config
|
||||
|
||||
# Build control plane auth config
|
||||
args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict)
|
||||
|
||||
@@ -278,6 +295,9 @@ class Router:
|
||||
"oracle_pool_timeout_secs",
|
||||
"postgres_db_url",
|
||||
"postgres_pool_max",
|
||||
"redis_url",
|
||||
"redis_pool_max",
|
||||
"redis_retention_days",
|
||||
# Control plane auth fields (converted to control_plane_auth)
|
||||
"control_plane_api_keys",
|
||||
"control_plane_audit_enabled",
|
||||
|
||||
@@ -117,6 +117,9 @@ class RouterArgs:
|
||||
oracle_pool_timeout_secs: int = 30
|
||||
postgres_db_url: Optional[str] = None
|
||||
postgres_pool_max: int = 16
|
||||
redis_url: Optional[str] = None
|
||||
redis_pool_max: int = 16
|
||||
redis_retention_days: int = 30
|
||||
# mTLS configuration for worker communication
|
||||
client_cert_path: Optional[str] = None
|
||||
client_key_path: Optional[str] = None
|
||||
@@ -200,6 +203,9 @@ class RouterArgs:
|
||||
postgres_group = parser.add_argument_group(
|
||||
"PostgreSQL Database", "PostgreSQL database backend configuration"
|
||||
)
|
||||
redis_group = parser.add_argument_group(
|
||||
"Redis Database", "Redis database backend configuration"
|
||||
)
|
||||
tls_group = parser.add_argument_group(
|
||||
"TLS/mTLS Security", "TLS certificates for server and worker communication"
|
||||
)
|
||||
@@ -663,7 +669,7 @@ class RouterArgs:
|
||||
f"--{prefix}history-backend",
|
||||
type=str,
|
||||
default=RouterArgs.history_backend,
|
||||
choices=["memory", "none", "oracle", "postgres"],
|
||||
choices=["memory", "none", "oracle", "postgres", "redis"],
|
||||
help="History storage backend for conversations and responses (default: memory)",
|
||||
)
|
||||
|
||||
@@ -733,6 +739,28 @@ class RouterArgs:
|
||||
help="Maximum PostgreSQL connection pool size (default: 16, env: POSTGRES_POOL_MAX)",
|
||||
)
|
||||
|
||||
# Redis configuration
|
||||
redis_group.add_argument(
|
||||
f"--{prefix}redis-url",
|
||||
type=str,
|
||||
default=os.getenv("REDIS_URL"),
|
||||
help="Redis connection URL (env: REDIS_URL)",
|
||||
)
|
||||
redis_group.add_argument(
|
||||
f"--{prefix}redis-pool-max",
|
||||
type=int,
|
||||
default=int(os.getenv("REDIS_POOL_MAX", RouterArgs.redis_pool_max)),
|
||||
help="Maximum Redis connection pool size (default: 16, env: REDIS_POOL_MAX)",
|
||||
)
|
||||
redis_group.add_argument(
|
||||
f"--{prefix}redis-retention-days",
|
||||
type=int,
|
||||
default=int(
|
||||
os.getenv("REDIS_RETENTION_DAYS", RouterArgs.redis_retention_days)
|
||||
),
|
||||
help="Redis data retention in days (-1 for persistent, default: 30, env: REDIS_RETENTION_DAYS)",
|
||||
)
|
||||
|
||||
# TLS/mTLS configuration
|
||||
tls_group.add_argument(
|
||||
f"--{prefix}client-cert-path",
|
||||
|
||||
Reference in New Issue
Block a user