Refactor pool candidate scheduling

This commit is contained in:
fawney19
2026-05-03 20:14:29 +08:00
parent 8ebee9922c
commit a24e4a793d
55 changed files with 4825 additions and 311 deletions

View File

@@ -146,6 +146,26 @@ impl MinimalCandidateSelectionRowSource for GatewayDataState {
.await
}
async fn read_minimal_candidate_selection_rows_for_api_format_and_requested_model(
&self,
api_format: &str,
requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.list_minimal_candidate_selection_rows_for_requested_model(
api_format,
requested_model_name,
)
.await
}
async fn read_minimal_candidate_selection_rows_for_api_format_and_requested_model_page(
&self,
query: &aether_data_contracts::repository::candidate_selection::StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.list_minimal_candidate_selection_rows_for_requested_model_page(query)
.await
}
async fn read_minimal_candidate_selection_rows_for_api_format(
&self,
api_format: &str,
@@ -153,6 +173,13 @@ impl MinimalCandidateSelectionRowSource for GatewayDataState {
self.list_minimal_candidate_selection_rows_for_api_format(api_format)
.await
}
async fn read_pool_key_candidate_rows_for_group(
&self,
query: &aether_data_contracts::repository::candidate_selection::StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.list_pool_key_candidate_rows_for_group(query).await
}
}
#[async_trait]

View File

@@ -77,6 +77,7 @@ use aether_data_contracts::repository::billing::{
};
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
use aether_data_contracts::repository::candidates::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,

View File

@@ -2,9 +2,10 @@ use super::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredProviderActiveGlobalModel,
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
impl GatewayDataState {
@@ -23,6 +24,35 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_minimal_candidate_selection_rows_for_requested_model(
&self,
api_format: &str,
requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => {
repository
.list_for_exact_api_format_and_requested_model(api_format, requested_model_name)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_minimal_candidate_selection_rows_for_requested_model_page(
&self,
query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => {
repository
.list_for_exact_api_format_and_requested_model_page(query)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format(
&self,
api_format: &str,
@@ -33,6 +63,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_pool_key_candidate_rows_for_group(
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => repository.list_pool_key_rows_for_group(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_public_global_models(
&self,
query: &PublicGlobalModelQuery,