fix(placement,fc): memory-aware sandbox placement and DeflateOnOom

packages/api/internal/orchestrator/nodemanager/mock.go

@@ -134,6 +134,13 @@ func WithAllocatedMemoryBytes(bytes uint64) TestOptions {
 	}
 }
 
+// WithTotalMemoryBytes sets the total physical memory reported by the test node
+func WithTotalMemoryBytes(bytes uint64) TestOptions {
+	return func(node *TestNode) {
+		node.metrics.MemoryTotalBytes = bytes
+	}
+}
+
 // MockSandboxClientCustom allows custom error logic per call
 type MockSandboxClientCustom struct {
 	orchestrator.SandboxServiceClient

packages/api/internal/orchestrator/orchestrator.go

@@ -315,13 +315,17 @@ func getBestOfKConfig(ctx context.Context, featureFlagsClient *featureflags.Clie
 
 	alphaPercent := featureFlagsClient.IntFlag(ctx, featureflags.BestOfKAlpha)
 
+	memOvercommitPercent := featureFlagsClient.IntFlag(ctx, featureflags.BestOfKMaxMemoryOvercommit)
+
 	// Convert percentage to decimal
 	alpha := float64(alphaPercent) / 100.0
 	maxOvercommit := float64(maxOvercommitPercent) / 100.0
+	memOvercommit := float64(memOvercommitPercent) / 100.0
 
 	return placement.BestOfKConfig{
 		R:     maxOvercommit,
 		K:     k,
 		Alpha: alpha,
+		M:     memOvercommit,
 	}
 }

packages/api/internal/orchestrator/placement/placement_best_of_K.go

@@ -20,6 +20,9 @@ type BestOfKConfig struct {
 	Alpha float64
 	// K is the number of candidate nodes sampled per placement ("power of K choices")
 	K int
+	// M is the memory overcommit ratio for the hard placement filter.
+	// 0 = disabled (current behaviour); 1.0 = no overcommit; 1.5 = 50% overcommit.
+	M float64
 }
 
 // DefaultBestOfKConfig returns the default placement configuration
@@ -95,7 +98,7 @@ func (b *BestOfK) chooseNode(_ context.Context, nodes []*nodemanager.Node, exclu
 	config := b.getConfig()
 
 	// Filter eligible nodes
-	candidates := b.sample(nodes, config, excludedNodes, buildMachineInfo, filterByLabels, requiredLabels)
+	candidates := b.sample(nodes, config, resources, excludedNodes, buildMachineInfo, filterByLabels, requiredLabels)
 
 	// Find the best node among candidates
 	bestScore := math.MaxFloat64
@@ -140,7 +143,7 @@ func (e FailedToPlaceSandboxError) Error() string {
 }
 
 // sample returns up to k items chosen uniformly from those passing ok.
-func (b *BestOfK) sample(items []*nodemanager.Node, config BestOfKConfig, excludedNodes map[string]struct{}, buildMachineInfo machineinfo.MachineInfo, filterByLabels bool, requiredLabels []string) []*nodemanager.Node {
+func (b *BestOfK) sample(items []*nodemanager.Node, config BestOfKConfig, resources nodemanager.SandboxResources, excludedNodes map[string]struct{}, buildMachineInfo machineinfo.MachineInfo, filterByLabels bool, requiredLabels []string) []*nodemanager.Node {
 	if config.K <= 0 || len(items) == 0 {
 		return nil
 	}
@@ -184,6 +187,19 @@ func (b *BestOfK) sample(items []*nodemanager.Node, config BestOfKConfig, exclud
 			continue
 		}
 
+		// Hard memory filter: skip when the node would exceed the configured
+		// overcommit ratio. Only active when M > 0 and the node has reported
+		// its total memory (MemoryTotalBytes > 0).
+		if config.M > 0 {
+			metrics := n.Metrics()
+			if metrics.MemoryTotalBytes > 0 {
+				requested := uint64(resources.MiBMemory) * 1024 * 1024
+				if float64(metrics.MemoryAllocatedBytes+requested) > config.M*float64(metrics.MemoryTotalBytes) {
+					continue
+				}
+			}
+		}
+
 		candidates = append(candidates, n)
 	}
 

packages/api/internal/orchestrator/placement/placement_best_of_K_test.go

@@ -321,7 +321,7 @@ func TestBestOfK_Sample(t *testing.T) {
 	excludedNodes := make(map[string]struct{})
 
 	// Test sampling fewer nodes than available
-	sampled := algo.sample(nodes, config, excludedNodes, machineinfo.MachineInfo{}, false, nil)
+	sampled := algo.sample(nodes, config, nodemanager.SandboxResources{}, excludedNodes, machineinfo.MachineInfo{}, false, nil)
 	assert.LessOrEqual(t, len(sampled), 3)
 
 	// Check all sampled nodes are unique
@@ -334,7 +334,7 @@ func TestBestOfK_Sample(t *testing.T) {
 	// Test sampling with exclusions
 	excludedNodes["a"] = struct{}{}
 	excludedNodes["b"] = struct{}{}
-	sampled = algo.sample(nodes, config, excludedNodes, machineinfo.MachineInfo{}, false, nil)
+	sampled = algo.sample(nodes, config, nodemanager.SandboxResources{}, excludedNodes, machineinfo.MachineInfo{}, false, nil)
 
 	for _, n := range sampled {
 		assert.NotEqual(t, "a", n.ID)

packages/api/internal/orchestrator/placement/placement_memory_filter_test.go

@@ -0,0 +1,126 @@
+package placement
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
+	"github.com/e2b-dev/infra/packages/shared/pkg/machineinfo"
+)
+
+const gib = uint64(1024 * 1024 * 1024)
+
+func TestBestOfK_MemoryFilter_Disabled_WhenMIsZero(t *testing.T) {
+	t.Parallel()
+
+	// Node has only 1 GiB total but is allocated 900 MiB; sandbox requests 512 MiB.
+	// Without a memory filter (M=0) this node should still be eligible.
+	node := nodemanager.NewTestNode("node-a", api.NodeStatusReady, 0, 8,
+		nodemanager.WithTotalMemoryBytes(gib),
+		nodemanager.WithAllocatedMemoryBytes(900*1024*1024),
+	)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 0}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	candidates := algo.sample([]*nodemanager.Node{node}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	require.Len(t, candidates, 1, "memory filter must be inactive when M=0")
+}
+
+func TestBestOfK_MemoryFilter_ExcludesFullNode(t *testing.T) {
+	t.Parallel()
+
+	// Node has 2 GiB total, 1.9 GiB allocated; sandbox requests 512 MiB.
+	// With M=1.0 (no overcommit), 1.9 GiB + 0.5 GiB = 2.4 GiB > 2 GiB → excluded.
+	node := nodemanager.NewTestNode("node-full", api.NodeStatusReady, 0, 8,
+		nodemanager.WithTotalMemoryBytes(2*gib),
+		nodemanager.WithAllocatedMemoryBytes(1900 * 1024 * 1024),
+	)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 1.0}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	candidates := algo.sample([]*nodemanager.Node{node}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	assert.Empty(t, candidates, "full node must be excluded when M=1.0")
+}
+
+func TestBestOfK_MemoryFilter_AllowsNodeWithHeadroom(t *testing.T) {
+	t.Parallel()
+
+	// Node has 4 GiB total, 1 GiB allocated; sandbox requests 512 MiB.
+	// 1 GiB + 0.5 GiB = 1.5 GiB < 4 GiB → eligible.
+	node := nodemanager.NewTestNode("node-ok", api.NodeStatusReady, 0, 8,
+		nodemanager.WithTotalMemoryBytes(4*gib),
+		nodemanager.WithAllocatedMemoryBytes(gib),
+	)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 1.0}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	candidates := algo.sample([]*nodemanager.Node{node}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	require.Len(t, candidates, 1, "node with headroom must be included")
+}
+
+func TestBestOfK_MemoryFilter_OvercommitAllowsExcess(t *testing.T) {
+	t.Parallel()
+
+	// Node has 2 GiB total, 1.9 GiB allocated; sandbox requests 512 MiB.
+	// With M=1.5 (50% overcommit), limit = 3 GiB; 1.9 + 0.5 = 2.4 GiB < 3 GiB → eligible.
+	node := nodemanager.NewTestNode("node-overcommit", api.NodeStatusReady, 0, 8,
+		nodemanager.WithTotalMemoryBytes(2*gib),
+		nodemanager.WithAllocatedMemoryBytes(1900 * 1024 * 1024),
+	)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 1.5}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	candidates := algo.sample([]*nodemanager.Node{node}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	require.Len(t, candidates, 1, "overcommit ratio M=1.5 must allow this allocation")
+}
+
+func TestBestOfK_MemoryFilter_SkipsWhenTotalBytesZero(t *testing.T) {
+	t.Parallel()
+
+	// Node has not yet reported MemoryTotalBytes (zero); filter should be skipped.
+	node := nodemanager.NewTestNode("node-unreported", api.NodeStatusReady, 0, 8)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 1.0}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	candidates := algo.sample([]*nodemanager.Node{node}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	require.Len(t, candidates, 1, "filter must be skipped when MemoryTotalBytes=0")
+}
+
+func TestBestOfK_MemoryFilter_SelectsNodeWithMoreHeadroom(t *testing.T) {
+	t.Parallel()
+
+	// Two nodes: tight has 200 MiB free, spacious has 2 GiB free.
+	// K=10 samples all nodes; Score() is CPU-only, so both are eligible.
+	// The memory filter only hard-excludes; this test verifies neither is excluded.
+	tight := nodemanager.NewTestNode("tight", api.NodeStatusReady, 4, 8,
+		nodemanager.WithTotalMemoryBytes(4*gib),
+		nodemanager.WithAllocatedMemoryBytes(4*gib-200*1024*1024),
+	)
+	spacious := nodemanager.NewTestNode("spacious", api.NodeStatusReady, 0, 8,
+		nodemanager.WithTotalMemoryBytes(4*gib),
+		nodemanager.WithAllocatedMemoryBytes(2*gib),
+	)
+
+	config := BestOfKConfig{R: 4, K: 10, Alpha: 0.5, M: 1.0}
+	algo := NewBestOfK(config).(*BestOfK)
+
+	resources := nodemanager.SandboxResources{CPUs: 1, MiBMemory: 512}
+	// tight: 3.8 GiB + 0.5 GiB = 4.3 GiB > 4 GiB → excluded
+	// spacious: 2 GiB + 0.5 GiB = 2.5 GiB < 4 GiB → eligible
+	candidates := algo.sample([]*nodemanager.Node{tight, spacious}, config, resources, map[string]struct{}{}, machineinfo.MachineInfo{}, false, nil)
+	require.Len(t, candidates, 1)
+	assert.Equal(t, "spacious", candidates[0].ID)
+}

packages/orchestrator/pkg/sandbox/fc/client.go

@@ -442,7 +442,7 @@ func (c *apiClient) installBalloon(ctx context.Context, freePageReporting, freeP
 	defer span.End()
 
 	amountMib := int64(0)
-	deflateOnOom := false
+	deflateOnOom := true
 
 	balloonConfig := operations.PutBalloonParams{
 		Context: ctx,

packages/shared/pkg/featureflags/flags.go

@@ -306,7 +306,8 @@ var (
 	ClickhouseBatcherQueueSize    = NewIntFlag("clickhouse-batcher-queue-size", 1000)
 	BestOfKSampleSize             = NewIntFlag("best-of-k-sample-size", 3)                           // Default K=3
 	BestOfKMaxOvercommit          = NewIntFlag("best-of-k-max-overcommit", 400)                      // Default R=4 (stored as percentage, max over-commit ratio)
-	BestOfKAlpha                  = NewIntFlag("best-of-k-alpha", 50)                                // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)
+	BestOfKAlpha                     = NewIntFlag("best-of-k-alpha", 50)                                    // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)
+	BestOfKMaxMemoryOvercommit       = NewIntFlag("best-of-k-max-memory-overcommit", 0)                      // Default 0 = disabled; set e.g. 100 for 1.0x (no overcommit) or 150 for 1.5x
 	EnvdInitTimeoutMilliseconds   = NewIntFlag("envd-init-request-timeout-milliseconds", 50)         // Timeout for envd init request in milliseconds
 	EnvdTimeoutMilliseconds       = NewIntFlag("envd-timeout-milliseconds", envdTimeoutFallbackMs()) // Timeout for waiting for envd on resume; falls back to ENVD_TIMEOUT env var (default 10s)
 	// GuestSyncTimeoutMs overrides the mandatory pre-pause guest-sync deadline