feat(function-autoscaler): add LLM gateway scaling

src/control-plane-services/function-autoscaler/crates/server/src/scaling/mod.rs

@@ -303,25 +303,19 @@ pub fn calculate_average_utilization(
 pub enum MetricSource {
     /// Worker-thread busy ratio. The default; emitted by managed worker pods.
     WorkerThreads,
+    /// Function-level request duration emitted by the LLM API gateway.
+    LlmGateway,
     /// Control-plane request-latency / instance metrics. Used for BYOC and any
     /// other function that never emits worker series.
     ControlPlane,
 }
 
-impl MetricSource {
-    /// Whether this source uses the control-plane utilization query.
-    pub fn uses_control_plane_metrics(self) -> bool {
-        matches!(self, MetricSource::ControlPlane)
-    }
-}
-
 /// Everything a scaling decision needs, gathered once from the timeseries DB.
 ///
 /// All NaN/parse sanitization happens while building this (see
 /// [`sanitize_utilization`]), so every downstream path operates on the same
 /// shape regardless of which [`MetricSource`] produced it. That is what keeps
-/// scale-to-zero behaving identically for worker-metric and control-plane
-/// functions.
+/// scale-to-zero behaving identically for every metric source.
 #[derive(Debug, Clone)]
 pub struct ScalingInputs {
     pub metric_source: MetricSource,

src/control-plane-services/function-autoscaler/crates/server/src/server.rs

@@ -25,7 +25,7 @@ use rs_autoscaler::{
     scaling::{policy_cache::PolicyCache, ScalingPolicy, ScalingSettings},
     secrets::secrets_file_watcher::SecretFileWatcher,
     settings, startup, timeseries_db, work,
-    work::new_function_state_cache,
+    work::{new_function_state_cache, new_metric_routing_cache},
 };
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -231,6 +231,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     );
 
     let function_state_cache = Arc::new(new_function_state_cache());
+    let metric_routing_cache = Arc::new(new_metric_routing_cache());
 
     tracing::info!("Initializing NVCF API client");
     let cassandra_service_manager_nvcf = Arc::clone(&cassandra_service_manager);
@@ -301,6 +302,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let bucket_manager_p0 = Arc::clone(&bucket_manager);
     let lock_manager_p0 = Arc::clone(&lock_manager);
     let function_state_cache_p0 = Arc::clone(&function_state_cache);
+    let metric_routing_cache_p0 = Arc::clone(&metric_routing_cache);
     let scaling_loop_interval = config.scaling.scaling_loop_interval;
     tokio::spawn(async move {
         let mut interval = tokio::time::interval(scaling_loop_interval);
@@ -315,6 +317,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 &bucket_manager_p0,
                 &lock_manager_p0,
                 Arc::clone(&function_state_cache_p0),
+                Arc::clone(&metric_routing_cache_p0),
             )
             .await
             {

src/control-plane-services/function-autoscaler/crates/server/src/work/discovery.rs

@@ -82,20 +82,23 @@ impl DiscoveryShard {
 enum InvocationMetricSource {
     InvocationService,
     GrpcProxy,
+    LlmGateway,
 }
 
 impl InvocationMetricSource {
     fn name(self) -> &'static str {
         match self {
             Self::InvocationService => "invocation_service",
             Self::GrpcProxy => "grpc_proxy",
+            Self::LlmGateway => "llm_gateway",
         }
     }
 
     fn order(self) -> u8 {
         match self {
             Self::InvocationService => 0,
             Self::GrpcProxy => 1,
+            Self::LlmGateway => 2,
         }
     }
 }
@@ -182,6 +185,24 @@ fn get_timeseries_db_query(
     template.replace("{env_filter}", &selector)
 }
 
+fn llm_gateway_discovery_query(env: &str, ignore_env: bool, shard: DiscoveryShard) -> String {
+    let function_id_regex = shard.function_id_regex();
+    let env_matcher = if ignore_env {
+        String::new()
+    } else {
+        format!(r#", aws_env="{}""#, env)
+    };
+    format!(
+        r#"(sum by(function_id) (
+            increase(llm_api_gateway_http_requests_total{{function_id=~"{function_id_regex}", function_id!="none"{env_matcher}}}[5m])
+        ) > 0)
+        * on(function_id) group_right(function_version_id, nca_id)
+        max by(function_id, function_version_id, nca_id) (
+            nvcf_function_info{{function_id=~"{function_id_regex}"{env_matcher}}}
+        )"#
+    )
+}
+
 fn recent_invocation_queries(
     env: &str,
     ignore_env: bool,
@@ -193,7 +214,7 @@ fn recent_invocation_queries(
         DiscoveryShard::ALL.into_iter().map(Some).collect()
     };
 
-    let mut queries = Vec::with_capacity(shards.len() * 2);
+    let mut queries = Vec::with_capacity(shards.len() * 3);
     for shard in shards {
         queries.push(RecentInvocationQuery {
             source: InvocationMetricSource::InvocationService,
@@ -217,6 +238,14 @@ fn recent_invocation_queries(
                 shard,
             ),
         });
+        if function_version_filter.is_none() {
+            let shard = shard.expect("discovery queries are sharded");
+            queries.push(RecentInvocationQuery {
+                source: InvocationMetricSource::LlmGateway,
+                shard: Some(shard),
+                query: llm_gateway_discovery_query(env, ignore_env, shard),
+            });
+        }
     }
     queries
 }
@@ -533,9 +562,10 @@ async fn get_recently_invoked_functions_with_semaphore(
         "Executing PromQL queries for recently invoked functions"
     );
 
-    // Discovery runs eight queries (two sources across four shards) through one
-    // shared concurrency bound. Per-function scaling remains two unsharded
-    // queries. Every query is polled even when another source or shard fails.
+    // Discovery runs twelve queries (three sources across four shards) through
+    // one shared concurrency bound. Per-version invocation checks remain two
+    // unsharded queries. Every query is polled even when another source or
+    // shard fails.
     let mut query_results = stream::iter(queries.into_iter().map(|query_spec| {
         let query_semaphore = query_semaphore.clone();
         async move {
@@ -952,24 +982,42 @@ mod tests {
     }
 
     #[test]
-    fn discovery_queries_cover_four_fixed_shards_for_both_sources() {
+    fn discovery_queries_cover_four_fixed_shards_for_all_sources() {
         let queries = recent_invocation_queries("prd", false, None);
-        assert_eq!(queries.len(), 8);
+        assert_eq!(queries.len(), 12);
 
         for shard in DiscoveryShard::ALL {
             let matcher = format!(r#"function_id=~"{}""#, shard.function_id_regex());
             let shard_queries: Vec<_> = queries
                 .iter()
                 .filter(|query| query.shard == Some(shard))
                 .collect();
-            assert_eq!(shard_queries.len(), 2);
+            assert_eq!(shard_queries.len(), 3);
+
             for query in shard_queries {
-                assert_eq!(query.query.matches(&matcher).count(), 4);
-                assert_eq!(query.query.matches(r#"aws_env="prd""#).count(), 4);
+                if matches!(query.source, InvocationMetricSource::LlmGateway) {
+                    assert_eq!(query.query.matches(&matcher).count(), 2);
+                    assert_eq!(query.query.matches(r#"aws_env="prd""#).count(), 2);
+                    assert!(query.query.contains("llm_api_gateway_http_requests_total"));
+                } else {
+                    assert_eq!(query.query.matches(&matcher).count(), 4);
+                    assert_eq!(query.query.matches(r#"aws_env="prd""#).count(), 4);
+                }
             }
         }
     }
 
+    #[test]
+    fn gateway_discovery_omits_environment_when_configured() {
+        let queries = recent_invocation_queries("stg", true, None);
+        for query in queries
+            .iter()
+            .filter(|query| matches!(query.source, InvocationMetricSource::LlmGateway))
+        {
+            assert!(!query.query.contains("aws_env"));
+        }
+    }
+
     #[test]
     fn per_function_recent_invocation_queries_are_not_sharded() {
         let function_version_id = Uuid::new_v4();