v3: machine config

.changeset/clean-pianos-listen.md

@@ -0,0 +1,8 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core-apps": patch
+"trigger.dev": patch
+"@trigger.dev/core": patch
+---
+
+add machine config and secure zod connection

apps/coordinator/.env.example

@@ -1,3 +1,4 @@
 HTTP_SERVER_PORT=8020
 PLATFORM_ENABLED=true
-PLATFORM_WS_PORT=3030
+PLATFORM_WS_PORT=3030
+SECURE_CONNECTION=false

apps/coordinator/src/index.ts

@@ -26,6 +26,7 @@ const PLATFORM_ENABLED = ["1", "true"].includes(process.env.PLATFORM_ENABLED ??
 const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1";
 const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030;
 const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret";
+const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false");
 
 const logger = new SimpleLogger(`[${NODE_NAME}]`);
 
@@ -365,6 +366,7 @@ class TaskCoordinator {
       namespace: "coordinator",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: CoordinatorToPlatformMessages,
       serverMessages: PlatformToCoordinatorMessages,
       authToken: PLATFORM_SECRET,

apps/docker-provider/.env.example

@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8050
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"

apps/docker-provider/src/index.ts

@@ -8,6 +8,7 @@ import {
   TaskOperationsIndexOptions,
 } from "@trigger.dev/core-apps";
 import { setTimeout } from "node:timers/promises";
+import { PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3";
 
 const MACHINE_NAME = process.env.MACHINE_NAME || "local";
 const COORDINATOR_PORT = process.env.COORDINATOR_PORT || 8020;
@@ -190,6 +191,9 @@ class DockerTaskOperations implements TaskOperations {
   async delete(opts: { runId: string }) {
     await this.#initialize();
 
+    const containerName = this.#getRunContainerName(opts.runId);
+    await this.#sendPreStop(containerName);
+
     logger.log("noop: delete");
   }
 
@@ -208,6 +212,26 @@ class DockerTaskOperations implements TaskOperations {
   }
 
   async #sendPostStart(containerName: string): Promise<void> {
+    try {
+      const port = await this.#getHttpServerPort(containerName);
+      logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore"));
+    } catch (error) {
+      logger.error("postStart error", { error });
+      throw new Error("postStart command failed");
+    }
+  }
+
+  async #sendPreStop(containerName: string): Promise<void> {
+    try {
+      const port = await this.#getHttpServerPort(containerName);
+      logger.debug(await this.#runLifecycleCommand(containerName, port, "preStop", "terminate"));
+    } catch (error) {
+      logger.error("preStop error", { error });
+      throw new Error("preStop command failed");
+    }
+  }
+
+  async #getHttpServerPort(containerName: string): Promise<number> {
     // We first get the correct port, which is random during dev as we run with host networking and need to avoid clashes
     // FIXME: Skip this in prod
     const logs = logger.debug(await $`docker logs ${containerName}`);
@@ -219,19 +243,14 @@ class DockerTaskOperations implements TaskOperations {
       throw new Error("failed to extract port from logs");
     }
 
-    try {
-      logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore"));
-    } catch (error) {
-      logger.error("postStart error", { error });
-      throw new Error("postStart command failed");
-    }
+    return port;
   }
 
-  async #runLifecycleCommand(
+  async #runLifecycleCommand<THookType extends "postStart" | "preStop">(
     containerName: string,
     port: number,
-    type: "postStart" | "preStop",
-    cause: "index" | "create" | "restore",
+    type: THookType,
+    cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses,
     retryCount = 0
   ): Promise<ExecaChildProcess> {
     try {
@@ -244,15 +263,15 @@ class DockerTaskOperations implements TaskOperations {
         `127.0.0.1:${port}/${type}?cause=${cause}`,
       ]);
     } catch (error: any) {
-      if (retryCount < 6) {
-        logger.debug("retriable postStart error", { retryCount, message: error?.message });
+      if (type === "postStart" && retryCount < 6) {
+        logger.debug(`retriable ${type} error`, { retryCount, message: error?.message });
         await setTimeout(exponentialBackoff(retryCount + 1, 2, 50, 1150, 50));
 
         return this.#runLifecycleCommand(containerName, port, type, cause, retryCount + 1);
       }
 
-      logger.error("final postStart error", { message: error?.message });
-      throw new Error(`postStart command failed after ${retryCount - 1} retries`);
+      logger.error(`final ${type} error`, { message: error?.message });
+      throw new Error(`${type} command failed after ${retryCount - 1} retries`);
     }
   }
 }

apps/kubernetes-provider/.env.example

@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8060
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"

apps/kubernetes-provider/src/index.ts

@@ -7,6 +7,7 @@ import {
   TaskOperationsIndexOptions,
   TaskOperationsRestoreOptions,
 } from "@trigger.dev/core-apps";
+import { Machine, PostStartCauses, PreStopCauses, EnvironmentType } from "@trigger.dev/core/v3";
 import { randomUUID } from "crypto";
 
 const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local";
@@ -55,6 +56,12 @@ class KubernetesTaskOperations implements TaskOperations {
             metadata: {
               labels: {
                 app: "task-index",
+                "app.kubernetes.io/part-of": "trigger-worker",
+                "app.kubernetes.io/component": "index",
+                env: opts.envId,
+                envtype: this.#envTypeToLabelValue(opts.envType),
+                org: opts.orgId,
+                project: opts.projectId,
               },
             },
             spec: {
@@ -64,6 +71,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   name: "registry-trigger",
                 },
               ],
+              nodeSelector: {
+                nodetype: "worker",
+              },
               containers: [
                 {
                   name: this.#getIndexContainerName(opts.shortCode),
@@ -79,6 +89,13 @@ class KubernetesTaskOperations implements TaskOperations {
                   //     memory: "50Mi",
                   //   },
                   // },
+                  lifecycle: {
+                    preStop: {
+                      exec: {
+                        command: this.#getLifecycleCommand("preStop", "terminate"),
+                      },
+                    },
+                  },
                   env: [
                     {
                       name: "DEBUG",
@@ -151,6 +168,13 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "create",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
           },
         },
         spec: {
@@ -160,6 +184,9 @@ class KubernetesTaskOperations implements TaskOperations {
               name: "registry-trigger",
             },
           ],
+          nodeSelector: {
+            nodetype: "worker",
+          },
           containers: [
             {
               name: this.#getRunContainerName(opts.runId),
@@ -169,9 +196,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   containerPort: 8000,
                 },
               ],
-              // resources: {
-              //   limits: opts.machine,
-              // },
+              resources: {
+                limits: this.#getResourcesFromMachineConfig(opts.machine),
+              },
               lifecycle: {
                 postStart: {
                   exec: {
@@ -180,7 +207,7 @@ class KubernetesTaskOperations implements TaskOperations {
                 },
                 preStop: {
                   exec: {
-                    command: this.#getLifecycleCommand("preStop", "create"),
+                    command: this.#getLifecycleCommand("preStop", "terminate"),
                   },
                 },
               },
@@ -262,6 +289,14 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "restore",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
+            checkpoint: opts.checkpointId,
           },
         },
         spec: {
@@ -271,6 +306,9 @@ class KubernetesTaskOperations implements TaskOperations {
               name: "registry-trigger",
             },
           ],
+          nodeSelector: {
+            nodetype: "worker",
+          },
           initContainers: [
             {
               name: "pull-base-image",
@@ -309,9 +347,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   containerPort: 8000,
                 },
               ],
-              // resources: {
-              //   limits: opts.machine,
-              // },
+              resources: {
+                limits: this.#getResourcesFromMachineConfig(opts.machine),
+              },
               lifecycle: {
                 postStart: {
                   exec: {
@@ -320,7 +358,7 @@ class KubernetesTaskOperations implements TaskOperations {
                 },
                 preStop: {
                   exec: {
-                    command: this.#getLifecycleCommand("preStop", "restore"),
+                    command: this.#getLifecycleCommand("preStop", "terminate"),
                   },
                 },
               },
@@ -355,7 +393,30 @@ class KubernetesTaskOperations implements TaskOperations {
     await this.#getPod(opts.runId, this.#namespace);
   }
 
-  #getLifecycleCommand(type: "postStart" | "preStop", cause: "index" | "create" | "restore") {
+  #envTypeToLabelValue(type: EnvironmentType) {
+    switch (type) {
+      case "PRODUCTION":
+        return "prod";
+      case "STAGING":
+        return "stg";
+      case "DEVELOPMENT":
+        return "dev";
+      case "PREVIEW":
+        return "preview";
+    }
+  }
+
+  #getResourcesFromMachineConfig(config: Machine) {
+    return {
+      cpu: `${config.cpu}`,
+      memory: `${config.memory}G`,
+    };
+  }
+
+  #getLifecycleCommand<THookType extends "postStart" | "preStop">(
+    type: THookType,
+    cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses
+  ) {
     return ["/bin/sh", "-c", `sleep 1; wget -q -O- 127.0.0.1:8000/${type}?cause=${cause}`];
   }
 

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -1,5 +1,6 @@
 import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api";
 import {
+  Machine,
   ProdTaskRunExecution,
   ProdTaskRunExecutionPayload,
   TaskRunError,
@@ -438,10 +439,27 @@ export class SharedQueueConsumer {
             queueId: queue.id,
             runtimeEnvironmentId: environment.id,
           },
+          include: {
+            backgroundWorkerTask: true,
+          },
         });
 
         const isRetry = taskRunAttempt.number > 1;
 
+        const { machineConfig } = taskRunAttempt.backgroundWorkerTask;
+        const machine = Machine.safeParse(machineConfig ?? {});
+
+        if (!machine.success) {
+          logger.error("Failed to parse machine config", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+            attemptId: taskRunAttempt.id,
+            machineConfig,
+          });
+
+          await this.#ackAndDoMoreWork(message.messageId);
+          return;
+        }
         try {
           if (messageBody.data.checkpointEventId) {
             const restoreService = new RestoreCheckpointService();
@@ -470,11 +488,16 @@ export class SharedQueueConsumer {
               backgroundWorkerId: deployment.worker.friendlyId,
               data: {
                 type: "SCHEDULE_ATTEMPT",
-                id: taskRunAttempt.id,
                 image: deployment.imageReference,
+                version: deployment.version,
+                machine: machine.data,
+                // identifiers
+                id: taskRunAttempt.id,
                 envId: environment.id,
+                envType: environment.type,
+                orgId: environment.organizationId,
+                projectId: environment.projectId,
                 runId: taskRunAttempt.taskRunId,
-                version: deployment.version,
               },
             });
           }

apps/webapp/app/v3/services/createBackgroundWorker.server.ts

@@ -101,6 +101,7 @@ export async function createBackgroundTasks(
           exportName: task.exportName,
           retryConfig: task.retry,
           queueConfig: task.queue,
+          machineConfig: task.machine,
         },
       });
 

apps/webapp/app/v3/services/indexDeployment.server.ts

@@ -49,9 +49,13 @@ export class IndexDeploymentService extends BaseService {
         version: "v1",
         shortCode: deployment.shortCode,
         imageTag: deployment.imageReference,
-        envId: deployment.environmentId,
         apiKey: deployment.environment.apiKey,
         apiUrl: env.APP_ORIGIN,
+        // identifiers
+        envId: deployment.environmentId,
+        envType: deployment.environment.type,
+        projectId: deployment.projectId,
+        orgId: deployment.environment.organizationId,
       });
 
       logger.debug("Index ACK received", { responses });