mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
Refactor pool candidate scheduling
This commit is contained in:
@@ -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]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user