Add pullpolicy when scheduling new workers (#238)

- Always pull an image when the tag is "latest" - Pull if not present when tag is anything else (default)
beam-cloud · May 21, 2024 · d8eb1a7 · d8eb1a7
1 parent eaa9044
commit d8eb1a7
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/pkg/scheduler/pool_external.go b/pkg/scheduler/pool_external.go
@@ -219,6 +219,11 @@ func (wpc *ExternalWorkerPoolController) createWorkerJob(workerId, machineId str
 		wpc.config.Worker.ImageTag,
 	)
 
+	pullPolicy := corev1.PullIfNotPresent
+	if wpc.config.Worker.ImageTag == "latest" {
+		pullPolicy = corev1.PullAlways
+	}
+
 	resources := corev1.ResourceRequirements{}
 	if workerGpuType != "" {
 		resources.Requests = corev1.ResourceList{
@@ -231,8 +236,9 @@ func (wpc *ExternalWorkerPoolController) createWorkerJob(workerId, machineId str
 
 	containers := []corev1.Container{
 		{
-			Name:  defaultContainerName,
-			Image: workerImage,
+			Name:            defaultContainerName,
+			Image:           workerImage,
+			ImagePullPolicy: pullPolicy,
 			Command: []string{
 				defaultWorkerEntrypoint,
 			},

diff --git a/pkg/scheduler/pool_local.go b/pkg/scheduler/pool_local.go
@@ -142,6 +142,11 @@ func (wpc *LocalKubernetesWorkerPoolController) createWorkerJob(workerId string,
 		wpc.config.Worker.ImageTag,
 	)
 
+	pullPolicy := corev1.PullIfNotPresent
+	if wpc.config.Worker.ImageTag == "latest" {
+		pullPolicy = corev1.PullAlways
+	}
+
 	resources := corev1.ResourceRequirements{}
 	if wpc.config.Worker.ResourcesEnforced {
 		resources.Requests = resourceRequests
@@ -150,8 +155,9 @@ func (wpc *LocalKubernetesWorkerPoolController) createWorkerJob(workerId string,
 
 	containers := []corev1.Container{
 		{
-			Name:  defaultContainerName,
-			Image: workerImage,
+			Name:            defaultContainerName,
+			Image:           workerImage,
+			ImagePullPolicy: pullPolicy,
 			Command: []string{
 				defaultWorkerEntrypoint,
 			},

diff --git a/sdk/src/beta9/__init__.py b/sdk/src/beta9/__init__.py
@@ -10,9 +10,7 @@
 from .abstractions.volume import Volume
 from .type import GpuType, PythonVersion
 
-__version__ = "0.0.1"
 __all__ = [
-    "__version__",
     "Map",
     "Image",
     "Queue",