add pre-emptible config

beam-cloud · Nov 4, 2024 · b763ca3 · b763ca3
1 parent 5b156a7
commit b763ca3
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 0 deletions.
diff --git a/pkg/scheduler/pool_external.go b/pkg/scheduler/pool_external.go
@@ -410,6 +410,10 @@ func (wpc *ExternalWorkerPoolController) getWorkerEnvironment(workerId, machineI
 			Name:  "NETWORK_PREFIX",
 			Value: machineId,
 		},
+		{
+			Name:  "PREEMPTABLE",
+			Value: strconv.FormatBool(wpc.workerPool.Preemptable),
+		},
 	}
 
 	remoteConfig, err := providers.GetRemoteConfig(wpc.config, wpc.tailscale)

diff --git a/pkg/scheduler/pool_local.go b/pkg/scheduler/pool_local.go
@@ -379,6 +379,10 @@ func (wpc *LocalKubernetesWorkerPoolController) getWorkerEnvironment(workerId st
 				},
 			},
 		},
+		{
+			Name:  "PREEMPTABLE",
+			Value: strconv.FormatBool(wpc.workerPool.Preemptable),
+		},
 	}
 
 	if wpc.config.Worker.UseGatewayServiceHostname {

diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
@@ -246,6 +246,7 @@ func (s *Scheduler) StartProcessingRequests() {
 							log.Printf("Unable to schedule request for container<%s>: %v\n", request.ContainerId, err)
 							s.addRequestToBacklog(request)
 						}
+
 						return
 					}
 				}

diff --git a/pkg/types/config.go b/pkg/types/config.go
@@ -213,6 +213,7 @@ type WorkerPoolConfig struct {
 	DefaultMachineCost   float64                           `key:"defaultMachineCost" json:"default_machine_cost"`
 	RequiresPoolSelector bool                              `key:"requiresPoolSelector" json:"requires_pool_selector"`
 	Priority             int32                             `key:"priority" json:"priority"`
+	Preemptable          bool                              `key:"preemptable" json:"preemptable"`
 }
 
 type WorkerPoolJobSpecConfig struct {