Expose usage info to the client

Author: ericallam

apps/webapp/app/env.server.ts

@@ -199,6 +199,7 @@ const EnvironmentSchema = z.object({
   CENTS_PER_HOUR_MEDIUM_1X: z.coerce.number().positive().default(0),
   CENTS_PER_HOUR_MEDIUM_2X: z.coerce.number().positive().default(0),
   CENTS_PER_HOUR_LARGE_1X: z.coerce.number().positive().default(0),
+  BASE_RUN_COST_IN_CENTS: z.coerce.number().positive().default(0),
 
   USAGE_OPEN_METER_API_KEY: z.string().optional(),
   USAGE_OPEN_METER_BASE_URL: z.string().optional(),

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

@@ -42,6 +42,7 @@ import { tracer } from "../tracer.server";
 import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { EnvironmentVariable } from "../environmentVariables/repository";
 import { machinePresetFromConfig } from "../machinePresets.server";
+import { env } from "~/env.server";
 
 const WithTraceContext = z.object({
   traceparent: z.string().optional(),
@@ -411,6 +412,7 @@ export class SharedQueueConsumer {
             lockedById: backgroundTask.id,
             lockedToVersionId: deployment.worker.id,
             startedAt: existingTaskRun.startedAt ?? new Date(),
+            baseCostInCents: env.BASE_RUN_COST_IN_CENTS,
           },
           include: {
             runtimeEnvironment: true,
@@ -1032,6 +1034,7 @@ class SharedQueueTasks {
         idempotencyKey: taskRun.idempotencyKey ?? undefined,
         durationMs: taskRun.usageDurationMs,
         costInCents: taskRun.costInCents,
+        baseCostInCents: taskRun.baseCostInCents,
       },
       queue: {
         id: queue.friendlyId,

packages/database/prisma/migrations/20240612171759_add_task_run_base_cost/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "TaskRun" ADD COLUMN     "baseCostInCents" DOUBLE PRECISION NOT NULL DEFAULT 0;

packages/database/prisma/schema.prisma

@@ -1648,6 +1648,7 @@ model TaskRun {
 
   usageDurationMs Int   @default(0)
   costInCents     Float @default(0)
+  baseCostInCents Float @default(0)
 
   lockedAt   DateTime?
   lockedBy   BackgroundWorkerTask? @relation(fields: [lockedById], references: [id])

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

@@ -12,14 +12,12 @@ export type CurrentUsage = {
     total: ComputeUsage;
   };
   baseCostInCents: number;
-  totalInCents: number;
+  totalCostInCents: number;
 };
 
 export const usage = {
   /**
-   * Get the current usage of this task run attempt.
-   *
-   * @returns The current usage of this task run attempt.
+   * Get the current running usage of this task run.
    *
    * @example
    *
@@ -33,35 +31,104 @@ export const usage = {
    *
    *   const currentUsage = usage.getCurrent();
    *
-   *   console.log("Current cost and duration", {
-   *     cost: currentUsage.costInCents,
-   *     duration: currentUsage.durationMs,
+   *   // You have access to the current compute cost and duration up to this point
+   *   console.log("Current attempt compute cost and duration", {
+   *     cost: currentUsage.compute.attempt.costInCents,
+   *     duration: currentUsage.compute.attempt.durationMs,
+   *   });
+   *
+   *   // You also can see the total compute cost and duration up to this point in the run, across all attempts
+   *   console.log("Current total compute cost and duration", {
+   *     cost: currentUsage.compute.total.costInCents,
+   *     duration: currentUsage.compute.total.durationMs,
    *   });
    *
-   *   // Use ctx to access the total run cost and duration
-   *   console.log("Total cost and duration", {
-   *     cost: ctx.run.costInCents + currentUsage.costInCents,
-   *     duration: ctx.run.durationMs + currentUsage.durationMs,
+   *   // You can see the base cost of the run, which is the cost of the run before any compute costs
+   *   console.log("Total cost", {
+   *     cost: currentUsage.totalCostInCents,
+   *     baseCost: currentUsage.baseCostInCents,
    *   });
    *  },
    * });
    * ```
    */
   getCurrent: (): CurrentUsage => {
     const sample = usageApi.sample();
+    const machine = taskContext.ctx?.machine;
+    const run = taskContext.ctx?.run;
 
     if (!sample) {
       return {
-        costInCents: 0,
-        durationMs: 0,
+        compute: {
+          attempt: {
+            costInCents: 0,
+            durationMs: 0,
+          },
+          total: {
+            costInCents: run?.costInCents ?? 0,
+            durationMs: run?.durationMs ?? 0,
+          },
+        },
+        baseCostInCents: run?.baseCostInCents ?? 0,
+        totalCostInCents: (run?.costInCents ?? 0) + (run?.baseCostInCents ?? 0),
       };
     }
 
+    const currentCostInCents = machine?.centsPerMs ? sample.cpuTime * machine.centsPerMs : 0;
+
+    return {
+      compute: {
+        attempt: {
+          costInCents: currentCostInCents,
+          durationMs: sample.cpuTime,
+        },
+        total: {
+          costInCents: (run?.costInCents ?? 0) + currentCostInCents,
+          durationMs: (run?.durationMs ?? 0) + sample.cpuTime,
+        },
+      },
+      baseCostInCents: run?.baseCostInCents ?? 0,
+      totalCostInCents: (run?.costInCents ?? 0) + currentCostInCents + (run?.baseCostInCents ?? 0),
+    };
+  },
+  /**
+   * Measure the cost and duration of a function.
+   *
+   * @example
+   *
+   * ```typescript
+   * import { usage } from "@trigger.dev/sdk/v3";
+   *
+   * export const myTask = task({
+   *  id: "my-task",
+   *  run: async (payload, { ctx }) => {
+   *    const { result, usage } = await usage.measure(async () => {
+   *      // Do some work
+   *      return "result";
+   *    });
+   *
+   *    console.log("Result", result);
+   *    console.log("Cost and duration", { cost: usage.costInCents, duration: usage.durationMs });
+   *  },
+   * });
+   * ```
+   */
+  measure: async <T>(cb: () => Promise<T>): Promise<{ result: T; usage: ComputeUsage }> => {
+    const measurement = usageApi.start();
+
+    const result = await cb();
+
+    const sample = usageApi.stop(measurement);
     const machine = taskContext.ctx?.machine;
 
+    const costInCents = machine?.centsPerMs ? sample.cpuTime * machine.centsPerMs : 0;
+
     return {
-      costInCents: machine?.centsPerMs ? sample.cpuTime * machine.centsPerMs : 0,
-      durationMs: sample.cpuTime,
+      result,
+      usage: {
+        costInCents,
+        durationMs: sample.cpuTime,
+      },
     };
   },
 };

references/v3-catalog/src/trigger/longRunning.ts

@@ -1,39 +1,14 @@
-import { logger, task, wait, usage } from "@trigger.dev/sdk/v3";
+import { logger, task, wait } from "@trigger.dev/sdk/v3";
 
 export const longRunning = task({
   id: "long-running",
-  machine: {
-    preset: "medium-2x",
-  },
   run: async (payload: { message: string }, { ctx }) => {
     logger.info("Long running payloadddd", { payload });
 
-    if (ctx.machine) {
-      logger.info("Machine preset", { preset: ctx.machine });
-    }
-
-    logger.info("Cost and duration", { cost: ctx.run.costInCents, duration: ctx.run.durationMs });
-
     // Wait for 3 minutes
     await new Promise((resolve) => setTimeout(resolve, 5000));
 
-    let currentUsage = usage.getCurrent();
-
-    logger.info("Current Cost and duration (before wait)", {
-      cost: currentUsage.costInCents,
-      duration: currentUsage.durationMs,
-    });
-
     await wait.for({ seconds: 5 });
-
-    currentUsage = usage.getCurrent();
-
-    logger.info("Current Cost and duration (after wait)", {
-      cost: currentUsage.costInCents,
-      duration: currentUsage.durationMs,
-    });
-
-    throw new Error(`This is an error at ${new Date().toISOString()}`);
   },
 });
 

references/v3-catalog/src/trigger/usage.ts

@@ -0,0 +1,29 @@
+import { logger, task, wait, usage } from "@trigger.dev/sdk/v3";
+
+export const usagePlayground = task({
+  id: "usage-playground",
+  machine: {
+    preset: "medium-2x",
+  },
+  run: async (payload: { duration: number }, { ctx }) => {
+    if (ctx.machine) {
+      logger.info("Machine preset", { preset: ctx.machine });
+    }
+
+    logger.info("Cost and duration", { cost: ctx.run.costInCents, duration: ctx.run.durationMs });
+
+    await new Promise((resolve) => setTimeout(resolve, payload.duration));
+
+    let currentUsage = usage.getCurrent();
+
+    logger.info("Current Cost and duration (before wait)", { currentUsage });
+
+    await wait.for({ seconds: 5 });
+
+    currentUsage = usage.getCurrent();
+
+    logger.info("Current Cost and duration (after wait)", { currentUsage });
+
+    throw new Error(`This is an error at ${new Date().toISOString()}`);
+  },
+});