Aether/apps/aether-gateway/src/scheduler/candidate/affinity.rs

use aether_scheduler_core::{
    build_scheduler_affinity_cache_key_for_api_key_id_with_client_session, candidate_affinity_hash,
    candidate_key, matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
};

use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;

use super::{
    SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState,
    SCHEDULER_AFFINITY_MAX_ENTRIES,
};

pub(super) fn build_scheduler_affinity_cache_key(
    auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
    api_format: &str,
    global_model_name: &str,
    client_session_affinity: Option<&ClientSessionAffinity>,
) -> Option<String> {
    let api_key_id = auth_snapshot
        .map(|snapshot| snapshot.api_key_id.trim())
        .filter(|value| !value.is_empty())?;
    build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
        api_key_id,
        api_format,
        global_model_name,
        client_session_affinity,
    )
}

pub(super) fn has_explicit_session_affinity(
    client_session_affinity: Option<&ClientSessionAffinity>,
) -> bool {
    client_session_affinity.is_some_and(ClientSessionAffinity::has_session_key)
}

pub(super) fn scheduler_candidate_affinity_hash(
    affinity_key: &str,
    candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> u64 {
    candidate_affinity_hash(affinity_key, candidate)
}

pub(super) fn scheduler_candidate_matches_affinity_target(
    candidate: &SchedulerMinimalCandidateSelectionCandidate,
    target: &SchedulerAffinityTarget,
) -> bool {
    matches_affinity_target(candidate, target)
}

#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn remember_scheduler_affinity(
    affinity_cache_key: Option<&str>,
    state: &(impl SchedulerRuntimeState + ?Sized),
    candidate: &SchedulerMinimalCandidateSelectionCandidate,
    expected_epoch: Option<u64>,
) {
    let Some(cache_key) = affinity_cache_key else {
        return;
    };
    let (provider_id, endpoint_id, key_id) = candidate_key(candidate);

    let _ = state.remember_scheduler_affinity_target_for_epoch(
        cache_key,
        SchedulerAffinityTarget {
            provider_id,
            endpoint_id,
            key_id,
        },
        SCHEDULER_AFFINITY_TTL,
        SCHEDULER_AFFINITY_MAX_ENTRIES,
        expected_epoch,
    );
}