feat(k8s): add placement tags for flexible node selection (#2390)

* add tier scheduling support to supervisor

* add billing info to dequeued message w/o cache

* add cache with best effort invalidation

* fix invalidate circular dep

* add changeset

* use new plan type on runs as fallback during dequeue

* tidy up

* be more explicit with plan type fallback

* remove additional billing check from hot path

* switch to placement tags

* update changeset

* update platform package

* start using new entitlement response

* ensure skipChecks optimization validates at batch level

* add optional items to add to queue manager limits

* make the bool env helper only accept boolean defaults

* remove redundant private field

* update placement tag helper to prevent unsupported tags

Author: nicktrn

.changeset/proud-nails-grin.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+Add optional placement tags to dequeued messages for targeted scheduling

apps/supervisor/src/env.ts

@@ -15,7 +15,7 @@ const Env = z.object({
   OTEL_EXPORTER_OTLP_ENDPOINT: z.string().url(), // set on the runners
 
   // Workload API settings (coordinator mode) - the workload API is what the run controller connects to
-  TRIGGER_WORKLOAD_API_ENABLED: BoolEnv.default("true"),
+  TRIGGER_WORKLOAD_API_ENABLED: BoolEnv.default(true),
   TRIGGER_WORKLOAD_API_PROTOCOL: z
     .string()
     .transform((s) => z.enum(["http", "https"]).parse(s.toLowerCase()))
@@ -32,7 +32,7 @@ const Env = z.object({
   RUNNER_PRETTY_LOGS: BoolEnv.default(false),
 
   // Dequeue settings (provider mode)
-  TRIGGER_DEQUEUE_ENABLED: BoolEnv.default("true"),
+  TRIGGER_DEQUEUE_ENABLED: BoolEnv.default(true),
   TRIGGER_DEQUEUE_INTERVAL_MS: z.coerce.number().int().default(250),
   TRIGGER_DEQUEUE_IDLE_INTERVAL_MS: z.coerce.number().int().default(1000),
   TRIGGER_DEQUEUE_MAX_RUN_COUNT: z.coerce.number().int().default(10),
@@ -77,6 +77,10 @@ const Env = z.object({
   KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
   KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
 
+  // Placement tags settings
+  PLACEMENT_TAGS_ENABLED: BoolEnv.default(false),
+  PLACEMENT_TAGS_PREFIX: z.string().default("node.cluster.x-k8s.io"),
+
   // Metrics
   METRICS_ENABLED: BoolEnv.default(true),
   METRICS_COLLECT_DEFAULTS: BoolEnv.default(true),

apps/supervisor/src/envUtil.ts

@@ -3,14 +3,19 @@ import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLog
 
 const logger = new SimpleStructuredLogger("env-util");
 
-export const BoolEnv = z.preprocess((val) => {
+const baseBoolEnv = z.preprocess((val) => {
   if (typeof val !== "string") {
     return val;
   }
 
   return ["true", "1"].includes(val.toLowerCase().trim());
 }, z.boolean());
 
+// Create a type-safe version that only accepts boolean defaults
+export const BoolEnv = baseBoolEnv as Omit<typeof baseBoolEnv, "default"> & {
+  default: (value: boolean) => z.ZodDefault<typeof baseBoolEnv>;
+};
+
 export const AdditionalEnvVars = z.preprocess((val) => {
   if (typeof val !== "string") {
     return val;

apps/supervisor/src/index.ts

@@ -247,6 +247,7 @@ class ManagedSupervisor {
           nextAttemptNumber: message.run.attemptNumber,
           snapshotId: message.snapshot.id,
           snapshotFriendlyId: message.snapshot.friendlyId,
+          placementTags: message.placementTags,
         });
 
         // Disabled for now

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -4,7 +4,7 @@ import {
   type WorkloadManagerCreateOptions,
   type WorkloadManagerOptions,
 } from "./types.js";
-import type { EnvironmentType, MachinePreset } from "@trigger.dev/core/v3";
+import type { EnvironmentType, MachinePreset, PlacementTag } from "@trigger.dev/core/v3";
 import { env } from "../env.js";
 import { type K8sApi, createK8sApi, type k8s } from "../clients/kubernetes.js";
 import { getRunnerId } from "../util.js";
@@ -13,6 +13,11 @@ type ResourceQuantities = {
   [K in "cpu" | "memory" | "ephemeral-storage"]?: string;
 };
 
+interface PlacementConfig {
+  enabled: boolean;
+  prefix: string;
+}
+
 export class KubernetesWorkloadManager implements WorkloadManager {
   private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
   private k8s: K8sApi;
@@ -28,6 +33,56 @@ export class KubernetesWorkloadManager implements WorkloadManager {
     }
   }
 
+  private get placementConfig(): PlacementConfig {
+    return {
+      enabled: env.PLACEMENT_TAGS_ENABLED,
+      prefix: env.PLACEMENT_TAGS_PREFIX,
+    };
+  }
+
+  private addPlacementTags(
+    podSpec: Omit<k8s.V1PodSpec, "containers">,
+    placementTags?: PlacementTag[]
+  ): Omit<k8s.V1PodSpec, "containers"> {
+    if (!this.placementConfig.enabled || !placementTags || placementTags.length === 0) {
+      return podSpec;
+    }
+
+    const nodeSelector: Record<string, string> = { ...podSpec.nodeSelector };
+
+    // Convert placement tags to nodeSelector labels
+    for (const tag of placementTags) {
+      const labelKey = `${this.placementConfig.prefix}/${tag.key}`;
+
+      // Print warnings (if any)
+      this.printTagWarnings(tag);
+
+      // For now we only support single values via nodeSelector
+      nodeSelector[labelKey] = tag.values?.[0] ?? "";
+    }
+
+    return {
+      ...podSpec,
+      nodeSelector,
+    };
+  }
+
+  private printTagWarnings(tag: PlacementTag) {
+    if (!tag.values || tag.values.length === 0) {
+      // No values provided
+      this.logger.warn(
+        "[KubernetesWorkloadManager] Placement tag has no values, using empty string",
+        tag
+      );
+    } else if (tag.values.length > 1) {
+      // Multiple values provided
+      this.logger.warn(
+        "[KubernetesWorkloadManager] Placement tag has multiple values, only using first one",
+        tag
+      );
+    }
+  }
+
   async create(opts: WorkloadManagerCreateOptions) {
     this.logger.log("[KubernetesWorkloadManager] Creating container", { opts });
 
@@ -48,7 +103,7 @@ export class KubernetesWorkloadManager implements WorkloadManager {
             },
           },
           spec: {
-            ...this.#defaultPodSpec,
+            ...this.addPlacementTags(this.#defaultPodSpec, opts.placementTags),
             terminationGracePeriodSeconds: 60 * 60,
             containers: [
               {

apps/supervisor/src/workloadManager/types.ts

@@ -1,4 +1,4 @@
-import { type EnvironmentType, type MachinePreset } from "@trigger.dev/core/v3";
+import type { EnvironmentType, MachinePreset, PlacementTag } from "@trigger.dev/core/v3";
 
 export interface WorkloadManagerOptions {
   workloadApiProtocol: "http" | "https";
@@ -23,6 +23,7 @@ export interface WorkloadManagerCreateOptions {
   version: string;
   nextAttemptNumber?: number;
   dequeuedAt: Date;
+  placementTags?: PlacementTag[];
   // identifiers
   envId: string;
   envType: EnvironmentType;

apps/webapp/app/env.server.ts

@@ -761,6 +761,8 @@ const EnvironmentSchema = z.object({
     .int()
     .default(60_000 * 5), // 5 minutes
 
+  BATCH_TRIGGER_CACHED_RUNS_CHECK_ENABLED: BoolEnv.default(false),
+
   BATCH_TRIGGER_WORKER_ENABLED: z.string().default(process.env.WORKER_ENABLED ?? "true"),
   BATCH_TRIGGER_WORKER_CONCURRENCY_WORKERS: z.coerce.number().int().default(2),
   BATCH_TRIGGER_WORKER_CONCURRENCY_TASKS_PER_WORKER: z.coerce.number().int().default(10),

apps/webapp/app/routes/resources.orgs.$organizationSlug.select-plan.tsx

@@ -42,6 +42,7 @@ import { redirectWithErrorMessage } from "~/models/message.server";
 import { logger } from "~/services/logger.server";
 import { setPlan } from "~/services/platform.v3.server";
 import { requireUser } from "~/services/session.server";
+import { engine } from "~/v3/runEngine.server";
 import { cn } from "~/utils/cn";
 import { sendToPlain } from "~/utils/plain.server";
 
@@ -152,7 +153,9 @@ export async function action({ request, params }: ActionFunctionArgs) {
     }
   }
 
-  return setPlan(organization, request, form.callerPath, payload);
+  return setPlan(organization, request, form.callerPath, payload, {
+    invalidateBillingCache: engine.invalidateBillingCache.bind(engine),
+  });
 }
 
 const pricingDefinitions = {

apps/webapp/app/runEngine/concerns/queues.server.ts

@@ -177,8 +177,11 @@ export class DefaultQueueManager implements QueueManager {
     return task.queue.name ?? defaultQueueName;
   }
 
-  async validateQueueLimits(environment: AuthenticatedEnvironment): Promise<QueueValidationResult> {
-    const queueSizeGuard = await guardQueueSizeLimitsForEnv(this.engine, environment);
+  async validateQueueLimits(
+    environment: AuthenticatedEnvironment,
+    itemsToAdd?: number
+  ): Promise<QueueValidationResult> {
+    const queueSizeGuard = await guardQueueSizeLimitsForEnv(this.engine, environment, itemsToAdd);
 
     logger.debug("Queue size guard result", {
       queueSizeGuard,