Skip to content

Add support for specifying machine preset at trigger time #1608

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jan 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions .changeset/gold-melons-fetch.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
---
"@trigger.dev/sdk": patch
"@trigger.dev/core": patch
---

Add support for specifying machine preset at trigger time. Works with any trigger function:

```ts
// Same as usual, will use the machine preset on childTask, defaults to "small-1x"
await childTask.trigger({ message: "Hello, world!" });

// This will override the task's machine preset and any defaults. Works with all trigger functions.
await childTask.trigger({ message: "Hello, world!" }, { machine: "small-2x" });
await childTask.triggerAndWait({ message: "Hello, world!" }, { machine: "small-2x" });

await childTask.batchTrigger([
{ payload: { message: "Hello, world!" }, options: { machine: "micro" } },
{ payload: { message: "Hello, world!" }, options: { machine: "large-1x" } },
]);
await childTask.batchTriggerAndWait([
{ payload: { message: "Hello, world!" }, options: { machine: "micro" } },
{ payload: { message: "Hello, world!" }, options: { machine: "large-1x" } },
]);

await tasks.trigger<typeof childTask>(
"child",
{ message: "Hello, world!" },
{ machine: "small-2x" }
);
await tasks.batchTrigger<typeof childTask>("child", [
{ payload: { message: "Hello, world!" }, options: { machine: "micro" } },
{ payload: { message: "Hello, world!" }, options: { machine: "large-1x" } },
]);
```
10 changes: 10 additions & 0 deletions apps/webapp/app/v3/machinePresets.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
25 changes: 18 additions & 7 deletions apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -1432,6 +1441,7 @@ class SharedQueueTasks {
machineConfig: true,
},
},
machinePreset: true,
},
});

Expand All @@ -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);

Expand Down
6 changes: 4 additions & 2 deletions apps/webapp/app/v3/services/createTaskRunAttempt.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 5 additions & 3 deletions apps/webapp/app/v3/services/restoreCheckpoint.server.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -24,6 +24,7 @@ export class RestoreCheckpointService extends BaseService {
run: {
select: {
status: true,
machinePreset: true,
},
},
attempt: {
Expand Down Expand Up @@ -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: {
Expand Down
1 change: 1 addition & 0 deletions apps/webapp/app/v3/services/triggerTask.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,7 @@ export class TriggerTaskService extends BaseService {
: undefined,
runTags: bodyTags,
oneTimeUseToken: options.oneTimeUseToken,
machinePreset: body.options?.machine,
},
});

Expand Down
12 changes: 12 additions & 0 deletions docs/v3-openapi.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -1595,6 +1595,18 @@ components:
We recommend prefixing tags with a namespace using an underscore or colon, like `user_1234567` or `org:9876543`. Stripe uses underscores.
items:
type: string
machine:
type: string
enum:
- micro
- small-1x
- small-2x
- medium-1x
- medium-2x
- large-1x
- large-2x
example: "small-2x"
description: The machine preset to use for this run. This will override the task's machine preset and any defaults.
TTL:
type:
- string
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/v3/schemas/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { z } from "zod";
import { DeserializedJsonSchema } from "../../schemas/json.js";
import {
FlushedRunMetadata,
MachinePresetName,
RunMetadataChangeOperation,
SerializedError,
TaskRunError,
Expand Down Expand Up @@ -91,6 +92,7 @@ export const TriggerTaskRequestBody = z.object({
metadata: z.any(),
metadataType: z.string().optional(),
maxDuration: z.number().optional(),
machine: MachinePresetName.optional(),
})
.optional(),
});
Expand Down Expand Up @@ -131,6 +133,7 @@ export const BatchTriggerTaskItem = z.object({
metadataType: z.string().optional(),
maxDuration: z.number().optional(),
parentAttempt: z.string().optional(),
machine: MachinePresetName.optional(),
})
.optional(),
});
Expand Down
73 changes: 36 additions & 37 deletions packages/core/src/v3/types/tasks.ts
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import type { Schema as AISchema } from "ai";
import { z } from "zod";
import { SerializableJson } from "../../schemas/json.js";
import { TriggerApiRequestOptions } from "../apiClient/index.js";
import { RunTags } from "../schemas/api.js";
import {
MachineCpu,
MachineMemory,
MachinePresetName,
RetryOptions,
TaskMetadata,
TaskRunContext,
Expand Down Expand Up @@ -220,41 +219,36 @@ type CommonTaskOptions<
});
* ```
*/
machine?: {
/** vCPUs. The default is 0.5.
*
* Possible values:
* - 0.25
* - 0.5
* - 1
* - 2
* - 4
* @deprecated use preset instead
*/
cpu?: MachineCpu;
/** In GBs of RAM. The default is 1.
*
* Possible values:
* - 0.25
* - 0.5
* - 1
* - 2
* - 4
* - 8
* * @deprecated use preset instead
*/
memory?: MachineMemory;

/** Preset to use for the machine. Defaults to small-1x */
preset?:
| "micro"
| "small-1x"
| "small-2x"
| "medium-1x"
| "medium-2x"
| "large-1x"
| "large-2x";
};
machine?:
| {
/** vCPUs. The default is 0.5.
*
* Possible values:
* - 0.25
* - 0.5
* - 1
* - 2
* - 4
* @deprecated use preset instead
*/
cpu?: MachineCpu;
/** In GBs of RAM. The default is 1.
*
* Possible values:
* - 0.25
* - 0.5
* - 1
* - 2
* - 4
* - 8
* * @deprecated use preset instead
*/
memory?: MachineMemory;

/** Preset to use for the machine. Defaults to small-1x */
preset?: MachinePresetName;
}
| MachinePresetName;

/**
* The maximum duration in compute-time seconds that a task run is allowed to run. If the task run exceeds this duration, it will be stopped.
Expand Down Expand Up @@ -775,6 +769,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.
*/
machine?: MachinePresetName;
};

export type TriggerAndWaitOptions = Omit<TriggerOptions, "idempotencyKey" | "idempotencyKeyTTL">;
Expand Down
Loading
Loading