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

The supervisor now supports routing an organization's runs to specific
nodes. `KUBERNETES_ORG_PLACEMENT_OVERRIDES` takes JSON keyed by the
internal org ID, adding node selector entries and tolerations to that
org's run pods, e.g. to route an org onto a dedicated, tainted node
pool:

```json
{"<orgId>": {"nodeSelector": {"pool": "dedicated"}, "tolerations": "dedicated=runs:NoSchedule"}}
```

The node selector merges over the defaults (the override wins on key
collision, with a warning logged). Tolerations append to the existing
runner and scheduled-run sets. Overrides are validated at startup
similar to `KUBERNETES_RUNNER_TOLERATIONS`.

Exposed in the Helm chart as
`supervisor.config.kubernetes.orgPlacementOverrides`, where tolerations
can also be given as a list.

Author: myftija

.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.

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):
+    // {"<orgId>": {"nodeSelector": {"<key>": "<value>"}, "tolerations": "<csv or array>"}}
+    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,

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

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<string, string> = {};
+
+    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):
+ * `{"<orgId>": {"nodeSelector": {"<key>": "<value>"}, "tolerations": "<csv>"}}`.
+ * 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;

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