Right-size actor sandboxes to declared ActorTemplate resource limits

cmd/ateapi/internal/controlapi/actor_sizing_test.go

@@ -0,0 +1,139 @@
+// 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 controlapi
+
+import (
+	"testing"
+
+	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+// TestActorResourceLimits covers the actor-side extraction: the CPU/memory limits
+// an ActorTemplate declares become the sandbox size and the scheduling floor.
+func TestActorResourceLimits(t *testing.T) {
+	tests := []struct {
+		name       string
+		res        *corev1.ResourceRequirements
+		wantCPU    int64
+		wantMemory int64
+	}{
+		{
+			name:       "nil resources yields zero",
+			res:        nil,
+			wantCPU:    0,
+			wantMemory: 0,
+		},
+		{
+			name: "cpu and memory limits are read",
+			res: &corev1.ResourceRequirements{
+				Limits: corev1.ResourceList{
+					corev1.ResourceCPU:    resource.MustParse("2"),
+					corev1.ResourceMemory: resource.MustParse("4Gi"),
+				},
+			},
+			wantCPU:    2000,
+			wantMemory: 4 << 30,
+		},
+		{
+			name: "millicpu is preserved",
+			res: &corev1.ResourceRequirements{
+				Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1500m")},
+			},
+			wantCPU:    1500,
+			wantMemory: 0,
+		},
+		{
+			name: "requests are ignored; only limits size the actor",
+			res: &corev1.ResourceRequirements{
+				Requests: corev1.ResourceList{
+					corev1.ResourceCPU:    resource.MustParse("1"),
+					corev1.ResourceMemory: resource.MustParse("1Gi"),
+				},
+			},
+			wantCPU:    0,
+			wantMemory: 0,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tmpl := &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{Resources: tc.res}}
+			cpu, mem := actorResourceLimits(tmpl)
+			if cpu != tc.wantCPU || mem != tc.wantMemory {
+				t.Fatalf("actorResourceLimits() = (%d, %d), want (%d, %d)", cpu, mem, tc.wantCPU, tc.wantMemory)
+			}
+		})
+	}
+}
+
+// TestWorkerCapacity covers the worker-side extraction: capacity comes from the
+// ateom container's limits, not the pod total, and other containers are ignored.
+func TestWorkerCapacity(t *testing.T) {
+	pod := func(ctrs ...corev1.Container) *corev1.Pod {
+		return &corev1.Pod{Spec: corev1.PodSpec{Containers: ctrs}}
+	}
+	limited := func(name, cpu, mem string) corev1.Container {
+		lim := corev1.ResourceList{}
+		if cpu != "" {
+			lim[corev1.ResourceCPU] = resource.MustParse(cpu)
+		}
+		if mem != "" {
+			lim[corev1.ResourceMemory] = resource.MustParse(mem)
+		}
+		return corev1.Container{Name: name, Resources: corev1.ResourceRequirements{Limits: lim}}
+	}
+
+	tests := []struct {
+		name       string
+		pod        *corev1.Pod
+		wantCPU    int64
+		wantMemory int64
+	}{
+		{
+			name:       "no ateom container yields zero",
+			pod:        pod(limited("sidecar", "1", "1Gi")),
+			wantCPU:    0,
+			wantMemory: 0,
+		},
+		{
+			name:       "ateom container limits become capacity",
+			pod:        pod(limited(ateomContainerName, "4", "8Gi")),
+			wantCPU:    4000,
+			wantMemory: 8 << 30,
+		},
+		{
+			name:       "only the ateom container counts, not the pod total",
+			pod:        pod(limited("sidecar", "16", "64Gi"), limited(ateomContainerName, "2", "2Gi")),
+			wantCPU:    2000,
+			wantMemory: 2 << 30,
+		},
+		{
+			name:       "unset dimension reports zero",
+			pod:        pod(limited(ateomContainerName, "2", "")),
+			wantCPU:    2000,
+			wantMemory: 0,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := workerCapacity(tc.pod)
+			if got.GetCpuMilli() != tc.wantCPU || got.GetMemoryBytes() != tc.wantMemory {
+				t.Fatalf("workerCapacity() = (%d, %d), want (%d, %d)",
+					got.GetCpuMilli(), got.GetMemoryBytes(), tc.wantCPU, tc.wantMemory)
+			}
+		})
+	}
+}

cmd/ateapi/internal/controlapi/syncer.go

@@ -216,6 +216,7 @@ func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerK
 			SandboxClass:    string(pool.Spec.SandboxClass),
 			Labels:          pool.GetLabels(),
 			State:           ateapipb.Worker_STATE_ACTIVE,
+			Capacity:        workerCapacity(pod),
 		}
 		// TODO(thockin): for now this is the only place Workers are
 		// created.  If/when this becomes a regular API, validation should
@@ -273,6 +274,38 @@ func isWorkerEligible(pod *corev1.Pod) bool {
 	return pod.Status.PodIP != ""
 }
 
+// ateomContainerName is the name of the container in a worker pod that hosts the
+// actor's sandbox; its resource limits bound what an actor placed here can use.
+const ateomContainerName = "ateom"
+
+// workerCapacity returns the worker pod's capacity for hosting an actor — CPU
+// in millicores and memory in bytes — taken from the ateom container's resource
+// limits. A dimension the pod does not limit reports 0, which the scheduler
+// treats as "unknown" (unconstrained); a pod that limits neither reports nil
+// rather than an all-zero message that says the same thing. The actor sandbox
+// runs nested in the ateom container's cgroup, so that container's limits — not
+// the pod total — are the relevant envelope.
+func workerCapacity(pod *corev1.Pod) *ateapipb.WorkerCapacity {
+	var capacity ateapipb.WorkerCapacity
+	for i := range pod.Spec.Containers {
+		c := &pod.Spec.Containers[i]
+		if c.Name != ateomContainerName {
+			continue
+		}
+		if v := c.Resources.Limits.Cpu(); v != nil {
+			capacity.CpuMilli = v.MilliValue()
+		}
+		if v := c.Resources.Limits.Memory(); v != nil {
+			capacity.MemoryBytes = v.Value()
+		}
+		break
+	}
+	if capacity.CpuMilli == 0 && capacity.MemoryBytes == 0 {
+		return nil
+	}
+	return &capacity
+}
+
 // markWorkerDraining transitions a worker to STATE_DRAINING so the scheduler
 // stops routing new actors to it while its pod is Terminating. If the worker is
 // already gone or already draining there is nothing more to do — the Pod

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -546,11 +546,32 @@ func workerAssignmentFrom(w *ateapipb.Worker) *ateapipb.WorkerAssignment {
 	}
 }
 
+// actorResourceLimits returns the actor's declared CPU (millicores) and memory
+// (bytes) limits from its ActorTemplate, or 0 for a dimension the template did
+// not set. These size the sandbox (supplied over the actor RPCs) and gate
+// scheduling (a worker must have >= capacity).
+func actorResourceLimits(tmpl *atev1alpha1.ActorTemplate) (cpuMilli, memBytes int64) {
+	res := tmpl.Spec.Resources
+	if res == nil {
+		return 0, 0
+	}
+	if c := res.Limits.Cpu(); c != nil {
+		cpuMilli = c.MilliValue()
+	}
+	if m := res.Limits.Memory(); m != nil {
+		memBytes = m.Value()
+	}
+	return cpuMilli, memBytes
+}
+
 func schedulingConstraints(actor *ateapipb.Actor, tmpl *atev1alpha1.ActorTemplate) (scheduling.Constraints, error) {
+	cpuMilli, memBytes := actorResourceLimits(tmpl)
 	c := scheduling.Constraints{
 		SandboxClass:  string(tmpl.Spec.SandboxClass),
 		ActorSelector: labels.SelectorFromSet(labels.Set(actor.GetWorkerSelector().GetMatchLabels())),
 		RequiredNodes: actor.GetLocalSnapshotInfo().GetNodeVmsWithLocalSnapshots(),
+		CPUMilli:      cpuMilli,
+		MemoryBytes:   memBytes,
 	}
 	if tmpl.Spec.WorkerSelector != nil {
 		sel, err := metav1.LabelSelectorAsSelector(tmpl.Spec.WorkerSelector)
@@ -613,6 +634,10 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 	}
 	egressGateway := w.egressGateway()
 
+	// The actor's declared limits ride the RPC down to the sandbox so it is sized
+	// to the actor (replacing the worker-pod downward-API approach).
+	cpuMilli, memBytes := actorResourceLimits(actorTemplate)
+
 	if local := actor.GetLocalSnapshotInfo(); local != nil {
 		slog.InfoContext(ctx, "Actor has snapshot; Restoring from snapshot")
 		tele.SnapshotKind = ateattr.SnapshotKindLocal
@@ -626,6 +651,8 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 			Spec:                   workloadSpec,
 			ActorUid:               actor.GetMetadata().Uid,
 			EgressGateway:          egressGateway,
+			CpuMilli:               cpuMilli,
+			MemoryBytes:            memBytes,
 		}
 		req.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
 		req.Config = &ateletpb.RestoreRequest_LocalConfig{
@@ -681,6 +708,8 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 			GoldenSnapshotUri: src.GoldenSnapshotURI.String(),
 			ActorUid:          actor.GetMetadata().Uid,
 			EgressGateway:     egressGateway,
+			CpuMilli:          cpuMilli,
+			MemoryBytes:       memBytes,
 		}
 		_, err = client.Restore(ctx, req)
 		return tele, maybeCrashActor(ctx, w.store, actorRef, err, "while restoring durable snapshot", ateattr.OperationResume)
@@ -706,6 +735,8 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 			Spec:                   workloadSpec,
 			ActorUid:               actor.GetMetadata().Uid,
 			EgressGateway:          egressGateway,
+			CpuMilli:               cpuMilli,
+			MemoryBytes:            memBytes,
 		}
 		_, err = client.Run(ctx, req)
 		return tele, maybeCrashActor(ctx, w.store, actorRef, err, "while creating workload from spec", ateattr.OperationResume)

cmd/ateapi/internal/scheduling/scheduling.go

@@ -41,6 +41,15 @@ type Constraints struct {
 	// on one of these nodes. Used when the actor's latest snapshot is local
 	// to specific node VMs.
 	RequiredNodes []string
+
+	// CPUMilli and MemoryBytes are the actor's declared resource limits, from
+	// the ActorTemplate. A worker is eligible only if its reported capacity is
+	// >= these. Zero means "unconstrained" for that dimension (the actor did not
+	// declare a limit), and a worker that reports zero capacity for a dimension
+	// is treated as unconstrained too, so placement is never blocked by missing
+	// data (matching the pre-capacity behavior).
+	CPUMilli    int64
+	MemoryBytes int64
 }
 
 // ErrNoCapacity is returned by Schedule when no free worker satisfies the
@@ -136,5 +145,17 @@ func (s *scheduler) Applies(worker *ateapipb.Worker, constraints Constraints) bo
 		return false
 	}
 
+	// The worker must be able to contain the actor's declared limits. A zero
+	// constraint (actor declared no limit) or zero worker capacity (capacity
+	// unknown) is treated as unconstrained, so placement is never blocked by
+	// missing data.
+	capacity := worker.GetCapacity()
+	if constraints.CPUMilli > 0 && capacity.GetCpuMilli() > 0 && capacity.GetCpuMilli() < constraints.CPUMilli {
+		return false
+	}
+	if constraints.MemoryBytes > 0 && capacity.GetMemoryBytes() > 0 && capacity.GetMemoryBytes() < constraints.MemoryBytes {
+		return false
+	}
+
 	return len(constraints.RequiredNodes) == 0 || slices.Contains(constraints.RequiredNodes, worker.GetNodeName())
 }

cmd/ateapi/internal/scheduling/scheduling_test.go

@@ -120,6 +120,47 @@ func TestSchedule(t *testing.T) {
 			constraints: Constraints{SandboxClass: "gvisor"},
 			wantPod:     "w-active",
 		},
+		{
+			name: "worker with too little cpu capacity is skipped",
+			fleet: fleet{
+				worker("w-small", "gvisor", "node-a", tierTwo, withCapacity(1000, 8<<30)),
+				worker("w-big", "gvisor", "node-a", tierTwo, withCapacity(4000, 8<<30)),
+			},
+			constraints: Constraints{SandboxClass: "gvisor", CPUMilli: 2000},
+			wantPod:     "w-big",
+		},
+		{
+			name: "worker with too little memory capacity is skipped",
+			fleet: fleet{
+				worker("w-small", "gvisor", "node-a", tierTwo, withCapacity(4000, 1<<30)),
+				worker("w-big", "gvisor", "node-a", tierTwo, withCapacity(4000, 4<<30)),
+			},
+			constraints: Constraints{SandboxClass: "gvisor", MemoryBytes: 2 << 30},
+			wantPod:     "w-big",
+		},
+		{
+			name: "no worker with enough capacity yields ErrNoCapacity",
+			fleet: fleet{
+				worker("w-small", "gvisor", "node-a", tierTwo, withCapacity(1000, 1<<30)),
+			},
+			constraints: Constraints{SandboxClass: "gvisor", CPUMilli: 2000},
+		},
+		{
+			name: "zero worker capacity is treated as unconstrained",
+			fleet: fleet{
+				worker("w-unknown", "gvisor", "node-a", tierTwo),
+			},
+			constraints: Constraints{SandboxClass: "gvisor", CPUMilli: 2000, MemoryBytes: 2 << 30},
+			wantPod:     "w-unknown",
+		},
+		{
+			name: "zero constraint ignores worker capacity",
+			fleet: fleet{
+				worker("w-tiny", "gvisor", "node-a", tierTwo, withCapacity(100, 1<<20)),
+			},
+			constraints: Constraints{SandboxClass: "gvisor"},
+			wantPod:     "w-tiny",
+		},
 		{
 			name:        "empty fleet",
 			fleet:       fleet{},
@@ -269,6 +310,12 @@ func assigned(atespace, name string) func(*ateapipb.Worker) {
 	}
 }
 
+func withCapacity(cpuMilli, memBytes int64) func(*ateapipb.Worker) {
+	return func(w *ateapipb.Worker) {
+		w.Capacity = &ateapipb.WorkerCapacity{CpuMilli: cpuMilli, MemoryBytes: memBytes}
+	}
+}
+
 // firstIntn always picks the first candidate, making Schedule deterministic.
 func firstIntn(int) int { return 0 }
 

cmd/atecontroller/internal/controllers/workerpool_apply_test.go

@@ -753,11 +753,13 @@ func expectedDeploymentApplyConfig(mutatePodSpec func(*corev1ac.PodSpecApplyConf
 					WithAdd(ateomGvisorCapabilities...)).
 				WithAppArmorProfile(corev1ac.AppArmorProfile().
 					WithType(corev1.AppArmorProfileTypeUnconfined))).
-			WithEnv(corev1ac.EnvVar().
-				WithName("POD_UID").
-				WithValueFrom(corev1ac.EnvVarSource().
-					WithFieldRef(corev1ac.ObjectFieldSelector().
-						WithFieldPath("metadata.uid")))).
+			WithEnv(
+				corev1ac.EnvVar().
+					WithName("POD_UID").
+					WithValueFrom(corev1ac.EnvVarSource().
+						WithFieldRef(corev1ac.ObjectFieldSelector().
+							WithFieldPath("metadata.uid"))),
+			).
 			WithVolumeMounts(
 				corev1ac.VolumeMount().
 					WithName("run-ateom").

cmd/atelet/main.go

@@ -454,6 +454,8 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (resp *
 		Spec:                   buildAteomWorkloadSpec(req.GetSpec()),
 		ActorUid:               actorUID,
 		EgressGateway:          toAteomEgressGateway(req.GetEgressGateway()),
+		CpuMilli:               req.GetCpuMilli(),
+		MemoryBytes:            req.GetMemoryBytes(),
 	}); err != nil {
 		return nil, fmt.Errorf("while calling ateom.RunWorkload: %w", err)
 	}
@@ -1091,6 +1093,8 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 		Scope:                  toAteomSnapshotScope(req.GetScope()),
 		ActorUid:               req.GetActorUid(),
 		EgressGateway:          toAteomEgressGateway(req.GetEgressGateway()),
+		CpuMilli:               req.GetCpuMilli(),
+		MemoryBytes:            req.GetMemoryBytes(),
 		// Informational: for DATA_ON_GOLDEN the golden snapshot's files are
 		// already staged into the restore dir by the combined download above;
 		// ateom restores from the shared dir and never fetches this URI.