adding a default docker image to the k8 backend (#42)

Author: ianhodge

README.md

@@ -61,6 +61,7 @@ backend:
   kubernetes:
     kubeconfig: "/path/to/kubeconfig"
     namespace: "agents"
+    default_image: "my-registry.io/dev-image:latest"
     unschedulable_timeout: "2m"
     pod_template:
       nodeSelector:
@@ -75,6 +76,7 @@ backend:
 
 Notes:
 
+- `default_image` sets the Docker image for task Jobs when no Warp environment is configured on the run; this lets you skip creating a Warp environment entirely if all your tasks use the same base image (precedence: Warp environment image > `default_image` > `ubuntu:22.04`)
 - `namespace` selects the namespace inside the chosen cluster; it does not choose the cluster itself, and defaults to `default` when omitted
 - `unschedulable_timeout` controls how long a Pod may remain unschedulable before the task is failed early; it defaults to `30s`, and `0s` disables that fail-fast behavior
 - `image_pull_policy` defaults to `IfNotPresent`

charts/oz-agent-worker/templates/configmap.yaml

@@ -15,6 +15,9 @@ data:
     backend:
       kubernetes:
         namespace: {{ default .Release.Namespace .Values.kubernetesBackend.namespace | quote }}
+        {{- if .Values.kubernetesBackend.defaultImage }}
+        default_image: {{ .Values.kubernetesBackend.defaultImage | quote }}
+        {{- end }}
         image_pull_policy: {{ .Values.kubernetesBackend.imagePullPolicy | quote }}
         unschedulable_timeout: {{ .Values.kubernetesBackend.unschedulableTimeout | quote }}
         {{- if .Values.kubernetesBackend.preflightImage }}

charts/oz-agent-worker/values.yaml

@@ -66,6 +66,7 @@ worker:
 
 kubernetesBackend:
   namespace: ""
+  defaultImage: ""
   imagePullPolicy: IfNotPresent
   preflightImage: ""
   setupCommand: ""

internal/config/config.go

@@ -56,6 +56,7 @@ type DirectConfig struct {
 type KubernetesConfig struct {
 	Namespace             string            `yaml:"namespace"`
 	Kubeconfig            string            `yaml:"kubeconfig"`
+	DefaultImage          string            `yaml:"default_image" validate:"omitempty,no_whitespace"`
 	ImagePullPolicy       string            `yaml:"image_pull_policy" validate:"omitempty,oneof=Always Never IfNotPresent"`
 	PreflightImage        string            `yaml:"preflight_image" validate:"omitempty,no_whitespace"`
 	SetupCommand          string            `yaml:"setup_command"`

internal/config/config_test.go

@@ -436,6 +436,56 @@ backend:
 	}
 }
 
+func TestLoadKubernetesDefaultImage(t *testing.T) {
+	t.Run("parses default_image when set", func(t *testing.T) {
+		path := writeTestConfig(t, `
+worker_id: "k8s-worker"
+backend:
+  kubernetes:
+    default_image: "my-registry.io/custom-image:latest"
+`)
+		cfg, err := Load(path)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if cfg.Backend.Kubernetes == nil {
+			t.Fatal("expected kubernetes backend to be set")
+		}
+		if cfg.Backend.Kubernetes.DefaultImage != "my-registry.io/custom-image:latest" {
+			t.Errorf("default_image = %q, want %q", cfg.Backend.Kubernetes.DefaultImage, "my-registry.io/custom-image:latest")
+		}
+	})
+
+	t.Run("default_image is empty when not set", func(t *testing.T) {
+		path := writeTestConfig(t, `
+worker_id: "k8s-worker"
+backend:
+  kubernetes:
+    namespace: "agents"
+`)
+		cfg, err := Load(path)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if cfg.Backend.Kubernetes.DefaultImage != "" {
+			t.Errorf("expected default_image to be empty, got %q", cfg.Backend.Kubernetes.DefaultImage)
+		}
+	})
+
+	t.Run("rejects default_image with whitespace", func(t *testing.T) {
+		path := writeTestConfig(t, `
+worker_id: "k8s-worker"
+backend:
+  kubernetes:
+    default_image: "my image:latest"
+`)
+		_, err := Load(path)
+		if err == nil {
+			t.Fatal("expected error for default_image with whitespace")
+		}
+	})
+}
+
 func TestLoadLegacyKubernetesFieldRejected(t *testing.T) {
 	tests := []string{
 		"image_pull_secret",

internal/worker/kubernetes.go

@@ -49,6 +49,7 @@ type KubernetesBackendConfig struct {
 	WorkerID              string
 	Namespace             string
 	Kubeconfig            string
+	DefaultImage          string
 	ImagePullPolicy       string
 	PreflightImage        string
 	SetupCommand          string

internal/worker/worker.go

@@ -347,17 +347,9 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) {
 func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *TaskParams {
 	task := assignment.Task
 
-	// Resolve Docker image with default fallback.
-	dockerImage := assignment.DockerImage
-	if dockerImage == "" {
-		dockerImage = "ubuntu:22.04"
-		if task.AgentConfigSnapshot != nil && task.AgentConfigSnapshot.EnvironmentID != nil {
-			log.Warnf(w.ctx, "Environment %s specified but no Docker image resolved. Using default: %s",
-				*task.AgentConfigSnapshot.EnvironmentID, dockerImage)
-		} else {
-			log.Infof(w.ctx, "No environment specified, using default image: %s", dockerImage)
-		}
-	}
+	// Resolve Docker image.
+	// Precedence: server-provided image (from environment) > worker config default_image > hardcoded ubuntu:22.04.
+	dockerImage := w.defaultImageForTask(assignment.DockerImage, task)
 
 	// Build common environment variables.
 	envVars := []string{
@@ -406,6 +398,26 @@ func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *Tas
 	}
 }
 
+// defaultImageForTask returns the Docker image to use for a task, applying the
+// precedence: server-provided > worker config default_image > hardcoded fallback.
+func (w *Worker) defaultImageForTask(assignmentImage string, task *types.Task) string {
+	if assignmentImage != "" {
+		return assignmentImage
+	}
+	if w.config.Kubernetes != nil && w.config.Kubernetes.DefaultImage != "" {
+		log.Infof(w.ctx, "Using worker-configured default image: %s", w.config.Kubernetes.DefaultImage)
+		return w.config.Kubernetes.DefaultImage
+	}
+	fallback := "ubuntu:22.04"
+	if task.AgentConfigSnapshot != nil && task.AgentConfigSnapshot.EnvironmentID != nil {
+		log.Warnf(w.ctx, "Environment %s specified but no Docker image resolved. Using default: %s",
+			*task.AgentConfigSnapshot.EnvironmentID, fallback)
+	} else {
+		log.Infof(w.ctx, "No environment specified, using default image: %s", fallback)
+	}
+	return fallback
+}
+
 func (w *Worker) executeTask(ctx context.Context, assignment *types.TaskAssignmentMessage) {
 	defer func() {
 		w.tasksMutex.Lock()

internal/worker/worker_test.go

@@ -3,6 +3,8 @@ package worker
 import (
 	"context"
 	"testing"
+
+	"github.com/warpdotdev/oz-agent-worker/internal/types"
 )
 
 type shutdownRecordingBackend struct {
@@ -19,6 +21,72 @@ func (b *shutdownRecordingBackend) Shutdown(ctx context.Context) {
 	b.shutdownCtxErr = ctx.Err()
 }
 
+func TestDefaultImageForTask(t *testing.T) {
+	newWorker := func(defaultImage string) *Worker {
+		ctx := context.Background()
+		var k8sConfig *KubernetesBackendConfig
+		if defaultImage != "" {
+			k8sConfig = &KubernetesBackendConfig{DefaultImage: defaultImage}
+		}
+		return &Worker{
+			ctx: ctx,
+			config: Config{
+				Kubernetes: k8sConfig,
+			},
+		}
+	}
+
+	envID := "env-123"
+
+	t.Run("server-provided image wins over default_image", func(t *testing.T) {
+		w := newWorker("my-registry.io/default:v1")
+		got := w.defaultImageForTask("server-image:latest", &types.Task{})
+		if got != "server-image:latest" {
+			t.Errorf("got %q, want %q", got, "server-image:latest")
+		}
+	})
+
+	t.Run("default_image used when server image empty", func(t *testing.T) {
+		w := newWorker("my-registry.io/default:v1")
+		got := w.defaultImageForTask("", &types.Task{})
+		if got != "my-registry.io/default:v1" {
+			t.Errorf("got %q, want %q", got, "my-registry.io/default:v1")
+		}
+	})
+
+	t.Run("hardcoded fallback when no default_image configured", func(t *testing.T) {
+		w := newWorker("")
+		got := w.defaultImageForTask("", &types.Task{})
+		if got != "ubuntu:22.04" {
+			t.Errorf("got %q, want %q", got, "ubuntu:22.04")
+		}
+	})
+
+	t.Run("hardcoded fallback when kubernetes config nil", func(t *testing.T) {
+		w := &Worker{
+			ctx:    context.Background(),
+			config: Config{},
+		}
+		got := w.defaultImageForTask("", &types.Task{})
+		if got != "ubuntu:22.04" {
+			t.Errorf("got %q, want %q", got, "ubuntu:22.04")
+		}
+	})
+
+	t.Run("hardcoded fallback with environment ID logs warning", func(t *testing.T) {
+		w := newWorker("")
+		task := &types.Task{
+			AgentConfigSnapshot: &types.AmbientAgentConfig{
+				EnvironmentID: &envID,
+			},
+		}
+		got := w.defaultImageForTask("", task)
+		if got != "ubuntu:22.04" {
+			t.Errorf("got %q, want %q", got, "ubuntu:22.04")
+		}
+	})
+}
+
 func TestWorkerShutdownUsesFreshContextForBackendCleanup(t *testing.T) {
 	workerCtx, cancel := context.WithCancel(context.Background())
 	backend := &shutdownRecordingBackend{}

main.go

@@ -154,6 +154,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
 		var (
 			namespace             string
 			kubeconfig            string
+			defaultImage          string
 			imagePullPolicy       string
 			preflightImage        string
 			setupCmd              string
@@ -170,6 +171,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
 			kc := fileConfig.Backend.Kubernetes
 			namespace = kc.Namespace
 			kubeconfig = kc.Kubeconfig
+			defaultImage = kc.DefaultImage
 			imagePullPolicy = kc.ImagePullPolicy
 			preflightImage = kc.PreflightImage
 			setupCmd = kc.SetupCommand
@@ -208,6 +210,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
 			WorkerID:              workerID,
 			Namespace:             namespace,
 			Kubeconfig:            kubeconfig,
+			DefaultImage:          defaultImage,
 			ImagePullPolicy:       imagePullPolicy,
 			PreflightImage:        preflightImage,
 			SetupCommand:          setupCmd,