feat: query dispatch gate (llm-d#194)

* feat: query dispatch gate

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

* address comments

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

---------

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

Author: evacchi

README.md

@@ -105,6 +105,7 @@ For more fine-grained control, configure gates per queue in your configuration f
 - `composite`: Combines multiple gates. Returns the minimum budget across all inner dispatch gates and acquires quota across all inner attribute gates (all or nothing).
 - `prometheus-saturation`: Queries Prometheus for pool saturation metric. The gate closes (returns `0.0`) when saturation ≥ threshold; when open it returns `(1 - saturation) - (1 - threshold)`, i.e. the margin below the threshold.
 - `prometheus-budget`: Computes a dispatch budget D using a cascade of two Prometheus metric sources. Both sources compute `max_SYS = ready_pods × max_concurrency` dynamically. The primary source uses the EPP flow control queue size: `D = 1 − (queue_size / max_SYS)`. If the primary is unavailable, it falls back to a secondary source using vLLM and pool metrics: `D = 1 − (running_requests / max_SYS)`. The gate closes when D ≤ B (baseline); callers compute `N = max_SYS × (D − B)`. See [docs/dispatch-budget.md](docs/dispatch-budget.md) for details.
+- `prometheus-query`: Evaluates an arbitrary user-supplied PromQL expression as the dispatch budget. The expression must resolve to an instant vector with a single sample whose value is in [0, 1]. Unlike `prometheus-saturation` and `prometheus-budget`, this gate does not construct queries internally — the user provides the complete PromQL expression. Values outside [0, 1] are clamped.
 
 **Example Configuration with Per-Queue Gates:**
 
@@ -147,6 +148,16 @@ For more fine-grained control, configure gates per queue in your configuration f
           "budget_key": "my-budget-key"
        }
     },
+    {
+       "queue_name": "custom_metric_queue",
+       "inference_objective": "custom-task",
+       "request_path_url": "/v1/inference",
+       "gate_type": "prometheus-query",
+       "gate_params": {
+          "query": "1 - (sum(rate(http_requests_total{job=\"inference\"}[5m])) / 100)",
+          "fallback": "0.0"
+       }
+    },
     {
        "queue_name": "composite_gated_queue",
        "inference_objective": "composite-task",
@@ -211,6 +222,15 @@ For more fine-grained control, configure gates per queue in your configuration f
       targetLabel: inference_pool
   ```
 
+- `prometheus-query`: Evaluates a user-supplied PromQL expression directly as the dispatch budget.
+  The expression must resolve to a Prometheus instant vector with a single sample whose value is in [0, 1].
+  Values outside this range are clamped.
+
+  - `query` (**required**): The PromQL expression to evaluate. This is sent to Prometheus exactly as provided.
+    The result is used directly as the dispatch budget (no transformation is applied).
+  - `fallback` (optional): Fallback budget value (0.0-1.0) returned when the query fails or returns no data.
+    Default is `0.0` (fail closed).
+
 ## Request Messages and Consumption
 
 The async processor expects request messages to have the following format:

pkg/async/inference/flowcontrol/gate_factory.go

@@ -87,6 +87,11 @@ func (f *GateFactory) Close() error {
 //     Gate closes when D ≤ B (baseline); returns D − B when open, so callers compute
 //     N = max_SYS × (D − B). Params: pool (required),
 //     max_concurrency (default 100), baseline (default 0.05), fallback (default 0.0)
+//   - "prometheus-query": Evaluates an arbitrary user-supplied PromQL expression as the dispatch
+//     budget. The expression must resolve to an instant vector with a single sample whose value
+//     is in [0, 1]. Unlike prometheus-saturation and prometheus-budget, this gate does not
+//     construct queries internally — the user provides the complete PromQL expression.
+//     Params: query (required), fallback (default 0.0)
 //
 // For unsupported or unknown gate types, returns ConstOpenGate as a safe default.
 func (f *GateFactory) CreateGate(gateType string, params map[string]string) (pipeline.DispatchGate, error) {
@@ -255,6 +260,28 @@ func (f *GateFactory) CreateGate(gateType string, params map[string]string) (pip
 		)
 		return NewBudgetDispatchGate(ms, baseline, fallback), nil
 
+	case "prometheus-query":
+		if f.prometheusURL == "" {
+			return nil, fmt.Errorf("prometheus-query gate type requires --prometheus-url flag to be set")
+		}
+
+		query := params["query"]
+		if query == "" {
+			return nil, fmt.Errorf("prometheus-query gate requires a 'query' parameter with a PromQL expression")
+		}
+
+		fallback, err := parseFloat("fallback", params["fallback"], 0.0)
+		if err != nil {
+			return nil, err
+		}
+
+		promConfig := promapi.Config{Address: f.prometheusURL}
+		source, err := NewPromQLMetricSource(promConfig, query)
+		if err != nil {
+			return nil, err
+		}
+		return NewMetricDispatchGate(cachedSource(source, f.cacheTTL), 0.0, fallback), nil
+
 	default:
 		// Unknown gate types default to open gate
 		return ConstOpenGate(), nil

pkg/async/inference/flowcontrol/gate_factory_test.go

@@ -249,3 +249,51 @@ func TestGateFactory_BudgetGateWithAllParams(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, gate)
 }
+
+func TestGateFactory_PrometheusQueryGateWithoutURL(t *testing.T) {
+	factory := NewGateFactory("")
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{
+		"query": "up",
+	})
+	assert.Error(t, err, "should return error when Prometheus URL is not set")
+	assert.Nil(t, gate)
+	assert.Contains(t, err.Error(), "prometheus-query gate type requires --prometheus-url flag to be set")
+}
+
+func TestGateFactory_PrometheusQueryGateMissingQuery(t *testing.T) {
+	factory := NewGateFactory("http://localhost:9090")
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{})
+	assert.Error(t, err, "should return error when query parameter is missing")
+	assert.Nil(t, gate)
+	assert.Contains(t, err.Error(), "requires a 'query' parameter")
+}
+
+func TestGateFactory_PrometheusQueryGateWithInvalidFallback(t *testing.T) {
+	factory := NewGateFactory("http://localhost:9090")
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{
+		"query":    "up",
+		"fallback": "not-a-number",
+	})
+	assert.Error(t, err, "should return error when fallback is not a valid float")
+	assert.Nil(t, gate)
+	assert.Contains(t, err.Error(), "invalid fallback value")
+}
+
+func TestGateFactory_PrometheusQueryGateWithDefaults(t *testing.T) {
+	factory := NewGateFactory("http://localhost:9090")
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{
+		"query": "up",
+	})
+	assert.NoError(t, err, "should create gate with default fallback")
+	assert.NotNil(t, gate)
+}
+
+func TestGateFactory_PrometheusQueryGateWithAllParams(t *testing.T) {
+	factory := NewGateFactory("http://localhost:9090")
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{
+		"query":    `1 - (sum(rate(http_requests_total[5m])) / 100)`,
+		"fallback": "0.5",
+	})
+	assert.NoError(t, err, "should create gate with all params specified")
+	assert.NotNil(t, gate)
+}

test/e2e/e2e_query_gate_test.go

@@ -0,0 +1,97 @@
+package e2e
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/onsi/ginkgo/v2"
+	"github.com/onsi/gomega"
+)
+
+const (
+	queryRequestQueue = "query-request-sortedset"
+	queryResultQueue  = "query-result-list"
+)
+
+// The prometheus-query gate evaluates an arbitrary PromQL expression as the
+// dispatch budget. The e2e Helm values configure it with the same saturation
+// metric used by the saturation gate tests:
+//
+//	1 - inference_extension_flow_control_pool_saturation{inference_pool="e2e-pool"}
+//
+// Because threshold=0 the raw PromQL result is used directly as the budget.
+var _ = ginkgo.Describe("Prometheus Query Dispatch Gate E2E", ginkgo.Ordered, func() {
+	var ctx context.Context
+
+	ginkgo.BeforeAll(func() {
+		redeployEPPWithFlowControl()
+	})
+
+	ginkgo.BeforeEach(func() {
+		ctx = context.Background()
+		rdb.Del(ctx, queryRequestQueue) //nolint:errcheck
+		rdb.Del(ctx, queryResultQueue)  //nolint:errcheck
+		setSimWaitingRequests(simAdminURL, 0)
+	})
+
+	ginkgo.It("processes a message when the query returns a positive budget", func() {
+		setSimWaitingRequests(simAdminURL, 0)
+
+		waitForSaturation(promURL, envoyURL, func(v float64) bool { return v < 0.5 })
+
+		msg := makeRequestMessage("query-positive", 5*time.Minute)
+		enqueueMessage(ctx, rdb, queryRequestQueue, msg)
+
+		gomega.Eventually(func() int64 {
+			return getResultCount(ctx, rdb, queryResultQueue)
+		}, 60*time.Second, 1*time.Second).Should(gomega.BeNumerically(">=", 1))
+
+		result := popResult(ctx, rdb, queryResultQueue)
+		gomega.Expect(result).NotTo(gomega.BeNil())
+		gomega.Expect(result.ID).To(gomega.Equal("query-positive"))
+	})
+
+	ginkgo.It("pauses processing when the query returns zero budget", func() {
+		setSimWaitingRequests(simAdminURL, 10)
+
+		waitForSaturation(promURL, envoyURL, func(v float64) bool { return v >= 0.7 })
+
+		msg := makeRequestMessage("query-blocked", 5*time.Minute)
+		enqueueMessage(ctx, rdb, queryRequestQueue, msg)
+
+		gomega.Consistently(func() int64 {
+			return getResultCount(ctx, rdb, queryResultQueue)
+		}, 10*time.Second, 1*time.Second).Should(gomega.Equal(int64(0)))
+
+		setSimWaitingRequests(simAdminURL, 0)
+
+		gomega.Eventually(func() int64 {
+			return getResultCount(ctx, rdb, queryResultQueue)
+		}, 60*time.Second, 1*time.Second).Should(gomega.BeNumerically(">=", 1))
+
+		result := popResult(ctx, rdb, queryResultQueue)
+		gomega.Expect(result).NotTo(gomega.BeNil())
+		gomega.Expect(result.ID).To(gomega.Equal("query-blocked"))
+	})
+
+	ginkgo.It("resumes processing when the query budget recovers", func() {
+		setSimWaitingRequests(simAdminURL, 10)
+		waitForSaturation(promURL, envoyURL, func(v float64) bool { return v >= 0.7 })
+
+		for i := 1; i <= 3; i++ {
+			msg := makeRequestMessage(fmt.Sprintf("query-resume-%d", i), 5*time.Minute)
+			enqueueMessage(ctx, rdb, queryRequestQueue, msg)
+		}
+
+		gomega.Consistently(func() int64 {
+			return getResultCount(ctx, rdb, queryResultQueue)
+		}, 5*time.Second, 1*time.Second).Should(gomega.Equal(int64(0)))
+
+		setSimWaitingRequests(simAdminURL, 0)
+
+		gomega.Eventually(func() int64 {
+			return getResultCount(ctx, rdb, queryResultQueue)
+		}, 60*time.Second, 1*time.Second).Should(gomega.BeNumerically(">=", 3))
+	})
+})

test/e2e/e2e_suite_test.go

@@ -319,6 +319,7 @@ func applyManifests() {
 		{"redis-gate", helmValuesDir + "/redis-gate.yaml"},
 		{"quota", helmValuesDir + "/quota.yaml"},
 		{"composite", helmValuesDir + "/composite.yaml"},
+		{"prometheus-query", helmValuesDir + "/prometheus-query.yaml"},
 	} {
 		helmInstall(r.name, r.values, map[string]string{
 			"ap.image.repository": imageRepo,
@@ -534,6 +535,7 @@ func doRedeployEPPWithFlowControl() {
 		"redis-gate-async-processor",
 		"quota-async-processor",
 		"composite-async-processor",
+		"prometheus-query-async-processor",
 	} {
 		cmd := exec.Command("kubectl", "--kubeconfig", kindKubeconfig,
 			"-n", nsName, "rollout", "restart", "deployment/"+deploy)
@@ -548,6 +550,7 @@ func doRedeployEPPWithFlowControl() {
 		"redis-gate-async-processor",
 		"quota-async-processor",
 		"composite-async-processor",
+		"prometheus-query-async-processor",
 	} {
 		cmd := exec.Command("kubectl", "--kubeconfig", kindKubeconfig,
 			"-n", nsName, "rollout", "status", "deployment/"+deploy, "--timeout=120s")

test/e2e/helm/prometheus-query.yaml

@@ -0,0 +1,25 @@
+ap:
+  imagePullPolicy: Never
+  messageQueueImpl: "redis-sortedset"
+  igwBaseURL: "http://envoy.e2e-integration.svc.cluster.local:8081"
+  concurrency: 1
+  prometheusURL: "http://prometheus.e2e-integration.svc.cluster.local:9090"
+  prometheusCacheTTL: "0s"
+  metrics:
+    enabled: true
+    port: 9090
+    secure: false
+  redis:
+    enabled: true
+    url: "redis://redis.e2e-integration.svc.cluster.local:6379"
+    requestPathURL: "/v1/completions"
+    requestQueueName: "query-request-sortedset"
+    resultQueueName: "query-result-list"
+    pollIntervalMs: 500
+    batchSize: 10
+    gateType: "prometheus-query"
+    gateParams:
+      query: "1 - inference_extension_flow_control_pool_saturation{inference_pool=\"e2e-pool\"}"
+      fallback: "0.0"
+  modelServerMonitor:
+    enabled: false

test/integration/prometheus_gate_factory_test.go

@@ -125,6 +125,55 @@ func TestGateFactory_PrometheusSaturation_MissingParams(t *testing.T) {
 	assert.Error(t, err, "Should fail when pool is missing")
 }
 
+// TestGateFactory_PrometheusQuery_EndToEnd validates the prometheus-query gate
+// wiring: GateFactory.CreateGate("prometheus-query", params) ->
+//
+//	NewPromQLMetricSource (user-supplied PromQL) ->
+//	CachedMetricSource -> MetricDispatchGate(threshold=0) -> Budget()
+func TestGateFactory_PrometheusQuery_EndToEnd(t *testing.T) {
+	var queryCount int64
+	body := &atomic.Value{}
+	body.Store(promVectorResponse("0.7"))
+
+	const promQL = "my_custom_metric"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt64(&queryCount, 1)
+		assert.Equal(t, promQL, r.FormValue("query"), "gate should forward the user-supplied PromQL expression")
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprint(w, body.Load().(string))
+	}))
+	defer server.Close()
+
+	factory := flowcontrol.NewGateFactoryWithCacheTTL(server.URL, 200*time.Millisecond)
+
+	gate, err := factory.CreateGate("prometheus-query", map[string]string{
+		"query":    promQL,
+		"fallback": "0.0",
+	})
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	// threshold=0 so Budget = value - 0 = 0.7.
+	budget := gate.Budget(ctx)
+	assert.InDelta(t, 0.7, budget, 0.01)
+	assert.Equal(t, int64(1), atomic.LoadInt64(&queryCount))
+
+	// Second call within cache TTL should NOT trigger another HTTP request.
+	budget2 := gate.Budget(ctx)
+	assert.InDelta(t, 0.7, budget2, 0.01)
+	assert.Equal(t, int64(1), atomic.LoadInt64(&queryCount), "Should use cached result")
+
+	// Wait for cache TTL to expire and update response.
+	time.Sleep(250 * time.Millisecond)
+	body.Store(promVectorResponse("0.0"))
+
+	// Query returns 0.0 => Budget = 0.0.
+	budget3 := gate.Budget(ctx)
+	assert.Equal(t, 0.0, budget3)
+	assert.Equal(t, int64(2), atomic.LoadInt64(&queryCount), "Should have queried Prometheus again after cache expiry")
+}
+
 func promVectorResponse(value string) string {
 	return fmt.Sprintf(`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"test","inference_pool":"my-pool"},"value":[1234567890,"%s"]}]}}`, value)
 }