From 3ba0238272a89f4b928f8629bdc232b1de0810f6 Mon Sep 17 00:00:00 2001
From: Brandon Duffany <brandon@buildbuddy.io>
Date: Wed, 18 Sep 2024 12:15:53 -0700
Subject: [PATCH] Enable task sizer for Firecracker CPU (#7487)

---
 enterprise/server/tasksize/tasksize.go | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/enterprise/server/tasksize/tasksize.go b/enterprise/server/tasksize/tasksize.go
index d93e7847d0d..81a73a3594b 100644
--- a/enterprise/server/tasksize/tasksize.go
+++ b/enterprise/server/tasksize/tasksize.go
@@ -139,11 +139,6 @@ func (s *taskSizer) Get(ctx context.Context, task *repb.ExecutionTask) *scpb.Tas
 	if props.DisableMeasuredTaskSize {
 		return nil
 	}
-	// Don't use measured task sizes for Firecracker tasks for now, since task
-	// sizes are used as hard limits on allowed resources.
-	if props.WorkloadIsolationType == string(platform.FirecrackerContainerType) {
-		return nil
-	}
 	statusLabel := "hit"
 	defer func() {
 		groupID, _ := s.groupKey(ctx)
@@ -167,10 +162,19 @@ func (s *taskSizer) Get(ctx context.Context, task *repb.ExecutionTask) *scpb.Tas
 		// executor run this task once to estimate the size.
 		return nil
 	}
-	return applyMinimums(task, &scpb.TaskSize{
+	size := applyMinimums(task, &scpb.TaskSize{
 		EstimatedMemoryBytes: recordedSize.EstimatedMemoryBytes,
 		EstimatedMilliCpu:    recordedSize.EstimatedMilliCpu,
 	})
+	// Don't size memory automatically for Firecracker VMs, since if we size it
+	// incorrectly, it can cause OOM errors.
+	// TODO(https://github.com/buildbuddy-io/buildbuddy-internal/issues/3824):
+	// use these sizes for determining how many VMs to schedule concurrently,
+	// but not for the upper limit on VM size.
+	if props.WorkloadIsolationType == string(platform.FirecrackerContainerType) {
+		size.EstimatedMemoryBytes = 0
+	}
+	return size
 }
 
 func (s *taskSizer) Predict(ctx context.Context, task *repb.ExecutionTask) *scpb.TaskSize {