switch to placement tags

Author: nicktrn

apps/supervisor/src/env.ts

@@ -77,11 +77,9 @@ const Env = z.object({
   KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
   KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
 
-  // Tier scheduling settings
-  ENABLE_TIER_SCHEDULING: BoolEnv.default(false),
-  TIER_LABEL_KEY: z.string().default("node.cluster.x-k8s.io/paid"),
-  TIER_LABEL_VALUE_FREE: z.string().default("false"),
-  TIER_LABEL_VALUE_PAID: z.string().default("true"),
+  // 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),

apps/supervisor/src/index.ts

@@ -247,7 +247,7 @@ class ManagedSupervisor {
           nextAttemptNumber: message.run.attemptNumber,
           snapshotId: message.snapshot.id,
           snapshotFriendlyId: message.snapshot.friendlyId,
-          isPaidTier: message.billing?.currentPlan.isPaying ?? false,
+          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,22 +13,20 @@ type ResourceQuantities = {
   [K in "cpu" | "memory" | "ephemeral-storage"]?: string;
 };
 
-interface TierConfig {
+interface PlacementConfig {
   enabled: boolean;
-  labelKey: string;
-  freeValue: string;
-  paidValue: string;
+  prefix: string;
 }
 
 export class KubernetesWorkloadManager implements WorkloadManager {
   private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
   private k8s: K8sApi;
   private namespace = env.KUBERNETES_NAMESPACE;
-  private tierConfig: TierConfig;
+  private placementConfig: PlacementConfig;
 
   constructor(private opts: WorkloadManagerOptions) {
     this.k8s = createK8sApi();
-    this.tierConfig = this.tierSchedulingConfig;
+    this.placementConfig = this.placementTagsConfig;
 
     if (opts.workloadApiDomain) {
       this.logger.warn("[KubernetesWorkloadManager] ⚠️ Custom workload API domain", {
@@ -37,34 +35,56 @@ export class KubernetesWorkloadManager implements WorkloadManager {
     }
   }
 
-  private get tierSchedulingConfig(): TierConfig {
+  private get placementTagsConfig(): PlacementConfig {
     return {
-      enabled: env.ENABLE_TIER_SCHEDULING,
-      labelKey: env.TIER_LABEL_KEY,
-      freeValue: env.TIER_LABEL_VALUE_FREE,
-      paidValue: env.TIER_LABEL_VALUE_PAID,
+      enabled: env.PLACEMENT_TAGS_ENABLED,
+      prefix: env.PLACEMENT_TAGS_PREFIX,
     };
   }
 
-  private addTierScheduling(
+  private addPlacementTags(
     podSpec: Omit<k8s.V1PodSpec, "containers">,
-    isPaidTier: boolean
+    placementTags?: PlacementTag[]
   ): Omit<k8s.V1PodSpec, "containers"> {
-    if (!this.tierConfig.enabled) {
+    if (!this.placementConfig.enabled || !placementTags || placementTags.length === 0) {
       return podSpec;
     }
 
-    const labelValue = isPaidTier ? this.tierConfig.paidValue : this.tierConfig.freeValue;
+    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: {
-        ...podSpec.nodeSelector,
-        [this.tierConfig.labelKey]: labelValue,
-      },
+      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 });
 
@@ -85,7 +105,7 @@ export class KubernetesWorkloadManager implements WorkloadManager {
             },
           },
           spec: {
-            ...this.addTierScheduling(this.#defaultPodSpec, opts.isPaidTier ?? false),
+            ...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;
@@ -32,6 +33,4 @@ export interface WorkloadManagerCreateOptions {
   runFriendlyId: string;
   snapshotId: string;
   snapshotFriendlyId: string;
-  // tier scheduling
-  isPaidTier?: boolean;
 }

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -1,7 +1,7 @@
 import type { BillingCache } from "../billingCache.js";
 import { startSpan } from "@internal/tracing";
 import { assertExhaustive } from "@trigger.dev/core";
-import { DequeuedMessage, RetryOptions } from "@trigger.dev/core/v3";
+import { DequeuedMessage, RetryOptions, placementTag } from "@trigger.dev/core/v3";
 import { getMaxDuration } from "@trigger.dev/core/v3/isomorphic";
 import { PrismaClientOrTransaction } from "@trigger.dev/database";
 import { getRunWithBackgroundWorkerTasks } from "../db/worker.js";
@@ -474,11 +474,7 @@ export class DequeueSystem {
                 project: {
                   id: lockedTaskRun.projectId,
                 },
-                billing: {
-                  currentPlan: {
-                    isPaying,
-                  },
-                },
+                placementTags: [placementTag("paid", isPaying ? "true" : "false")],
               } satisfies DequeuedMessage;
             }
           );

packages/core/src/v3/schemas/runEngine.ts

@@ -224,6 +224,17 @@ export const DequeueMessageCheckpoint = z.object({
 });
 export type DequeueMessageCheckpoint = z.infer<typeof DequeueMessageCheckpoint>;
 
+export const PlacementTag = z.object({
+  key: z.string(),
+  values: z.array(z.string()).optional(),
+});
+export type PlacementTag = z.infer<typeof PlacementTag>;
+
+/** Helper functions for placement tags. Accepts a single value or an array of values. */
+export function placementTag(key: string, valueOrValues: string | string[]): PlacementTag {
+  return { key, values: Array.isArray(valueOrValues) ? valueOrValues : [valueOrValues] };
+}
+
 /** This is sent to a Worker when a run is dequeued (a new run or continuing run) */
 export const DequeuedMessage = z.object({
   version: z.literal("1"),
@@ -261,12 +272,6 @@ export const DequeuedMessage = z.object({
   project: z.object({
     id: z.string(),
   }),
-  billing: z
-    .object({
-      currentPlan: z.object({
-        isPaying: z.boolean(),
-      }),
-    })
-    .optional(),
+  placementTags: z.array(PlacementTag).optional(),
 });
 export type DequeuedMessage = z.infer<typeof DequeuedMessage>;