feat(supervisor): custom tolerations for scheduled runs (#3297)

Adds support for taint tolerations for scheduled runs. Useful for
selectively tolerating taints on dedicated node pools.

The new `KUBERNETES_SCHEDULED_RUN_TOLERATIONS` env variable accepts a
comma-separated list in the format key=value:effect (or key:effect for
the Exists operator).

Drive-by: renames all `KUBERNETES_SCHEDULE_*` affinity env vars to
KUBERNETES_SCHEDULED_RUN_* for clarity — this feature isn't used in
production yet or published in a tagged image; the name change is fine.
This commit is contained in:
Saadi Myftija
2026-03-30 13:04:13 +02:00
committed by GitHub
parent 66b70101e9
commit 7210bdee9b
2 changed files with 102 additions and 14 deletions
+86 -7
View File
@@ -123,8 +123,16 @@ const Env = z
KUBERNETES_SCHEDULER_NAME: z.string().optional(), // Custom scheduler name for pods
// Large machine affinity settings - large-* presets prefer a dedicated pool
KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED: BoolEnv.default(false),
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY: z.string().trim().min(1).default("node.cluster.x-k8s.io/machinepool"),
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE: z.string().trim().min(1).default("large-machines"),
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY: z
.string()
.trim()
.min(1)
.default("node.cluster.x-k8s.io/machinepool"),
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE: z
.string()
.trim()
.min(1)
.default("large-machines"),
KUBERNETES_LARGE_MACHINE_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(100),
// Project affinity settings - pods from the same project prefer the same node
@@ -137,11 +145,82 @@ const Env = z
.default("kubernetes.io/hostname"),
// Schedule affinity settings - runs from schedule trees prefer a dedicated pool
KUBERNETES_SCHEDULE_AFFINITY_ENABLED: BoolEnv.default(false),
KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_KEY: z.string().trim().min(1).default("node.cluster.x-k8s.io/machinepool"),
KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_VALUE: z.string().trim().min(1).default("scheduled-runs"),
KUBERNETES_SCHEDULE_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(80),
KUBERNETES_SCHEDULE_ANTI_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(20),
KUBERNETES_SCHEDULED_RUN_AFFINITY_ENABLED: BoolEnv.default(false),
KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY: z
.string()
.trim()
.min(1)
.default("node.cluster.x-k8s.io/machinepool"),
KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE: z
.string()
.trim()
.min(1)
.default("scheduled-runs"),
KUBERNETES_SCHEDULED_RUN_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(80),
KUBERNETES_SCHEDULED_RUN_ANTI_AFFINITY_WEIGHT: z.coerce
.number()
.int()
.min(1)
.max(100)
.default(20),
// Schedule toleration settings - scheduled runs tolerate taints on the dedicated pool
// Comma-separated list of tolerations in the format: key=value:effect
// For Exists operator (no value): key:effect
KUBERNETES_SCHEDULED_RUN_TOLERATIONS: z
.string()
.transform((val, ctx) => {
const tolerations = val
.split(",")
.map((entry) => entry.trim())
.filter((entry) => entry.length > 0)
.map((entry) => {
const colonIdx = entry.lastIndexOf(":");
if (colonIdx === -1) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
message: `Invalid toleration format (missing effect): "${entry}"`,
});
return z.NEVER;
}
const effect = entry.slice(colonIdx + 1);
const validEffects = ["NoSchedule", "NoExecute", "PreferNoSchedule"];
if (!validEffects.includes(effect)) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
message: `Invalid toleration effect "${effect}" in "${entry}". Must be one of: ${validEffects.join(", ")}`,
});
return z.NEVER;
}
const keyValue = entry.slice(0, colonIdx);
const eqIdx = keyValue.indexOf("=");
const key = eqIdx === -1 ? keyValue : keyValue.slice(0, eqIdx);
if (!key) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
message: `Invalid toleration format (empty key): "${entry}"`,
});
return z.NEVER;
}
if (eqIdx === -1) {
return { key, operator: "Exists" as const, effect };
}
return {
key,
operator: "Equal" as const,
value: keyValue.slice(eqIdx + 1),
effect,
};
});
return tolerations;
})
.optional(),
// Placement tags settings
PLACEMENT_TAGS_ENABLED: BoolEnv.default(false),
@@ -121,6 +121,7 @@ export class KubernetesWorkloadManager implements WorkloadManager {
spec: {
...this.addPlacementTags(this.#defaultPodSpec, opts.placementTags),
affinity: this.#getAffinity(opts),
tolerations: this.#getScheduleTolerations(this.#isScheduledRun(opts)),
terminationGracePeriodSeconds: 60 * 60,
containers: [
{
@@ -485,7 +486,7 @@ export class KubernetesWorkloadManager implements WorkloadManager {
}
#getScheduleNodeAffinityRules(isScheduledRun: boolean): k8s.V1NodeAffinity | undefined {
if (!env.KUBERNETES_SCHEDULE_AFFINITY_ENABLED || !env.KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_VALUE) {
if (!env.KUBERNETES_SCHEDULED_RUN_AFFINITY_ENABLED || !env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE) {
return undefined;
}
@@ -494,13 +495,13 @@ export class KubernetesWorkloadManager implements WorkloadManager {
return {
preferredDuringSchedulingIgnoredDuringExecution: [
{
weight: env.KUBERNETES_SCHEDULE_AFFINITY_WEIGHT,
weight: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_WEIGHT,
preference: {
matchExpressions: [
{
key: env.KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_KEY,
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
operator: "In",
values: [env.KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_VALUE],
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
},
],
},
@@ -513,13 +514,13 @@ export class KubernetesWorkloadManager implements WorkloadManager {
return {
preferredDuringSchedulingIgnoredDuringExecution: [
{
weight: env.KUBERNETES_SCHEDULE_ANTI_AFFINITY_WEIGHT,
weight: env.KUBERNETES_SCHEDULED_RUN_ANTI_AFFINITY_WEIGHT,
preference: {
matchExpressions: [
{
key: env.KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_KEY,
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
operator: "NotIn",
values: [env.KUBERNETES_SCHEDULE_AFFINITY_POOL_LABEL_VALUE],
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
},
],
},
@@ -528,6 +529,14 @@ export class KubernetesWorkloadManager implements WorkloadManager {
};
}
#getScheduleTolerations(isScheduledRun: boolean): k8s.V1Toleration[] | undefined {
if (!isScheduledRun || !env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS?.length) {
return undefined;
}
return env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS;
}
#getProjectPodAffinity(projectId: string): k8s.V1PodAffinity | undefined {
if (!env.KUBERNETES_PROJECT_AFFINITY_ENABLED) {
return undefined;