Refactor pool candidate scheduling

2026-10-12 20:19:51 +08:00 · 2026-05-03 20:14:29 +08:00
parent 8ebee9922c
commit a24e4a793d
55 changed files with 4825 additions and 311 deletions
@@ -1,7 +1,9 @@
 use tracing::warn;

 use crate::ai_serving::planner::candidate_materialization::{
+    build_local_execution_candidate_attempt_source_with_serving,
    materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
+    LocalExecutionCandidateAttemptSource,
 };
 use crate::ai_serving::planner::candidate_metadata::{
    build_local_execution_candidate_contract_metadata,
@@ -169,3 +171,96 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(

    Ok((outcome.attempts, outcome.candidate_count))
 }
+
+pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a>(
+    state: &'a AppState,
+    trace_id: &str,
+    input: &LocalSameFormatProviderDecisionInput,
+    body_json: &serde_json::Value,
+    spec: LocalSameFormatProviderSpec,
+) -> Result<(LocalExecutionCandidateAttemptSource<'a>, usize), GatewayError> {
+    let spec_metadata = local_same_format_provider_spec_metadata(spec);
+    let planner_state = PlannerAppState::new(state);
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
+    let persistence_policy = build_local_candidate_persistence_policy(
+        &input.auth_context,
+        input.required_capabilities.as_ref(),
+        LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
+    );
+    let (candidates, preselection_skipped) = planner_state
+        .list_selectable_candidates_with_skip_reasons(
+            spec_metadata.api_format,
+            &input.requested_model,
+            spec_metadata.require_streaming,
+            input.required_capabilities.as_ref(),
+            Some(&input.auth_snapshot),
+            current_unix_secs(),
+        )
+        .await?;
+
+    Ok(build_local_execution_candidate_attempt_source_with_serving(
+        planner_state,
+        trace_id,
+        spec_metadata.api_format,
+        Some(&input.requested_model),
+        Some(&input.auth_snapshot),
+        input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
+        input.request_auth_channel.as_deref(),
+        persistence_policy,
+        candidates,
+        preselection_skipped
+            .into_iter()
+            .map(|item| SkippedLocalExecutionCandidate {
+                candidate: item.candidate,
+                skip_reason: item.skip_reason,
+                transport: None,
+                ranking: None,
+                extra_data: None,
+            })
+            .collect(),
+        LocalCandidateResolutionMode::Standard,
+        |eligible| {
+            let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
+                spec_metadata.api_format,
+                spec_metadata.api_format,
+            );
+            Some(build_local_execution_candidate_contract_metadata(
+                LocalExecutionCandidateMetadataParts {
+                    eligible,
+                    provider_api_format: spec_metadata.api_format,
+                    client_api_format: spec_metadata.api_format,
+                    extra_fields: serde_json::Map::new(),
+                },
+                execution_strategy,
+                conversion_mode,
+                spec_metadata.api_format,
+            ))
+        },
+        |mut skipped_candidate| {
+            let provider_api_format = skipped_candidate
+                .transport
+                .as_ref()
+                .map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
+                .unwrap_or_else(|| spec_metadata.api_format.to_string());
+            let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
+                spec_metadata.api_format,
+                provider_api_format.as_str(),
+            );
+            skipped_candidate.extra_data = Some(
+                build_local_execution_candidate_contract_metadata_for_candidate(
+                    &skipped_candidate.candidate,
+                    skipped_candidate.transport_ref(),
+                    provider_api_format.as_str(),
+                    spec_metadata.api_format,
+                    serde_json::Map::new(),
+                    execution_strategy,
+                    conversion_mode,
+                    provider_api_format.as_str(),
+                ),
+            );
+            skipped_candidate
+        },
+    )
+    .await)
+}
@@ -8,10 +8,12 @@ pub(crate) use self::build::{
    maybe_build_sync_local_same_format_provider_decision_payload,
 };
 pub(crate) use self::candidates::{
+    build_local_same_format_provider_candidate_attempt_source,
    materialize_local_same_format_provider_candidate_attempts,
    resolve_local_same_format_provider_decision_input,
 };
 pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
 pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
+pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalSameFormatProviderCandidateAttemptSource;
 pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
 pub(crate) use crate::ai_serving::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};