mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 21:20:20 +08:00
Invalidate model routing cache after admin model writes
This commit is contained in:
Novick Yuan
@@ -284,20 +284,30 @@ impl AppState {
|
||||
&self,
|
||||
record: &global_models::UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<global_models::StoredAdminProviderModel>, GatewayError> {
|
||||
self.data
|
||||
let created = self
|
||||
.data
|
||||
.create_admin_provider_model(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if created.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_admin_provider_model(
|
||||
&self,
|
||||
record: &global_models::UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<global_models::StoredAdminProviderModel>, GatewayError> {
|
||||
self.data
|
||||
let updated = self
|
||||
.data
|
||||
.update_admin_provider_model(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if updated.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_admin_provider_model(
|
||||
@@ -305,40 +315,60 @@ impl AppState {
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
let deleted = self
|
||||
.data
|
||||
.delete_admin_provider_model(provider_id, model_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if deleted {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn create_admin_global_model(
|
||||
&self,
|
||||
record: &global_models::CreateAdminGlobalModelRecord,
|
||||
) -> Result<Option<global_models::StoredAdminGlobalModel>, GatewayError> {
|
||||
self.data
|
||||
let created = self
|
||||
.data
|
||||
.create_admin_global_model(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if created.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_admin_global_model(
|
||||
&self,
|
||||
record: &global_models::UpdateAdminGlobalModelRecord,
|
||||
) -> Result<Option<global_models::StoredAdminGlobalModel>, GatewayError> {
|
||||
self.data
|
||||
let updated = self
|
||||
.data
|
||||
.update_admin_global_model(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if updated.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_admin_global_model(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
let deleted = self
|
||||
.data
|
||||
.delete_admin_global_model(global_model_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if deleted {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_model_stats(
|
||||
@@ -840,13 +870,30 @@ impl AppState {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
use std::sync::{
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::{
|
||||
global_models::InMemoryGlobalModelReadRepository,
|
||||
provider_catalog::InMemoryProviderCatalogReadRepository,
|
||||
};
|
||||
use aether_data::DataLayerError;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
CreateAdminGlobalModelRecord, StoredAdminGlobalModel, UpdateAdminGlobalModelRecord,
|
||||
UpsertAdminProviderModelRecord,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::cache::SchedulerAffinityTarget;
|
||||
use crate::data::GatewayDataState;
|
||||
@@ -897,6 +944,195 @@ mod tests {
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
fn sample_admin_global_model() -> StoredAdminGlobalModel {
|
||||
StoredAdminGlobalModel::new(
|
||||
"global-1".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
"GPT 5".to_string(),
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000),
|
||||
)
|
||||
.expect("global model should build")
|
||||
}
|
||||
|
||||
fn sample_provider_model_record(
|
||||
id: &str,
|
||||
global_model_id: &str,
|
||||
is_active: bool,
|
||||
) -> UpsertAdminProviderModelRecord {
|
||||
UpsertAdminProviderModelRecord::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
global_model_id.to_string(),
|
||||
"gpt-5-upstream".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(true),
|
||||
None,
|
||||
None,
|
||||
is_active,
|
||||
true,
|
||||
None,
|
||||
)
|
||||
.expect("provider model record should build")
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct ClearCountingCandidateSelectionReadRepository {
|
||||
clear_count: AtomicUsize,
|
||||
}
|
||||
|
||||
impl ClearCountingCandidateSelectionReadRepository {
|
||||
fn clear_count(&self) -> usize {
|
||||
self.clear_count.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MinimalCandidateSelectionReadRepository for ClearCountingCandidateSelectionReadRepository {
|
||||
fn clear_local_cache(&self) {
|
||||
self.clear_count.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format(
|
||||
&self,
|
||||
_api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
_api_format: &str,
|
||||
_global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_requested_model(
|
||||
&self,
|
||||
_api_format: &str,
|
||||
_requested_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_requested_model_page(
|
||||
&self,
|
||||
_query: &StoredRequestedModelCandidateRowsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_pool_key_rows_for_group(
|
||||
&self,
|
||||
_query: &StoredPoolKeyCandidateRowsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_pool_key_rows_for_group_key_ids(
|
||||
&self,
|
||||
_query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_model_writes_invalidate_candidate_selection_cache() {
|
||||
let candidate_repository =
|
||||
Arc::new(ClearCountingCandidateSelectionReadRepository::default());
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![sample_admin_global_model()]),
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_minimal_candidate_selection_reader_for_tests(
|
||||
candidate_repository.clone(),
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
);
|
||||
|
||||
assert_eq!(candidate_repository.clear_count(), 0);
|
||||
|
||||
let provider_model = sample_provider_model_record("model-1", "global-1", true);
|
||||
state
|
||||
.create_admin_provider_model(&provider_model)
|
||||
.await
|
||||
.expect("provider model create should succeed")
|
||||
.expect("provider model should create");
|
||||
assert_eq!(candidate_repository.clear_count(), 1);
|
||||
|
||||
let disabled_provider_model = sample_provider_model_record("model-1", "global-1", false);
|
||||
state
|
||||
.update_admin_provider_model(&disabled_provider_model)
|
||||
.await
|
||||
.expect("provider model update should succeed")
|
||||
.expect("provider model should update");
|
||||
assert_eq!(candidate_repository.clear_count(), 2);
|
||||
|
||||
assert!(state
|
||||
.delete_admin_provider_model("provider-1", "model-1")
|
||||
.await
|
||||
.expect("provider model delete should succeed"));
|
||||
assert_eq!(candidate_repository.clear_count(), 3);
|
||||
|
||||
let created_global_model = CreateAdminGlobalModelRecord::new(
|
||||
"global-2".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
"GPT 4.1".to_string(),
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("global model create record should build");
|
||||
state
|
||||
.create_admin_global_model(&created_global_model)
|
||||
.await
|
||||
.expect("global model create should succeed")
|
||||
.expect("global model should create");
|
||||
assert_eq!(candidate_repository.clear_count(), 4);
|
||||
|
||||
let disabled_global_model = UpdateAdminGlobalModelRecord::new(
|
||||
"global-1".to_string(),
|
||||
"GPT 5".to_string(),
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("global model update record should build");
|
||||
state
|
||||
.update_admin_global_model(&disabled_global_model)
|
||||
.await
|
||||
.expect("global model update should succeed")
|
||||
.expect("global model should update");
|
||||
assert_eq!(candidate_repository.clear_count(), 5);
|
||||
|
||||
assert!(state
|
||||
.delete_admin_global_model("global-2")
|
||||
.await
|
||||
.expect("global model delete should succeed"));
|
||||
assert_eq!(candidate_repository.clear_count(), 6);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_catalog_update_invalidates_scheduler_affinity_and_transport_snapshot_cache() {
|
||||
let provider = sample_provider();
|
||||
|
||||
@@ -10,15 +10,21 @@ use crate::tests::{
|
||||
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
|
||||
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
};
|
||||
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::DataLayerError;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
StoredAdminGlobalModel, UpdateAdminGlobalModelRecord,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use axum::response::IntoResponse;
|
||||
use std::collections::HashMap;
|
||||
use std::future::pending;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
|
||||
match status {
|
||||
@@ -183,6 +189,120 @@ impl MinimalCandidateSelectionReadRepository for PendingMinimalCandidateSelectio
|
||||
}
|
||||
}
|
||||
|
||||
struct CachedToggleMinimalCandidateSelectionReadRepository {
|
||||
row: StoredMinimalCandidateSelectionRow,
|
||||
active: AtomicBool,
|
||||
cached_rows_by_api_format: Mutex<HashMap<String, Vec<StoredMinimalCandidateSelectionRow>>>,
|
||||
}
|
||||
|
||||
impl CachedToggleMinimalCandidateSelectionReadRepository {
|
||||
fn new(row: StoredMinimalCandidateSelectionRow) -> Self {
|
||||
Self {
|
||||
row,
|
||||
active: AtomicBool::new(true),
|
||||
cached_rows_by_api_format: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_active(&self, active: bool) {
|
||||
self.active.store(active, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
fn rows_for_api_format(&self, api_format: &str) -> Vec<StoredMinimalCandidateSelectionRow> {
|
||||
let api_format = api_format.trim().to_string();
|
||||
let mut cached = self
|
||||
.cached_rows_by_api_format
|
||||
.lock()
|
||||
.expect("candidate row cache lock");
|
||||
if let Some(rows) = cached.get(&api_format) {
|
||||
return rows.clone();
|
||||
}
|
||||
|
||||
let rows = if self.active.load(Ordering::SeqCst)
|
||||
&& self
|
||||
.row
|
||||
.endpoint_api_format
|
||||
.eq_ignore_ascii_case(&api_format)
|
||||
{
|
||||
vec![self.row.clone()]
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
cached.insert(api_format, rows.clone());
|
||||
rows
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MinimalCandidateSelectionReadRepository
|
||||
for CachedToggleMinimalCandidateSelectionReadRepository
|
||||
{
|
||||
fn clear_local_cache(&self) {
|
||||
self.cached_rows_by_api_format
|
||||
.lock()
|
||||
.expect("candidate row cache lock")
|
||||
.clear();
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(self.rows_for_api_format(api_format))
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.rows_for_api_format(api_format)
|
||||
.into_iter()
|
||||
.filter(|row| row.global_model_name == global_model_name)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_requested_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
requested_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.rows_for_api_format(api_format)
|
||||
.into_iter()
|
||||
.filter(|row| row.global_model_name == requested_model_name)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_for_exact_api_format_and_requested_model_page(
|
||||
&self,
|
||||
query: &StoredRequestedModelCandidateRowsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(self
|
||||
.rows_for_api_format(&query.api_format)
|
||||
.into_iter()
|
||||
.filter(|row| row.global_model_name == query.requested_model_name)
|
||||
.skip(query.offset as usize)
|
||||
.take(query.limit as usize)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_pool_key_rows_for_group(
|
||||
&self,
|
||||
_query: &StoredPoolKeyCandidateRowsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_pool_key_rows_for_group_key_ids(
|
||||
&self,
|
||||
_query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
|
||||
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -240,6 +360,102 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
|
||||
fallback_probe_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_openai_models_list_drops_disabled_global_model_after_cache_invalidation() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-openai-models-cache")),
|
||||
unrestricted_models_snapshot("key-models-cache", "user-models-cache"),
|
||||
)]));
|
||||
let row = sample_models_candidate_row(
|
||||
"provider-openai-cache",
|
||||
"openai",
|
||||
"openai:chat",
|
||||
"gpt-5",
|
||||
10,
|
||||
);
|
||||
let global_model_id = row.global_model_id.clone();
|
||||
let candidate_repository = Arc::new(CachedToggleMinimalCandidateSelectionReadRepository::new(
|
||||
row.clone(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
|
||||
StoredAdminGlobalModel::new(
|
||||
global_model_id.clone(),
|
||||
row.global_model_name.clone(),
|
||||
"GPT 5".to_string(),
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
1,
|
||||
0,
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000),
|
||||
)
|
||||
.expect("global model should build"),
|
||||
]),
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
|
||||
candidate_repository.clone(),
|
||||
auth_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
);
|
||||
let gateway = build_router_with_state(state.clone());
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let response = client
|
||||
.get(format!("{gateway_url}/v1/models"))
|
||||
.header("authorization", "Bearer sk-openai-models-cache")
|
||||
.send()
|
||||
.await
|
||||
.expect("initial models request should succeed");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["data"][0]["id"], "gpt-5");
|
||||
|
||||
candidate_repository.set_active(false);
|
||||
let disabled_global_model = UpdateAdminGlobalModelRecord::new(
|
||||
global_model_id,
|
||||
"GPT 5".to_string(),
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("global model update record should build");
|
||||
state
|
||||
.update_admin_global_model(&disabled_global_model)
|
||||
.await
|
||||
.expect("global model update should succeed")
|
||||
.expect("global model should update");
|
||||
|
||||
let response = client
|
||||
.get(format!("{gateway_url}/v1/models"))
|
||||
.header("authorization", "Bearer sk-openai-models-cache")
|
||||
.send()
|
||||
.await
|
||||
.expect("models request after disable should succeed");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(
|
||||
payload["data"]
|
||||
.as_array()
|
||||
.expect("data should be an array")
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() {
|
||||
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
Reference in New Issue
Block a user