From b359b0fe55d536ce9135981f2932d9b0a372c09c Mon Sep 17 00:00:00 2001 From: Ian Hodge Date: Tue, 7 Apr 2026 13:01:19 -0400 Subject: [PATCH] ian/adding_sidecar_image_param --- README.md | 1 + .../oz-agent-worker/templates/configmap.yaml | 3 + charts/oz-agent-worker/values.yaml | 1 + internal/config/config.go | 1 + internal/config/config_test.go | 50 ++++++++++++++++ internal/worker/kubernetes.go | 1 + internal/worker/worker.go | 7 ++- internal/worker/worker_test.go | 60 +++++++++++++++++++ main.go | 3 + 9 files changed, 126 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 386fdf8..1eab26d 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,7 @@ Notes: - `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` +- `sidecar_image` overrides the warp-agent sidecar image reference sent by the server (e.g. `docker.io/warpdotdev/warp-agent:latest`); set this when cluster nodes cannot pull directly from Docker Hub and must use an internal registry mirror or pull-through cache instead. This only affects the warp-agent sidecar (mounted at `/agent`), not any additional sidecars. When using this override, you are responsible for keeping your mirror in sync with `docker.io/warpdotdev/warp-agent` — the server normally sends the correct version-matched image per task, so a stale mirror may cause version incompatibility - by default, the Kubernetes backend materializes sidecars with root init containers into `emptyDir` volumes, matching the existing behavior - set `use_image_volumes: true` to opt into native image volumes for sidecars; in that mode, sidecar mounts are read-only and Kubernetes/runtime support for the built-in `ImageVolume` Pod volume source is required - Kubernetes `1.35+` is the recommended and tested target for `use_image_volumes: true`; Kubernetes `1.33`-`1.34` may work if `ImageVolume` is enabled and the container runtime supports image volumes diff --git a/charts/oz-agent-worker/templates/configmap.yaml b/charts/oz-agent-worker/templates/configmap.yaml index 25172f0..994ff52 100644 --- a/charts/oz-agent-worker/templates/configmap.yaml +++ b/charts/oz-agent-worker/templates/configmap.yaml @@ -24,6 +24,9 @@ data: {{- if .Values.kubernetesBackend.preflightImage }} preflight_image: {{ .Values.kubernetesBackend.preflightImage | quote }} {{- end }} + {{- if .Values.kubernetesBackend.sidecarImage }} + sidecar_image: {{ .Values.kubernetesBackend.sidecarImage | quote }} + {{- end }} {{- if .Values.kubernetesBackend.setupCommand }} setup_command: |- {{ .Values.kubernetesBackend.setupCommand | indent 10 }} diff --git a/charts/oz-agent-worker/values.yaml b/charts/oz-agent-worker/values.yaml index a8b1e05..d5b6eb4 100644 --- a/charts/oz-agent-worker/values.yaml +++ b/charts/oz-agent-worker/values.yaml @@ -70,6 +70,7 @@ kubernetesBackend: imagePullPolicy: IfNotPresent useImageVolumes: false preflightImage: "" + sidecarImage: "" setupCommand: "" teardownCommand: "" extraLabels: {} diff --git a/internal/config/config.go b/internal/config/config.go index 3f5e98c..e3a13da 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -60,6 +60,7 @@ type KubernetesConfig struct { ImagePullPolicy string `yaml:"image_pull_policy" validate:"omitempty,oneof=Always Never IfNotPresent"` UseImageVolumes bool `yaml:"use_image_volumes"` PreflightImage string `yaml:"preflight_image" validate:"omitempty,no_whitespace"` + SidecarImage string `yaml:"sidecar_image" validate:"omitempty,no_whitespace"` SetupCommand string `yaml:"setup_command"` TeardownCommand string `yaml:"teardown_command"` ExtraLabels map[string]string `yaml:"extra_labels"` diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 3ad27f2..8c43506 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -490,6 +490,56 @@ backend: }) } +func TestLoadKubernetesSidecarImage(t *testing.T) { + t.Run("parses sidecar_image when set", func(t *testing.T) { + path := writeTestConfig(t, ` +worker_id: "k8s-worker" +backend: + kubernetes: + sidecar_image: "my-registry.io/warpdotdev/warp-agent: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.SidecarImage != "my-registry.io/warpdotdev/warp-agent:latest" { + t.Errorf("sidecar_image = %q, want %q", cfg.Backend.Kubernetes.SidecarImage, "my-registry.io/warpdotdev/warp-agent:latest") + } + }) + + t.Run("sidecar_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.SidecarImage != "" { + t.Errorf("expected sidecar_image to be empty, got %q", cfg.Backend.Kubernetes.SidecarImage) + } + }) + + t.Run("rejects sidecar_image with whitespace", func(t *testing.T) { + path := writeTestConfig(t, ` +worker_id: "k8s-worker" +backend: + kubernetes: + sidecar_image: "my image:latest" +`) + _, err := Load(path) + if err == nil { + t.Fatal("expected error for sidecar_image with whitespace") + } + }) +} + func TestLoadLegacyKubernetesFieldRejected(t *testing.T) { tests := []string{ "image_pull_secret", diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index 726a736..670c87f 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -55,6 +55,7 @@ type KubernetesBackendConfig struct { ImagePullPolicy string UseImageVolumes bool PreflightImage string + SidecarImage string SetupCommand string TeardownCommand string NoCleanup bool diff --git a/internal/worker/worker.go b/internal/worker/worker.go index d8ced90..fa85a51 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -381,8 +381,13 @@ func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *Tas // entrypoint.sh lives) comes first, followed by any additional sidecars. var sidecars []types.SidecarMount if assignment.SidecarImage != "" { + sidecarImage := assignment.SidecarImage + if w.config.Kubernetes != nil && w.config.Kubernetes.SidecarImage != "" { + log.Infof(w.ctx, "Overriding server sidecar image %s with configured sidecar image %s", assignment.SidecarImage, w.config.Kubernetes.SidecarImage) + sidecarImage = w.config.Kubernetes.SidecarImage + } sidecars = append(sidecars, types.SidecarMount{ - Image: assignment.SidecarImage, + Image: sidecarImage, MountPath: "/agent", }) } diff --git a/internal/worker/worker_test.go b/internal/worker/worker_test.go index 22b6bab..1ba1c86 100644 --- a/internal/worker/worker_test.go +++ b/internal/worker/worker_test.go @@ -87,6 +87,66 @@ func TestDefaultImageForTask(t *testing.T) { }) } +func TestPrepareTaskParamsSidecarImageOverride(t *testing.T) { + newWorker := func(sidecarImage string) *Worker { + ctx := context.Background() + var k8sConfig *KubernetesBackendConfig + if sidecarImage != "" { + k8sConfig = &KubernetesBackendConfig{SidecarImage: sidecarImage} + } else { + k8sConfig = &KubernetesBackendConfig{} + } + return &Worker{ + ctx: ctx, + config: Config{ + Kubernetes: k8sConfig, + }, + } + } + + t.Run("config sidecar_image overrides server-provided image", func(t *testing.T) { + w := newWorker("my-registry.io/warpdotdev/warp-agent:latest") + params := w.prepareTaskParams(&types.TaskAssignmentMessage{ + TaskID: "task-1", + Task: &types.Task{ID: "task-1"}, + SidecarImage: "docker.io/warpdotdev/warp-agent:latest", + }) + if len(params.Sidecars) == 0 { + t.Fatal("expected at least one sidecar") + } + if params.Sidecars[0].Image != "my-registry.io/warpdotdev/warp-agent:latest" { + t.Errorf("sidecar image = %q, want %q", params.Sidecars[0].Image, "my-registry.io/warpdotdev/warp-agent:latest") + } + }) + + t.Run("server-provided image used when config sidecar_image empty", func(t *testing.T) { + w := newWorker("") + params := w.prepareTaskParams(&types.TaskAssignmentMessage{ + TaskID: "task-1", + Task: &types.Task{ID: "task-1"}, + SidecarImage: "docker.io/warpdotdev/warp-agent:latest", + }) + if len(params.Sidecars) == 0 { + t.Fatal("expected at least one sidecar") + } + if params.Sidecars[0].Image != "docker.io/warpdotdev/warp-agent:latest" { + t.Errorf("sidecar image = %q, want %q", params.Sidecars[0].Image, "docker.io/warpdotdev/warp-agent:latest") + } + }) + + t.Run("no sidecar when server provides empty sidecar image", func(t *testing.T) { + w := newWorker("my-registry.io/warpdotdev/warp-agent:latest") + params := w.prepareTaskParams(&types.TaskAssignmentMessage{ + TaskID: "task-1", + Task: &types.Task{ID: "task-1"}, + SidecarImage: "", + }) + if len(params.Sidecars) != 0 { + t.Errorf("expected no sidecars when server sidecar image is empty, got %d", len(params.Sidecars)) + } + }) +} + func TestWorkerShutdownUsesFreshContextForBackendCleanup(t *testing.T) { workerCtx, cancel := context.WithCancel(context.Background()) backend := &shutdownRecordingBackend{} diff --git a/main.go b/main.go index 4d4fba5..9f16757 100644 --- a/main.go +++ b/main.go @@ -158,6 +158,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { imagePullPolicy string useImageVolumes bool preflightImage string + sidecarImage string setupCmd string teardownCmd string extraLabels map[string]string @@ -176,6 +177,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { imagePullPolicy = kc.ImagePullPolicy useImageVolumes = kc.UseImageVolumes preflightImage = kc.PreflightImage + sidecarImage = kc.SidecarImage setupCmd = kc.SetupCommand teardownCmd = kc.TeardownCommand extraLabels = copyStringMap(kc.ExtraLabels) @@ -216,6 +218,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { ImagePullPolicy: imagePullPolicy, UseImageVolumes: useImageVolumes, PreflightImage: preflightImage, + SidecarImage: sidecarImage, SetupCommand: setupCmd, TeardownCommand: teardownCmd, NoCleanup: noCleanup,