add tier scheduling support to supervisor

Author: nicktrn

apps/supervisor/src/env.ts

@@ -77,6 +77,12 @@ 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"),
+
   // Metrics
   METRICS_ENABLED: BoolEnv.default(true),
   METRICS_COLLECT_DEFAULTS: BoolEnv.default(true),

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -13,13 +13,22 @@ type ResourceQuantities = {
   [K in "cpu" | "memory" | "ephemeral-storage"]?: string;
 };
 
+interface TierConfig {
+  enabled: boolean;
+  labelKey: string;
+  freeValue: string;
+  paidValue: 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;
 
   constructor(private opts: WorkloadManagerOptions) {
     this.k8s = createK8sApi();
+    this.tierConfig = this.tierSchedulingConfig;
 
     if (opts.workloadApiDomain) {
       this.logger.warn("[KubernetesWorkloadManager] ⚠️ Custom workload API domain", {
@@ -28,6 +37,34 @@ export class KubernetesWorkloadManager implements WorkloadManager {
     }
   }
 
+  private get tierSchedulingConfig(): TierConfig {
+    return {
+      enabled: env.ENABLE_TIER_SCHEDULING,
+      labelKey: env.TIER_LABEL_KEY,
+      freeValue: env.TIER_LABEL_VALUE_FREE,
+      paidValue: env.TIER_LABEL_VALUE_PAID,
+    };
+  }
+
+  private addTierScheduling(
+    podSpec: Omit<k8s.V1PodSpec, "containers">,
+    isPaidTier: boolean
+  ): Omit<k8s.V1PodSpec, "containers"> {
+    if (!this.tierConfig.enabled) {
+      return podSpec;
+    }
+
+    const labelValue = isPaidTier ? this.tierConfig.paidValue : this.tierConfig.freeValue;
+
+    return {
+      ...podSpec,
+      nodeSelector: {
+        ...podSpec.nodeSelector,
+        [this.tierConfig.labelKey]: labelValue,
+      },
+    };
+  }
+
   async create(opts: WorkloadManagerCreateOptions) {
     this.logger.log("[KubernetesWorkloadManager] Creating container", { opts });
 
@@ -48,7 +85,7 @@ export class KubernetesWorkloadManager implements WorkloadManager {
             },
           },
           spec: {
-            ...this.#defaultPodSpec,
+            ...this.addTierScheduling(this.#defaultPodSpec, opts.isPaidTier ?? false),
             terminationGracePeriodSeconds: 60 * 60,
             containers: [
               {

apps/supervisor/src/workloadManager/types.ts

@@ -32,4 +32,6 @@ export interface WorkloadManagerCreateOptions {
   runFriendlyId: string;
   snapshotId: string;
   snapshotFriendlyId: string;
+  // tier scheduling
+  isPaidTier?: boolean;
 }