github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

restore dependencies from events

Browse Source
This commit is contained in:
nicktrn
2024-03-22 18:54:07 +00:00
parent 30c74f3c4b
commit 6ababeae55
15 changed files with 267 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/coordinator/src/index.ts
+5 -2
View File
@@ -515,6 +515,8 @@ class TaskCoordinator {
socket.on("READY_FOR_RESUME", async (message) => {
logger.log("[READY_FOR_RESUME]", message);
socket.data.attemptFriendlyId = message.attemptFriendlyId;
this.#platformSocket?.send("READY_FOR_RESUME", message);
});
@@ -723,7 +725,7 @@ class TaskCoordinator {
location: checkpoint.location,
reason: {
type: "WAIT_FOR_TASK",
id: message.id,
friendlyId: message.friendlyId,
},
});
});
@@ -765,7 +767,8 @@ class TaskCoordinator {
location: checkpoint.location,
reason: {
type: "WAIT_FOR_BATCH",
id: message.id,
batchFriendlyId: message.batchFriendlyId,
runFriendlyIds: message.runFriendlyIds,
},
});
});
apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts
-33
View File
@@ -614,39 +614,6 @@ export class SharedQueueConsumer {
return;
}
if (messageBody.data.checkpointEventId) {
try {
const restoreService = new RestoreCheckpointService();
const checkpoint = await restoreService.call({
eventId: messageBody.data.checkpointEventId,
});
if (!checkpoint) {
logger.error("Failed to restore checkpoint", {
queueMessage: message.data,
messageId: message.messageId,
});
await marqs?.acknowledgeMessage(message.messageId);
return;
}
} catch (e) {
if (e instanceof Error) {
this._currentSpan?.recordException(e);
} else {
this._currentSpan?.recordException(new Error(String(e)));
}
this._endSpanInNextIteration = true;
// Finally we need to nack the message so it can be retried
await marqs?.nackMessage(message.messageId);
return;
} finally {
setTimeout(() => this.#doWork(), this._options.interval);
}
}
const completions: TaskRunExecutionResult[] = [];
const executions: TaskRunExecution[] = [];
apps/webapp/app/v3/services/createCheckpoint.server.ts
+52 -25
View File
@@ -1,5 +1,9 @@
import { CoordinatorToPlatformMessages, InferSocketMessageSchema } from "@trigger.dev/core/v3";
import type { TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database";
import type {
CheckpointRestoreEvent,
TaskRunAttemptStatus,
TaskRunStatus,
} from "@trigger.dev/database";
import { logger } from "~/services/logger.server";
import { generateFriendlyId } from "../friendlyIdentifiers";
import { marqs } from "../marqs.server";
@@ -85,18 +89,6 @@ export class CreateCheckpointService extends BaseService {
});
const eventService = new CreateCheckpointRestoreEventService(this._prisma);
const checkpointEvent = await eventService.call({
checkpointId: checkpoint.id,
type: "CHECKPOINT",
});
if (!checkpointEvent) {
logger.error("No checkpoint event", {
attemptId: attempt.id,
checkpointId: checkpoint.id,
});
return;
}
await this._prisma.taskRunAttempt.update({
where: {
@@ -112,26 +104,40 @@ export class CreateCheckpointService extends BaseService {
},
});
switch (params.reason.type) {
const { reason } = params;
let checkpointEvent: CheckpointRestoreEvent | undefined;
switch (reason.type) {
case "WAIT_FOR_DURATION": {
await marqs?.replaceMessage(
attempt.taskRunId,
{
type: "RESUME_AFTER_DURATION",
resumableAttemptId: attempt.id,
checkpointEventId: checkpointEvent.id,
},
params.reason.now + params.reason.ms
);
checkpointEvent = await eventService.checkpoint({
checkpointId: checkpoint.id,
});
break;
}
case "WAIT_FOR_TASK": {
checkpointEvent = await eventService.checkpoint({
checkpointId: checkpoint.id,
dependencyFriendlyRunId: reason.friendlyId,
});
await marqs?.acknowledgeMessage(attempt.taskRunId);
break;
}
// TODO: Attach the checkpoint event ID to in-progress dependencies
case "WAIT_FOR_TASK":
case "WAIT_FOR_BATCH": {
checkpointEvent = await eventService.checkpoint({
checkpointId: checkpoint.id,
batchDependencyFriendlyId: reason.batchFriendlyId,
});
await marqs?.acknowledgeMessage(attempt.taskRunId);
break;
}
case "RETRYING_AFTER_FAILURE": {
checkpointEvent = await eventService.checkpoint({
checkpointId: checkpoint.id,
});
// ACK is already handled by attempt completion
break;
}
@@ -140,6 +146,27 @@ export class CreateCheckpointService extends BaseService {
}
}
if (!checkpointEvent) {
logger.error("No checkpoint event", {
attemptId: attempt.id,
checkpointId: checkpoint.id,
});
await marqs?.acknowledgeMessage(attempt.taskRunId);
return;
}
if (reason.type === "WAIT_FOR_DURATION") {
await marqs?.replaceMessage(
attempt.taskRunId,
{
type: "RESUME_AFTER_DURATION",
resumableAttemptId: attempt.id,
checkpointEventId: checkpointEvent.id,
},
reason.now + reason.ms
);
}
return {
checkpoint,
event: checkpointEvent,
apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts
+69 -5
View File
@@ -2,11 +2,32 @@ import type { CheckpointRestoreEvent, CheckpointRestoreEventType } from "@trigge
import { logger } from "~/services/logger.server";
import { BaseService } from "./baseService.server";
interface CheckpointRestoreEventCallParams {
checkpointId: string;
type: CheckpointRestoreEventType;
dependencyFriendlyRunId?: string;
batchDependencyFriendlyId?: string;
}
type CheckpointRestoreEventParams = Omit<CheckpointRestoreEventCallParams, "type">;
export class CreateCheckpointRestoreEventService extends BaseService {
public async call(params: {
checkpointId: string;
type: CheckpointRestoreEventType;
}): Promise<CheckpointRestoreEvent | undefined> {
async checkpoint(params: CheckpointRestoreEventParams) {
return this.#call({ ...params, type: "CHECKPOINT" });
}
async restore(params: CheckpointRestoreEventParams) {
return this.#call({ ...params, type: "RESTORE" });
}
async #call(
params: CheckpointRestoreEventCallParams
): Promise<CheckpointRestoreEvent | undefined> {
if (params.dependencyFriendlyRunId && params.batchDependencyFriendlyId) {
logger.error("Only one dependency can be set", { params });
return;
}
const checkpoint = await this._prisma.checkpoint.findUnique({
where: {
id: params.checkpointId,
@@ -18,7 +39,32 @@ export class CreateCheckpointRestoreEventService extends BaseService {
return;
}
logger.debug(`Creating checkpoint/restore event`, params);
logger.debug(`Creating checkpoint/restore event`, { params });
let taskRunDependencyId: string | undefined;
if (params.dependencyFriendlyRunId) {
const run = await this._prisma.taskRun.findUnique({
where: {
friendlyId: params.dependencyFriendlyRunId,
},
select: {
id: true,
dependency: {
select: {
id: true,
},
},
},
});
taskRunDependencyId = run?.dependency?.id;
if (!taskRunDependencyId) {
logger.error("Dependency or run not found", { runId: params.dependencyFriendlyRunId });
return;
}
}
const checkpointEvent = await this._prisma.checkpointRestoreEvent.create({
data: {
@@ -30,6 +76,24 @@ export class CreateCheckpointRestoreEventService extends BaseService {
type: params.type,
reason: checkpoint.reason,
metadata: checkpoint.metadata,
...(taskRunDependencyId
? {
taskRunDependency: {
connect: {
id: taskRunDependencyId,
},
},
}
: undefined),
...(params.batchDependencyFriendlyId
? {
batchTaskRunDependency: {
connect: {
friendlyId: params.batchDependencyFriendlyId,
},
},
}
: undefined),
},
});
apps/webapp/app/v3/services/restoreCheckpoint.server.ts
+1 -1
View File
@@ -64,7 +64,7 @@ export class RestoreCheckpointService extends BaseService {
}
const eventService = new CreateCheckpointRestoreEventService(this._prisma);
await eventService.call({ checkpointId: checkpoint.id, type: "RESTORE" });
await eventService.restore({ checkpointId: checkpoint.id });
socketIo.providerNamespace.emit("RESTORE", {
version: "v1",
apps/webapp/app/v3/services/resumeBatchRun.server.ts
+12 -2
View File
@@ -2,6 +2,7 @@ import { PrismaClientOrTransaction } from "~/db.server";
import { workerQueue } from "~/services/worker.server";
import { marqs } from "../marqs.server";
import { BaseService } from "./baseService.server";
import { logger } from "~/services/logger.server";
export class ResumeBatchRunService extends BaseService {
public async call(batchRunId: string, sourceTaskAttemptId: string) {
@@ -40,6 +41,7 @@ export class ResumeBatchRunService extends BaseService {
return;
}
// We need to update the batchRun status so we don't resume it again
await this._prisma.batchTaskRun.update({
where: {
id: batchRun.id,
@@ -49,8 +51,6 @@ export class ResumeBatchRunService extends BaseService {
},
});
// We need to update the batchRun status so we don't resume it again
// This batch has a dependent attempt and just finalized, we should resume that attempt
const environment = batchRun.dependentTaskAttempt.runtimeEnvironment;
@@ -62,6 +62,15 @@ export class ResumeBatchRunService extends BaseService {
const dependentRun = batchRun.dependentTaskAttempt.taskRun;
if (batchRun.dependentTaskAttempt.status === "PAUSED") {
if (!batchRun.checkpointEventId) {
logger.error("Can't resume paused attempt without checkpoint event", {
batchRunId: batchRun.id,
});
await marqs?.acknowledgeMessage(dependentRun.id);
return;
}
await marqs?.enqueueMessage(
environment,
dependentRun.queue,
@@ -70,6 +79,7 @@ export class ResumeBatchRunService extends BaseService {
type: "RESUME",
completedAttemptIds: [sourceTaskAttemptId],
resumableAttemptId: batchRun.dependentTaskAttempt.id,
checkpointEventId: batchRun.checkpointEventId,
},
dependentRun.concurrencyKey ?? undefined
);
apps/webapp/app/v3/services/resumeTaskDependency.server.ts
+12
View File
@@ -2,6 +2,7 @@ import { PrismaClientOrTransaction } from "~/db.server";
import { workerQueue } from "~/services/worker.server";
import { marqs } from "../marqs.server";
import { BaseService } from "./baseService.server";
import { logger } from "~/services/logger.server";
export class ResumeTaskDependencyService extends BaseService {
public async call(dependencyId: string, sourceTaskAttemptId: string) {
@@ -34,9 +35,19 @@ export class ResumeTaskDependencyService extends BaseService {
if (dependency.taskRun.runtimeEnvironment.type === "DEVELOPMENT") {
return;
}
const dependentRun = dependency.dependentAttempt.taskRun;
if (dependency.dependentAttempt.status === "PAUSED") {
if (!dependency.checkpointEventId) {
logger.error("Can't resume paused attempt without checkpoint event", {
attemptId: dependency.id,
});
await marqs?.acknowledgeMessage(dependentRun.id);
return;
}
await marqs?.enqueueMessage(
dependency.taskRun.runtimeEnvironment,
dependentRun.queue,
@@ -45,6 +56,7 @@ export class ResumeTaskDependencyService extends BaseService {
type: "RESUME",
completedAttemptIds: [sourceTaskAttemptId],
resumableAttemptId: dependency.dependentAttempt.id,
checkpointEventId: dependency.checkpointEventId,
},
dependentRun.concurrencyKey ?? undefined
);
packages/cli-v3/src/workers/prod/entry-point.ts
-2
View File
@@ -28,7 +28,6 @@ class ProdWorker {
private projectRef = process.env.TRIGGER_PROJECT_REF!;
private envId = process.env.TRIGGER_ENV_ID!;
private runId = process.env.TRIGGER_RUN_ID || "index-only";
private attemptId = process.env.TRIGGER_ATTEMPT_ID || "index-only";
private deploymentId = process.env.TRIGGER_DEPLOYMENT_ID!;
private deploymentVersion = process.env.TRIGGER_DEPLOYMENT_VERSION!;
@@ -187,7 +186,6 @@ class ProdWorker {
"x-pod-name": POD_NAME,
"x-trigger-content-hash": this.contentHash,
"x-trigger-project-ref": this.projectRef,
"x-trigger-attempt-id": this.attemptId,
"x-trigger-env-id": this.envId,
"x-trigger-deployment-id": this.deploymentId,
"x-trigger-run-id": this.runId,
packages/core/src/v3/runtime/prodRuntimeManager.ts
+3 -3
View File
@@ -104,7 +104,7 @@ export class ProdRuntimeManager implements RuntimeManager {
});
await this.ipc.send("WAIT_FOR_TASK", {
id: params.id,
friendlyId: params.id,
});
return await promise;
@@ -128,8 +128,8 @@ export class ProdRuntimeManager implements RuntimeManager {
);
await this.ipc.send("WAIT_FOR_BATCH", {
id: params.id,
runs: params.runs,
batchFriendlyId: params.id,
runFriendlyIds: params.runs,
});
const results = await promise;
packages/core/src/v3/schemas/messages.ts
+3 -3
View File
@@ -281,14 +281,14 @@ export const ProdChildToWorkerMessages = {
WAIT_FOR_TASK: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
friendlyId: z.string(),
}),
},
WAIT_FOR_BATCH: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
runs: z.string().array(),
batchFriendlyId: z.string(),
runFriendlyIds: z.string().array(),
}),
},
UNCAUGHT_EXCEPTION: {
packages/core/src/v3/schemas/schemas.ts
+8 -5
View File
@@ -207,11 +207,12 @@ export const CoordinatorToPlatformMessages = {
}),
z.object({
type: z.literal("WAIT_FOR_BATCH"),
id: z.string(),
batchFriendlyId: z.string(),
runFriendlyIds: z.string().array(),
}),
z.object({
type: z.literal("WAIT_FOR_TASK"),
id: z.string(),
friendlyId: z.string(),
}),
z.object({
type: z.literal("RETRYING_AFTER_FAILURE"),
@@ -376,7 +377,8 @@ export const ProdWorkerToCoordinatorMessages = {
WAIT_FOR_TASK: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
friendlyId: z.string(),
// This is the attempt that is waiting
attemptFriendlyId: z.string(),
}),
callback: z.object({
@@ -386,8 +388,9 @@ export const ProdWorkerToCoordinatorMessages = {
WAIT_FOR_BATCH: {
message: z.object({
version: z.literal("v1").default("v1"),
id: z.string(),
runs: z.string().array(),
batchFriendlyId: z.string(),
runFriendlyIds: z.string().array(),
// This is the attempt that is waiting
attemptFriendlyId: z.string(),
}),
callback: z.object({
packages/database/prisma/migrations/20240322172035_add_checkpoint_event_to_dependencies/migration.sql
+24
View File
@@ -0,0 +1,24 @@
/*
Warnings:
- A unique constraint covering the columns `[checkpointEventId]` on the table `BatchTaskRun` will be added. If there are existing duplicate values, this will fail.
- A unique constraint covering the columns `[checkpointEventId]` on the table `TaskRunDependency` will be added. If there are existing duplicate values, this will fail.
*/
-- AlterTable
ALTER TABLE "BatchTaskRun" ADD COLUMN "checkpointEventId" TEXT;
-- AlterTable
ALTER TABLE "TaskRunDependency" ADD COLUMN "checkpointEventId" TEXT;
-- CreateIndex
CREATE UNIQUE INDEX "BatchTaskRun_checkpointEventId_key" ON "BatchTaskRun"("checkpointEventId");
-- CreateIndex
CREATE UNIQUE INDEX "TaskRunDependency_checkpointEventId_key" ON "TaskRunDependency"("checkpointEventId");
-- AddForeignKey
ALTER TABLE "TaskRunDependency" ADD CONSTRAINT "TaskRunDependency_checkpointEventId_fkey" FOREIGN KEY ("checkpointEventId") REFERENCES "CheckpointRestoreEvent"("id") ON DELETE CASCADE ON UPDATE CASCADE;
-- AddForeignKey
ALTER TABLE "BatchTaskRun" ADD CONSTRAINT "BatchTaskRun_checkpointEventId_fkey" FOREIGN KEY ("checkpointEventId") REFERENCES "CheckpointRestoreEvent"("id") ON DELETE CASCADE ON UPDATE CASCADE;
packages/database/prisma/schema.prisma
+9
View File
@@ -1651,6 +1651,9 @@ model TaskRunDependency {
taskRun TaskRun @relation(fields: [taskRunId], references: [id], onDelete: Cascade, onUpdate: Cascade)
taskRunId String @unique
checkpointEvent CheckpointRestoreEvent? @relation(fields: [checkpointEventId], references: [id], onDelete: Cascade, onUpdate: Cascade)
checkpointEventId String? @unique
/// An attempt that is dependent on this task run.
dependentAttempt TaskRunAttempt? @relation("dependentAttempt", fields: [dependentAttemptId], references: [id])
dependentAttemptId String? @unique
@@ -1876,6 +1879,9 @@ model BatchTaskRun {
idempotencyKey String
taskIdentifier String
checkpointEvent CheckpointRestoreEvent? @relation(fields: [checkpointEventId], references: [id], onDelete: Cascade, onUpdate: Cascade)
checkpointEventId String? @unique
runtimeEnvironment RuntimeEnvironment @relation(fields: [runtimeEnvironmentId], references: [id], onDelete: Cascade, onUpdate: Cascade)
runtimeEnvironmentId String
@@ -2006,6 +2012,9 @@ model CheckpointRestoreEvent {
runtimeEnvironment RuntimeEnvironment @relation(fields: [runtimeEnvironmentId], references: [id], onDelete: Cascade, onUpdate: Cascade)
runtimeEnvironmentId String
taskRunDependency TaskRunDependency?
batchTaskRunDependency BatchTaskRun?
createdAt DateTime @default(now())
updatedAt DateTime @updatedAt
}
references/v3-catalog/src/trigger/retries.ts
+1 -1
View File
@@ -1,4 +1,4 @@
import { logger, retry, task } from "@trigger.dev/sdk/v3";
import { logger, retry, task, wait } from "@trigger.dev/sdk/v3";
import { cache } from "./utils/cache";
import { interceptor } from "./utils/interceptor";
references/v3-catalog/src/trigger/subtasks.ts
+68 -1
View File
@@ -1,4 +1,5 @@
import { Context, logger, task } from "@trigger.dev/sdk/v3";
import { logger, task } from "@trigger.dev/sdk/v3";
import { taskWithRetries } from "./retries";
export const simpleParentTask = task({
id: "simple-parent-task",
@@ -47,3 +48,69 @@ export const simpleChildTask = task({
logger.log("Simple child task payload", { payload, ctx });
},
});
export const subtasksWithRetries = task({
id: "subtasks-with-retries",
run: async (payload: { message: string }) => {
await taskWithRetries.triggerAndWait({
payload: {
message: `${payload.message} - 2.b`,
},
});
await taskWithRetries.batchTrigger({
items: [
{
payload: {
message: `${payload.message} - 2.c`,
},
},
{
payload: {
message: `${payload.message} - 2.cc`,
},
},
],
});
await taskWithRetries.batchTriggerAndWait({
items: [
{
payload: {
message: `${payload.message} - 2.d`,
},
},
{
payload: {
message: `${payload.message} - 2.dd`,
},
},
],
});
await taskWithRetries.triggerAndWait({
payload: {
message: `${payload.message} - 2.e`,
},
});
await taskWithRetries.batchTriggerAndWait({
items: [
{
payload: {
message: `${payload.message} - 2.f`,
},
},
{
payload: {
message: `${payload.message} - 2.ff`,
},
},
],
});
return {
hello: "world",
};
},
});