From 03d6f5d3e04bf73369043f1eeb2ea8d9dbf302bc Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Fri, 25 Sep 2026 06:11:20 +0000 Subject: [PATCH 1/3] feat(controller): apply Task resource limits to the ActorTemplate TaskSpec.resources is documented and accepted by ax-server, but the controller never read it, so spec.resources.limits had no effect on the sandbox. Translate limits.cpu and limits.memory into the Substrate ActorTemplate resources block (Substrate's {name, quantity} list) when provisioning the per-task template. Substrate sizes sandboxes by limits only, so requests are documented as not applied. A limits change already yields a new template because the launch spec is part of the template digest via AX_TASK_YAML; a test now pins that behaviour. Fixes #369 Co-authored-by: Matt Van Horn --- docs/manifests.md | 6 ++ internal/controller/reconciler.go | 5 +- internal/controller/reconciler_test.go | 87 ++++++++++++++++++++++++++ internal/substrate/client.go | 30 +++++++-- internal/substrate/client_test.go | 85 +++++++++++++++++++++++++ 5 files changed, 208 insertions(+), 5 deletions(-) create mode 100644 internal/substrate/client_test.go diff --git a/docs/manifests.md b/docs/manifests.md index c3bdcb12..22f43431 100644 --- a/docs/manifests.md +++ b/docs/manifests.md @@ -48,6 +48,12 @@ spec: Each entry is set up independently at its own path, in order. Every entry needs a `name`; without a `path` it lands at `/workspace/`, and paths must be unique. The first entry is the working directory of `spec.command`, and the task reports `WorkspaceReady` only once all of them are prepared. See [`examples/multi-workspace.yaml`](../examples/multi-workspace.yaml) for a complete set. +### Sizing the sandbox + +`spec.resources.limits` caps the CPU and memory of the task's sandbox. The controller copies the limits onto the Substrate `ActorTemplate` it provisions for the task, using Kubernetes quantity syntax (`500m`, `2`, `4Gi`). Only `cpu` and `memory` are supported, each quantity must be greater than zero, and the CPU limit must be below 1000 cores. A task without limits is sized by its worker's defaults. + +Substrate sizes sandboxes by limits alone, so `spec.resources.requests` is stored on the task but not applied. + ## Workspace ```yaml diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index 8a3cde17..da7d1ad2 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -173,7 +173,10 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, wor slog.Info("ensuring custom ActorTemplate for task", "image", task.Spec.Image) customTemplateName := taskTemplateName(task.Metadata.Name, task.Spec.Image, extraEnv) - tmpl, err := r.client.EnsureActorTemplateWithImage(ctx, templateAtespace, templateName, atespace, customTemplateName, task.Spec.Image, extraEnv) + // spec.resources rides along in AX_TASK_YAML, so a limits change is already + // part of the template digest and re-provisions the template. + resources := substrate.ResourceLimits(task.Spec.Resources) + tmpl, err := r.client.EnsureActorTemplateWithImage(ctx, templateAtespace, templateName, atespace, customTemplateName, task.Spec.Image, resources, extraEnv) if err != nil { slog.Warn("could not create custom ActorTemplate, falling back to default template", "error", err) } else if tmpl != nil && tmpl.Metadata != nil { diff --git a/internal/controller/reconciler_test.go b/internal/controller/reconciler_test.go index 94396024..0226fdd1 100644 --- a/internal/controller/reconciler_test.go +++ b/internal/controller/reconciler_test.go @@ -29,6 +29,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" ) type mockControlServer struct { @@ -40,6 +41,7 @@ type mockControlServer struct { suspendedActors []string deletedActors []string actorTemplates map[string]bool + createdTemplates []*ateapipb.ActorTemplate deletedTemplates []string } @@ -62,6 +64,7 @@ func (m *mockControlServer) CreateActorTemplate(_ context.Context, req *ateapipb } tmpl := req.GetActorTemplate() m.actorTemplates[tmpl.GetMetadata().GetName()] = true + m.createdTemplates = append(m.createdTemplates, tmpl) return tmpl, nil } @@ -387,6 +390,90 @@ func TestTaskReconciler_WorkspaceReady(t *testing.T) { } } +func TestTaskReconciler_ResourceLimits(t *testing.T) { + ctx := context.Background() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + defer lis.Close() + + mockSrv := &mockControlServer{} + grpcServer := grpc.NewServer() + ateapipb.RegisterControlServer(grpcServer, mockSrv) + go grpcServer.Serve(lis) + defer grpcServer.Stop() + + client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to create substrate client: %v", err) + } + defer client.Close() + + reconciler := controller.NewTaskReconciler(client, "test-template", "ax-system") + reconciler.SecretResolver = noSecrets + reconciler.WorkspaceReadyTimeout = 200 * time.Millisecond + + task := &v1alpha1.Task{ + ApiVersion: v1alpha1.APIVersion, + Kind: v1alpha1.KindTask, + Metadata: &v1alpha1.ObjectMeta{ + Name: "sized-task", + Atespace: "default", + }, + Spec: &v1alpha1.TaskSpec{ + Image: "ghrc.io/my-org/my-image", + Resources: &v1alpha1.ResourceReqs{ + Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}, + Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, + }, + }, + } + + if _, err := reconciler.Reconcile(ctx, task); err != nil { + t.Fatalf("Reconcile failed: %v", err) + } + if len(mockSrv.createdTemplates) != 1 { + t.Fatalf("created %d templates, want 1", len(mockSrv.createdTemplates)) + } + want := &ateapipb.Resources{Limits: []*ateapipb.Limits{ + {Name: "cpu", Quantity: "2"}, + {Name: "memory", Quantity: "4Gi"}, + }} + if got := mockSrv.createdTemplates[0].GetResources(); !proto.Equal(got, want) { + t.Errorf("template resources = %v, want %v", got, want) + } + + // Raising a limit is a launch configuration change and must provision a new + // template carrying the new value. + task.Spec.Resources.Limits.Memory = "8Gi" + if _, err := reconciler.Reconcile(ctx, task); err != nil { + t.Fatalf("Reconcile with changed limits failed: %v", err) + } + if len(mockSrv.createdTemplates) != 2 { + t.Fatalf("limits change left %d templates, want 2", len(mockSrv.createdTemplates)) + } + want.Limits[1].Quantity = "8Gi" + if got := mockSrv.createdTemplates[1].GetResources(); !proto.Equal(got, want) { + t.Errorf("template resources after change = %v, want %v", got, want) + } + + // A task without limits inherits the worker defaults: no resources block. + plain := &v1alpha1.Task{ + ApiVersion: v1alpha1.APIVersion, + Kind: v1alpha1.KindTask, + Metadata: &v1alpha1.ObjectMeta{Name: "plain-task", Atespace: "default"}, + Spec: &v1alpha1.TaskSpec{Image: "ghrc.io/my-org/my-image"}, + } + if _, err := reconciler.Reconcile(ctx, plain); err != nil { + t.Fatalf("Reconcile of task without limits failed: %v", err) + } + if got := mockSrv.createdTemplates[len(mockSrv.createdTemplates)-1].GetResources(); got != nil { + t.Errorf("template for task without limits has resources %v, want none", got) + } +} + // assertCondition fails the test unless the task has a condition of the given type with // the expected status and reason. func assertCondition(t *testing.T, task *v1alpha1.Task, condType, wantStatus, wantReason string) { diff --git a/internal/substrate/client.go b/internal/substrate/client.go index e3a32d29..92d1a006 100644 --- a/internal/substrate/client.go +++ b/internal/substrate/client.go @@ -208,8 +208,28 @@ const ( DefaultSnapshotsBucket = "gs://dberkov-gke-dev3/ate-env/" ) +// ResourceLimits translates a Task's resource requirements into the Substrate +// ActorTemplate resources block. Substrate sizes a sandbox by limits alone, so +// only spec.resources.limits is carried over; requests are not applied. It +// returns nil when no limit is set so the template inherits the worker defaults. +func ResourceLimits(reqs *v1alpha1.ResourceReqs) *ateapipb.Resources { + limits := reqs.GetLimits() + var out []*ateapipb.Limits + if cpu := limits.GetCpu(); cpu != "" { + out = append(out, &ateapipb.Limits{Name: "cpu", Quantity: cpu}) + } + if memory := limits.GetMemory(); memory != "" { + out = append(out, &ateapipb.Limits{Name: "memory", Quantity: memory}) + } + if len(out) == 0 { + return nil + } + return &ateapipb.Resources{Limits: out} +} + // BuildActorTemplate constructs a Substrate ActorTemplate based on the standard ate-env specification. -func BuildActorTemplate(atespace, name, image string, envMap map[string]string, command []string, snapshotsBucket string) *ateapipb.ActorTemplate { +// A nil resources leaves the sandbox sized by the worker defaults. +func BuildActorTemplate(atespace, name, image string, envMap map[string]string, command []string, snapshotsBucket string, resources *ateapipb.Resources) *ateapipb.ActorTemplate { if atespace == "" { atespace = "default" } @@ -272,11 +292,13 @@ func BuildActorTemplate(atespace, name, image string, envMap map[string]string, SandboxClass: ateapipb.SandboxClass_SANDBOX_CLASS_GVISOR, ConfigName: "gvisor-default", }, + Resources: resources, } } -// EnsureActorTemplateWithImage creates an ActorTemplate using the specified container image and optional environment variables. -func (c *Client) EnsureActorTemplateWithImage(ctx context.Context, baseAtespace, baseTemplate, targetAtespace, targetTemplate, image string, extraEnv ...map[string]string) (*ateapipb.ActorTemplate, error) { +// EnsureActorTemplateWithImage creates an ActorTemplate using the specified container image, +// resource limits (nil for worker defaults), and optional environment variables. +func (c *Client) EnsureActorTemplateWithImage(ctx context.Context, baseAtespace, baseTemplate, targetAtespace, targetTemplate, image string, resources *ateapipb.Resources, extraEnv ...map[string]string) (*ateapipb.ActorTemplate, error) { existing, err := c.GetActorTemplate(ctx, targetAtespace, targetTemplate) if err == nil && existing != nil { return existing, nil @@ -289,7 +311,7 @@ func (c *Client) EnsureActorTemplateWithImage(ctx context.Context, baseAtespace, } } - tmpl := BuildActorTemplate(targetAtespace, targetTemplate, image, envMap, nil, "") + tmpl := BuildActorTemplate(targetAtespace, targetTemplate, image, envMap, nil, "", resources) req := &ateapipb.CreateActorTemplateRequest{ ActorTemplate: tmpl, } diff --git a/internal/substrate/client_test.go b/internal/substrate/client_test.go new file mode 100644 index 00000000..5205ba98 --- /dev/null +++ b/internal/substrate/client_test.go @@ -0,0 +1,85 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package substrate + +import ( + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/ax/pkg/apis/v1alpha1" + "google.golang.org/protobuf/proto" +) + +func TestResourceLimits(t *testing.T) { + tests := []struct { + name string + reqs *v1alpha1.ResourceReqs + want *ateapipb.Resources + }{ + {name: "nil"}, + {name: "empty", reqs: &v1alpha1.ResourceReqs{}}, + { + name: "requests only are not applied", + reqs: &v1alpha1.ResourceReqs{Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}}, + }, + { + name: "cpu limit", + reqs: &v1alpha1.ResourceReqs{Limits: &v1alpha1.ResourceList{Cpu: "2"}}, + want: &ateapipb.Resources{Limits: []*ateapipb.Limits{{Name: "cpu", Quantity: "2"}}}, + }, + { + name: "memory limit", + reqs: &v1alpha1.ResourceReqs{Limits: &v1alpha1.ResourceList{Memory: "4Gi"}}, + want: &ateapipb.Resources{Limits: []*ateapipb.Limits{{Name: "memory", Quantity: "4Gi"}}}, + }, + { + name: "cpu and memory limits", + reqs: &v1alpha1.ResourceReqs{ + Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}, + Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, + }, + want: &ateapipb.Resources{Limits: []*ateapipb.Limits{ + {Name: "cpu", Quantity: "2"}, + {Name: "memory", Quantity: "4Gi"}, + }}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ResourceLimits(tt.reqs) + if !proto.Equal(got, tt.want) { + t.Fatalf("ResourceLimits() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestBuildActorTemplate_Resources(t *testing.T) { + limits := &ateapipb.Resources{Limits: []*ateapipb.Limits{ + {Name: "cpu", Quantity: "2"}, + {Name: "memory", Quantity: "4Gi"}, + }} + tmpl := BuildActorTemplate("default", "task-tmpl-01234567", "ghcr.io/my-org/agent@sha256:abc", nil, nil, "", limits) + if !proto.Equal(tmpl.GetResources(), limits) { + t.Fatalf("template resources = %v, want %v", tmpl.GetResources(), limits) + } + + // Without limits the template must not carry a resources block, so the + // worker defaults keep applying. + tmpl = BuildActorTemplate("default", "task-tmpl-01234567", "ghcr.io/my-org/agent@sha256:abc", nil, nil, "", nil) + if tmpl.GetResources() != nil { + t.Fatalf("template without limits has resources %v, want none", tmpl.GetResources()) + } +} From c648889d9a1d16f71c506682ce5b2800c1fe76ac Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Fri, 25 Sep 2026 06:45:50 +0000 Subject: [PATCH 2/3] controller: validate resource limits and never fall back without them Address review on #405: - Validate spec.resources before provisioning, at apply time via ValidateTask and again in the reconciler. Quantities follow the Kubernetes quantity grammar, must be greater than zero, and the cpu limit must be below 1000 cores, matching what Substrate enforces on the ActorTemplate. Invalid values fail the reconcile with Ready=False, reason InvalidResources. - When a task has limits and Substrate rejects the per-task template, fail the reconcile (reason TemplateCreationFailed) instead of falling back to the default template, which would run the task without the limits it asked for. Tasks without limits keep the existing fallback. - Reject spec.resources.requests with a clear message: Substrate sizes sandboxes by limits alone, so accepting and ignoring it was misleading. The docs and examples/task.yaml drop the requests block. - Tests: failure paths for invalid limits and for a rejected template with and without limits, plus an explicit assertion that a limits change yields a differently named template. Co-authored-by: Matt Van Horn --- docs/concepts.md | 2 +- docs/manifests.md | 7 +- examples/task.yaml | 3 - internal/controller/reconciler.go | 17 +++- internal/controller/reconciler_test.go | 118 ++++++++++++++++++++++- internal/substrate/client.go | 8 +- internal/substrate/client_test.go | 8 +- pkg/apis/v1alpha1/resources.go | 126 +++++++++++++++++++++++++ pkg/apis/v1alpha1/resources_test.go | 95 +++++++++++++++++++ pkg/apis/v1alpha1/types.go | 2 +- 10 files changed, 362 insertions(+), 24 deletions(-) create mode 100644 pkg/apis/v1alpha1/resources.go create mode 100644 pkg/apis/v1alpha1/resources_test.go diff --git a/docs/concepts.md b/docs/concepts.md index 81380c23..a165149c 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -4,7 +4,7 @@ Every AX resource lives in an **atespace**. The default atespace is `default`. ## Task -The smallest unit of isolated execution. A `Task` declares the container image and command, compute requests and limits, environment variables, and references to one or more `Workspace`s under `spec.workspaces`. Each workspace is mounted at its own path, and the first serves as the command's working directory. +The smallest unit of isolated execution. A `Task` declares the container image and command, compute limits, environment variables, and references to one or more `Workspace`s under `spec.workspaces`. Each workspace is mounted at its own path, and the first serves as the command's working directory. The unit is deliberately small. An agent is not one process that runs to completion; over its lifetime it plans, delegates, retries, and fans work out. AX does not try to model that shape. It gives you one primitive that is cheap to create, isolate, suspend, and throw away, and lets the agent compose as many of them as its work demands. A single task may be the whole job, or it may be the root of a large tree of tasks spawned as the agent breaks the problem down. Either way each node gets the same sandbox, the same lifecycle, and the same tooling. diff --git a/docs/manifests.md b/docs/manifests.md index 22f43431..2086722d 100644 --- a/docs/manifests.md +++ b/docs/manifests.md @@ -18,9 +18,6 @@ spec: value: "production" resources: - requests: - cpu: "500m" - memory: "1Gi" limits: cpu: "2" memory: "4Gi" @@ -50,9 +47,9 @@ Each entry is set up independently at its own path, in order. Every entry needs ### Sizing the sandbox -`spec.resources.limits` caps the CPU and memory of the task's sandbox. The controller copies the limits onto the Substrate `ActorTemplate` it provisions for the task, using Kubernetes quantity syntax (`500m`, `2`, `4Gi`). Only `cpu` and `memory` are supported, each quantity must be greater than zero, and the CPU limit must be below 1000 cores. A task without limits is sized by its worker's defaults. +`spec.resources.limits` caps the CPU and memory of the task's sandbox. The controller copies the limits onto the Substrate `ActorTemplate` it provisions for the task, using Kubernetes quantity syntax (`500m`, `2`, `4Gi`). Only `cpu` and `memory` are supported, each quantity must be greater than zero, and the CPU limit must be below 1000 cores. `ax apply` rejects values outside these rules, and a task whose limits Substrate refuses is marked `Failed` with reason `TemplateCreationFailed` rather than run without them. A task without limits is sized by its worker's defaults. -Substrate sizes sandboxes by limits alone, so `spec.resources.requests` is stored on the task but not applied. +Substrate sizes sandboxes by limits alone, so `spec.resources.requests` is not supported and `ax apply` rejects a manifest that sets it. ## Workspace diff --git a/examples/task.yaml b/examples/task.yaml index bce5e875..c9d5467a 100644 --- a/examples/task.yaml +++ b/examples/task.yaml @@ -24,9 +24,6 @@ spec: image: "gcr.io/ax-substrate/ate-images/ax-task-runner@sha256:3a0dea6ad8b55278685db58aca6e37dc4ba04056831d45bef3aaeafdca43cac6" resources: - requests: - cpu: "500m" - memory: "1Gi" limits: cpu: "2" memory: "4Gi" diff --git a/internal/controller/reconciler.go b/internal/controller/reconciler.go index da7d1ad2..750cd310 100644 --- a/internal/controller/reconciler.go +++ b/internal/controller/reconciler.go @@ -168,6 +168,15 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, wor extraEnv["AX_WORKSPACES_YAML"] = wsYAML } + // Resource limits are enforced by the per-task template, so they are checked + // here as well as at apply time: a task must never run without limits it asked for. + if err := v1alpha1.ValidateResources(task.Spec.Resources); err != nil { + r.setNotReady(task, "InvalidResources", err.Error(), now) + task.Status.Phase = "Failed" + return task, fmt.Errorf("validating resources: %w", err) + } + resources := substrate.ResourceLimits(task.Spec.Resources) + // If a custom image, workspace, or extra environment is specified, provision or use a dedicated ActorTemplate if task.Spec != nil && (task.Spec.Image != "" || len(extraEnv) > 0) { slog.Info("ensuring custom ActorTemplate for task", "image", task.Spec.Image) @@ -175,8 +184,14 @@ func (r *TaskReconciler) Reconcile(ctx context.Context, task *v1alpha1.Task, wor // spec.resources rides along in AX_TASK_YAML, so a limits change is already // part of the template digest and re-provisions the template. - resources := substrate.ResourceLimits(task.Spec.Resources) tmpl, err := r.client.EnsureActorTemplateWithImage(ctx, templateAtespace, templateName, atespace, customTemplateName, task.Spec.Image, resources, extraEnv) + if err != nil && resources != nil { + // The default template does not carry the task's limits, so falling back + // would silently run the task unconstrained. + r.setNotReady(task, "TemplateCreationFailed", err.Error(), now) + task.Status.Phase = "Failed" + return task, fmt.Errorf("ensuring actor template: %w", err) + } if err != nil { slog.Warn("could not create custom ActorTemplate, falling back to default template", "error", err) } else if tmpl != nil && tmpl.Metadata != nil { diff --git a/internal/controller/reconciler_test.go b/internal/controller/reconciler_test.go index 0226fdd1..1e4ec7e3 100644 --- a/internal/controller/reconciler_test.go +++ b/internal/controller/reconciler_test.go @@ -43,6 +43,9 @@ type mockControlServer struct { actorTemplates map[string]bool createdTemplates []*ateapipb.ActorTemplate deletedTemplates []string + // createTemplateErr, when set, is returned by CreateActorTemplate to stand in + // for Substrate rejecting a template (for example an invalid quantity). + createTemplateErr error } // noSecrets is a SecretResolver for tests: it never finds a key and never touches a cluster. @@ -59,6 +62,9 @@ func (m *mockControlServer) GetActorTemplate(_ context.Context, req *ateapipb.Ge } func (m *mockControlServer) CreateActorTemplate(_ context.Context, req *ateapipb.CreateActorTemplateRequest) (*ateapipb.ActorTemplate, error) { + if m.createTemplateErr != nil { + return nil, m.createTemplateErr + } if m.actorTemplates == nil { m.actorTemplates = make(map[string]bool) } @@ -425,8 +431,7 @@ func TestTaskReconciler_ResourceLimits(t *testing.T) { Spec: &v1alpha1.TaskSpec{ Image: "ghrc.io/my-org/my-image", Resources: &v1alpha1.ResourceReqs{ - Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}, - Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, + Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, }, }, } @@ -446,7 +451,7 @@ func TestTaskReconciler_ResourceLimits(t *testing.T) { } // Raising a limit is a launch configuration change and must provision a new - // template carrying the new value. + // template, under a new name, carrying the new value. task.Spec.Resources.Limits.Memory = "8Gi" if _, err := reconciler.Reconcile(ctx, task); err != nil { t.Fatalf("Reconcile with changed limits failed: %v", err) @@ -454,6 +459,10 @@ func TestTaskReconciler_ResourceLimits(t *testing.T) { if len(mockSrv.createdTemplates) != 2 { t.Fatalf("limits change left %d templates, want 2", len(mockSrv.createdTemplates)) } + first, second := mockSrv.createdTemplates[0].GetMetadata().GetName(), mockSrv.createdTemplates[1].GetMetadata().GetName() + if first == second { + t.Errorf("limits change reused template name %q, want a distinct name", first) + } want.Limits[1].Quantity = "8Gi" if got := mockSrv.createdTemplates[1].GetResources(); !proto.Equal(got, want) { t.Errorf("template resources after change = %v, want %v", got, want) @@ -474,6 +483,109 @@ func TestTaskReconciler_ResourceLimits(t *testing.T) { } } +// A task that asks for limits must never run without them: invalid limits and a +// Substrate rejection of the template both fail the reconcile instead of falling +// back to the default template. Without limits the fallback contract is unchanged. +func TestTaskReconciler_ResourceLimitsFailurePaths(t *testing.T) { + ctx := context.Background() + + newTask := func(name string, resources *v1alpha1.ResourceReqs) *v1alpha1.Task { + return &v1alpha1.Task{ + ApiVersion: v1alpha1.APIVersion, + Kind: v1alpha1.KindTask, + Metadata: &v1alpha1.ObjectMeta{Name: name, Atespace: "default"}, + Spec: &v1alpha1.TaskSpec{Image: "ghrc.io/my-org/my-image", Resources: resources}, + } + } + setup := func(t *testing.T, mockSrv *mockControlServer) *controller.TaskReconciler { + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + grpcServer := grpc.NewServer() + ateapipb.RegisterControlServer(grpcServer, mockSrv) + go grpcServer.Serve(lis) + client, err := substrate.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to create substrate client: %v", err) + } + t.Cleanup(func() { + client.Close() + grpcServer.Stop() + lis.Close() + }) + reconciler := controller.NewTaskReconciler(client, "test-template", "ax-system") + reconciler.SecretResolver = noSecrets + reconciler.WorkspaceReadyTimeout = 200 * time.Millisecond + return reconciler + } + + t.Run("invalid limits fail before provisioning", func(t *testing.T) { + mockSrv := &mockControlServer{} + reconciler := setup(t, mockSrv) + + for _, reqs := range []*v1alpha1.ResourceReqs{ + {Limits: &v1alpha1.ResourceList{Cpu: "two"}}, + {Limits: &v1alpha1.ResourceList{Cpu: "0"}}, + {Limits: &v1alpha1.ResourceList{Cpu: "1000"}}, + {Limits: &v1alpha1.ResourceList{Memory: "-4Gi"}}, + {Requests: &v1alpha1.ResourceList{Cpu: "500m"}, Limits: &v1alpha1.ResourceList{Cpu: "2"}}, + } { + reconciled, err := reconciler.Reconcile(ctx, newTask("bad-limits", reqs)) + if err == nil { + t.Errorf("Reconcile(%v) succeeded, want error", reqs) + } + if reconciled.Status.Phase != "Failed" { + t.Errorf("Reconcile(%v) phase = %q, want Failed", reqs, reconciled.Status.Phase) + } + assertCondition(t, reconciled, "Ready", "False", "InvalidResources") + } + if len(mockSrv.createdTemplates) != 0 || len(mockSrv.createdActors) != 0 || len(mockSrv.resumedActors) != 0 { + t.Errorf("invalid limits reached Substrate: templates=%d actors=%v resumed=%v", + len(mockSrv.createdTemplates), mockSrv.createdActors, mockSrv.resumedActors) + } + }) + + t.Run("rejected template with limits does not fall back", func(t *testing.T) { + mockSrv := &mockControlServer{ + createTemplateErr: status.Error(codes.InvalidArgument, "actor_template.resources.limits[0].quantity: Invalid value"), + } + reconciler := setup(t, mockSrv) + + reconciled, err := reconciler.Reconcile(ctx, newTask("rejected-limits", &v1alpha1.ResourceReqs{ + Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, + })) + if err == nil { + t.Fatal("Reconcile succeeded although the template with limits was rejected") + } + if reconciled.Status.Phase != "Failed" { + t.Errorf("phase = %q, want Failed", reconciled.Status.Phase) + } + assertCondition(t, reconciled, "Ready", "False", "TemplateCreationFailed") + if len(mockSrv.createdActors) != 0 || len(mockSrv.resumedActors) != 0 { + t.Errorf("task ran on a fallback template without its limits: actors=%v resumed=%v", mockSrv.createdActors, mockSrv.resumedActors) + } + }) + + t.Run("rejected template without limits still falls back", func(t *testing.T) { + mockSrv := &mockControlServer{ + createTemplateErr: status.Error(codes.InvalidArgument, "image must be pinned by digest"), + } + reconciler := setup(t, mockSrv) + + reconciled, err := reconciler.Reconcile(ctx, newTask("no-limits", nil)) + if err != nil { + t.Fatalf("Reconcile without limits failed: %v", err) + } + if reconciled.Status.Phase != "Running" { + t.Errorf("phase = %q, want Running", reconciled.Status.Phase) + } + if len(mockSrv.createdActors) != 1 { + t.Errorf("fallback did not create the actor: %v", mockSrv.createdActors) + } + }) +} + // assertCondition fails the test unless the task has a condition of the given type with // the expected status and reason. func assertCondition(t *testing.T, task *v1alpha1.Task, condType, wantStatus, wantReason string) { diff --git a/internal/substrate/client.go b/internal/substrate/client.go index 92d1a006..68c7fe20 100644 --- a/internal/substrate/client.go +++ b/internal/substrate/client.go @@ -208,10 +208,10 @@ const ( DefaultSnapshotsBucket = "gs://dberkov-gke-dev3/ate-env/" ) -// ResourceLimits translates a Task's resource requirements into the Substrate -// ActorTemplate resources block. Substrate sizes a sandbox by limits alone, so -// only spec.resources.limits is carried over; requests are not applied. It -// returns nil when no limit is set so the template inherits the worker defaults. +// ResourceLimits translates a Task's resource limits into the Substrate +// ActorTemplate resources block. Substrate sizes a sandbox by limits alone +// (requests are rejected by v1alpha1.ValidateResources). It returns nil when no +// limit is set so the template inherits the worker defaults. func ResourceLimits(reqs *v1alpha1.ResourceReqs) *ateapipb.Resources { limits := reqs.GetLimits() var out []*ateapipb.Limits diff --git a/internal/substrate/client_test.go b/internal/substrate/client_test.go index 5205ba98..b1142484 100644 --- a/internal/substrate/client_test.go +++ b/internal/substrate/client_test.go @@ -30,10 +30,7 @@ func TestResourceLimits(t *testing.T) { }{ {name: "nil"}, {name: "empty", reqs: &v1alpha1.ResourceReqs{}}, - { - name: "requests only are not applied", - reqs: &v1alpha1.ResourceReqs{Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}}, - }, + {name: "empty limits", reqs: &v1alpha1.ResourceReqs{Limits: &v1alpha1.ResourceList{}}}, { name: "cpu limit", reqs: &v1alpha1.ResourceReqs{Limits: &v1alpha1.ResourceList{Cpu: "2"}}, @@ -47,8 +44,7 @@ func TestResourceLimits(t *testing.T) { { name: "cpu and memory limits", reqs: &v1alpha1.ResourceReqs{ - Requests: &v1alpha1.ResourceList{Cpu: "500m", Memory: "1Gi"}, - Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, + Limits: &v1alpha1.ResourceList{Cpu: "2", Memory: "4Gi"}, }, want: &ateapipb.Resources{Limits: []*ateapipb.Limits{ {Name: "cpu", Quantity: "2"}, diff --git a/pkg/apis/v1alpha1/resources.go b/pkg/apis/v1alpha1/resources.go new file mode 100644 index 00000000..195e3f22 --- /dev/null +++ b/pkg/apis/v1alpha1/resources.go @@ -0,0 +1,126 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package v1alpha1 + +import ( + "errors" + "fmt" + "math/big" + "regexp" + "strconv" +) + +// Resource requirements. +// +// spec.resources.limits is copied onto the Substrate ActorTemplate, which +// enforces the same rules as below: cpu and memory only, each quantity greater +// than zero, and a cpu limit below 1000 cores. Checking them here means a bad +// value fails at apply time instead of when the template is created. +// Substrate sizes a sandbox by limits alone, so requests are rejected rather +// than accepted and ignored. + +// MaxCPULimitCores is the exclusive upper bound Substrate places on a cpu limit. +const MaxCPULimitCores = 1000 + +// ValidateResources reports the first problem with a task's resource +// requirements. A nil or empty value is valid: the sandbox is then sized by its +// worker's defaults. +func ValidateResources(reqs *ResourceReqs) error { + if req := reqs.GetRequests(); req.GetCpu() != "" || req.GetMemory() != "" { + return errors.New("spec.resources.requests: not supported, only spec.resources.limits is applied to the sandbox") + } + limits := reqs.GetLimits() + if cpu := limits.GetCpu(); cpu != "" { + q, err := parseQuantity(cpu) + if err != nil { + return fmt.Errorf("spec.resources.limits.cpu: invalid quantity %q: %w", cpu, err) + } + if q.Sign() <= 0 { + return fmt.Errorf("spec.resources.limits.cpu: %q must be greater than zero", cpu) + } + if q.Cmp(big.NewRat(MaxCPULimitCores, 1)) >= 0 { + return fmt.Errorf("spec.resources.limits.cpu: %q must be less than %d cores", cpu, MaxCPULimitCores) + } + } + if memory := limits.GetMemory(); memory != "" { + q, err := parseQuantity(memory) + if err != nil { + return fmt.Errorf("spec.resources.limits.memory: invalid quantity %q: %w", memory, err) + } + if q.Sign() <= 0 { + return fmt.Errorf("spec.resources.limits.memory: %q must be greater than zero", memory) + } + } + return nil +} + +// quantityRegexp follows the Kubernetes resource.Quantity grammar: a decimal +// number followed by an optional binary SI suffix (Ki..Ei), decimal SI suffix +// (n, u, m, k, M, G, T, P, E) or decimal exponent (e3, E-2). +var quantityRegexp = regexp.MustCompile(`^([+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+))([eE][+-]?[0-9]+|Ki|Mi|Gi|Ti|Pi|Ei|[numkMGTPE])?$`) + +var quantitySuffixes = map[string]*big.Rat{ + "n": big.NewRat(1, 1_000_000_000), + "u": big.NewRat(1, 1_000_000), + "m": big.NewRat(1, 1000), + "k": big.NewRat(1000, 1), + "M": big.NewRat(1_000_000, 1), + "G": big.NewRat(1_000_000_000, 1), + "T": big.NewRat(1_000_000_000_000, 1), + "P": big.NewRat(1_000_000_000_000_000, 1), + "E": big.NewRat(1_000_000_000_000_000_000, 1), + "Ki": big.NewRat(1<<10, 1), + "Mi": big.NewRat(1<<20, 1), + "Gi": big.NewRat(1<<30, 1), + "Ti": big.NewRat(1<<40, 1), + "Pi": big.NewRat(1<<50, 1), + "Ei": big.NewRat(1<<60, 1), +} + +// parseQuantity evaluates a Kubernetes quantity string to its exact value in +// base units (cores for cpu, bytes for memory). +func parseQuantity(s string) (*big.Rat, error) { + m := quantityRegexp.FindStringSubmatch(s) + if m == nil { + return nil, errors.New("must be a number with an optional suffix such as 500m, 2, 4Gi or 1e3") + } + value, ok := new(big.Rat).SetString(m[1]) + if !ok { + return nil, errors.New("not a number") + } + switch suffix := m[2]; { + case suffix == "": + case suffix[0] == 'e' || suffix[0] == 'E': + exp, err := strconv.Atoi(suffix[1:]) + if err != nil || exp > 64 || exp < -64 { + return nil, errors.New("exponent out of range") + } + pow := new(big.Rat).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(abs(exp))), nil)) + if exp < 0 { + pow.Inv(pow) + } + value.Mul(value, pow) + default: + value.Mul(value, quantitySuffixes[suffix]) + } + return value, nil +} + +func abs(n int) int { + if n < 0 { + return -n + } + return n +} diff --git a/pkg/apis/v1alpha1/resources_test.go b/pkg/apis/v1alpha1/resources_test.go new file mode 100644 index 00000000..e396ac7d --- /dev/null +++ b/pkg/apis/v1alpha1/resources_test.go @@ -0,0 +1,95 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package v1alpha1_test + +import ( + "strings" + "testing" + + "github.com/google/ax/pkg/apis/v1alpha1" +) + +func TestValidateResources(t *testing.T) { + limits := func(cpu, memory string) *v1alpha1.ResourceReqs { + return &v1alpha1.ResourceReqs{Limits: &v1alpha1.ResourceList{Cpu: cpu, Memory: memory}} + } + tests := []struct { + name string + reqs *v1alpha1.ResourceReqs + wantErr string + }{ + {name: "nil"}, + {name: "empty", reqs: &v1alpha1.ResourceReqs{}}, + {name: "empty limits", reqs: limits("", "")}, + {name: "docs example", reqs: limits("2", "4Gi")}, + {name: "millicores", reqs: limits("500m", "")}, + {name: "fractional cores", reqs: limits("1.5", "")}, + {name: "leading dot", reqs: limits(".5", "")}, + {name: "exponent", reqs: limits("2e0", "1e9")}, + {name: "largest cpu", reqs: limits("999999m", "")}, + {name: "decimal SI memory", reqs: limits("", "512M")}, + {name: "binary SI memory", reqs: limits("", "1Ti")}, + {name: "small units", reqs: limits("1n", "1u")}, + { + name: "requests are not supported", + reqs: &v1alpha1.ResourceReqs{Requests: &v1alpha1.ResourceList{Cpu: "500m"}, Limits: &v1alpha1.ResourceList{Cpu: "2"}}, + wantErr: "spec.resources.requests: not supported", + }, + { + name: "memory request alone", + reqs: &v1alpha1.ResourceReqs{Requests: &v1alpha1.ResourceList{Memory: "1Gi"}}, + wantErr: "spec.resources.requests: not supported", + }, + {name: "cpu zero", reqs: limits("0", ""), wantErr: `spec.resources.limits.cpu: "0" must be greater than zero`}, + {name: "cpu zero millicores", reqs: limits("0m", ""), wantErr: "must be greater than zero"}, + {name: "cpu negative", reqs: limits("-1", ""), wantErr: "must be greater than zero"}, + {name: "cpu at bound", reqs: limits("1000", ""), wantErr: `spec.resources.limits.cpu: "1000" must be less than 1000 cores`}, + {name: "cpu over bound via suffix", reqs: limits("1k", ""), wantErr: "must be less than 1000 cores"}, + {name: "cpu over bound via exponent", reqs: limits("1e3", ""), wantErr: "must be less than 1000 cores"}, + {name: "cpu over bound via binary suffix", reqs: limits("1Ki", ""), wantErr: "must be less than 1000 cores"}, + {name: "cpu not a quantity", reqs: limits("two", ""), wantErr: `spec.resources.limits.cpu: invalid quantity "two"`}, + {name: "cpu bad suffix", reqs: limits("2cores", ""), wantErr: "invalid quantity"}, + {name: "cpu double dot", reqs: limits("1.2.3", ""), wantErr: "invalid quantity"}, + {name: "cpu whitespace", reqs: limits(" 2", ""), wantErr: "invalid quantity"}, + {name: "memory zero", reqs: limits("", "0Gi"), wantErr: `spec.resources.limits.memory: "0Gi" must be greater than zero`}, + {name: "memory negative", reqs: limits("", "-4Gi"), wantErr: "must be greater than zero"}, + {name: "memory lowercase suffix", reqs: limits("", "4gi"), wantErr: `spec.resources.limits.memory: invalid quantity "4gi"`}, + {name: "memory bytes suffix", reqs: limits("", "4GiB"), wantErr: "invalid quantity"}, + {name: "memory exponent out of range", reqs: limits("", "1e100"), wantErr: "exponent out of range"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := v1alpha1.ValidateResources(tt.reqs) + if tt.wantErr == "" { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %v, want containing %q", err, tt.wantErr) + } + }) + } + + // ValidateTask runs the same check, so ax apply rejects bad limits up front. + err := v1alpha1.ValidateTask(&v1alpha1.Task{Spec: &v1alpha1.TaskSpec{Resources: limits("abc", "")}}) + if err == nil || !strings.Contains(err.Error(), "spec.resources.limits.cpu") { + t.Fatalf("ValidateTask error = %v, want cpu limit error", err) + } + if err := v1alpha1.ValidateTask(&v1alpha1.Task{Spec: &v1alpha1.TaskSpec{Resources: limits("2", "4Gi")}}); err != nil { + t.Fatalf("ValidateTask rejected valid limits: %v", err) + } +} diff --git a/pkg/apis/v1alpha1/types.go b/pkg/apis/v1alpha1/types.go index bb1b02c7..63ada834 100644 --- a/pkg/apis/v1alpha1/types.go +++ b/pkg/apis/v1alpha1/types.go @@ -273,5 +273,5 @@ func ValidateTask(t *Task) error { } seen[p] = r.GetName() } - return nil + return ValidateResources(spec.GetResources()) } From c1194dfe2856b1f4b0e81bf09bf64665fc433785 Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Fri, 25 Sep 2026 06:46:12 +0000 Subject: [PATCH 3/3] test: name the ValidateTask fixtures in TestValidateResources Keeps the test valid alongside metadata.name validation (#406). Co-authored-by: Matt Van Horn --- pkg/apis/v1alpha1/resources_test.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pkg/apis/v1alpha1/resources_test.go b/pkg/apis/v1alpha1/resources_test.go index e396ac7d..60cee744 100644 --- a/pkg/apis/v1alpha1/resources_test.go +++ b/pkg/apis/v1alpha1/resources_test.go @@ -85,11 +85,12 @@ func TestValidateResources(t *testing.T) { } // ValidateTask runs the same check, so ax apply rejects bad limits up front. - err := v1alpha1.ValidateTask(&v1alpha1.Task{Spec: &v1alpha1.TaskSpec{Resources: limits("abc", "")}}) + meta := &v1alpha1.ObjectMeta{Name: "task", Atespace: "default"} + err := v1alpha1.ValidateTask(&v1alpha1.Task{Metadata: meta, Spec: &v1alpha1.TaskSpec{Resources: limits("abc", "")}}) if err == nil || !strings.Contains(err.Error(), "spec.resources.limits.cpu") { t.Fatalf("ValidateTask error = %v, want cpu limit error", err) } - if err := v1alpha1.ValidateTask(&v1alpha1.Task{Spec: &v1alpha1.TaskSpec{Resources: limits("2", "4Gi")}}); err != nil { + if err := v1alpha1.ValidateTask(&v1alpha1.Task{Metadata: meta, Spec: &v1alpha1.TaskSpec{Resources: limits("2", "4Gi")}}); err != nil { t.Fatalf("ValidateTask rejected valid limits: %v", err) } }