Skip to content

Commit 2e2b14e

Browse files
feat: add TenantQuota watch for agent deployments and implement Prometheus query client with unit tests
1 parent fa63118 commit 2e2b14e

7 files changed

Lines changed: 348 additions & 36 deletions

File tree

internal/controller/agentdeployment_controller.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ func (r *AgentDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Requ
181181
// Return the shorter of the two requeue intervals.
182182
statusResult, err := r.updateStatus(ctx, ad, logger, qs)
183183
if err != nil {
184-
return statusResult, err
184+
return statusResult, fmt.Errorf("updating status: %w", err)
185185
}
186186
if hpaResult.RequeueAfter > 0 {
187187
if statusResult.RequeueAfter == 0 || hpaResult.RequeueAfter < statusResult.RequeueAfter {
@@ -688,7 +688,11 @@ func (r *AgentDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) error {
688688
For(&agentraxv1alpha1.AgentDeployment{}).
689689
Owns(&appsv1.Deployment{}).
690690
Owns(&corev1.Service{}).
691-
Owns(&autoscalingv2.HorizontalPodAutoscaler{})
691+
Owns(&autoscalingv2.HorizontalPodAutoscaler{}).
692+
Watches(
693+
&agentraxv1alpha1.TenantQuota{},
694+
enqueueAgentDeploymentsForTenantQuota(mgr.GetClient()),
695+
)
692696

693697
if r.hasServiceMonitorCRD {
694698
bldr = bldr.Owns(&monitoringv1.ServiceMonitor{})

internal/controller/agentdeployment_controller_test.go

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -755,6 +755,11 @@ var _ = Describe("AgentDeployment HPA lifecycle", func() {
755755
})
756756

757757
It("sets the HPA owner reference to the AgentDeployment", func() {
758+
parent := &agentraxv1alpha1.AgentDeployment{}
759+
Eventually(func() error {
760+
return k8sClient.Get(ctx, key, parent)
761+
}, testTimeout, testInterval).Should(Succeed())
762+
758763
hpa := &autoscalingv2.HorizontalPodAutoscaler{}
759764
Eventually(func() error {
760765
return k8sClient.Get(ctx, key, hpa)
@@ -763,6 +768,7 @@ var _ = Describe("AgentDeployment HPA lifecycle", func() {
763768
Expect(hpa.OwnerReferences).To(HaveLen(1))
764769
Expect(hpa.OwnerReferences[0].Kind).To(Equal("AgentDeployment"))
765770
Expect(hpa.OwnerReferences[0].Name).To(Equal(key.Name))
771+
Expect(hpa.OwnerReferences[0].UID).To(Equal(parent.UID))
766772
Expect(hpa.OwnerReferences[0].Controller).NotTo(BeNil())
767773
Expect(*hpa.OwnerReferences[0].Controller).To(BeTrue())
768774
})
@@ -952,15 +958,13 @@ var _ = Describe("AgentDeployment HPA lifecycle", func() {
952958

953959
// Verify QuotaLimited condition is NOT True — max=5 is within headroom of 10.
954960
// Use Consistently so a transient True that later settles does not go undetected.
955-
Consistently(func() bool {
961+
Consistently(func(g Gomega) {
956962
latest := &agentraxv1alpha1.AgentDeployment{}
957-
if err := k8sClient.Get(ctx, key, latest); err != nil {
958-
return false // treat Get error as not-True; outer Eventually guards timing
959-
}
963+
g.Expect(k8sClient.Get(ctx, key, latest)).To(Succeed())
960964
c := apimeta.FindStatusCondition(latest.Status.Conditions, agentraxv1alpha1.ConditionQuotaLimited)
961-
return c != nil && c.Status == metav1.ConditionTrue
962-
}, 3*time.Second, testInterval).Should(BeFalse(),
963-
"QuotaLimited should never be True when max (5) <= headroom (10)")
965+
g.Expect(c != nil && c.Status == metav1.ConditionTrue).To(BeFalse(),
966+
"QuotaLimited should never be True when max (5) <= headroom (10)")
967+
}, 3*time.Second, testInterval).Should(Succeed())
964968
})
965969
})
966970

@@ -1037,16 +1041,6 @@ var _ = Describe("AgentDeployment HPA lifecycle", func() {
10371041
patch.Spec.MaxTotalReplicas = 2
10381042
Expect(k8sClient.Patch(ctx, patch, client.MergeFrom(tq))).To(Succeed())
10391043

1040-
// Trigger a reconcile by patching the AD (no-op label change).
1041-
ad := &agentraxv1alpha1.AgentDeployment{}
1042-
Expect(k8sClient.Get(ctx, key, ad)).To(Succeed())
1043-
adPatch := ad.DeepCopy()
1044-
if adPatch.Labels == nil {
1045-
adPatch.Labels = make(map[string]string)
1046-
}
1047-
adPatch.Labels["agentrax.io/reconcile-trigger"] = "quota-lower"
1048-
Expect(k8sClient.Patch(ctx, adPatch, client.MergeFrom(ad))).To(Succeed())
1049-
10501044
// HPA maxReplicas must be reduced to the new quota ceiling (2).
10511045
Eventually(func() int32 {
10521046
hpa := &autoscalingv2.HorizontalPodAutoscaler{}

internal/controller/enqueue_handlers.go

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,33 @@ func enqueueTenantQuota() handler.EventHandler {
4747
}
4848
})
4949
}
50+
51+
// enqueueAgentDeploymentsForTenantQuota returns an EventHandler that maps every
52+
// TenantQuota event to reconcile requests for all AgentDeployments in the same
53+
// namespace that reference that TenantQuota. This ensures that lowering or
54+
// raising quota ceilings updates HPA maxReplicas and QuotaLimited conditions
55+
// across all affected agents without requiring out-of-band edits to each AD.
56+
func enqueueAgentDeploymentsForTenantQuota(c client.Client) handler.EventHandler {
57+
return handler.EnqueueRequestsFromMapFunc(func(ctx context.Context, obj client.Object) []reconcile.Request {
58+
tq, ok := obj.(*agentraxv1alpha1.TenantQuota)
59+
if !ok {
60+
return nil
61+
}
62+
list := &agentraxv1alpha1.AgentDeploymentList{}
63+
if err := c.List(ctx, list, client.InNamespace(tq.Namespace)); err != nil {
64+
return nil
65+
}
66+
var reqs []reconcile.Request
67+
for _, ad := range list.Items {
68+
if ad.Spec.TenantRef == tq.Name {
69+
reqs = append(reqs, reconcile.Request{
70+
NamespacedName: types.NamespacedName{
71+
Namespace: ad.Namespace,
72+
Name: ad.Name,
73+
},
74+
})
75+
}
76+
}
77+
return reqs
78+
})
79+
}

internal/metrics/prometheus.go

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -172,18 +172,14 @@ func parseScalarFromQueryResponse(body []byte) (float64, error) {
172172

173173
switch r.Data.ResultType {
174174
case "scalar":
175-
// Scalar result: Data.Result is a [timestamp, "value"] pair at index 0.
175+
// Scalar result: Data.Result is a [timestamp, "value"] pair.
176176
// Require exactly two elements to guard against a malformed payload;
177177
// we read the value string from index 1 (index 0 is the Unix timestamp).
178-
if len(r.Data.Result) == 0 {
179-
return 0, fmt.Errorf("Prometheus scalar result is empty")
180-
}
181-
var pair [2]json.RawMessage
182-
if err := json.Unmarshal(r.Data.Result[0], &pair); err != nil {
183-
return 0, fmt.Errorf("decoding scalar value pair: %w", err)
178+
if len(r.Data.Result) != 2 {
179+
return 0, fmt.Errorf("Prometheus scalar result must contain exactly 2 elements, got %d", len(r.Data.Result))
184180
}
185181
var valStr string
186-
if err := json.Unmarshal(pair[1], &valStr); err != nil {
182+
if err := json.Unmarshal(r.Data.Result[1], &valStr); err != nil {
187183
return 0, fmt.Errorf("decoding scalar value string: %w", err)
188184
}
189185
return strconv.ParseFloat(valStr, 64)
Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
/*
2+
Copyright 2026.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package metrics
18+
19+
import (
20+
"context"
21+
"net/http"
22+
"net/http/httptest"
23+
"testing"
24+
"time"
25+
)
26+
27+
func TestParseScalarFromQueryResponse(t *testing.T) {
28+
t.Parallel()
29+
30+
tests := []struct {
31+
name string
32+
body string
33+
wantVal float64
34+
wantErr bool
35+
}{
36+
{
37+
name: "valid scalar response",
38+
body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"42.5"]}}`,
39+
wantVal: 42.5,
40+
wantErr: false,
41+
},
42+
{
43+
name: "scalar with fewer than 2 elements",
44+
body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781]}}`,
45+
wantErr: true,
46+
},
47+
{
48+
name: "scalar with more than 2 elements",
49+
body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"42.5","extra"]}}`,
50+
wantErr: true,
51+
},
52+
{
53+
name: "scalar with non-string value",
54+
body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,42.5]}}`,
55+
wantErr: true,
56+
},
57+
{
58+
name: "scalar with non-float string value",
59+
body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"invalid"]}}`,
60+
wantErr: true,
61+
},
62+
{
63+
name: "valid vector response",
64+
body: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"http_requests_total"},"value":[1435781451.781,"100.5"]}]}}`,
65+
wantVal: 100.5,
66+
wantErr: false,
67+
},
68+
{
69+
name: "empty vector response",
70+
body: `{"status":"success","data":{"resultType":"vector","result":[]}}`,
71+
wantErr: true,
72+
},
73+
{
74+
name: "unsupported resultType",
75+
body: `{"status":"success","data":{"resultType":"matrix","result":[]}}`,
76+
wantErr: true,
77+
},
78+
{
79+
name: "status error",
80+
body: `{"status":"error","error":"bad query"}`,
81+
wantErr: true,
82+
},
83+
{
84+
name: "malformed JSON",
85+
body: `{"status":`,
86+
wantErr: true,
87+
},
88+
}
89+
90+
for _, tc := range tests {
91+
tc := tc
92+
t.Run(tc.name, func(t *testing.T) {
93+
t.Parallel()
94+
got, err := parseScalarFromQueryResponse([]byte(tc.body))
95+
if tc.wantErr {
96+
if err == nil {
97+
t.Errorf("expected error for %s, got nil", tc.name)
98+
}
99+
return
100+
}
101+
if err != nil {
102+
t.Fatalf("unexpected error for %s: %v", tc.name, err)
103+
}
104+
if got != tc.wantVal {
105+
t.Errorf("got %v, want %v", got, tc.wantVal)
106+
}
107+
})
108+
}
109+
}
110+
111+
func TestParseLastValueFromRangeResponse(t *testing.T) {
112+
t.Parallel()
113+
114+
tests := []struct {
115+
name string
116+
body string
117+
wantVal float64
118+
wantErr bool
119+
}{
120+
{
121+
name: "valid range response",
122+
body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[[1435781430.781,"10"],[1435781451.781,"20.5"]]}]}}`,
123+
wantVal: 20.5,
124+
wantErr: false,
125+
},
126+
{
127+
name: "empty result",
128+
body: `{"status":"success","data":{"resultType":"matrix","result":[]}}`,
129+
wantErr: true,
130+
},
131+
{
132+
name: "empty values in series",
133+
body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[]}]}}`,
134+
wantErr: true,
135+
},
136+
{
137+
name: "status error",
138+
body: `{"status":"error","error":"bad query"}`,
139+
wantErr: true,
140+
},
141+
{
142+
name: "malformed json",
143+
body: `invalid json`,
144+
wantErr: true,
145+
},
146+
}
147+
148+
for _, tc := range tests {
149+
tc := tc
150+
t.Run(tc.name, func(t *testing.T) {
151+
t.Parallel()
152+
got, err := parseLastValueFromRangeResponse([]byte(tc.body))
153+
if tc.wantErr {
154+
if err == nil {
155+
t.Errorf("expected error for %s, got nil", tc.name)
156+
}
157+
return
158+
}
159+
if err != nil {
160+
t.Fatalf("unexpected error for %s: %v", tc.name, err)
161+
}
162+
if got != tc.wantVal {
163+
t.Errorf("got %v, want %v", got, tc.wantVal)
164+
}
165+
})
166+
}
167+
}
168+
169+
func TestClient_QueryScalar(t *testing.T) {
170+
t.Parallel()
171+
172+
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
173+
if r.URL.Path != "/api/v1/query" {
174+
http.NotFound(w, r)
175+
return
176+
}
177+
q := r.URL.Query().Get("query")
178+
if q == "scalar_metric" {
179+
w.Header().Set("Content-Type", "application/json")
180+
w.WriteHeader(http.StatusOK)
181+
_, _ = w.Write([]byte(`{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"12.34"]}}`))
182+
return
183+
}
184+
if q == "error_metric" {
185+
w.WriteHeader(http.StatusBadRequest)
186+
_, _ = w.Write([]byte(`bad query`))
187+
return
188+
}
189+
http.Error(w, "unknown query", http.StatusInternalServerError)
190+
}))
191+
defer ts.Close()
192+
193+
c := NewClient(ts.URL, WithTimeout(2*time.Second))
194+
val, err := c.QueryScalar(context.Background(), "scalar_metric")
195+
if err != nil {
196+
t.Fatalf("unexpected error: %v", err)
197+
}
198+
if val != 12.34 {
199+
t.Errorf("expected 12.34, got %v", val)
200+
}
201+
202+
_, err = c.QueryScalar(context.Background(), "error_metric")
203+
if err == nil {
204+
t.Error("expected error for error_metric, got nil")
205+
}
206+
}
207+
208+
func TestClient_QueryRange(t *testing.T) {
209+
t.Parallel()
210+
211+
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
212+
if r.URL.Path != "/api/v1/query_range" {
213+
http.NotFound(w, r)
214+
return
215+
}
216+
w.Header().Set("Content-Type", "application/json")
217+
w.WriteHeader(http.StatusOK)
218+
_, _ = w.Write([]byte(`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[[1000,"5.5"],[2000,"8.5"]]}]}}`))
219+
}))
220+
defer ts.Close()
221+
222+
c := NewClient(ts.URL)
223+
now := time.Now()
224+
val, err := c.QueryRange(context.Background(), "range_query", now.Add(-10*time.Minute), now, time.Minute)
225+
if err != nil {
226+
t.Fatalf("unexpected error: %v", err)
227+
}
228+
if val != 8.5 {
229+
t.Errorf("expected 8.5, got %v", val)
230+
}
231+
}

internal/quota/enforcer.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ func evalQuotaRules(
238238
quota.MaxAgents, committedUsage.UsedAgents, inFlight.agents, deltaAgents,
239239
)
240240
}
241-
if quota.MaxGPUs > 0 && projGPUs > quota.MaxGPUs && (!isUpdate || deltaGPUs > 0) {
241+
if projGPUs > quota.MaxGPUs && (!isUpdate || deltaGPUs > 0) {
242242
return false, fmt.Sprintf(
243243
"would exceed maxGPUs (%d): current=%d in-flight=%d delta=%d",
244244
quota.MaxGPUs, committedUsage.UsedGPUs, inFlight.gpus, deltaGPUs,
@@ -367,7 +367,7 @@ func (e *Enforcer) IsOverQuota(
367367
if usage.UsedAgents > quota.MaxAgents {
368368
return true, fmt.Sprintf("usedAgents (%d) exceeds maxAgents (%d)", usage.UsedAgents, quota.MaxAgents)
369369
}
370-
if quota.MaxGPUs > 0 && usage.UsedGPUs > quota.MaxGPUs {
370+
if usage.UsedGPUs > quota.MaxGPUs {
371371
return true, fmt.Sprintf("usedGPUs (%d) exceeds maxGPUs (%d)", usage.UsedGPUs, quota.MaxGPUs)
372372
}
373373
if usage.UsedTotalReplicas > quota.MaxTotalReplicas {

0 commit comments

Comments
 (0)