Skip to content

Commit d50f1b3

Browse files
nicktrnnicktrn
committed
add machine preset trigger option
1 parent ca51f65 commit d50f1b3

File tree

8 files changed

+55
-12
lines changed

8 files changed

+55
-12
lines changed

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,16 @@ export function machinePresetFromName(name: MachinePresetName): MachinePreset {
3131
};
3232
}
3333

34+
export function machinePresetFromRun(run: { machinePreset: string | null }): MachinePreset | null {
35+
const presetName = MachinePresetName.safeParse(run.machinePreset).data;
36+
37+
if (!presetName) {
38+
return null;
39+
}
40+
41+
return machinePresetFromName(presetName);
42+
}
43+
3444
// Finds the smallest machine preset name that satisfies the given CPU and memory requirements
3545
function derivePresetNameFromValues(cpu: number, memory: number): MachinePresetName {
3646
for (const [name, preset] of Object.entries(machines)) {

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
4343
import { SEMINTATTRS_FORCE_RECORDING, tracer } from "../tracer.server";
4444
import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
4545
import { EnvironmentVariable } from "../environmentVariables/repository";
46-
import { machinePresetFromConfig } from "../machinePresets.server";
46+
import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
4747
import { env } from "~/env.server";
4848
import {
4949
FINAL_ATTEMPT_STATUSES,
@@ -413,7 +413,9 @@ export class SharedQueueConsumer {
413413
cliVersion: deployment.worker.cliVersion,
414414
startedAt: existingTaskRun.startedAt ?? new Date(),
415415
baseCostInCents: env.CENTS_PER_RUN,
416-
machinePreset: machinePresetFromConfig(backgroundTask.machineConfig ?? {}).name,
416+
machinePreset:
417+
existingTaskRun.machinePreset ??
418+
machinePresetFromConfig(backgroundTask.machineConfig ?? {}).name,
417419
maxDurationInSeconds: getMaxDuration(
418420
existingTaskRun.maxDurationInSeconds,
419421
backgroundTask.maxDurationInSeconds
@@ -542,8 +544,9 @@ export class SharedQueueConsumer {
542544

543545
// Retries for workers with disabled retry checkpoints will be handled just like normal attempts
544546
} else {
545-
const machineConfig = lockedTaskRun.lockedBy?.machineConfig;
546-
const machine = machinePresetFromConfig(machineConfig ?? {});
547+
const machine =
548+
machinePresetFromRun(lockedTaskRun) ??
549+
machinePresetFromConfig(lockedTaskRun.lockedBy?.machineConfig ?? {});
547550

548551
await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
549552
backgroundWorkerId: deployment.worker.friendlyId,
@@ -1077,7 +1080,9 @@ class SharedQueueTasks {
10771080
const { backgroundWorkerTask, taskRun, queue } = attempt;
10781081

10791082
if (!machinePreset) {
1080-
machinePreset = machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
1083+
machinePreset =
1084+
machinePresetFromRun(attempt.taskRun) ??
1085+
machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
10811086
}
10821087

10831088
const metadata = await parsePacket({
@@ -1294,9 +1299,13 @@ class SharedQueueTasks {
12941299
},
12951300
});
12961301
}
1302+
12971303
const { backgroundWorkerTask, taskRun } = attempt;
12981304

1299-
const machinePreset = machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
1305+
const machinePreset =
1306+
machinePresetFromRun(attempt.taskRun) ??
1307+
machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
1308+
13001309
const execution = await this._executionFromAttempt(attempt, machinePreset);
13011310
const variables = await this.#buildEnvironmentVariables(
13021311
attempt.runtimeEnvironment,
@@ -1432,6 +1441,7 @@ class SharedQueueTasks {
14321441
machineConfig: true,
14331442
},
14341443
},
1444+
machinePreset: true,
14351445
},
14361446
});
14371447

@@ -1451,7 +1461,8 @@ class SharedQueueTasks {
14511461
attemptCount,
14521462
});
14531463

1454-
const machinePreset = machinePresetFromConfig(run.lockedBy?.machineConfig ?? {});
1464+
const machinePreset =
1465+
machinePresetFromRun(run) ?? machinePresetFromConfig(run.lockedBy?.machineConfig ?? {});
14551466

14561467
const variables = await this.#buildEnvironmentVariables(environment, run.id, machinePreset);
14571468

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
66
import { logger } from "~/services/logger.server";
77
import { reportInvocationUsage } from "~/services/platform.v3.server";
88
import { generateFriendlyId } from "../friendlyIdentifiers";
9-
import { machinePresetFromConfig } from "../machinePresets.server";
9+
import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
1010
import { BaseService, ServiceValidationError } from "./baseService.server";
1111
import { CrashTaskRunService } from "./crashTaskRun.server";
1212
import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
@@ -173,7 +173,9 @@ export class CreateTaskRunAttemptService extends BaseService {
173173
});
174174
}
175175

176-
const machinePreset = machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
176+
const machinePreset =
177+
machinePresetFromRun(taskRun) ??
178+
machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
177179

178180
const metadata = await parsePacket({
179181
data: taskRun.metadata ?? undefined,

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { type Checkpoint } from "@trigger.dev/database";
22
import { logger } from "~/services/logger.server";
33
import { socketIo } from "../handleSocketIo.server";
4-
import { machinePresetFromConfig } from "../machinePresets.server";
4+
import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
55
import { BaseService } from "./baseService.server";
66
import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
77
import { isRestorableAttemptStatus, isRestorableRunStatus } from "../taskStatus";
@@ -24,6 +24,7 @@ export class RestoreCheckpointService extends BaseService {
2424
run: {
2525
select: {
2626
status: true,
27+
machinePreset: true,
2728
},
2829
},
2930
attempt: {
@@ -69,8 +70,9 @@ export class RestoreCheckpointService extends BaseService {
6970
return;
7071
}
7172

72-
const { machineConfig } = checkpoint.attempt.backgroundWorkerTask;
73-
const machine = machinePresetFromConfig(machineConfig ?? {});
73+
const machine =
74+
machinePresetFromRun(checkpoint.run) ??
75+
machinePresetFromConfig(checkpoint.attempt.backgroundWorkerTask.machineConfig ?? {});
7476

7577
const restoreEvent = await this._prisma.checkpointRestoreEvent.findFirst({
7678
where: {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,7 @@ export class TriggerTaskService extends BaseService {
415415
: undefined,
416416
runTags: bodyTags,
417417
oneTimeUseToken: options.oneTimeUseToken,
418+
machinePreset: body.options?.machinePreset,
418419
},
419420
});
420421

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import { DeserializedJsonSchema } from "../../schemas/json.js";
33
import {
44
FlushedRunMetadata,
5+
MachinePresetName,
56
RunMetadataChangeOperation,
67
SerializedError,
78
TaskRunError,
@@ -91,6 +92,7 @@ export const TriggerTaskRequestBody = z.object({
9192
metadata: z.any(),
9293
metadataType: z.string().optional(),
9394
maxDuration: z.number().optional(),
95+
machinePreset: MachinePresetName.optional(),
9496
})
9597
.optional(),
9698
});
@@ -131,6 +133,7 @@ export const BatchTriggerTaskItem = z.object({
131133
metadataType: z.string().optional(),
132134
maxDuration: z.number().optional(),
133135
parentAttempt: z.string().optional(),
136+
machinePreset: MachinePresetName.optional(),
134137
})
135138
.optional(),
136139
});

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { RunTags } from "../schemas/api.js";
66
import {
77
MachineCpu,
88
MachineMemory,
9+
MachinePresetName,
910
RetryOptions,
1011
TaskMetadata,
1112
TaskRunContext,
@@ -775,6 +776,11 @@ export type TriggerOptions = {
775776
* Minimum value is 5 seconds
776777
*/
777778
maxDuration?: number;
779+
780+
/**
781+
* The machine preset to use for this run. This will override the task's machine preset and any defaults.
782+
*/
783+
machinePreset?: MachinePresetName;
778784
};
779785

780786
export type TriggerAndWaitOptions = Omit<TriggerOptions, "idempotencyKey" | "idempotencyKeyTTL">;

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,7 @@ export async function batchTriggerById<TTask extends AnyTask>(
613613
parentAttempt: taskContext.ctx?.attempt.id,
614614
metadata: item.options?.metadata,
615615
maxDuration: item.options?.maxDuration,
616+
machinePreset: item.options?.machinePreset,
616617
},
617618
};
618619
})
@@ -786,6 +787,7 @@ export async function batchTriggerByIdAndWait<TTask extends AnyTask>(
786787
maxAttempts: item.options?.maxAttempts,
787788
metadata: item.options?.metadata,
788789
maxDuration: item.options?.maxDuration,
790+
machinePreset: item.options?.machinePreset,
789791
},
790792
};
791793
})
@@ -947,6 +949,7 @@ export async function batchTriggerTasks<TTasks extends readonly AnyTask[]>(
947949
parentAttempt: taskContext.ctx?.attempt.id,
948950
metadata: item.options?.metadata,
949951
maxDuration: item.options?.maxDuration,
952+
machinePreset: item.options?.machinePreset,
950953
},
951954
};
952955
})
@@ -1122,6 +1125,7 @@ export async function batchTriggerAndWaitTasks<TTasks extends readonly AnyTask[]
11221125
maxAttempts: item.options?.maxAttempts,
11231126
metadata: item.options?.metadata,
11241127
maxDuration: item.options?.maxDuration,
1128+
machinePreset: item.options?.machinePreset,
11251129
},
11261130
};
11271131
})
@@ -1200,6 +1204,7 @@ async function trigger_internal<TRunTypes extends AnyRunTypes>(
12001204
parentAttempt: taskContext.ctx?.attempt.id,
12011205
metadata: options?.metadata,
12021206
maxDuration: options?.maxDuration,
1207+
machinePreset: options?.machinePreset,
12031208
},
12041209
},
12051210
{
@@ -1259,6 +1264,7 @@ async function batchTrigger_internal<TRunTypes extends AnyRunTypes>(
12591264
parentAttempt: taskContext.ctx?.attempt.id,
12601265
metadata: item.options?.metadata,
12611266
maxDuration: item.options?.maxDuration,
1267+
machinePreset: item.options?.machinePreset,
12621268
},
12631269
};
12641270
})
@@ -1352,6 +1358,7 @@ async function triggerAndWait_internal<TIdentifier extends string, TPayload, TOu
13521358
maxAttempts: options?.maxAttempts,
13531359
metadata: options?.metadata,
13541360
maxDuration: options?.maxDuration,
1361+
machinePreset: options?.machinePreset,
13551362
},
13561363
},
13571364
{},
@@ -1428,6 +1435,7 @@ async function batchTriggerAndWait_internal<TIdentifier extends string, TPayload
14281435
maxAttempts: item.options?.maxAttempts,
14291436
metadata: item.options?.metadata,
14301437
maxDuration: item.options?.maxDuration,
1438+
machinePreset: item.options?.machinePreset,
14311439
},
14321440
};
14331441
})

0 commit comments

Comments
 (0)