feat(supervisor): per-org placement overrides for run pods

KUBERNETES_ORG_PLACEMENT_OVERRIDES takes JSON keyed by org id, adding
node selector entries and tolerations to that org's run pods so an
operator can pin an org onto a dedicated node pool. Validated at
startup so a typo fails fast instead of rejecting every pod create.

Author: myftija

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,10 @@ 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 org id:
+    // {"<orgId>": {"nodeSelector": {"<key>": "<value>"}, "tolerations": "<csv>"}}
+    KUBERNETES_ORG_PLACEMENT_OVERRIDES: OrgPlacementOverrides.optional(),
+
     // Placement tags settings
     PLACEMENT_TAGS_ENABLED: BoolEnv.default(false),
     PLACEMENT_TAGS_PREFIX: z.string().default("node.cluster.x-k8s.io"),

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,72 @@ 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 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);
+  });
+});

apps/supervisor/src/envUtil.ts

@@ -146,6 +146,56 @@ export const Tolerations = z.string().transform((val, ctx) => {
     });
 });
 
+const NodeSelector = z.record(z.string(), z.string()).superRefine((selector, ctx) => {
+  for (const [key, value] of Object.entries(selector)) {
+    if (!isQualifiedName(key)) {
+      ctx.addIssue({
+        code: z.ZodIssueCode.custom,
+        message: `Invalid node selector key "${key}". Must be a Kubernetes label key, optionally prefixed with a DNS subdomain.`,
+      });
+    }
+
+    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.`,
+      });
+    }
+  }
+});
+
+/**
+ * Per-organization placement overrides for run pods, as JSON keyed by org id:
+ * `{"<orgId>": {"nodeSelector": {"<key>": "<value>"}, "tolerations": "<csv>"}}`.
+ * Tolerations use the same CSV format as `Tolerations`. 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.
+ */
+export const OrgPlacementOverrides = z
+  .string()
+  .transform((val, ctx) => {
+    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),
+      z
+        .object({
+          nodeSelector: NodeSelector.optional(),
+          tolerations: Tolerations.optional(),
+        })
+        .strict()
+    )
+  );
+
 export const AdditionalEnvVars = z.preprocess((val) => {
   if (typeof val !== "string") {
     return val;

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", () => {

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -18,6 +18,7 @@ import {
   nodetypeNodeSelector,
   runPodTolerations,
   withBlockIoUringSeccompProfile,
+  withNodeSelector,
 } from "./kubernetesPodSpec.js";
 
 type ResourceQuantities = {
@@ -110,7 +111,11 @@ 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 basePodSpec = withNodeSelector(
+        this.addPlacementTags(this.#defaultPodSpec, opts.placementTags),
+        orgOverride?.nodeSelector
+      );
       const podSpec = this.opts.checkpointsEnabled
         ? withBlockIoUringSeccompProfile(basePodSpec, opts.runtime)
         : basePodSpec;
@@ -131,7 +136,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 +560,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
     );
   }
 

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<k8s.V1PodSpec, "containers">,
+  nodeSelector: Record<string, string> | undefined
+): Omit<k8s.V1PodSpec, "containers"> {
+  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,