optional labels with sample rates

Author: nicktrn

apps/kubernetes-provider/src/index.ts

@@ -18,6 +18,7 @@ import { PodCleaner } from "./podCleaner";
 import { TaskMonitor } from "./taskMonitor";
 import { UptimeHeartbeat } from "./uptimeHeartbeat";
 import { assertExhaustive } from "@trigger.dev/core";
+import { CustomLabelHelper } from "./labelHelper";
 
 const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local";
 const NODE_NAME = process.env.NODE_NAME || "local";
@@ -73,6 +74,8 @@ class KubernetesTaskOperations implements TaskOperations {
     apps: k8s.AppsV1Api;
   };
 
+  #labelHelper = new CustomLabelHelper();
+
   constructor(opts: { namespace?: string } = {}) {
     if (opts.namespace) {
       this.#namespace.metadata.name = opts.namespace;
@@ -165,6 +168,7 @@ class KubernetesTaskOperations implements TaskOperations {
           name: containerName,
           namespace: this.#namespace.metadata.name,
           labels: {
+            ...this.#labelHelper.getAdditionalLabels("create"),
             ...this.#getSharedLabels(opts),
             app: "task-run",
             "app.kubernetes.io/part-of": "trigger-worker",
@@ -226,6 +230,7 @@ class KubernetesTaskOperations implements TaskOperations {
           name: `${this.#getRunContainerName(opts.runId)}-${opts.checkpointId.slice(-8)}`,
           namespace: this.#namespace.metadata.name,
           labels: {
+            ...this.#labelHelper.getAdditionalLabels("restore"),
             ...this.#getSharedLabels(opts),
             app: "task-run",
             "app.kubernetes.io/part-of": "trigger-worker",

apps/kubernetes-provider/src/labelHelper.ts

@@ -0,0 +1,160 @@
+import { assertExhaustive } from "@trigger.dev/core";
+
+const CREATE_LABEL_ENV_VAR_PREFIX = "DEPLOYMENT_LABEL_";
+const RESTORE_LABEL_ENV_VAR_PREFIX = "RESTORE_LABEL_";
+const LABEL_SAMPLE_RATE_POSTFIX = "_SAMPLE_RATE";
+
+type OperationType = "create" | "restore";
+
+type CustomLabel = {
+  key: string;
+  value: string;
+  sampleRate: number;
+};
+
+export class CustomLabelHelper {
+  // Labels and sample rates are defined in environment variables so only need to be computed once
+  private createLabels?: CustomLabel[];
+  private restoreLabels?: CustomLabel[];
+
+  private getLabelPrefix(type: OperationType) {
+    const prefix = type === "create" ? CREATE_LABEL_ENV_VAR_PREFIX : RESTORE_LABEL_ENV_VAR_PREFIX;
+    return prefix.toLowerCase();
+  }
+
+  private getLabelSampleRatePostfix() {
+    return LABEL_SAMPLE_RATE_POSTFIX.toLowerCase();
+  }
+
+  // Can only range from 0 to 1
+  private fractionFromPercent(percent: number) {
+    return Math.min(1, Math.max(0, percent / 100));
+  }
+
+  private isLabelSampleRateEnvVar(key: string) {
+    return key.toLowerCase().endsWith(this.getLabelSampleRatePostfix());
+  }
+
+  private isLabelEnvVar(type: OperationType, key: string) {
+    const prefix = this.getLabelPrefix(type);
+    return key.toLowerCase().startsWith(prefix) && !this.isLabelSampleRateEnvVar(key);
+  }
+
+  private getSampleRateEnvVarKey(type: OperationType, envKey: string) {
+    return `${envKey.toLowerCase()}${this.getLabelSampleRatePostfix()}`;
+  }
+
+  private getLabelNameFromEnvVarKey(type: OperationType, key: string) {
+    return key
+      .slice(this.getLabelPrefix(type).length)
+      .toLowerCase()
+      .replace(/___/g, ".")
+      .replace(/__/g, "/")
+      .replace(/_/g, "-");
+  }
+
+  private getCaseInsensitiveEnvValue(key: string) {
+    for (const [envKey, value] of Object.entries(process.env)) {
+      if (envKey.toLowerCase() === key.toLowerCase()) {
+        return value;
+      }
+    }
+  }
+
+  /** Returns the sample rate for a given label as fraction of 100 */
+  private getSampleRateFromEnvVarKey(type: OperationType, envKey: string) {
+    // Apply default: always sample
+    const DEFAULT_SAMPLE_RATE_PERCENT = 100;
+    const defaultSampleRateFraction = this.fractionFromPercent(DEFAULT_SAMPLE_RATE_PERCENT);
+
+    const value = this.getCaseInsensitiveEnvValue(this.getSampleRateEnvVarKey(type, envKey));
+
+    if (!value) {
+      return defaultSampleRateFraction;
+    }
+
+    const sampleRatePercent = parseFloat(value || String(DEFAULT_SAMPLE_RATE_PERCENT));
+
+    if (isNaN(sampleRatePercent)) {
+      return defaultSampleRateFraction;
+    }
+
+    const fractionalSampleRate = this.fractionFromPercent(sampleRatePercent);
+
+    return fractionalSampleRate;
+  }
+
+  private getCustomLabels(type: OperationType): CustomLabel[] {
+    switch (type) {
+      case "create":
+        if (this.createLabels) {
+          return this.createLabels;
+        }
+        break;
+      case "restore":
+        if (this.restoreLabels) {
+          return this.restoreLabels;
+        }
+        break;
+      default:
+        assertExhaustive(type);
+    }
+
+    const customLabels: CustomLabel[] = [];
+
+    for (const [envKey, value] of Object.entries(process.env)) {
+      const key = envKey.toLowerCase();
+
+      // Only process env vars that start with the expected prefix
+      if (!this.isLabelEnvVar(type, key)) {
+        continue;
+      }
+
+      // Skip sample rates - deal with them separately
+      if (this.isLabelSampleRateEnvVar(key)) {
+        continue;
+      }
+
+      const labelName = this.getLabelNameFromEnvVarKey(type, key);
+      const sampleRate = this.getSampleRateFromEnvVarKey(type, key);
+
+      console.log({ key, value, sampleRate, labelName });
+
+      const label = {
+        key: labelName,
+        value: value || "",
+        sampleRate,
+      } satisfies CustomLabel;
+
+      customLabels.push(label);
+    }
+
+    console.log({
+      type,
+      customLabels,
+    });
+
+    return customLabels;
+  }
+
+  getAdditionalLabels(type: OperationType): Record<string, string> {
+    const labels = this.getCustomLabels(type);
+
+    const additionalLabels: Record<string, string> = {};
+
+    for (const { key, value, sampleRate } of labels) {
+      // Always apply label if sample rate is 1
+      if (sampleRate === 1) {
+        additionalLabels[key] = value;
+        continue;
+      }
+
+      if (Math.random() <= sampleRate) {
+        additionalLabels[key] = value;
+        continue;
+      }
+    }
+
+    return additionalLabels;
+  }
+}