diff --git a/apps/supervisor/src/env.ts b/apps/supervisor/src/env.ts index fc18b7f09d8..3f2cc8b46e9 100644 --- a/apps/supervisor/src/env.ts +++ b/apps/supervisor/src/env.ts @@ -185,6 +185,8 @@ export const Env = z KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"), KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"), KUBERNETES_STRIP_IMAGE_DIGEST: BoolEnv.default(false), + KUBERNETES_IMAGE_REGISTRY_REWRITE_FROM: z.string().optional(), + KUBERNETES_IMAGE_REGISTRY_REWRITE_TO: z.string().optional(), KUBERNETES_CPU_REQUEST_MIN_CORES: z.coerce.number().min(0).default(0), KUBERNETES_CPU_REQUEST_RATIO: z.coerce.number().min(0).max(1).default(0.75), // Ratio of CPU limit, so 0.75 = 75% of CPU limit KUBERNETES_MEMORY_REQUEST_MIN_GB: z.coerce.number().min(0).default(0), diff --git a/apps/supervisor/src/workloadManager/imageRegistry.test.ts b/apps/supervisor/src/workloadManager/imageRegistry.test.ts new file mode 100644 index 00000000000..66f829564c3 --- /dev/null +++ b/apps/supervisor/src/workloadManager/imageRegistry.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from "vitest"; +import { rewriteImageRegistry } from "./imageRegistry.js"; + +const FROM = "123456789012.dkr.ecr.us-east-1.amazonaws.com"; +const TO = "123456789012.dkr.ecr.eu-central-1.amazonaws.com"; + +describe("rewriteImageRegistry", () => { + it("rewrites the registry host and keeps the rest of the reference", () => { + expect(rewriteImageRegistry(`${FROM}/deployments/proj_abc:20260818.1`, FROM, TO)).toBe( + `${TO}/deployments/proj_abc:20260818.1` + ); + }); + + it("preserves a digest", () => { + expect(rewriteImageRegistry(`${FROM}/deployments/proj_abc@sha256:abc123`, FROM, TO)).toBe( + `${TO}/deployments/proj_abc@sha256:abc123` + ); + }); + + it("is a no-op unless both ends are configured", () => { + const ref = `${FROM}/deployments/proj_abc:tag`; + + expect(rewriteImageRegistry(ref, undefined, TO)).toBe(ref); + expect(rewriteImageRegistry(ref, FROM, undefined)).toBe(ref); + expect(rewriteImageRegistry(ref, undefined, undefined)).toBe(ref); + }); + + it("leaves other registries alone", () => { + const ref = "ghcr.io/triggerdotdev/something:tag"; + expect(rewriteImageRegistry(ref, FROM, TO)).toBe(ref); + }); + + it("only matches on a host boundary", () => { + const lookalike = `${FROM}.evil.example.com/deployments/proj_abc:tag`; + expect(rewriteImageRegistry(lookalike, FROM, TO)).toBe(lookalike); + }); + + it("does not rewrite a host that merely contains the source", () => { + const ref = `registry.example.com/${FROM}/proj_abc:tag`; + expect(rewriteImageRegistry(ref, FROM, TO)).toBe(ref); + }); +}); diff --git a/apps/supervisor/src/workloadManager/imageRegistry.ts b/apps/supervisor/src/workloadManager/imageRegistry.ts new file mode 100644 index 00000000000..a49da7f85f3 --- /dev/null +++ b/apps/supervisor/src/workloadManager/imageRegistry.ts @@ -0,0 +1,15 @@ +export function rewriteImageRegistry( + imageRef: string, + from: string | undefined, + to: string | undefined +): string { + if (!from || !to) { + return imageRef; + } + + if (!imageRef.startsWith(`${from}/`)) { + return imageRef; + } + + return `${to}${imageRef.slice(from.length)}`; +} diff --git a/apps/supervisor/src/workloadManager/kubernetes.ts b/apps/supervisor/src/workloadManager/kubernetes.ts index d09c86cf9cc..54a9428efeb 100644 --- a/apps/supervisor/src/workloadManager/kubernetes.ts +++ b/apps/supervisor/src/workloadManager/kubernetes.ts @@ -20,6 +20,7 @@ import { withRunnerSeccompProfile, withNodeSelector, } from "./kubernetesPodSpec.js"; +import { rewriteImageRegistry } from "./imageRegistry.js"; type ResourceQuantities = { [K in "cpu" | "memory" | "ephemeral-storage"]?: string; @@ -163,7 +164,11 @@ export class KubernetesWorkloadManager implements WorkloadManager { containers: [ { name: "run-controller", - image: this.stripImageDigest(opts.image), + image: rewriteImageRegistry( + this.stripImageDigest(opts.image), + env.KUBERNETES_IMAGE_REGISTRY_REWRITE_FROM, + env.KUBERNETES_IMAGE_REGISTRY_REWRITE_TO + ), ports: [ { containerPort: 8000,