add machine preset trigger option

Author: nicktrn

apps/webapp/app/v3/machinePresets.server.ts

@@ -31,6 +31,16 @@ export function machinePresetFromName(name: MachinePresetName): MachinePreset {
   };
 }
 
+export function machinePresetFromRun(run: { machinePreset: string | null }): MachinePreset | null {
+  const presetName = MachinePresetName.safeParse(run.machinePreset).data;
+
+  if (!presetName) {
+    return null;
+  }
+
+  return machinePresetFromName(presetName);
+}
+
 // Finds the smallest machine preset name that satisfies the given CPU and memory requirements
 function derivePresetNameFromValues(cpu: number, memory: number): MachinePresetName {
   for (const [name, preset] of Object.entries(machines)) {

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

@@ -43,7 +43,7 @@ import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
 import { SEMINTATTRS_FORCE_RECORDING, tracer } from "../tracer.server";
 import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { EnvironmentVariable } from "../environmentVariables/repository";
-import { machinePresetFromConfig } from "../machinePresets.server";
+import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
 import { env } from "~/env.server";
 import {
   FINAL_ATTEMPT_STATUSES,
@@ -413,7 +413,9 @@ export class SharedQueueConsumer {
             cliVersion: deployment.worker.cliVersion,
             startedAt: existingTaskRun.startedAt ?? new Date(),
             baseCostInCents: env.CENTS_PER_RUN,
-            machinePreset: machinePresetFromConfig(backgroundTask.machineConfig ?? {}).name,
+            machinePreset:
+              existingTaskRun.machinePreset ??
+              machinePresetFromConfig(backgroundTask.machineConfig ?? {}).name,
             maxDurationInSeconds: getMaxDuration(
               existingTaskRun.maxDurationInSeconds,
               backgroundTask.maxDurationInSeconds
@@ -542,8 +544,9 @@ export class SharedQueueConsumer {
 
             // Retries for workers with disabled retry checkpoints will be handled just like normal attempts
           } else {
-            const machineConfig = lockedTaskRun.lockedBy?.machineConfig;
-            const machine = machinePresetFromConfig(machineConfig ?? {});
+            const machine =
+              machinePresetFromRun(lockedTaskRun) ??
+              machinePresetFromConfig(lockedTaskRun.lockedBy?.machineConfig ?? {});
 
             await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
               backgroundWorkerId: deployment.worker.friendlyId,
@@ -1077,7 +1080,9 @@ class SharedQueueTasks {
     const { backgroundWorkerTask, taskRun, queue } = attempt;
 
     if (!machinePreset) {
-      machinePreset = machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
+      machinePreset =
+        machinePresetFromRun(attempt.taskRun) ??
+        machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
     }
 
     const metadata = await parsePacket({
@@ -1294,9 +1299,13 @@ class SharedQueueTasks {
         },
       });
     }
+
     const { backgroundWorkerTask, taskRun } = attempt;
 
-    const machinePreset = machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
+    const machinePreset =
+      machinePresetFromRun(attempt.taskRun) ??
+      machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
+
     const execution = await this._executionFromAttempt(attempt, machinePreset);
     const variables = await this.#buildEnvironmentVariables(
       attempt.runtimeEnvironment,
@@ -1432,6 +1441,7 @@ class SharedQueueTasks {
             machineConfig: true,
           },
         },
+        machinePreset: true,
       },
     });
 
@@ -1451,7 +1461,8 @@ class SharedQueueTasks {
       attemptCount,
     });
 
-    const machinePreset = machinePresetFromConfig(run.lockedBy?.machineConfig ?? {});
+    const machinePreset =
+      machinePresetFromRun(run) ?? machinePresetFromConfig(run.lockedBy?.machineConfig ?? {});
 
     const variables = await this.#buildEnvironmentVariables(environment, run.id, machinePreset);
 

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

@@ -6,7 +6,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { reportInvocationUsage } from "~/services/platform.v3.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
-import { machinePresetFromConfig } from "../machinePresets.server";
+import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
 import { BaseService, ServiceValidationError } from "./baseService.server";
 import { CrashTaskRunService } from "./crashTaskRun.server";
 import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
@@ -173,7 +173,9 @@ export class CreateTaskRunAttemptService extends BaseService {
         });
       }
 
-      const machinePreset = machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
+      const machinePreset =
+        machinePresetFromRun(taskRun) ??
+        machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
 
       const metadata = await parsePacket({
         data: taskRun.metadata ?? undefined,

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

@@ -1,7 +1,7 @@
 import { type Checkpoint } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
 import { socketIo } from "../handleSocketIo.server";
-import { machinePresetFromConfig } from "../machinePresets.server";
+import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
 import { BaseService } from "./baseService.server";
 import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
 import { isRestorableAttemptStatus, isRestorableRunStatus } from "../taskStatus";
@@ -24,6 +24,7 @@ export class RestoreCheckpointService extends BaseService {
             run: {
               select: {
                 status: true,
+                machinePreset: true,
               },
             },
             attempt: {
@@ -69,8 +70,9 @@ export class RestoreCheckpointService extends BaseService {
       return;
     }
 
-    const { machineConfig } = checkpoint.attempt.backgroundWorkerTask;
-    const machine = machinePresetFromConfig(machineConfig ?? {});
+    const machine =
+      machinePresetFromRun(checkpoint.run) ??
+      machinePresetFromConfig(checkpoint.attempt.backgroundWorkerTask.machineConfig ?? {});
 
     const restoreEvent = await this._prisma.checkpointRestoreEvent.findFirst({
       where: {

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

@@ -415,6 +415,7 @@ export class TriggerTaskService extends BaseService {
                       : undefined,
                     runTags: bodyTags,
                     oneTimeUseToken: options.oneTimeUseToken,
+                    machinePreset: body.options?.machinePreset,
                   },
                 });
 

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

@@ -2,6 +2,7 @@ import { z } from "zod";
 import { DeserializedJsonSchema } from "../../schemas/json.js";
 import {
   FlushedRunMetadata,
+  MachinePresetName,
   RunMetadataChangeOperation,
   SerializedError,
   TaskRunError,
@@ -91,6 +92,7 @@ export const TriggerTaskRequestBody = z.object({
       metadata: z.any(),
       metadataType: z.string().optional(),
       maxDuration: z.number().optional(),
+      machinePreset: MachinePresetName.optional(),
     })
     .optional(),
 });
@@ -131,6 +133,7 @@ export const BatchTriggerTaskItem = z.object({
       metadataType: z.string().optional(),
       maxDuration: z.number().optional(),
       parentAttempt: z.string().optional(),
+      machinePreset: MachinePresetName.optional(),
     })
     .optional(),
 });

packages/core/src/v3/types/tasks.ts

@@ -6,6 +6,7 @@ import { RunTags } from "../schemas/api.js";
 import {
   MachineCpu,
   MachineMemory,
+  MachinePresetName,
   RetryOptions,
   TaskMetadata,
   TaskRunContext,
@@ -775,6 +776,11 @@ export type TriggerOptions = {
    * Minimum value is 5 seconds
    */
   maxDuration?: number;
+
+  /**
+   * The machine preset to use for this run. This will override the task's machine preset and any defaults.
+   */
+  machinePreset?: MachinePresetName;
 };
 
 export type TriggerAndWaitOptions = Omit<TriggerOptions, "idempotencyKey" | "idempotencyKeyTTL">;

packages/trigger-sdk/src/v3/shared.ts

@@ -613,6 +613,7 @@ export async function batchTriggerById<TTask extends AnyTask>(
               parentAttempt: taskContext.ctx?.attempt.id,
               metadata: item.options?.metadata,
               maxDuration: item.options?.maxDuration,
+              machinePreset: item.options?.machinePreset,
             },
           };
         })
@@ -786,6 +787,7 @@ export async function batchTriggerByIdAndWait<TTask extends AnyTask>(
                   maxAttempts: item.options?.maxAttempts,
                   metadata: item.options?.metadata,
                   maxDuration: item.options?.maxDuration,
+                  machinePreset: item.options?.machinePreset,
                 },
               };
             })
@@ -947,6 +949,7 @@ export async function batchTriggerTasks<TTasks extends readonly AnyTask[]>(
               parentAttempt: taskContext.ctx?.attempt.id,
               metadata: item.options?.metadata,
               maxDuration: item.options?.maxDuration,
+              machinePreset: item.options?.machinePreset,
             },
           };
         })
@@ -1122,6 +1125,7 @@ export async function batchTriggerAndWaitTasks<TTasks extends readonly AnyTask[]
                   maxAttempts: item.options?.maxAttempts,
                   metadata: item.options?.metadata,
                   maxDuration: item.options?.maxDuration,
+                  machinePreset: item.options?.machinePreset,
                 },
               };
             })
@@ -1200,6 +1204,7 @@ async function trigger_internal<TRunTypes extends AnyRunTypes>(
         parentAttempt: taskContext.ctx?.attempt.id,
         metadata: options?.metadata,
         maxDuration: options?.maxDuration,
+        machinePreset: options?.machinePreset,
       },
     },
     {
@@ -1259,6 +1264,7 @@ async function batchTrigger_internal<TRunTypes extends AnyRunTypes>(
               parentAttempt: taskContext.ctx?.attempt.id,
               metadata: item.options?.metadata,
               maxDuration: item.options?.maxDuration,
+              machinePreset: item.options?.machinePreset,
             },
           };
         })
@@ -1352,6 +1358,7 @@ async function triggerAndWait_internal<TIdentifier extends string, TPayload, TOu
             maxAttempts: options?.maxAttempts,
             metadata: options?.metadata,
             maxDuration: options?.maxDuration,
+            machinePreset: options?.machinePreset,
           },
         },
         {},
@@ -1428,6 +1435,7 @@ async function batchTriggerAndWait_internal<TIdentifier extends string, TPayload
                   maxAttempts: item.options?.maxAttempts,
                   metadata: item.options?.metadata,
                   maxDuration: item.options?.maxDuration,
+                  machinePreset: item.options?.machinePreset,
                 },
               };
             })