Report start run usage

Author: ericallam

apps/webapp/app/routes/api.v1.usage.ingest.ts

@@ -4,6 +4,7 @@ import { z } from "zod";
 import { prisma } from "~/db.server";
 import { validateJWTToken } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
+import { workerQueue } from "~/services/worker.server";
 import { machinePresetFromName } from "~/v3/machinePresets.server";
 import { reportUsageEvent } from "~/v3/openMeter.server";
 
@@ -66,15 +67,29 @@ export async function action({ request }: ActionFunctionArgs) {
       },
     });
 
-    await reportUsageEvent({
-      source: "webapp",
-      type: "usage",
-      subject: jwtPayload.org_id,
-      data: {
-        durationMs: json.data.durationMs,
-        costInCents: String(costInCents),
-      },
-    });
+    try {
+      await reportUsageEvent({
+        source: "webapp",
+        type: "usage",
+        subject: jwtPayload.org_id,
+        data: {
+          durationMs: json.data.durationMs,
+          costInCents: String(costInCents),
+        },
+      });
+    } catch (e) {
+      logger.error("Failed to report usage event, enqueing v3.reportUsage", { error: e });
+
+      await workerQueue.enqueue("v3.reportUsage", {
+        orgId: jwtPayload.org_id,
+        data: {
+          costInCents: String(costInCents),
+        },
+        additionalData: {
+          durationMs: json.data.durationMs,
+        },
+      });
+    }
   }
 
   return new Response(null, {

apps/webapp/app/services/worker.server.ts

@@ -46,6 +46,7 @@ import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout.se
 import { ResumeTaskService } from "./tasks/resumeTask.server";
 import { RequeueV2Message } from "~/v3/marqs/requeueV2Message.server";
 import { MarqsConcurrencyMonitor } from "~/v3/marqs/concurrencyMonitor.server";
+import { reportUsageEvent } from "~/v3/openMeter.server";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -169,6 +170,13 @@ const workerCatalog = {
   "v2.requeueMessage": z.object({
     runId: z.string(),
   }),
+  "v3.reportUsage": z.object({
+    orgId: z.string(),
+    data: z.object({
+      costInCents: z.string(),
+    }),
+    additionalData: z.record(z.any()).optional(),
+  }),
 };
 
 const executionWorkerCatalog = {
@@ -649,6 +657,21 @@ function getWorkerQueue() {
           await service.call(payload.runId);
         },
       },
+      "v3.reportUsage": {
+        priority: 0,
+        maxAttempts: 8,
+        handler: async (payload, job) => {
+          await reportUsageEvent({
+            source: "webapp",
+            type: "usage",
+            subject: payload.orgId,
+            data: {
+              costInCents: payload.data.costInCents,
+              ...payload.additionalData,
+            },
+          });
+        },
+      },
     },
   });
 }

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

@@ -5,18 +5,21 @@ import { logger } from "~/services/logger.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import { BaseService, ServiceValidationError } from "./baseService.server";
 import { TaskRun, TaskRunAttempt } from "@trigger.dev/database";
+import { machinePresetFromConfig } from "../machinePresets.server";
+import { workerQueue } from "~/services/worker.server";
 
 export class CreateTaskRunAttemptService extends BaseService {
   public async call(
     runId: string,
-    env?: AuthenticatedEnvironment,
+    authenticatedEnv?: AuthenticatedEnvironment,
     setToExecuting = true
   ): Promise<{
     execution: TaskRunExecution;
     run: TaskRun;
     attempt: TaskRunAttempt;
   }> {
-    const environment = env ?? (await getAuthenticatedEnvironmentFromRun(runId, this._prisma));
+    const environment =
+      authenticatedEnv ?? (await getAuthenticatedEnvironmentFromRun(runId, this._prisma));
 
     if (!environment) {
       throw new ServiceValidationError("Environment not found", 404);
@@ -128,6 +131,20 @@ export class CreateTaskRunAttemptService extends BaseService {
         throw new ServiceValidationError("Failed to create task run attempt", 500);
       }
 
+      if (taskRunAttempt.number === 1 && taskRun.baseCostInCents > 0) {
+        await workerQueue.enqueue("v3.reportUsage", {
+          orgId: environment.organizationId,
+          data: {
+            costInCents: String(taskRun.baseCostInCents),
+          },
+          additionalData: {
+            runId: taskRun.id,
+          },
+        });
+      }
+
+      const machinePreset = machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
+
       const execution: TaskRunExecution = {
         task: {
           id: taskRun.lockedBy.slug,
@@ -151,6 +168,10 @@ export class CreateTaskRunAttemptService extends BaseService {
           tags: taskRun.tags.map((tag) => tag.name),
           isTest: taskRun.isTest,
           idempotencyKey: taskRun.idempotencyKey ?? undefined,
+          startedAt: taskRun.startedAt ?? taskRun.createdAt,
+          durationMs: taskRun.usageDurationMs,
+          costInCents: taskRun.costInCents,
+          baseCostInCents: taskRun.baseCostInCents,
         },
         queue: {
           id: queue.friendlyId,
@@ -176,6 +197,7 @@ export class CreateTaskRunAttemptService extends BaseService {
           taskRun.batchItems[0] && taskRun.batchItems[0].batchTaskRun
             ? { id: taskRun.batchItems[0].batchTaskRun.friendlyId }
             : undefined,
+        machine: machinePreset,
       };
 
       return {

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

@@ -102,18 +102,18 @@ export const usage = {
    * export const myTask = task({
    *  id: "my-task",
    *  run: async (payload, { ctx }) => {
-   *    const { result, usage } = await usage.measure(async () => {
+   *    const { result, compute } = await usage.measure(async () => {
    *      // Do some work
    *      return "result";
    *    });
    *
    *    console.log("Result", result);
-   *    console.log("Cost and duration", { cost: usage.costInCents, duration: usage.durationMs });
+   *    console.log("Cost and duration", { cost: compute.costInCents, duration: compute.durationMs });
    *  },
    * });
    * ```
    */
-  measure: async <T>(cb: () => Promise<T>): Promise<{ result: T; usage: ComputeUsage }> => {
+  measure: async <T>(cb: () => Promise<T>): Promise<{ result: T; compute: ComputeUsage }> => {
     const measurement = usageApi.start();
 
     const result = await cb();
@@ -125,7 +125,7 @@ export const usage = {
 
     return {
       result,
-      usage: {
+      compute: {
         costInCents,
         durationMs: sample.cpuTime,
       },