Fix: prevent C/R during serves (#755)

beam-cloud · Dec 4, 2024 · 70a0e7d · 70a0e7d
1 parent 35df058
commit 70a0e7d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/pkg/abstractions/endpoint/instance.go b/pkg/abstractions/endpoint/instance.go
@@ -73,6 +73,11 @@ func (i *endpointInstance) startContainers(containersToRun int) error {
 		gpuCount = 1
 	}
 
+	checkpointEnabled := i.StubConfig.CheckpointEnabled
+	if i.Stub.Type.IsServe() {
+		checkpointEnabled = false
+	}
+
 	for c := 0; c < containersToRun; c++ {
 		containerId := i.genContainerId()
 
@@ -90,7 +95,7 @@ func (i *endpointInstance) startContainers(containersToRun int) error {
 			EntryPoint:        i.EntryPoint,
 			Mounts:            mounts,
 			Stub:              *i.Stub,
-			CheckpointEnabled: i.StubConfig.CheckpointEnabled,
+			CheckpointEnabled: checkpointEnabled,
 		}
 
 		// Set initial keepwarm to prevent rapid spin-up/spin-down of containers

diff --git a/pkg/abstractions/taskqueue/instance.go b/pkg/abstractions/taskqueue/instance.go
@@ -70,6 +70,11 @@ func (i *taskQueueInstance) startContainers(containersToRun int) error {
 		gpuCount = 1
 	}
 
+	checkpointEnabled := i.StubConfig.CheckpointEnabled
+	if i.Stub.Type.IsServe() {
+		checkpointEnabled = false
+	}
+
 	for c := 0; c < containersToRun; c++ {
 		runRequest := &types.ContainerRequest{
 			ContainerId:       i.genContainerId(),
@@ -85,7 +90,7 @@ func (i *taskQueueInstance) startContainers(containersToRun int) error {
 			EntryPoint:        i.EntryPoint,
 			Mounts:            mounts,
 			Stub:              *i.Stub,
-			CheckpointEnabled: i.StubConfig.CheckpointEnabled,
+			CheckpointEnabled: checkpointEnabled,
 		}
 
 		err := i.Scheduler.Run(runRequest)