diff --git a/.github/workflows/build-docker-images.yml b/.github/workflows/build-docker-images.yml new file mode 100644 index 000000000..a1bd3457d --- /dev/null +++ b/.github/workflows/build-docker-images.yml @@ -0,0 +1,211 @@ +name: Build Docker Images + +on: + push: + branches: ["main", "test"] + workflow_dispatch: + +env: + REGISTRY: ghcr.io + IMAGE_PREFIX: ${{ github.repository_owner }}/1xtoken + CLUSTER_GOPROXY: https://proxy.golang.org + # 阿里云镜像仓库：国内服务器拉 ghcr.io 太慢，构建时同步推一份到阿里云 + ALIYUN_REGISTRY: crpi-w0utadukedw6a9ld.cn-hangzhou.personal.cr.aliyuncs.com + ALIYUN_NAMESPACE: 1xtoken + +jobs: + # ------------------------------------------------------------------ + # 阶段 1: 智能分析 (Configuration Job) + # ------------------------------------------------------------------ + configure: + name: Configure Matrix + runs-on: astro-web3-arc-runner-set + outputs: + has-changes: ${{ steps.detect.outputs.has-changes }} + matrix: ${{ steps.matrix.outputs.matrix }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Detect changes to build inputs + id: detect + run: | + # 只在影响镜像构建的文件变更时才触发构建，避免无关改动白白排队。 + # workflow_dispatch 始终构建。 + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "has-changes=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + BEFORE="${{ github.event.before }}" + AFTER="${{ github.event.after }}" + if [ -z "$BEFORE" ] || [ "$BEFORE" = "0000000000000000000000000000000000000000" ]; then + echo "has-changes=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + CHANGED=$(git diff --name-only "$BEFORE" "$AFTER" || true) + PATTERNS='^(docker/Dockerfile|\.cargo/config\.toml|bindings/python/pyproject\.toml|model_gateway/|crates/|Cargo\.lock)' + if printf '%s\n' "$CHANGED" | grep -qE "$PATTERNS"; then + echo "has-changes=true" >> "$GITHUB_OUTPUT" + else + echo "has-changes=false" >> "$GITHUB_OUTPUT" + fi + + - name: Build matrix + id: matrix + run: | + # 单镜像仓库，matrix 只有一行；保留 include 结构便于后续多 app 扩展。 + MATRIX=$(cat <<'JSON' + {"include":[{"app":"smg","dockerfile":"docker/Dockerfile","context":"."}]} + JSON + ) + # jq -c 压成单行 JSON 给 fromJSON 用 + echo "matrix=$(printf '%s' "$MATRIX" | jq -c .)" >> "$GITHUB_OUTPUT" + + # ------------------------------------------------------------------ + # 阶段 2: 构建 (Build Job) + # ------------------------------------------------------------------ + build-and-push: + name: Build ${{ matrix.app }} + needs: configure + if: needs.configure.outputs.has-changes == 'true' + runs-on: astro-web3-arc-runner-set + permissions: + contents: read + packages: write + strategy: + fail-fast: false + matrix: + include: ${{ fromJSON(needs.configure.outputs.matrix) }} + steps: + - name: Export pod env to GitHub env + run: | + echo "ZOT_REGISTRY=${ZOT_REGISTRY}" >> $GITHUB_ENV + echo "CLUSTER_GOPROXY=${CLUSTER_GOPROXY}" >> $GITHUB_ENV + echo "BUILDKIT_HOST=${BUILDKIT_HOST}" >> $GITHUB_ENV + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + id: setup-buildx + uses: docker/setup-buildx-action@v3 + with: + driver: remote + endpoint: ${{ env.BUILDKIT_HOST }} + + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + # 阿里云镜像仓库偶发 TLS 握手超时（自建 runner 到杭州个人版实例网络抖动）， + # docker/login-action 内部只 login 一次无重试，这里改成带退避的重试。 + - name: Log in to Aliyun Registry + env: + ALIYUN_REGISTRY: ${{ env.ALIYUN_REGISTRY }} + ALIYUN_USERNAME: ${{ secrets.ALIYUN_USERNAME }} + ALIYUN_PASSWORD: ${{ secrets.ALIYUN_PASSWORD }} + run: | + max_attempts=5 + for attempt in $(seq 1 "$max_attempts"); do + if printf '%s' "$ALIYUN_PASSWORD" | docker login "$ALIYUN_REGISTRY" -u "$ALIYUN_USERNAME" --password-stdin; then + echo "Aliyun login succeeded on attempt $attempt" + exit 0 + fi + echo "Aliyun login attempt $attempt failed, backing off before retry..." + sleep $((attempt * 5)) + done + echo "::error::Aliyun registry login failed after $max_attempts attempts" + exit 1 + + - name: Metadata + id: meta + uses: docker/metadata-action@v5 + with: + images: | + ${{ env.REGISTRY }}/${{ env.IMAGE_PREFIX }}-${{ matrix.app }} + ${{ env.ZOT_REGISTRY }}/${{ env.IMAGE_PREFIX }}-${{ matrix.app }} + ${{ env.ALIYUN_REGISTRY }}/${{ env.ALIYUN_NAMESPACE }}/1xtoken-${{ matrix.app }} + tags: | + type=sha,prefix=${{ github.ref_name }}-,suffix=-${{ github.run_number }} + + # 构建与推送拆分：构建是确定性故障（pnpm install、Dockerfile 语法、编译错误）， + # 重试无意义，失败直接在这一步标红并阻断 job，报错落在真实出错位置； + # 推送才可能因阿里云个人版仓库鉴权/网络抖动失败，值得退避重试。 + - name: Build image + id: build + uses: docker/build-push-action@v6 + with: + context: ${{ matrix.context || '.' }} + file: ${{ matrix.dockerfile }} + push: false + load: false + cache-from: type=registry,ref=${{ env.ZOT_REGISTRY }}/cache/${{ env.IMAGE_PREFIX }}-${{ matrix.app }}:buildcache + cache-to: type=registry,ref=${{ env.ZOT_REGISTRY }}/cache/${{ env.IMAGE_PREFIX }}-${{ matrix.app }}:buildcache,mode=min + + # 阿里云鉴权端点偶发 connection reset by peer（push 时拉 oauth token 网络抖动）， + # docker/build-push-action 内部不重试推送，这里用「首尝试 + 2 次退避重试 + 状态校验」兜底。 + # 构建产物已在 BuildKit 缓存里，重试几乎只重推 manifest，成本可忽略。 + - name: Push image + id: push-1 + continue-on-error: true + uses: docker/build-push-action@v6 + with: + context: ${{ matrix.context || '.' }} + file: ${{ matrix.dockerfile }} + push: true + tags: ${{ steps.meta.outputs.tags }} + cache-from: type=registry,ref=${{ env.ZOT_REGISTRY }}/cache/${{ env.IMAGE_PREFIX }}-${{ matrix.app }}:buildcache + + - name: Retry push (attempt 2) + if: steps.push-1.outcome == 'failure' + run: | + echo "::warning::首次推送失败（可能阿里云鉴权端点网络抖动），退避 5s 后重试..." + sleep 5 + + - name: Push image (retry 2) + id: push-2 + continue-on-error: true + if: steps.push-1.outcome == 'failure' + uses: docker/build-push-action@v6 + with: + context: ${{ matrix.context || '.' }} + file: ${{ matrix.dockerfile }} + push: true + tags: ${{ steps.meta.outputs.tags }} + cache-from: type=registry,ref=${{ env.ZOT_REGISTRY }}/cache/${{ env.IMAGE_PREFIX }}-${{ matrix.app }}:buildcache + + - name: Retry push (attempt 3) + if: steps.push-2.outcome == 'failure' + run: | + echo "::warning::第二次推送失败，退避 10s 后最后一次重试..." + sleep 10 + + - name: Push image (retry 3) + id: push-3 + continue-on-error: true + if: steps.push-2.outcome == 'failure' + uses: docker/build-push-action@v6 + with: + context: ${{ matrix.context || '.' }} + file: ${{ matrix.dockerfile }} + push: true + tags: ${{ steps.meta.outputs.tags }} + cache-from: type=registry,ref=${{ env.ZOT_REGISTRY }}/cache/${{ env.IMAGE_PREFIX }}-${{ matrix.app }}:buildcache + + # 三次推送任一成功即通过；全败则标红并阻断 job。 + - name: Verify push result + if: always() + run: | + if [ "${{ steps.push-1.outcome }}" = "success" ] \ + || [ "${{ steps.push-2.outcome }}" = "success" ] \ + || [ "${{ steps.push-3.outcome }}" = "success" ]; then + echo "镜像推送成功" + else + echo "::error::镜像推送在 3 次尝试后均失败（阿里云个人版仓库鉴权/网络问题）" + exit 1 + fi diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index fc17c72b2..ebb748d89 100755 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -644,6 +644,45 @@ pub enum PolicyConfig { cache_boundaries: Vec, }, + /// Cache-aware length policy: cache affinity with a long/short pool split + /// driven by the `pool` worker label (`pool=long` → long pool, otherwise + /// short pool). Step 1-3 mirror `cache_aware` (string tree only); step 4 + /// routes by uncached prefill tokens. See `policies/cache_aware_length.rs`. + #[serde(rename = "cache_aware_length")] + CacheAwareLength { + /// Minimum matched-prefix share before a request pins to a holder. + #[serde(alias = "cache_match_threshold")] + #[serde(default = "default_cal_cache_threshold")] + cache_threshold: f32, + /// Spill gate, absolute part: the global imbalance fires when the + /// healthy-fleet load spread exceeds this. + #[serde(alias = "spill_abs_threshold")] + #[serde(default = "default_cal_balance_abs_threshold")] + balance_abs_threshold: usize, + /// Spill gate, relative part (multiple of the healthy-fleet min load); + /// fires only together with `balance_abs_threshold`. + #[serde(alias = "spill_rel_threshold")] + #[serde(default = "default_cal_balance_rel_threshold")] + balance_rel_threshold: f32, + #[serde(default = "default_cal_eviction_interval_secs")] + eviction_interval_secs: u64, + #[serde(default = "default_cal_max_tree_size")] + max_tree_size: usize, + /// Divisor for char-level token estimation when `X-Prompt-Tokens` is + /// absent (default 4). + #[serde(default = "default_cal_chars_per_token")] + chars_per_token: usize, + /// Uncached-prefill-token boundary between long and short requests. + #[serde(default = "default_cal_long_prefill_threshold")] + long_prefill_threshold: usize, + /// Load ceiling for the long pool (`pool=long` workers). + #[serde(default = "default_cal_long_pool_max_load")] + long_pool_max_load: usize, + /// Load ceiling for the short pool (remaining workers). + #[serde(default = "default_cal_short_pool_max_load")] + short_pool_max_load: usize, + }, + /// Power-of-two choices policy: samples two workers and routes to the one /// with the lower expected wait, scored like `least_load` /// (`(queued_tokens + inflight_tokens) / throughput + kv_pressure_weight * k/(1-k)`). @@ -777,6 +816,35 @@ fn default_cache_ttl_secs() -> u64 { 180 } +// cache_aware_length defaults (kept aligned with CacheAwareLengthConfig::default). +fn default_cal_cache_threshold() -> f32 { + 0.3 +} +fn default_cal_balance_abs_threshold() -> usize { + 32 +} +fn default_cal_balance_rel_threshold() -> f32 { + 1.1 +} +fn default_cal_eviction_interval_secs() -> u64 { + 30 +} +fn default_cal_max_tree_size() -> usize { + 10000 +} +fn default_cal_chars_per_token() -> usize { + 4 +} +fn default_cal_long_prefill_threshold() -> usize { + 100_000 +} +fn default_cal_long_pool_max_load() -> usize { + 4 +} +fn default_cal_short_pool_max_load() -> usize { + 32 +} + fn default_prefix_token_count() -> usize { 256 } @@ -828,6 +896,7 @@ impl PolicyConfig { PolicyConfig::RoundRobin => "round_robin", PolicyConfig::Passthrough => "passthrough", PolicyConfig::CacheAware { .. } => "cache_aware", + PolicyConfig::CacheAwareLength { .. } => "cache_aware_length", PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::LeastLoad { .. } => "least_load", PolicyConfig::Bucket { .. } => "bucket", diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index ad70151f6..474a69235 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -565,6 +565,81 @@ impl ConfigValidator { }); } } + PolicyConfig::CacheAwareLength { + cache_threshold, + balance_abs_threshold: _, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + chars_per_token, + long_prefill_threshold, + long_pool_max_load, + short_pool_max_load, + } => { + if !(0.0..=1.0).contains(cache_threshold) { + return Err(ConfigError::InvalidValue { + field: "cache_threshold".to_string(), + value: cache_threshold.to_string(), + reason: "Must be between 0.0 and 1.0".to_string(), + }); + } + + if *balance_rel_threshold < 1.0 { + return Err(ConfigError::InvalidValue { + field: "balance_rel_threshold".to_string(), + value: balance_rel_threshold.to_string(), + reason: "Must be >= 1.0".to_string(), + }); + } + + if *eviction_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "eviction_interval_secs".to_string(), + value: eviction_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *max_tree_size == 0 { + return Err(ConfigError::InvalidValue { + field: "max_tree_size".to_string(), + value: max_tree_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *chars_per_token == 0 { + return Err(ConfigError::InvalidValue { + field: "chars_per_token".to_string(), + value: chars_per_token.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *long_prefill_threshold == 0 { + return Err(ConfigError::InvalidValue { + field: "long_prefill_threshold".to_string(), + value: long_prefill_threshold.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *long_pool_max_load == 0 { + return Err(ConfigError::InvalidValue { + field: "long_pool_max_load".to_string(), + value: long_pool_max_load.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *short_pool_max_load == 0 { + return Err(ConfigError::InvalidValue { + field: "short_pool_max_load".to_string(), + value: short_pool_max_load.to_string(), + reason: "Must be > 0".to_string(), + }); + } + } PolicyConfig::PowerOfTwo { load_check_interval_secs, } => { diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 29aea9f88..85c543e87 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -230,7 +230,7 @@ struct CliArgs { // ==================== Routing Policy ==================== /// Load balancing policy to use - #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "passthrough", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")] + #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "passthrough", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")] policy: String, /// Minimum matched-prefix share (0.0-1.0) before cache-aware routing @@ -369,6 +369,27 @@ struct CliArgs { #[arg(long, default_value_t = 180, value_parser = clap::value_parser!(u64).range(1..), help_heading = "Routing Policy")] cache_ttl_secs: u64, + // ---- cache_aware_length policy ---- + /// Divisor for char-level token estimation when X-Prompt-Tokens is absent + /// (cache_aware_length policy). Default 4. + #[arg(long, default_value_t = 4, help_heading = "Routing Policy")] + chars_per_token: usize, + + /// Uncached-prefill-token boundary between long and short requests + /// (cache_aware_length policy). Default 100000. + #[arg(long, default_value_t = 100_000, help_heading = "Routing Policy")] + long_prefill_threshold: usize, + + /// Load ceiling for the long pool (pool=long workers) in the + /// cache_aware_length policy. Default 4. + #[arg(long, default_value_t = 4, help_heading = "Routing Policy")] + long_pool_max_load: usize, + + /// Load ceiling for the short pool (remaining workers) in the + /// cache_aware_length policy. Default 32. + #[arg(long, default_value_t = 32, help_heading = "Routing Policy")] + short_pool_max_load: usize, + /// How long an unused sticky routing key stays pinned: keys idle beyond /// this many seconds are evicted from the manual-policy / sticky-session /// map @@ -470,11 +491,11 @@ struct CliArgs { decode: Vec, /// Specific policy for prefill nodes in PD mode - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] prefill_policy: Option, /// Specific policy for decode nodes in PD mode - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] decode_policy: Option, /// Specific policy for encode nodes in EPD mode. Defaults to consistent_hashing. @@ -1382,6 +1403,17 @@ impl CliArgs { cache_ttl_secs: self.cache_ttl_secs, cache_boundaries: self.cache_boundaries.clone(), }, + "cache_aware_length" => PolicyConfig::CacheAwareLength { + cache_threshold: self.cache_threshold, + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + eviction_interval_secs: self.eviction_interval, + max_tree_size: self.max_tree_size, + chars_per_token: self.chars_per_token, + long_prefill_threshold: self.long_prefill_threshold, + long_pool_max_load: self.long_pool_max_load, + short_pool_max_load: self.short_pool_max_load, + }, "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, }, diff --git a/model_gateway/src/policies/cache_aware_length.rs b/model_gateway/src/policies/cache_aware_length.rs new file mode 100644 index 000000000..2def2b78b --- /dev/null +++ b/model_gateway/src/policies/cache_aware_length.rs @@ -0,0 +1,845 @@ +/* + Cache-Aware Length Load Balancing Router (cache_aware_length) + + Splits healthy workers into a "long pool" and a "short pool" by the + `pool` worker label (`pool=long` → long pool, otherwise short pool), then + applies cache-affinity routing on top of the split. Designed for P/D + disaggregated prefill fleets and regular single-node fleets alike — pool + membership is label-driven and independent of `WorkerType`. + + Routing pipeline (5 steps), mirroring `cache_aware` for steps 1-3 and + adding a long/short split at step 4: + + Step 1 - Health filter + is_available()==false workers are dropped; empty → 503 (None), no tree. + Step 2 - Global imbalance check (same formula as cache_aware) + (max_load - min_load) > abs_threshold AND + max_load > min_load * rel_threshold + YES → route to healthy min-load, record tree, return. + NO → continue. + Step 3 - Cache hit check (approximate string tree, char-level) + tree missing → random healthy worker, no tree (init race). + match_rate = matched_chars / input_chars + match_rate > cache_threshold + YES → hit branch: route to the highest-matching worker regardless of + pool; if that worker is unhealthy, clean its stale tenant and + fall back to the first healthy worker. Record tree. + NO → continue to step 4. + Step 4 - No-cache branch: split by uncached prefill tokens + token source (priority): + 1. X-Prompt-Tokens header (exact, supplied by an upstream gateway). + 2. (input_chars - matched_chars) / chars_per_token (char estimate). + 3. neither computable → all-healthy min-load, record tree. + long pool = healthy workers with labels["pool"] == "long" + short pool = remaining healthy workers + uncached >= long_prefill_threshold (long request): + long pool has free worker (load < long_pool_max_load) + → long pool min-load + else short pool has an idle worker (load == 0) + → that worker (long→short overflow) + else long pool has a healthy worker + → long pool min-load (queue) + else → all-healthy min-load + uncached < long_prefill_threshold (short request): + short pool has free worker (load < short_pool_max_load) + → short pool min-load + else long pool has free worker + → long pool min-load (short→long overflow) + else short pool has a worker + → short pool min-load (fallback queue) + else long pool has a worker + → long pool min-load + else (both pools empty) → all-healthy min-load + Step 5 - Record tree + return + tree.insert_text(text, selected worker url) and increment_processed() + for every selection except the 503 and the no-tree random fallback. + + Configuration Parameters: + ------------------------ + cache_threshold: Min prefix match ratio for hit routing (0.0-1.0) + balance_abs_threshold: Absolute load diff for global imbalance detection + balance_rel_threshold: Relative load ratio for global imbalance detection + eviction_interval_secs: Interval between LRU eviction cycles + max_tree_size: Max total chars of each model's approximate tree + chars_per_token: Divisor for char-level token estimation (default 4) + long_prefill_threshold: Uncached-token boundary between long and short + long_pool_max_load: Load ceiling for the long pool + short_pool_max_load: Load ceiling for the short pool +*/ + +use std::sync::Arc; + +use dashmap::DashMap; +use kv_index::{PrefixMatchResult, TenantId, Tree}; +use rand::RngExt; +use tracing::debug; + +use super::{ + normalize_model_key, utils::PeriodicTask, CacheAwareLengthConfig, LoadBalancingPolicy, + SelectWorkerInfo, +}; +use crate::{observability::metrics::Metrics, worker::Worker}; + +/// HTTP header carrying the exact prompt token count, supplied by an +/// upstream gateway that has already tokenized the request. Case-insensitive +/// per the `http` crate's `HeaderName` equality. +const HEADER_PROMPT_TOKENS: &str = "x-prompt-tokens"; + +/// Cache-aware length routing policy. +/// +/// Routes requests based on cache affinity and a long/short pool split. The +/// split is driven by the `pool` worker label (`pool=long`); workers without +/// the label form the short pool. Self-contained: maintains only a +/// per-model string radix tree and a background eviction task. Does not +/// participate in mesh tree sync, KV-event monitoring, or the hash index — +/// those are handled by `CacheAwarePolicy`. +#[derive(Debug)] +pub struct CacheAwareLengthPolicy { + config: CacheAwareLengthConfig, + /// String-based trees for HTTP connections (text input), keyed by model. + string_trees: Arc>>, + _eviction_task: Option, +} + +impl Default for CacheAwareLengthPolicy { + fn default() -> Self { + Self::new() + } +} + +impl CacheAwareLengthPolicy { + pub fn new() -> Self { + Self::with_config(CacheAwareLengthConfig::default()) + } + + pub fn with_config(config: CacheAwareLengthConfig) -> Self { + let string_trees = Arc::new(DashMap::>::new()); + + let eviction_task = (config.eviction_interval_secs > 0).then(|| { + let trees_clone = Arc::clone(&string_trees); + let max_tree_size = config.max_tree_size; + PeriodicTask::spawn( + config.eviction_interval_secs, + "LengthTreeEviction", + move || { + let mut total_chars = 0usize; + for tree_ref in trees_clone.iter() { + let model_id = tree_ref.key(); + let tree = tree_ref.value(); + tree.evict_tenant_by_size(max_tree_size); + + let counts = tree.get_tenant_char_count(); + let chars: usize = counts.values().sum(); + let tenants = counts.len(); + total_chars += chars; + Metrics::set_cache_tree_chars(model_id, chars); + Metrics::set_cache_tree_tenants(model_id, "string", tenants); + + debug!( + "String tree eviction completed for model {}, max_size: {}", + model_id, max_tree_size + ); + } + tracing::info!( + "Length tree memory: string_trees={} models / {} chars", + trees_clone.len(), + total_chars, + ); + }, + ) + }); + + Self { + config, + string_trees, + _eviction_task: eviction_task, + } + } + + /// Initialize trees for a set of workers (seed each worker as a tenant of + /// the root so it is always a cache-hit candidate). Call after workers are + /// registered. + pub fn init_workers(&self, workers: &[Arc]) { + // Group workers by model + let mut model_workers: std::collections::HashMap>> = + std::collections::HashMap::new(); + for worker in workers { + let tree_key = normalize_model_key(worker.model_id()); + model_workers + .entry(tree_key.to_string()) + .or_default() + .push(worker); + } + + for (tree_key, model_workers) in model_workers { + let string_tree = self + .string_trees + .entry(tree_key) + .or_insert_with(|| Arc::new(Tree::new())); + for worker in model_workers { + string_tree.insert_text("", worker.url()); + } + } + } + + /// Add a single worker to the trees (incremental update). + pub fn add_worker(&self, worker: &dyn Worker) { + let tree_key = normalize_model_key(worker.model_id()).to_string(); + let string_tree = self + .string_trees + .entry(tree_key) + .or_insert_with(|| Arc::new(Tree::new())); + string_tree.insert_text("", worker.url()); + } + + /// Remove a worker from the trees, purging its tenant from every model's + /// string tree. A removed worker's tenant count never grows again, so + /// size-based eviction alone would retain its subtree forever. + pub fn remove_worker_by_url(&self, url: &str) { + let tenant: TenantId = Arc::from(url); + for tree_ref in self.string_trees.iter() { + tree_ref.value().remove_tenant_all(&tenant); + } + } +} + +impl LoadBalancingPolicy for CacheAwareLengthPolicy { + fn select_worker(&self, workers: &[Arc], info: &SelectWorkerInfo) -> Option { + let text = info.request_text.unwrap_or(""); + + // Step 1: health filter — single O(workers) gather reading each worker + // once via routing_state() (health + load + processed under one guard). + let mut healthy_indices: Vec = Vec::with_capacity(workers.len()); + let mut min_key: Option<(usize, usize, usize)> = None; + let mut min_load_idx: Option = None; + for (idx, worker) in workers.iter().enumerate() { + let state = worker.routing_state(); + if state.eligible() { + healthy_indices.push(idx); + let key = (state.load, state.processed, idx); + match min_key { + Some(best) if key >= best => {} + _ => { + min_key = Some(key); + min_load_idx = Some(idx); + } + } + } + } + + if healthy_indices.is_empty() { + return None; // 503, do not record tree + } + + let model_id = normalize_model_key(workers[healthy_indices[0]].model_id()).to_string(); + + // Step 2: global imbalance check (same formula as cache_aware). + // The min/max are over the healthy fleet only. + let healthy_min = min_key.map(|(load, _, _)| load).unwrap_or(0); + let healthy_max = healthy_indices + .iter() + .map(|&i| workers[i].routing_state().load) + .max() + .unwrap_or(0); + let abs_diff = healthy_max.saturating_sub(healthy_min); + let rel_threshold = self.config.balance_rel_threshold * healthy_min as f32; + if abs_diff > self.config.balance_abs_threshold && healthy_max as f32 > rel_threshold { + // min_load_idx is guaranteed Some here (healthy_indices non-empty + // populates it in the same loop), but `?` keeps the analyzer happy + // without a deny-listed unwrap. + let selected = min_load_idx?; + self.record_tree(&model_id, text, workers[selected].url()); + workers[selected].increment_processed(); + debug!( + branch = "global_imbalance_min_load", + worker = workers[selected].url(), + model_id = model_id, + "cache_aware_length selection" + ); + return Some(selected); + } + + // Step 3: cache hit check (approximate string tree, char-level). + let tree = self + .string_trees + .get(&model_id) + .map(|entry| entry.value().clone()); + + let Some(tree) = tree else { + // tree missing: init race — random healthy worker, do not record. + let idx = healthy_indices[rand::rng().random_range(0..healthy_indices.len())]; + debug!( + branch = "no_tree_random", + worker = workers[idx].url(), + model_id = model_id, + "cache_aware_length selection" + ); + return Some(idx); + }; + + let result = tree.match_prefix_with_counts(text); + let match_rate = if result.input_char_count == 0 { + 0.0 + } else { + result.matched_char_count as f32 / result.input_char_count as f32 + }; + + if match_rate > self.config.cache_threshold { + // Cache hit: route to the highest-matching worker regardless of pool. + if let Some(idx) = Self::select_matched_candidate(workers, &healthy_indices, &result) { + self.record_tree(&model_id, text, workers[idx].url()); + workers[idx].increment_processed(); + debug!( + branch = "cache_hit", + worker = workers[idx].url(), + match_rate, + model_id = model_id, + "cache_aware_length selection" + ); + return Some(idx); + } + // Hit but the matched worker is unhealthy: clean stale tenant and + // fall back to the first healthy worker (record tree). + if let Some(tenant) = result.matched_tenants.first() { + tree.remove_tenant_all(tenant); + } + let idx = healthy_indices[0]; + self.record_tree(&model_id, text, workers[idx].url()); + workers[idx].increment_processed(); + debug!( + branch = "hit_unhealthy_first_healthy", + worker = workers[idx].url(), + model_id = model_id, + "cache_aware_length selection" + ); + return Some(idx); + } + + // Step 4: no-cache branch — split by uncached prefill tokens. + let uncached_tokens = self.compute_uncached_tokens(info, &result); + let Some(uncached) = uncached_tokens else { + // Neither source computable → all-healthy min-load. + let selected = min_load_idx?; + self.record_tree(&model_id, text, workers[selected].url()); + workers[selected].increment_processed(); + debug!( + branch = "uncached_unknown_min_load", + worker = workers[selected].url(), + model_id = model_id, + "cache_aware_length selection" + ); + return Some(selected); + }; + + let long_indices: Vec = healthy_indices + .iter() + .copied() + .filter(|&i| is_long_pool(&*workers[i])) + .collect(); + let short_indices: Vec = healthy_indices + .iter() + .copied() + .filter(|&i| !is_long_pool(&*workers[i])) + .collect(); + + let selected = if uncached >= self.config.long_prefill_threshold { + self.select_long_request( + workers, + &long_indices, + &short_indices, + &healthy_indices, + min_load_idx, + ) + } else { + self.select_short_request( + workers, + &long_indices, + &short_indices, + &healthy_indices, + min_load_idx, + ) + }; + + // Step 5: record tree + return. + let selected = selected.unwrap_or_else(|| min_load_idx.unwrap_or(healthy_indices[0])); + self.record_tree(&model_id, text, workers[selected].url()); + workers[selected].increment_processed(); + debug!( + branch = "pool_split", + worker = workers[selected].url(), + uncached_tokens = uncached, + is_long = uncached >= self.config.long_prefill_threshold, + model_id = model_id, + "cache_aware_length selection" + ); + Some(selected) + } + + fn name(&self) -> &'static str { + "cache_aware_length" + } + + fn needs_request_text(&self) -> bool { + true + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +// Private helper methods +impl CacheAwareLengthPolicy { + /// Record a routing decision into the model's string tree. + fn record_tree(&self, model_id: &str, text: &str, worker_url: &str) { + if let Some(tree) = self.string_trees.get(model_id).map(|e| e.value().clone()) { + tree.insert_text(text, worker_url); + } + } + + /// Pressure-select among the tenants holding the matched prefix. The + /// match is on the raw prefix, so every matched tenant holds the same + /// prefix; select the least-loaded healthy matched tenant. + fn select_matched_candidate( + workers: &[Arc], + healthy_indices: &[usize], + result: &PrefixMatchResult, + ) -> Option { + let mut best: Option = None; + let mut best_key: Option<(usize, usize, usize)> = None; + for &idx in healthy_indices { + let url = workers[idx].url(); + if result.matched_tenants.iter().any(|t| t.as_ref() == url) { + let state = workers[idx].routing_state(); + let key = (state.load, state.processed, idx); + match best_key { + Some(b) if key >= b => {} + _ => { + best = Some(idx); + best_key = Some(key); + } + } + } + } + best + } + + /// Compute uncached prefill tokens by priority: + /// 1. X-Prompt-Tokens header (exact). + /// 2. (input_chars - matched_chars) / chars_per_token (char estimate). + /// 3. None when neither is computable. + fn compute_uncached_tokens( + &self, + info: &SelectWorkerInfo, + result: &PrefixMatchResult, + ) -> Option { + // 1. Exact header value. + if let Some(n) = parse_prompt_tokens_header(info.headers) { + return Some(n); + } + // 2. Char-level estimate from the match result. + let uncached_chars = result + .input_char_count + .saturating_sub(result.matched_char_count); + if uncached_chars > 0 && self.config.chars_per_token > 0 { + // Ceiling so a fractional block still counts as one token. + let est = uncached_chars.div_ceil(self.config.chars_per_token); + return Some(est); + } + None + } + + /// Long request (uncached >= long_prefill_threshold). Does not overflow to + /// a short-pool worker that already has load. + fn select_long_request( + &self, + workers: &[Arc], + long_indices: &[usize], + short_indices: &[usize], + _healthy_indices: &[usize], + min_load_idx: Option, + ) -> Option { + let long_has_free = pool_has_free(workers, long_indices, self.config.long_pool_max_load); + if long_has_free { + return pool_min_load_worker(workers, long_indices); + } + // Long pool full/unhealthy: overflow to an idle short-pool worker only. + if let Some(idx) = pool_idle_worker(workers, short_indices) { + return Some(idx); + } + // Short pool all busy: queue on long pool if it still has a worker. + if let Some(idx) = pool_min_load_worker(workers, long_indices) { + return Some(idx); + } + // Long pool fully unhealthy and short pool busy: all-healthy min-load. + min_load_idx + } + + /// Short request (uncached < long_prefill_threshold). May overflow to an + /// idle long-pool worker. + fn select_short_request( + &self, + workers: &[Arc], + long_indices: &[usize], + short_indices: &[usize], + _healthy_indices: &[usize], + min_load_idx: Option, + ) -> Option { + let short_has_free = pool_has_free(workers, short_indices, self.config.short_pool_max_load); + if short_has_free { + return pool_min_load_worker(workers, short_indices); + } + // Short pool full: overflow to long pool if it has a free worker. + let long_has_free = pool_has_free(workers, long_indices, self.config.long_pool_max_load); + if long_has_free { + return pool_min_load_worker(workers, long_indices); + } + // Both full: queue on short pool if it has a worker. + if let Some(idx) = pool_min_load_worker(workers, short_indices) { + return Some(idx); + } + // Short pool empty: queue on long pool if it has a worker. + if let Some(idx) = pool_min_load_worker(workers, long_indices) { + return Some(idx); + } + // Both pools empty: all-healthy min-load. + min_load_idx + } +} + +/// Whether a worker belongs to the long pool (`labels["pool"] == "long"`). +fn is_long_pool(worker: &dyn Worker) -> bool { + worker + .metadata() + .spec + .labels + .get("pool") + .is_some_and(|v| v == "long") +} + +/// Does any worker in `pool` have `load() < max_load`? +fn pool_has_free(workers: &[Arc], pool: &[usize], max_load: usize) -> bool { + pool.iter() + .any(|&i| workers[i].routing_state().load < max_load) +} + +/// Return the index of an idle (`load == 0`) worker in `pool`, if any. +fn pool_idle_worker(workers: &[Arc], pool: &[usize]) -> Option { + pool.iter() + .copied() + .find(|&i| workers[i].routing_state().load == 0) +} + +/// Lowest-load worker in `pool` with the `(load, processed, idx)` tie-break. +/// Returns `None` when `pool` is empty. +fn pool_min_load_worker(workers: &[Arc], pool: &[usize]) -> Option { + let mut best: Option = None; + let mut best_key: Option<(usize, usize, usize)> = None; + for &idx in pool { + let state = workers[idx].routing_state(); + let key = (state.load, state.processed, idx); + match best_key { + Some(b) if key >= b => {} + _ => { + best = Some(idx); + best_key = Some(key); + } + } + } + best +} + +/// Parse the `X-Prompt-Tokens` header into a token count. Returns `None` on +/// missing/unparseable values. Header lookup is case-insensitive. +fn parse_prompt_tokens_header(headers: Option<&http::HeaderMap>) -> Option { + let headers = headers?; + headers + .get(HEADER_PROMPT_TOKENS) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.trim().parse::().ok()) + .filter(|n| *n > 0) +} + +#[cfg(test)] +mod tests { + use openai_protocol::worker::{HealthCheckConfig, WorkerStatus}; + + use super::*; + use crate::worker::{BasicWorkerBuilder, WorkerType}; + + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + + /// Build a worker with an optional `pool` label and a pre-set load. + fn make_worker(url: &str, pool: Option<&str>, load: usize) -> Arc { + let mut builder = BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .health_config(no_health_check()); + if let Some(p) = pool { + builder = builder.label("pool", p); + } + let worker: Arc = Arc::new(builder.build()); + for _ in 0..load { + // Leak the guard so the load stays elevated for the test without + // holding a handle; the process tears down on test exit. + std::mem::forget(crate::worker::WorkerLoadGuard::new( + Arc::clone(&worker), + None, + )); + } + worker + } + + fn info_with_text(text: &str) -> SelectWorkerInfo<'_> { + SelectWorkerInfo { + request_text: Some(text), + ..Default::default() + } + } + + /// Build a `SelectWorkerInfo` with an `X-Prompt-Tokens` header. The header + /// map must outlive the returned info — callers hold it in the same scope. + fn info_with_header<'a>(headers: &'a http::HeaderMap, text: &'a str) -> SelectWorkerInfo<'a> { + SelectWorkerInfo { + request_text: Some(text), + headers: Some(headers), + ..Default::default() + } + } + + fn tokens_headers(tokens: usize) -> http::HeaderMap { + let mut headers = http::HeaderMap::new(); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_str(&tokens.to_string()).unwrap(), + ); + headers + } + + fn test_config() -> CacheAwareLengthConfig { + CacheAwareLengthConfig { + eviction_interval_secs: 0, // disable eviction thread in tests + long_prefill_threshold: 100_000, + long_pool_max_load: 2, + short_pool_max_load: 2, + chars_per_token: 4, + ..Default::default() + } + } + + #[test] + fn step1_returns_none_when_all_unhealthy() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("k") + .health_config(no_health_check()) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("k") + .health_config(no_health_check()) + .build(), + ), + ]; + for w in &workers { + w.set_status(WorkerStatus::NotReady); + } + assert!(policy + .select_worker(&workers, &info_with_text("hello")) + .is_none()); + } + + #[test] + fn step3_cache_hit_pins_to_same_worker() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", None, 0), + ]; + policy.init_workers(&workers); + + let prompt = "shared long prompt prefix that both workers could cache"; + let idx1 = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + // Same prompt again → cache hit, same worker. + let idx2 = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + assert_eq!(idx1, idx2); + } + + #[test] + fn step3_tree_missing_falls_back_random() { + // No init_workers → no tree for the model → random healthy, no panic. + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", None, 0), + ]; + let idx = policy + .select_worker(&workers, &info_with_text("novel prompt")) + .unwrap(); + assert!(idx < workers.len()); + } + + #[test] + fn step4_long_request_uses_long_pool_when_free() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), // short pool, idle + make_worker("http://w2:8000", Some("long"), 0), // long pool, free + ]; + policy.init_workers(&workers); + // novel prompt → no hit → long request via header. + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_long_request_overflows_to_idle_short() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), // short, idle (load 0) + make_worker("http://w2:8000", Some("long"), 2), // long, full + ]; + policy.init_workers(&workers); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_long_request_queues_on_long_when_short_busy() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 1), // short, busy (load>0) + make_worker("http://w2:8000", Some("long"), 2), // long, full but healthy + ]; + policy.init_workers(&workers); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_short_request_uses_short_pool_when_free() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), // short, free + make_worker("http://w2:8000", Some("long"), 0), // long, free + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); // short request + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_short_request_overflows_to_long_when_short_full() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 2), // short, full + make_worker("http://w2:8000", Some("long"), 0), // long, free + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_short_request_falls_back_to_short_when_both_full() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 2), // short, full + make_worker("http://w2:8000", Some("long"), 2), // long, full + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + // short pool exists → queue on short pool min-load. + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_short_request_uses_long_when_short_pool_empty() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![make_worker("http://w2:8000", Some("long"), 0)]; // only long + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn char_estimate_falls_back_when_no_header() { + // No header: uncached derived from (input - matched) / chars_per_token. + // A novel prompt with 400 chars / 4 = 100 tokens < threshold → short. + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let prompt = "a".repeat(400); + let idx = policy + .select_worker(&workers, &info_with_text(&prompt)) + .unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn is_long_pool_reads_label() { + let w_long = make_worker("http://w:1", Some("long"), 0); + let w_short = make_worker("http://w:2", None, 0); + let w_other = make_worker("http://w:3", Some("short"), 0); + assert!(is_long_pool(&*w_long)); + assert!(!is_long_pool(&*w_short)); + assert!(!is_long_pool(&*w_other)); + } + + #[test] + fn parse_prompt_tokens_header_works() { + let mut headers = http::HeaderMap::new(); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_static("12345"), + ); + assert_eq!(parse_prompt_tokens_header(Some(&headers)), Some(12345)); + assert_eq!(parse_prompt_tokens_header(None), None); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_static("notanum"), + ); + assert_eq!(parse_prompt_tokens_header(Some(&headers)), None); + } + + #[test] + fn labels_map_is_accessible() { + // Smoke-test the label access chain the policy depends on. + let w = make_worker("http://w:1", Some("long"), 0); + assert_eq!( + w.metadata().spec.labels.get("pool").map(|s| s.as_str()), + Some("long") + ); + } +} diff --git a/model_gateway/src/policies/factory.rs b/model_gateway/src/policies/factory.rs index e18213870..866cecb5e 100755 --- a/model_gateway/src/policies/factory.rs +++ b/model_gateway/src/policies/factory.rs @@ -3,9 +3,10 @@ use std::sync::Arc; use super::{ - BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, - LeastLoadPolicy, LoadBalancingPolicy, ManualConfig, ManualPolicy, PassthroughPolicy, - PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, RandomPolicy, RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwareLengthConfig, CacheAwareLengthPolicy, + CacheAwarePolicy, ConsistentHashingPolicy, LeastLoadPolicy, LoadBalancingPolicy, ManualConfig, + ManualPolicy, PassthroughPolicy, PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, + RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -72,6 +73,30 @@ impl PolicyFactory { }; Arc::new(CacheAwarePolicy::with_config(config)) } + PolicyConfig::CacheAwareLength { + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + chars_per_token, + long_prefill_threshold, + long_pool_max_load, + short_pool_max_load, + } => { + let config = CacheAwareLengthConfig { + cache_threshold: *cache_threshold, + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + eviction_interval_secs: *eviction_interval_secs, + max_tree_size: *max_tree_size, + chars_per_token: *chars_per_token, + long_prefill_threshold: *long_prefill_threshold, + long_pool_max_load: *long_pool_max_load, + short_pool_max_load: *short_pool_max_load, + }; + Arc::new(CacheAwareLengthPolicy::with_config(config)) + } PolicyConfig::Bucket { balance_abs_threshold, balance_rel_threshold, @@ -123,6 +148,9 @@ impl PolicyFactory { "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())), "least_load" | "leastload" => Some(Arc::new(LeastLoadPolicy::new())), "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), + "cache_aware_length" | "cacheawarelength" => { + Some(Arc::new(CacheAwareLengthPolicy::new())) + } "bucket" => Some(Arc::new(BucketPolicy::new())), "manual" => Some(Arc::new(ManualPolicy::new())), "consistent_hashing" | "consistenthashing" => { @@ -171,6 +199,19 @@ mod tests { }); assert_eq!(policy.name(), "cache_aware"); + let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAwareLength { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 30, + max_tree_size: 10000, + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 256, + short_pool_max_load: 256, + }); + assert_eq!(policy.name(), "cache_aware_length"); + let policy = PolicyFactory::create_from_config(&PolicyConfig::Bucket { balance_abs_threshold: 10, balance_rel_threshold: 1.5, @@ -212,6 +253,8 @@ mod tests { assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some()); assert!(PolicyFactory::create_by_name("cache_aware").is_some()); assert!(PolicyFactory::create_by_name("CacheAware").is_some()); + assert!(PolicyFactory::create_by_name("cache_aware_length").is_some()); + assert!(PolicyFactory::create_by_name("CacheAwareLength").is_some()); assert!(PolicyFactory::create_by_name("bucket").is_some()); assert!(PolicyFactory::create_by_name("Bucket").is_some()); assert!(PolicyFactory::create_by_name("manual").is_some()); diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index d6b192592..e863c39ba 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -14,6 +14,7 @@ use crate::{ mod bucket; mod cache_aware; +mod cache_aware_length; mod consistent_hashing; mod dp_min_token; mod factory; @@ -29,6 +30,7 @@ pub(crate) mod utils; pub use bucket::BucketPolicy; pub use cache_aware::{CacheAwarePolicy, TreeHandle, TreeKind}; +pub use cache_aware_length::CacheAwareLengthPolicy; pub use consistent_hashing::ConsistentHashingPolicy; pub use dp_min_token::MinimumTokensPolicy; pub use factory::PolicyFactory; @@ -198,6 +200,51 @@ impl Default for CacheAwareConfig { } } +/// Configuration for the cache_aware_length policy (long/short pool split). +#[derive(Debug, Clone)] +pub struct CacheAwareLengthConfig { + /// Min matched-prefix share before a request pins to a holder (0.0-1.0). + pub cache_threshold: f32, + /// Absolute load diff for the global imbalance check (step 2). + pub balance_abs_threshold: usize, + /// Relative load ratio for the global imbalance check (step 2); fires + /// only together with `balance_abs_threshold`. + pub balance_rel_threshold: f32, + /// Interval between LRU eviction cycles (seconds). `0` disables. + pub eviction_interval_secs: u64, + /// Max total chars of each model's approximate string tree, shared across + /// all workers; enforced by eviction. + pub max_tree_size: usize, + /// Divisor for char-level token estimation when `X-Prompt-Tokens` is + /// absent: `uncached_tokens = (input_chars - matched_chars) / + /// chars_per_token`. + pub chars_per_token: usize, + /// Uncached-prefill-token boundary between long and short requests. + pub long_prefill_threshold: usize, + /// Load ceiling for the long pool (`pool=long` workers). + pub long_pool_max_load: usize, + /// Load ceiling for the short pool (remaining workers). + pub short_pool_max_load: usize, +} + +impl Default for CacheAwareLengthConfig { + fn default() -> Self { + Self { + cache_threshold: 0.3, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 30, + max_tree_size: 10000, + chars_per_token: 4, + long_prefill_threshold: 100_000, + // 0 would reject every worker as non-free; pick generous defaults so + // the policy is usable out of the box before operators tune. + long_pool_max_load: 4, + short_pool_max_load: 32, + } + } +} + #[derive(Debug, Clone)] pub struct BucketConfig { pub balance_abs_threshold: usize, diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 42777507e..d8c4200eb 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -17,8 +17,8 @@ use tracing::{debug, info, warn}; use super::{ get_healthy_worker_indices, manual::{ExecutionBranch, PinState}, - BucketPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, ManualConfig, - ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, + BucketPolicy, CacheAwareLengthPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, + ManualConfig, ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, }; use crate::{ config::types::{ManualAssignmentMode, PolicyConfig, RoutingKeyOverrideConfig}, @@ -807,6 +807,15 @@ impl PolicyRegistry { ); cache_aware.init_workers(workers); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + debug!( + "Initializing cache_aware_length policy with {} workers for model {}", + workers.len(), + model_id + ); + cal.init_workers(workers); + } } } } @@ -824,6 +833,14 @@ impl PolicyRegistry { worker_url, model_id ); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + cal.remove_worker_by_url(worker_url); + debug!( + "Removed worker {} from cache_aware_length policy for model {}", + worker_url, model_id + ); + } } } } @@ -845,6 +862,14 @@ impl PolicyRegistry { worker_url, worker_type ); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + cal.remove_worker_by_url(worker_url); + debug!( + "Removed worker {} from {} cache_aware_length policy", + worker_url, worker_type + ); + } } } } @@ -882,6 +907,19 @@ impl PolicyRegistry { cache_aware.init_workers(prefill_workers); } } + } else if prefill_policy.name() == "cache_aware_length" { + if let Some(cal) = prefill_policy + .as_any() + .downcast_ref::() + { + if !prefill_workers.is_empty() { + debug!( + "Initializing prefill cache_aware_length policy with {} workers", + prefill_workers.len() + ); + cal.init_workers(prefill_workers); + } + } } } @@ -898,6 +936,19 @@ impl PolicyRegistry { cache_aware.init_workers(decode_workers); } } + } else if decode_policy.name() == "cache_aware_length" { + if let Some(cal) = decode_policy + .as_any() + .downcast_ref::() + { + if !decode_workers.is_empty() { + debug!( + "Initializing decode cache_aware_length policy with {} workers", + decode_workers.len() + ); + cal.init_workers(decode_workers); + } + } } } } diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 21c3d50e7..e14169442 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -2122,7 +2122,7 @@ mod tests { use super::*; use crate::{ config::types::{PolicyConfig, RoutingKeyOverrideConfig}, - policies::CacheAwarePolicy, + policies::{CacheAwareLengthPolicy, CacheAwarePolicy}, routers::common::request_lease::test_probe::{spawn_release_gated_stub, DropProbeRequest}, worker::BasicWorkerBuilder, }; @@ -2475,6 +2475,32 @@ mod tests { .build() } + /// Like `plain_worker` but attaches a `pool=` label so the + /// cache_aware_length policy can split workers into long/short pools. + fn labeled_worker(url: &str, pool: Option<&str>) -> crate::worker::BasicWorker { + let mut builder = BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .health_config(no_health_check()); + if let Some(p) = pool { + builder = builder.label("pool", p); + } + builder.build() + } + + fn cache_aware_length_policy() -> PolicyConfig { + PolicyConfig::CacheAwareLength { + cache_threshold: 0.3, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 0, + max_tree_size: 4096, + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 2, + short_pool_max_load: 2, + } + } + type CapturedUpstreamRequest = Arc>>; /// Loopback engine stub: captures the forwarded `/generate` request and @@ -2987,4 +3013,331 @@ mod tests { let router = streaming_router(least_load_policy(), 1024 * 1024, vec![]); assert_falls_back_with_body_intact(&router).await; } + + // ===== cache_aware_length E2E: all 16 decision-table scenarios ===== + // + // Each test drives the real routing decision point + // (`select_worker_for_model`), which builds the exact `SelectWorkerInfo` + // (request_text + headers) the policy receives in production — the same + // path the buffered HTTP request takes. Workers carry `pool=long` labels; + // pre-set loads use WorkerLoadGuard + mem::forget (same as the unit tests). + + /// Pin a worker's in-flight load at `load` for the lifetime of the test by + /// leaking the RAII guard (the process tears down on test exit). + fn pin_load(worker: &crate::worker::BasicWorker, load: usize) { + let w: Arc = Arc::new(worker.clone()); + for _ in 0..load { + std::mem::forget(WorkerLoadGuard::new(Arc::clone(&w), None)); + } + } + + /// Spawn a two-pool router: short workers (no label) + long workers + /// (`pool=long`), with cache_aware_length as the default policy and trees + /// seeded. Returns the router so tests can assert routed worker. + async fn length_router(short_urls: &[&str], long_urls: &[&str]) -> Router { + let mut workers: Vec = Vec::new(); + for u in short_urls { + workers.push(labeled_worker(u, None)); + } + for u in long_urls { + workers.push(labeled_worker(u, Some("long"))); + } + let router = streaming_router(cache_aware_length_policy(), 1024 * 1024, workers); + let live = router.worker_registry.get_all(); + router + .policy_registry + .get_default_policy() + .as_any() + .downcast_ref::() + .unwrap() + .init_workers(&live); + router + } + + /// Route `prompt` through the real `select_worker_for_model` (the exact + /// path a buffered HTTP request takes), returning the selected worker's + /// URL — the value the router writes to `x-smg-routed-worker-id`. + /// Returns `None` when the fleet rejects (503). Panics if the selection + /// unexpectedly fails (use `route_or_none` for the reject case). + fn route_to_url(router: &Router, prompt: &str, headers: Option<&HeaderMap>) -> String { + router + .select_worker_for_model( + crate::worker::UNKNOWN_MODEL_ID, + Some(prompt), + None, + headers, + None, + ) + .map(|w| w.url().to_string()) + .expect("selection should succeed with healthy workers") + } + + /// Like `route_to_url` but returns `None` when the fleet rejects (503). + fn route_or_none(router: &Router, prompt: &str, headers: Option<&HeaderMap>) -> Option { + router + .select_worker_for_model( + crate::worker::UNKNOWN_MODEL_ID, + Some(prompt), + None, + headers, + None, + ) + .map(|w| w.url().to_string()) + } + + /// Build a HeaderMap with `X-Prompt-Tokens: `. + fn tokens_header(tokens: usize) -> HeaderMap { + let mut h = HeaderMap::new(); + h.insert( + "x-prompt-tokens", + http::HeaderValue::from_str(&tokens.to_string()).unwrap(), + ); + h + } + + /// Find a live worker by URL in the router's registry and pin its load. + fn pin_worker(router: &Router, url: &str, load: usize) { + pin_load( + router + .worker_registry + .get_all() + .iter() + .find(|w| w.url() == url) + .cloned() + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(), + load, + ); + } + + /// Mark a worker unhealthy by URL. + fn mark_unhealthy(router: &Router, url: &str) { + router + .worker_registry + .get_all() + .iter() + .find(|w| w.url() == url) + .cloned() + .unwrap() + .set_status(openai_protocol::worker::WorkerStatus::NotReady); + } + + // --- Step 1: health filter --- + + #[tokio::test] + async fn cal_step1_all_unhealthy_returns_503() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_s); + mark_unhealthy(&router, &url_l); + let h = tokens_header(200_000); + let routed = route_or_none(&router, "hello", Some(&h)); + assert!( + routed.is_none(), + "all unhealthy → 503 (None), got {routed:?}" + ); + } + + // --- Step 2: global imbalance --- + + #[tokio::test] + async fn cal_step2_global_imbalance_picks_min_load() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + // Pin the long-pool worker high so the fleet is imbalanced (100 vs 0). + pin_worker(&router, &url_l, 100); + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "imbalanced fleet → healthy min-load worker"); + } + + // --- Step 3: cache hit --- + + #[tokio::test] + async fn cal_step3_cache_hit_sticks_regardless_of_pool() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "shared long prompt prefix that builds cache affinity"; + let first = route_to_url(&router, prompt, None); + let second = route_to_url(&router, prompt, None); + assert_eq!(first, second, "cache hit pins to the same worker"); + } + + #[tokio::test] + async fn cal_step3_no_tree_random_healthy_does_not_panic() { + // Do NOT call init_workers → no tree → random healthy fallback. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = streaming_router( + cache_aware_length_policy(), + 1024 * 1024, + vec![ + labeled_worker(&url_s, None), + labeled_worker(&url_l, Some("long")), + ], + ); + let routed = route_to_url(&router, "novel", None); + assert!( + routed == url_s || routed == url_l, + "random fallback picks a healthy worker: {routed}" + ); + } + + // --- Step 4 long request (uncached ≥ 100K): 4 paths --- + + #[tokio::test] + async fn cal_step4_long_uses_long_pool_when_free() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "long request, free long pool → long pool"); + } + + #[tokio::test] + async fn cal_step4_long_overflows_to_idle_short() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_l, 2); // long pool full (load = max) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "full long pool → idle (load 0) short worker"); + } + + #[tokio::test] + async fn cal_step4_long_queues_on_long_when_short_busy() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_l, 2); // long full + pin_worker(&router, &url_s, 1); // short busy (load > 0, not idle) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "long full + short busy → queue on long pool"); + } + + // Step 4: long pool全不健康 + 短池都 load>0 → 全部健康 worker min-load. + #[tokio::test] + async fn cal_step4_long_unhealthy_short_busy_all_healthy_min_load() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_l); // long pool unhealthy + pin_worker(&router, &url_s, 1); // short busy (load > 0) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!( + routed, url_s, + "long unhealthy + short busy → all-healthy min-load (short worker)" + ); + } + + // --- Step 4 short request (uncached < 100K): 5 paths --- + + #[tokio::test] + async fn cal_step4_short_uses_short_pool_when_free() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "short request, free short pool → short pool"); + } + + #[tokio::test] + async fn cal_step4_short_overflows_to_long_when_short_full() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_s, 2); // short pool full + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "short pool full → overflow to long pool"); + } + + #[tokio::test] + async fn cal_step4_short_falls_back_to_short_when_both_full() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_s, 2); // short full + pin_worker(&router, &url_l, 2); // long full + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "both full → queue on short pool min-load"); + } + + #[tokio::test] + async fn cal_step4_short_uses_long_when_short_pool_empty() { + // Only long-pool workers configured; short pool is empty. + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[], &[&url_l]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "no short pool → long pool min-load"); + } + + #[tokio::test] + async fn cal_step4_short_both_pools_empty_all_healthy_min_load() { + // Only one healthy short worker (long pool empty); it is the sole + // all-healthy min-load candidate. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!( + routed, url_s, + "single healthy worker → all-healthy min-load" + ); + } + + // --- Step 4: token source priority (header vs char estimate vs none) --- + + #[tokio::test] + async fn cal_step4_char_estimate_fallback_when_no_header() { + // No X-Prompt-Tokens header → char-level estimate. A 400-char novel + // prompt → 100 tokens < 100K threshold → short request → short pool. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "a".repeat(400); + let routed = route_to_url(&router, &prompt, None); + assert_eq!(routed, url_s, "char estimate < threshold → short pool"); + } + + #[tokio::test] + async fn cal_step4_uncached_unknown_all_healthy_min_load() { + // Empty prompt, no header → uncached not computable → all-healthy + // min-load (the first healthy worker, i.e. the short worker). + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let routed = route_to_url(&router, "", None); + assert!( + routed == url_s || routed == url_l, + "uncached unknown → all-healthy min-load: {routed}" + ); + } + + // --- Step 5: tree recording (cache hit after pool routing) --- + + #[tokio::test] + async fn cal_step5_pool_routing_records_tree_for_future_hit() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "novel pool-routed prompt that will be recorded"; + // First request: novel → pool split (short, since <100K by char estimate) + let first = route_to_url(&router, prompt, None); + // Second request: same prompt → now a cache hit → same worker. + let second = route_to_url(&router, prompt, None); + assert_eq!(first, second, "pool routing recorded tree → subsequent hit"); + } }