Merge pull request #364 from diggerhq/fix/upsize-default-builder

make multistage built, add adjustable mb

Author: breardon2011

docs/reference/python-sdk/image.mdx

@@ -103,6 +103,18 @@ image = (
 
 **Returns:** `Image`
 
+### `image.builder_memory(mb)`
+
+Sets the RAM (in MB) used during the build phase. Raise this when a build OOMs
+(heavy `apt`/`pip`/`npm`). Defaults to **4096**. Does **not** affect the resulting
+sandbox's memory — size that at create time via `memory_mb`.
+
+<ParamField body="mb" type="int" required>
+  Build-phase memory in MB
+</ParamField>
+
+**Returns:** `Image`
+
 ---
 
 ## Utility Methods

docs/reference/typescript-sdk/image.mdx

@@ -115,6 +115,18 @@ const image = Image.base()
 
 **Returns:** `Image`
 
+### `image.builderMemory(mb)`
+
+Sets the RAM (MB) used during the build phase. Raise this when a build OOMs
+(heavy `apt`/`pip`/`npm`). Defaults to **4096**. Does **not** affect the resulting
+sandbox's memory — size that at create time via `memoryMB`.
+
+<ParamField body="mb" type="number" required>
+  Build-phase memory in MB
+</ParamField>
+
+**Returns:** `Image`
+
 ---
 
 ## Utility Methods

docs/sandboxes/templates.mdx

@@ -90,6 +90,39 @@ When you pass an `image` to `Sandbox.create()`, the server:
 2. If cached, creates the sandbox from the existing checkpoint instantly
 3. If not cached, boots a build sandbox, executes each step, checkpoints the result, then creates your sandbox from it
 
+### Build memory
+
+Images build in a **4 GB** sandbox by default. If a build runs out of memory (heavy `apt`/`pip`/`npm`, compiling a large toolchain), raise the build-phase RAM with **`.builderMemory(mb)` / `.builder_memory(mb)`**.
+
+This only affects the build. The resulting image is unchanged — you size the actual sandbox when you create it, via `memoryMB`:
+
+<CodeGroup>
+
+```typescript TypeScript
+// 8 GB to build…
+const image = Image.base()
+  .aptInstall(['build-essential', 'cmake'])
+  .runCommands('make -j')
+  .builderMemory(8192);
+
+// …but the sandbox runs at whatever you ask for
+const sandbox = await Sandbox.create({ image, memoryMB: 4096 });
+```
+
+```python Python
+image = (
+    Image.base()
+    .apt_install(["build-essential", "cmake"])
+    .run_commands("make -j")
+    .builder_memory(8192)
+)
+sandbox = await Sandbox.create(image=image)  # size via the HTTP API's memoryMB
+```
+
+</CodeGroup>
+
+`builderMemory` doesn't change the cache key — it's a build resource, not image content.
+
 ## Creating pre-built snapshots
 
 Create named snapshots that persist permanently and can be shared across sandboxes. Snapshots are visible in the dashboard and don't need to be rebuilt.
@@ -285,6 +318,7 @@ The `Image` class provides a fluent, immutable API for defining sandbox environm
 | `.addFile(path, content)` / `.add_file(path, content)` | Embed a file with inline content |
 | `.addLocalFile(local, remote)` / `.add_local_file(local, remote)` | Read a local file into the image |
 | `.addLocalDir(local, remote)` / `.add_local_dir(local, remote)` | Read a local directory into the image |
+| `.builderMemory(mb)` / `.builder_memory(mb)` | RAM for the build phase (default 4 GB; doesn't affect the resulting sandbox) |
 | `.toJSON()` / `.to_dict()` | Return the image manifest |
 | `.cacheKey()` / `.cache_key()` | Compute SHA-256 content hash |
 

internal/api/image_builder.go

@@ -23,8 +23,34 @@ type ImageManifest struct {
 	Base  string      `json:"base"`
 	Steps []ImageStep `json:"steps"`
 	Name  string      `json:"name,omitempty"` // optional — makes image addressable as a snapshot (for patches, etc.)
+
+	// BuilderMemoryMB is the RAM given to the build VM while it runs the steps
+	// (apt/pip can be memory-hungry; the old 1 GB default OOM'd large builds).
+	// It does NOT pin the output: the finalize pass re-snapshots at
+	// RuntimeMemoryMB. Default defaultBuilderMemoryMB. Exposed in the SDKs as
+	// .builderMemory().
+	BuilderMemoryMB int `json:"builderMemoryMB,omitempty"`
+
+	// RuntimeMemoryMB is the memory floor of the OUTPUT image — the RAM the
+	// finalize VM is cold-booted + savevm'd at. Forks can't run below it (a warm
+	// savevm can't shrink past its running size) but can hotplug up. Defaults to
+	// defaultRuntimeMemoryMB (1 GB); callers size the actual sandbox at create
+	// time via memoryMB. NOT exposed in the SDKs — an advanced/raw-manifest knob
+	// for images whose services auto-start heavy at boot.
+	RuntimeMemoryMB int `json:"runtimeMemoryMB,omitempty"`
 }
 
+const (
+	// defaultBuilderMemoryMB is the build-phase RAM. Enough that common
+	// apt/pip/npm builds don't OOM; safe because it never pins the output (see
+	// the finalize pass in buildImage).
+	defaultBuilderMemoryMB = 4096
+
+	// defaultRuntimeMemoryMB is the output image's memory floor — the historical
+	// 1 GB, so every image forks down to 1 GB and scales up on demand.
+	defaultRuntimeMemoryMB = 1024
+)
+
 // ImageStep is a single build step in an image manifest.
 type ImageStep struct {
 	Type string                 `json:"type"`
@@ -380,15 +406,36 @@ func (s *Server) buildImage(ctx context.Context, orgID uuid.UUID, manifest *Imag
 		base = "base"
 	}
 
+	// Build-phase RAM (apt/pip), and the output image's memory floor. Both
+	// settable per-manifest; defaults preserve "builds get headroom, floor stays
+	// at the historical 1 GB" (the finalize pass below re-snapshots at the floor).
+	builderMem := manifest.BuilderMemoryMB
+	if builderMem <= 0 {
+		builderMem = defaultBuilderMemoryMB
+	}
+	runtimeMem := manifest.RuntimeMemoryMB
+	if runtimeMem <= 0 {
+		runtimeMem = defaultRuntimeMemoryMB
+	}
+	// Only run the cold-boot finalize pass when it actually lowers the floor.
+	// If runtime >= builder there's nothing to gain (the in-place savevm already
+	// floors at builderMem ≤ runtimeMem), so snapshot in place (finalizeMem = 0).
+	finalizeMem := 0
+	if runtimeMem < builderMem {
+		finalizeMem = runtimeMem
+	}
+
 	// Create a throwaway sandbox
 	buildSandboxID := "sb-build-" + uuid.New().String()[:8]
 	cfg := types.SandboxConfig{
 		Template:  base,
 		Timeout:   600, // 10 min max for builds
 		SandboxID: buildSandboxID,
+		MemoryMB:  builderMem,
 	}
 
-	log.Printf("image-builder: creating build sandbox %s (base=%s, steps=%d)", buildSandboxID, base, len(manifest.Steps))
+	log.Printf("image-builder: creating build sandbox %s (base=%s, steps=%d, builderMem=%dMB, runtimeMem=%dMB)",
+		buildSandboxID, base, len(manifest.Steps), builderMem, runtimeMem)
 
 	var grpcClient pb.SandboxWorkerClient
 	var workerID string
@@ -414,6 +461,7 @@ func (s *Server) buildImage(ctx context.Context, orgID uuid.UUID, manifest *Imag
 			Timeout:        int32(cfg.Timeout),
 			NetworkEnabled: true, // Need network for apt/pip
 			SandboxId:      buildSandboxID,
+			MemoryMb:       int32(builderMem),
 		})
 		if err != nil {
 			return uuid.Nil, fmt.Errorf("failed to create build sandbox: %w", err)
@@ -512,10 +560,10 @@ func (s *Server) buildImage(ctx context.Context, orgID uuid.UUID, manifest *Imag
 	if s.store != nil {
 		cfgJSON, _ := json.Marshal(cfg)
 		cp := &db.Checkpoint{
-			ID:           checkpointID,
-			SandboxID:    buildSandboxID,
-			OrgID:        orgID,
-			Name:         fmt.Sprintf("_image_build_%s", checkpointID.String()[:8]),
+			ID:            checkpointID,
+			SandboxID:     buildSandboxID,
+			OrgID:         orgID,
+			Name:          fmt.Sprintf("_image_build_%s", checkpointID.String()[:8]),
 			SandboxConfig: cfgJSON,
 		}
 		if err := s.store.CreateCheckpoint(ctx, cp); err != nil {
@@ -533,9 +581,10 @@ func (s *Server) buildImage(ctx context.Context, orgID uuid.UUID, manifest *Imag
 		defer cancel()
 
 		resp, err := grpcClient.CreateCheckpoint(cpCtx, &pb.CreateCheckpointRequest{
-			SandboxId:     buildSandboxID,
-			CheckpointId:  checkpointID.String(),
-			PrepareGolden: true, // prepare golden snapshot for instant template creates
+			SandboxId:        buildSandboxID,
+			CheckpointId:     checkpointID.String(),
+			PrepareGolden:    true,               // prepare golden snapshot for instant template creates
+			FinalizeMemoryMb: int32(finalizeMem), // 0 = snapshot in place; >0 = cold-boot finalize at this floor
 		})
 		if err != nil {
 			return uuid.Nil, fmt.Errorf("failed to checkpoint build sandbox: %w", err)
@@ -557,7 +606,17 @@ func (s *Server) buildImage(ctx context.Context, orgID uuid.UUID, manifest *Imag
 			return uuid.Nil, fmt.Errorf("manager does not support checkpoints")
 		}
 
-		rootfsKey, workspaceKey, sizeBytes, err := cpMgr.CreateCheckpoint(ctx, buildSandboxID, checkpointID.String(), s.checkpointStore, func() {})
+		var rootfsKey, workspaceKey string
+		var sizeBytes int64
+		var err error
+		type finalizer interface {
+			CreateCheckpointFinalized(ctx context.Context, buildSandboxID, checkpointID string, store *storage.CheckpointStore, finalizeMemMB int, onReady func()) (string, string, int64, error)
+		}
+		if fz, ok := s.manager.(finalizer); ok && finalizeMem > 0 {
+			rootfsKey, workspaceKey, sizeBytes, err = fz.CreateCheckpointFinalized(ctx, buildSandboxID, checkpointID.String(), s.checkpointStore, finalizeMem, func() {})
+		} else {
+			rootfsKey, workspaceKey, sizeBytes, err = cpMgr.CreateCheckpoint(ctx, buildSandboxID, checkpointID.String(), s.checkpointStore, func() {})
+		}
 		if err != nil {
 			return uuid.Nil, fmt.Errorf("failed to checkpoint build sandbox: %w", err)
 		}
@@ -693,4 +752,3 @@ func stepDescription(step ImageStep) string {
 		return step.Type
 	}
 }
-

internal/qemu/manager.go

@@ -332,7 +332,7 @@ type Manager struct {
 	goldenVersion string // hash of base image — used for overlay-based migration
 
 	// Metadata service callbacks (set via SetMetadataCallbacks)
-	onSandboxReady   func(sandboxID, guestIP, template string, startedAt time.Time)
+	onSandboxReady      func(sandboxID, guestIP, template string, startedAt time.Time)
 	onSandboxDestroy    func(sandboxID string)
 	onMigrationOutgoing func(sandboxID string)
 
@@ -2845,6 +2845,67 @@ func (m *Manager) CheckpointCachePath(checkpointID, filename string) string {
 // checkpoint actually is. Upload failures now propagate as an error rather
 // than being silently logged — the control plane gets the reason and
 // persists it via SetCheckpointFailed (migration 039 added error_msg).
+// CreateCheckpointFinalized produces an image checkpoint whose memory floor is
+// finalizeMemMB, decoupled from the build VM's (larger) build-phase RAM. The
+// image builder uses this so a build can run at, say, 8 GB (apt/pip don't OOM)
+// while the resulting image still forks down to a 1 GB floor.
+//
+// Why two VMs: a savevm captures the *running* memory, and ForkFromCheckpoint
+// floors every fork at that size (shrinking past it would OOM restored
+// processes). So the only way to get a low floor from a high-memory build is to
+// cold-boot a fresh VM from the built disks at the target floor and snapshot
+// THAT. The cold-boot machinery already exists (Create with TemplateRootfsKey).
+//
+// Sequence: sync the build guest's FS → cold-boot a finalize VM from a copy of
+// the build disks at finalizeMemMB → snapshot the finalize VM (the output) →
+// tear the finalize VM down. The build VM is left running for its caller to
+// destroy (Kill would delete the disks we copy from). Cold-boot replays the
+// ext4 journal, so a post-sync disk copy restores cleanly.
+//
+// NOTE: the disk-consistency of the live-disk copy and the finalize cold-boot's
+// agent bring-up are the things to validate on dev before prod.
+func (m *Manager) CreateCheckpointFinalized(ctx context.Context, buildSandboxID, checkpointID string, checkpointStore *storage.CheckpointStore, finalizeMemMB int, onReady func()) (rootfsKey, workspaceKey string, sizeBytes int64, err error) {
+	// 1. Flush the build guest's filesystem so the on-disk qcow2 is consistent
+	//    before we copy it. Best-effort: the cold-boot journal-replays regardless.
+	if syncErr := m.SyncFS(ctx, buildSandboxID); syncErr != nil {
+		log.Printf("qemu: finalize %s: SyncFS warning: %v (continuing)", buildSandboxID, syncErr)
+	}
+
+	buildDir := filepath.Join(m.cfg.DataDir, "sandboxes", buildSandboxID)
+	buildRootfs := filepath.Join(buildDir, "rootfs.qcow2")
+	buildWorkspace := filepath.Join(buildDir, "workspace.qcow2")
+	if !fileExists(buildRootfs) || !fileExists(buildWorkspace) {
+		return "", "", 0, fmt.Errorf("finalize: build disks not found for %s", buildSandboxID)
+	}
+
+	// 2. Cold-boot a fresh finalize VM from a copy of the build disks at the
+	//    target floor. Create() copies the disks (reflink) via TemplateRootfsKey
+	//    and cold-boots — no savevm restore, so no memory floor inherited.
+	finalizeID := buildSandboxID + "-fin"
+	netEnabled := true
+	finCfg := types.SandboxConfig{
+		SandboxID:            finalizeID,
+		MemoryMB:             finalizeMemMB,
+		NetworkEnabled:       &netEnabled,
+		TemplateRootfsKey:    "local://" + buildRootfs,
+		TemplateWorkspaceKey: "local://" + buildWorkspace,
+	}
+	log.Printf("qemu: finalize: cold-booting %s from %s disks at %dMB", finalizeID, buildSandboxID, finalizeMemMB)
+	if _, err := m.Create(ctx, finCfg); err != nil {
+		return "", "", 0, fmt.Errorf("finalize: cold-boot at %dMB: %w", finalizeMemMB, err)
+	}
+	// Tear down the ephemeral finalize VM after we snapshot it.
+	defer func() {
+		if kErr := m.Kill(context.Background(), finalizeID); kErr != nil {
+			log.Printf("qemu: finalize %s: cleanup Kill failed: %v", finalizeID, kErr)
+		}
+	}()
+
+	// 3. Snapshot the finalize VM → the output checkpoint (floor = finalizeMemMB).
+	log.Printf("qemu: finalize %s → checkpoint %s (floor %dMB)", finalizeID, checkpointID, finalizeMemMB)
+	return m.CreateCheckpoint(ctx, finalizeID, checkpointID, checkpointStore, onReady)
+}
+
 func (m *Manager) CreateCheckpoint(ctx context.Context, sandboxID, checkpointID string, checkpointStore *storage.CheckpointStore, onReady func()) (rootfsKey, workspaceKey string, sizeBytes int64, err error) {
 	tStart := time.Now()
 	// failureReason is updated at each error site below so the defer can attribute

internal/worker/grpc_server.go

@@ -882,7 +882,25 @@ func (s *GRPCServer) CreateCheckpoint(ctx context.Context, req *pb.CreateCheckpo
 	// it's marked "ready" — forks poll for "ready" before downloading.
 	// The gRPC call returns immediately with the S3 keys — the CP's fork path
 	// polls for checkpoint readiness and blocks until onReady fires.
-	rootfsKey, workspaceKey, sizeBytes, err := s.manager.CreateCheckpoint(ctx, req.SandboxId, checkpointID, s.checkpointStore, onReady)
+	// Finalize: when finalize_memory_mb is set (image builder), produce the
+	// checkpoint at that memory floor by cold-booting a fresh VM from the build
+	// disks and snapshotting it — decoupling the image's floor from the build's
+	// (larger) RAM. Otherwise snapshot the sandbox in place.
+	var rootfsKey, workspaceKey string
+	var sizeBytes int64
+	var err error
+	if req.FinalizeMemoryMb > 0 {
+		type finalizer interface {
+			CreateCheckpointFinalized(ctx context.Context, buildSandboxID, checkpointID string, store *storage.CheckpointStore, finalizeMemMB int, onReady func()) (string, string, int64, error)
+		}
+		fz, ok := s.manager.(finalizer)
+		if !ok {
+			return nil, fmt.Errorf("manager does not support finalized checkpoints")
+		}
+		rootfsKey, workspaceKey, sizeBytes, err = fz.CreateCheckpointFinalized(ctx, req.SandboxId, checkpointID, s.checkpointStore, int(req.FinalizeMemoryMb), onReady)
+	} else {
+		rootfsKey, workspaceKey, sizeBytes, err = s.manager.CreateCheckpoint(ctx, req.SandboxId, checkpointID, s.checkpointStore, onReady)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("create checkpoint failed: %w", err)
 	}

proto/worker/worker.pb.go

@@ -338,6 +338,11 @@ message CreateCheckpointRequest {
   string sandbox_id = 1;
   string checkpoint_id = 2;  // UUID assigned by control plane
   bool prepare_golden = 3;   // If true, prepare a golden snapshot from this checkpoint
+  // If > 0, finalize the checkpoint at this memory: the worker stops the build VM,
+  // cold-boots a fresh VM from its disks at finalize_memory_mb, and snapshots THAT
+  // — so the image's memory floor is finalize_memory_mb, decoupled from the
+  // (larger) build-phase RAM. Used by the image builder. 0 = snapshot in place.
+  int32 finalize_memory_mb = 4;
 }
 
 message CreateCheckpointResponse {