diff --git a/.server-changes/supervisor-org-placement-overrides.md b/.server-changes/supervisor-org-placement-overrides.md new file mode 100644 index 00000000000..532c4a73e2a --- /dev/null +++ b/.server-changes/supervisor-org-placement-overrides.md @@ -0,0 +1,6 @@ +--- +area: supervisor +type: feature +--- + +Operators can now route an organization's runs to specific Kubernetes node pools. diff --git a/apps/supervisor/src/env.ts b/apps/supervisor/src/env.ts index 6830d5b8642..c5aaf70e1fe 100644 --- a/apps/supervisor/src/env.ts +++ b/apps/supervisor/src/env.ts @@ -1,7 +1,13 @@ import { randomUUID } from "crypto"; import { env as stdEnv } from "std-env"; import { z } from "zod"; -import { AdditionalEnvVars, BoolEnv, NodeLabelValue, Tolerations } from "./envUtil.js"; +import { + AdditionalEnvVars, + BoolEnv, + NodeLabelValue, + OrgPlacementOverrides, + Tolerations, +} from "./envUtil.js"; export const Env = z .object({ @@ -260,6 +266,11 @@ export const Env = z KUBERNETES_RUNNER_TOLERATIONS: Tolerations.optional(), // every run pod KUBERNETES_SCHEDULED_RUN_TOLERATIONS: Tolerations.optional(), // schedule-tree runs only + // Per-org placement overrides, JSON keyed by the internal org id + // (the `org` label on run pods): + // {"": {"nodeSelector": {"": ""}, "tolerations": ""}} + KUBERNETES_ORG_PLACEMENT_OVERRIDES: OrgPlacementOverrides, + // Placement tags settings PLACEMENT_TAGS_ENABLED: BoolEnv.default(false), PLACEMENT_TAGS_PREFIX: z.string().default("node.cluster.x-k8s.io"), @@ -305,6 +316,22 @@ export const Env = z path: ["TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE"], }); } + if (data.KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED && data.KUBERNETES_ORG_PLACEMENT_OVERRIDES) { + // Non-large presets carry a hard NotIn on the large-machine pool, so an org + // pinned to that pool could never schedule its non-large runs. + for (const [orgId, override] of Object.entries(data.KUBERNETES_ORG_PLACEMENT_OVERRIDES)) { + const pinnedPool = + override.nodeSelector?.[data.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY]; + + if (pinnedPool === data.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: `Org "${orgId}" pins run pods to the large-machine pool, but non-large presets are required to stay off it, so those runs would never schedule. Use a different pool or disable KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED.`, + path: ["KUBERNETES_ORG_PLACEMENT_OVERRIDES"], + }); + } + } + } if (data.COMPUTE_SNAPSHOTS_ENABLED && !data.TRIGGER_METADATA_URL) { ctx.addIssue({ code: z.ZodIssueCode.custom, diff --git a/apps/supervisor/src/envUtil.test.ts b/apps/supervisor/src/envUtil.test.ts index 378830f8ab0..9231925303f 100644 --- a/apps/supervisor/src/envUtil.test.ts +++ b/apps/supervisor/src/envUtil.test.ts @@ -1,5 +1,11 @@ import { describe, it, expect } from "vitest"; -import { BoolEnv, AdditionalEnvVars, NodeLabelValue, Tolerations } from "./envUtil.js"; +import { + BoolEnv, + AdditionalEnvVars, + NodeLabelValue, + OrgPlacementOverrides, + Tolerations, +} from "./envUtil.js"; describe("BoolEnv", () => { it("should parse string 'true' as true", () => { @@ -203,3 +209,125 @@ describe("Tolerations", () => { expect(Tolerations.safeParse("dedicated=runs:NoSchedule:NoExecute").success).toBe(false); }); }); + +describe("OrgPlacementOverrides", () => { + it("should parse a full override with nodeSelector and tolerations", () => { + expect( + OrgPlacementOverrides.parse( + JSON.stringify({ + org_123: { + nodeSelector: { "node.cluster.x-k8s.io/machinepool": "dedicated-pool" }, + tolerations: "dedicated=pool:NoSchedule", + }, + }) + ) + ).toEqual({ + org_123: { + nodeSelector: { "node.cluster.x-k8s.io/machinepool": "dedicated-pool" }, + tolerations: [{ key: "dedicated", operator: "Equal", value: "pool", effect: "NoSchedule" }], + }, + }); + }); + + it("should allow either half to be omitted", () => { + expect( + OrgPlacementOverrides.parse(JSON.stringify({ org_123: { nodeSelector: { pool: "a" } } })) + ).toEqual({ org_123: { nodeSelector: { pool: "a" } } }); + + expect( + OrgPlacementOverrides.parse(JSON.stringify({ org_123: { tolerations: "spot:NoExecute" } })) + ).toEqual({ + org_123: { tolerations: [{ key: "spot", operator: "Exists", effect: "NoExecute" }] }, + }); + + expect(OrgPlacementOverrides.parse(JSON.stringify({ org_123: {} }))).toEqual({ org_123: {} }); + }); + + it("should reject invalid JSON at startup rather than silently skipping the override", () => { + for (const invalid of ["not json", "[]", '"org_123"', "{"]) { + expect(OrgPlacementOverrides.safeParse(invalid).success).toBe(false); + } + }); + + it("should treat a blank or missing value as no overrides, like the sibling settings", () => { + expect(OrgPlacementOverrides.parse(undefined)).toBeUndefined(); + expect(OrgPlacementOverrides.parse("")).toBeUndefined(); + expect(OrgPlacementOverrides.parse(" ")).toBeUndefined(); + }); + + it("should accept tolerations as an array of entries, matching the Helm list shape", () => { + expect( + OrgPlacementOverrides.parse( + JSON.stringify({ + org_123: { tolerations: ["dedicated=pool:NoSchedule", "spot:NoExecute"] }, + }) + ) + ).toEqual({ + org_123: { + tolerations: [ + { key: "dedicated", operator: "Equal", value: "pool", effect: "NoSchedule" }, + { key: "spot", operator: "Exists", effect: "NoExecute" }, + ], + }, + }); + }); + + it("should coerce scalar node selector values to strings, as Kubernetes labels are", () => { + expect( + OrgPlacementOverrides.parse( + JSON.stringify({ org_123: { nodeSelector: { paid: true, replicas: 3 } } }) + ) + ).toEqual({ org_123: { nodeSelector: { paid: "true", replicas: "3" } } }); + }); + + it("should trim whitespace around node selector keys and values", () => { + expect( + OrgPlacementOverrides.parse( + JSON.stringify({ org_123: { nodeSelector: { " pool ": " a " } } }) + ) + ).toEqual({ org_123: { nodeSelector: { pool: "a" } } }); + }); + + it("should reject blank or padded org keys, since the lookup is exact", () => { + for (const key of [" ", " org_123", "org_123 "]) { + expect(OrgPlacementOverrides.safeParse(JSON.stringify({ [key]: {} })).success).toBe(false); + } + }); + + it("should reject an empty node selector value instead of pinning the org to nothing", () => { + for (const value of ["", " "]) { + expect( + OrgPlacementOverrides.safeParse( + JSON.stringify({ org_123: { nodeSelector: { pool: value } } }) + ).success + ).toBe(false); + } + }); + + it("should reject an unknown field, so a typo cannot silently drop an override", () => { + expect( + OrgPlacementOverrides.safeParse( + JSON.stringify({ org_123: { toleration: "dedicated=pool:NoSchedule" } }) + ).success + ).toBe(false); + }); + + it("should reject a node selector key or value Kubernetes would reject", () => { + for (const invalid of [ + { org_123: { nodeSelector: { "bad key": "a" } } }, + { org_123: { nodeSelector: { pool: "bad value" } } }, + { org_123: { nodeSelector: { "a/b/c": "a" } } }, + { org_123: { nodeSelector: { pool: "v".repeat(64) } } }, + ]) { + expect(OrgPlacementOverrides.safeParse(JSON.stringify(invalid)).success).toBe(false); + } + }); + + it("should reject an invalid toleration inside an override", () => { + expect( + OrgPlacementOverrides.safeParse( + JSON.stringify({ org_123: { tolerations: "dedicated=pool:Nope" } }) + ).success + ).toBe(false); + }); +}); diff --git a/apps/supervisor/src/envUtil.ts b/apps/supervisor/src/envUtil.ts index 67811f76fcb..e905141a43e 100644 --- a/apps/supervisor/src/envUtil.ts +++ b/apps/supervisor/src/envUtil.ts @@ -146,6 +146,102 @@ export const Tolerations = z.string().transform((val, ctx) => { }); }); +/** + * Scalar values are coerced: YAML/JSON easily produce `true` or `3` where a label + * value is meant, and Kubernetes label values are always strings. An empty value + * is rejected rather than passed through - as a selector it matches only nodes + * carrying a literal empty-valued label, which pins the org to nothing. + */ +const NodeSelector = z + .record(z.string(), z.union([z.string(), z.number(), z.boolean()])) + .transform((selector, ctx) => { + const result: Record = {}; + + for (const [rawKey, rawValue] of Object.entries(selector)) { + const key = rawKey.trim(); + const value = String(rawValue).trim(); + + if (!isQualifiedName(key)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: `Invalid node selector key "${rawKey}". Must be a Kubernetes label key, optionally prefixed with a DNS subdomain.`, + }); + continue; + } + + if (!value) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: `Empty node selector value for key "${key}". Remove the key instead of blanking the value.`, + }); + continue; + } + + if (!isLabelValue(value)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: `Invalid node selector value "${value}" for key "${key}". Must be a Kubernetes label value: alphanumeric, with dashes, underscores and dots inside, at most 63 characters.`, + }); + continue; + } + + result[key] = value; + } + + return result; + }); + +/** + * Per-organization placement overrides for run pods, as JSON keyed by the + * internal org id (the `org` label on run pods): + * `{"": {"nodeSelector": {"": ""}, "tolerations": ""}}`. + * Tolerations use the same CSV format as `Tolerations`, or an array of such + * entries. Everything is validated at startup for the same reason as + * tolerations above: a typo would otherwise reject every pod create for that + * org, with the cause buried in API errors. A blank value means no overrides. + */ +export const OrgPlacementOverrides = z + .string() + .optional() + .transform((val, ctx) => { + if (val === undefined || val.trim() === "") { + return undefined; + } + + try { + return JSON.parse(val) as unknown; + } catch { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Invalid org placement overrides: not valid JSON", + }); + return z.NEVER; + } + }) + .pipe( + z + .record( + z + .string() + .min(1) + .refine((key) => key === key.trim() && key.trim().length > 0, { + message: + "Org override keys must not be blank or padded with whitespace; the lookup is exact", + }), + z + .object({ + nodeSelector: NodeSelector.optional(), + tolerations: z + .union([z.string(), z.array(z.string())]) + .transform((val) => (Array.isArray(val) ? val.join(",") : val)) + .pipe(Tolerations) + .optional(), + }) + .strict() + ) + .optional() + ); + export const AdditionalEnvVars = z.preprocess((val) => { if (typeof val !== "string") { return val; diff --git a/apps/supervisor/src/workloadManager/kubernetes.test.ts b/apps/supervisor/src/workloadManager/kubernetes.test.ts index bb15c23e9f4..3c6419e9c6f 100644 --- a/apps/supervisor/src/workloadManager/kubernetes.test.ts +++ b/apps/supervisor/src/workloadManager/kubernetes.test.ts @@ -4,6 +4,7 @@ import { nodetypeNodeSelector, runPodTolerations, withBlockIoUringSeccompProfile, + withNodeSelector, } from "./kubernetesPodSpec.js"; const basePodSpec = { @@ -54,6 +55,49 @@ describe("runPodTolerations", () => { expect(runPodTolerations(worker, [], true)).toEqual(worker); expect(runPodTolerations(worker, scheduled, true)).toEqual([...worker, ...scheduled]); }); + + it("appends the org tolerations regardless of run type", () => { + const org = [{ key: "dedicated", operator: "Equal", value: "org-pool", effect: "NoSchedule" }]; + + expect(runPodTolerations(undefined, undefined, false, org)).toEqual(org); + expect(runPodTolerations(worker, undefined, false, org)).toEqual([...worker, ...org]); + expect(runPodTolerations(worker, scheduled, true, org)).toEqual([ + ...worker, + ...scheduled, + ...org, + ]); + expect(runPodTolerations(undefined, undefined, false, [])).toBeUndefined(); + }); +}); + +describe("withNodeSelector", () => { + const podSpec = { ...basePodSpec, nodeSelector: { nodetype: "v4-worker", paid: "true" } }; + + it("returns the pod spec untouched when there is nothing to merge", () => { + expect(withNodeSelector(podSpec, undefined)).toBe(podSpec); + expect(withNodeSelector(podSpec, {})).toBe(podSpec); + }); + + it("merges extra entries with existing ones", () => { + expect(withNodeSelector(podSpec, { machinepool: "dedicated-pool" })).toEqual({ + ...podSpec, + nodeSelector: { nodetype: "v4-worker", paid: "true", machinepool: "dedicated-pool" }, + }); + }); + + it("lets the extra entries win on key collision", () => { + expect(withNodeSelector(podSpec, { nodetype: "other" }).nodeSelector).toEqual({ + nodetype: "other", + paid: "true", + }); + }); + + it("adds a nodeSelector to a spec that had none", () => { + expect(withNodeSelector(basePodSpec, { machinepool: "dedicated-pool" })).toEqual({ + ...basePodSpec, + nodeSelector: { machinepool: "dedicated-pool" }, + }); + }); }); describe("withBlockIoUringSeccompProfile", () => { diff --git a/apps/supervisor/src/workloadManager/kubernetes.ts b/apps/supervisor/src/workloadManager/kubernetes.ts index 1b88bafbc28..e0bf01a050f 100644 --- a/apps/supervisor/src/workloadManager/kubernetes.ts +++ b/apps/supervisor/src/workloadManager/kubernetes.ts @@ -18,6 +18,7 @@ import { nodetypeNodeSelector, runPodTolerations, withBlockIoUringSeccompProfile, + withNodeSelector, } from "./kubernetesPodSpec.js"; type ResourceQuantities = { @@ -69,6 +70,12 @@ export class KubernetesWorkloadManager implements WorkloadManager { domain: opts.workloadApiDomain, }); } + + if (env.KUBERNETES_ORG_PLACEMENT_OVERRIDES) { + this.logger.info("[KubernetesWorkloadManager] Org placement overrides enabled", { + orgIds: Object.keys(env.KUBERNETES_ORG_PLACEMENT_OVERRIDES), + }); + } } private addPlacementTags( @@ -110,7 +117,24 @@ export class KubernetesWorkloadManager implements WorkloadManager { const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber); try { - const basePodSpec = this.addPlacementTags(this.#defaultPodSpec, opts.placementTags); + const orgOverride = env.KUBERNETES_ORG_PLACEMENT_OVERRIDES?.[opts.orgId]; + const taggedPodSpec = this.addPlacementTags(this.#defaultPodSpec, opts.placementTags); + const basePodSpec = withNodeSelector(taggedPodSpec, orgOverride?.nodeSelector); + + if (orgOverride?.nodeSelector) { + const replacedKeys = Object.keys(orgOverride.nodeSelector).filter( + (key) => + taggedPodSpec.nodeSelector?.[key] !== undefined && + taggedPodSpec.nodeSelector[key] !== orgOverride.nodeSelector?.[key] + ); + + if (replacedKeys.length > 0) { + this.logger.warn( + "[KubernetesWorkloadManager] Org placement override replaces node selector keys", + { orgId: opts.orgId, replacedKeys } + ); + } + } const podSpec = this.opts.checkpointsEnabled ? withBlockIoUringSeccompProfile(basePodSpec, opts.runtime) : basePodSpec; @@ -131,7 +155,7 @@ export class KubernetesWorkloadManager implements WorkloadManager { spec: { ...podSpec, affinity: this.#getAffinity(opts), - tolerations: this.#getTolerations(this.#isScheduledRun(opts)), + tolerations: this.#getTolerations(this.#isScheduledRun(opts), orgOverride?.tolerations), terminationGracePeriodSeconds: 60 * 60, containers: [ { @@ -555,11 +579,15 @@ export class KubernetesWorkloadManager implements WorkloadManager { }; } - #getTolerations(isScheduledRun: boolean): k8s.V1Toleration[] | undefined { + #getTolerations( + isScheduledRun: boolean, + orgTolerations?: k8s.V1Toleration[] + ): k8s.V1Toleration[] | undefined { return runPodTolerations( env.KUBERNETES_RUNNER_TOLERATIONS, env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS, - isScheduledRun + isScheduledRun, + orgTolerations ); } diff --git a/apps/supervisor/src/workloadManager/kubernetesPodSpec.ts b/apps/supervisor/src/workloadManager/kubernetesPodSpec.ts index ba15e563f6d..f521369f082 100644 --- a/apps/supervisor/src/workloadManager/kubernetesPodSpec.ts +++ b/apps/supervisor/src/workloadManager/kubernetesPodSpec.ts @@ -19,23 +19,47 @@ export function nodetypeNodeSelector( /** * Tolerations for a run pod: the cluster-wide set, plus the scheduled-run set when the - * run came from a schedule tree. Not reconciled - Kubernetes matches tolerations as an - * any-match set, so a broad entry in one set can subsume a narrower one in the other. + * run came from a schedule tree, plus the org's own set when a placement override + * matches. Not reconciled - Kubernetes matches tolerations as an any-match set, so a + * broad entry in one set can subsume a narrower one in another. * Returns undefined rather than an empty array to leave the field unset. */ export function runPodTolerations( runnerTolerations: k8s.V1Toleration[] | undefined, scheduledRunTolerations: k8s.V1Toleration[] | undefined, - isScheduledRun: boolean + isScheduledRun: boolean, + orgTolerations?: k8s.V1Toleration[] ): k8s.V1Toleration[] | undefined { const tolerations = [ ...(runnerTolerations ?? []), ...(isScheduledRun ? (scheduledRunTolerations ?? []) : []), + ...(orgTolerations ?? []), ]; return tolerations.length > 0 ? tolerations : undefined; } +/** + * Merges extra node selector entries into a pod spec. Later entries win on key + * collision, so an override can retarget a key set by an earlier stage. + */ +export function withNodeSelector( + podSpec: Omit, + nodeSelector: Record | undefined +): Omit { + if (!nodeSelector || Object.keys(nodeSelector).length === 0) { + return podSpec; + } + + return { + ...podSpec, + nodeSelector: { + ...podSpec.nodeSelector, + ...nodeSelector, + }, + }; +} + /** * Node >= 24 always creates io_uring fds, which can't be checkpointed. Blocking * io_uring_setup makes libuv fall back to epoll. Other runtimes don't need this, diff --git a/docs/self-hosting/env/supervisor.mdx b/docs/self-hosting/env/supervisor.mdx index a7e4ef96692..92cfe100d65 100644 --- a/docs/self-hosting/env/supervisor.mdx +++ b/docs/self-hosting/env/supervisor.mdx @@ -48,6 +48,7 @@ mode: "wide" | `KUBERNETES_NAMESPACE` | No | default | The namespace that runs should be in. | | `KUBERNETES_WORKER_NODETYPE_LABEL` | No | v4-worker | Nodes for runs need `nodetype=`. Empty: any node. | | `KUBERNETES_RUNNER_TOLERATIONS` | No | — | Run pod tolerations. CSV: `key=value:effect`/`key:effect`. | +| `KUBERNETES_ORG_PLACEMENT_OVERRIDES` | No | — | Per-org run pod placement. JSON keyed by internal org ID. | | `KUBERNETES_IMAGE_PULL_SECRETS` | No | — | Image pull secrets (CSV). | | `KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT` | No | 10Gi | Ephemeral storage size limit. Applies to all runs. | | `KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST` | No | 2Gi | Ephemeral storage size request. Applies to all runs. | diff --git a/hosting/k8s/helm/templates/supervisor.yaml b/hosting/k8s/helm/templates/supervisor.yaml index 84a6350f035..59b797c53af 100644 --- a/hosting/k8s/helm/templates/supervisor.yaml +++ b/hosting/k8s/helm/templates/supervisor.yaml @@ -174,6 +174,10 @@ spec: - name: KUBERNETES_RUNNER_TOLERATIONS value: {{ join "," . | quote }} {{- end }} + {{- with .Values.supervisor.config.kubernetes.orgPlacementOverrides }} + - name: KUBERNETES_ORG_PLACEMENT_OVERRIDES + value: {{ toJson . | quote }} + {{- end }} {{- $registryAuthEnabled := false }} {{- if .Values.registry.deploy }} {{- $registryAuthEnabled = .Values.registry.auth.enabled }} diff --git a/hosting/k8s/helm/values.yaml b/hosting/k8s/helm/values.yaml index 354d8e55ba1..8e7229bf9a3 100644 --- a/hosting/k8s/helm/values.yaml +++ b/hosting/k8s/helm/values.yaml @@ -297,6 +297,9 @@ supervisor: namespace: "" # Default: uses release namespace workerNodetypeLabel: "" # When set, runs will only be scheduled on nodes with "nodetype=