v3: usage tracking (#1163)

* Starting to measure wall time and cpu time in the workers, and reporting that via otel and to completed task run attempts

* Move usage tracking outside of the executor

* WIP prod usage tracking

* WIP

* WIP custom fetch to openmeter

* Create a usage client

* WIP

* WIP

* Implement new machine preset stuff and send usage reports to OpenMeter from webapp

* WIP

* Expose usage info to the client

* Add usage and cost to TaskEvent

* Add ability to globally configure the task machine preset

* Report start run usage

* Change the machine docs to use presets

* setExpirationTime to 24h

* Removed logs

* Update machines.mdx

* Removed console.logs

* Handle revalidating JWT tokens

* Couple tweaks

---------

Co-authored-by: Matt Aitken <[email protected]>

Author: ericallam

apps/kubernetes-provider/src/index.ts

@@ -7,7 +7,12 @@ import {
   TaskOperationsIndexOptions,
   TaskOperationsRestoreOptions,
 } from "@trigger.dev/core-apps";
-import { Machine, PostStartCauses, PreStopCauses, EnvironmentType } from "@trigger.dev/core/v3";
+import {
+  MachinePreset,
+  PostStartCauses,
+  PreStopCauses,
+  EnvironmentType,
+} from "@trigger.dev/core/v3";
 import { randomUUID } from "crypto";
 import { TaskMonitor } from "./taskMonitor";
 import { PodCleaner } from "./podCleaner";
@@ -399,10 +404,10 @@ class KubernetesTaskOperations implements TaskOperations {
     };
   }
 
-  #getResourcesFromMachineConfig(config: Machine): ComputeResources {
+  #getResourcesFromMachineConfig(preset: MachinePreset): ComputeResources {
     return {
-      cpu: `${config.cpu}`,
-      memory: `${config.memory}G`,
+      cpu: `${preset.cpu}`,
+      memory: `${preset.memory}G`,
     };
   }
 

apps/webapp/app/env.server.ts

@@ -1,5 +1,5 @@
-import { z } from "zod";
 import { SecretStoreOptionsSchema } from "./services/secrets/secretStoreOptionsSchema.server";
+import { z } from "zod";
 import { isValidRegex } from "./utils/regex";
 import { isValidDatabaseUrl } from "./utils/db";
 
@@ -189,6 +189,21 @@ const EnvironmentSchema = z.object({
   V2_MARQS_VERBOSE: z.string().default("0"),
   V3_MARQS_CONCURRENCY_MONITOR_ENABLED: z.string().default("0"),
   V2_MARQS_CONCURRENCY_MONITOR_ENABLED: z.string().default("0"),
+  /* Usage settings */
+  USAGE_EVENT_URL: z.string().optional(),
+  PROD_USAGE_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(),
+
+  CENTS_PER_HOUR_MICRO: z.coerce.number().default(0),
+  CENTS_PER_HOUR_SMALL_1X: z.coerce.number().default(0),
+  CENTS_PER_HOUR_SMALL_2X: z.coerce.number().default(0),
+  CENTS_PER_HOUR_MEDIUM_1X: z.coerce.number().default(0),
+  CENTS_PER_HOUR_MEDIUM_2X: z.coerce.number().default(0),
+  CENTS_PER_HOUR_LARGE_1X: z.coerce.number().default(0),
+  CENTS_PER_HOUR_LARGE_2X: z.coerce.number().default(0),
+  BASE_RUN_COST_IN_CENTS: z.coerce.number().default(0),
+
+  USAGE_OPEN_METER_API_KEY: z.string().optional(),
+  USAGE_OPEN_METER_BASE_URL: z.string().optional(),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;

apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts

@@ -80,7 +80,7 @@ export class ApiRetrieveRunPresenter extends BasePresenter {
         version: taskRun.lockedToVersion ? taskRun.lockedToVersion.version : undefined,
         createdAt: taskRun.createdAt ?? undefined,
         updatedAt: taskRun.updatedAt ?? undefined,
-        startedAt: taskRun.lockedAt ?? undefined,
+        startedAt: taskRun.startedAt ?? undefined,
         finishedAt: ApiRetrieveRunPresenter.isStatusFinished(apiStatus)
           ? taskRun.updatedAt
           : undefined,

apps/webapp/app/presenters/v3/RunListPresenter.server.ts

@@ -156,7 +156,7 @@ export class RunListPresenter extends BasePresenter {
         runtimeEnvironmentId: string;
         status: TaskRunStatus;
         createdAt: Date;
-        lockedAt: Date | null;
+        startedAt: Date | null;
         updatedAt: Date;
         isTest: boolean;
         spanId: string;
@@ -172,7 +172,7 @@ export class RunListPresenter extends BasePresenter {
     tr."runtimeEnvironmentId" AS "runtimeEnvironmentId",
     tr.status AS status,
     tr."createdAt" AS "createdAt",
-    tr."lockedAt" AS "lockedAt",
+    tr."startedAt" AS "startedAt",
     tr."updatedAt" AS "updatedAt",
     tr."isTest" AS "isTest",
     tr."spanId" AS "spanId",
@@ -278,7 +278,7 @@ export class RunListPresenter extends BasePresenter {
           number: Number(run.number),
           createdAt: run.createdAt.toISOString(),
           updatedAt: run.updatedAt.toISOString(),
-          startedAt: run.lockedAt ? run.lockedAt.toISOString() : undefined,
+          startedAt: run.startedAt ? run.startedAt.toISOString() : undefined,
           hasFinished,
           finishedAt: hasFinished ? run.updatedAt.toISOString() : undefined,
           isTest: run.isTest,

apps/webapp/app/presenters/v3/TaskListPresenter.server.ts

@@ -311,7 +311,7 @@ export class TaskListPresenter extends BasePresenter {
     >`    
     SELECT 
       tr."taskIdentifier", 
-      AVG(EXTRACT(EPOCH FROM (tr."updatedAt" - tr."lockedAt"))) as duration
+      AVG(EXTRACT(EPOCH FROM (tr."updatedAt" - tr."startedAt"))) as duration
       FROM 
       ${sqlDatabaseSchema}."TaskRun" as tr
     WHERE 

apps/webapp/app/routes/api.v1.deployments.ts

@@ -1,5 +1,8 @@
 import { ActionFunctionArgs, json } from "@remix-run/server-runtime";
-import { InitializeDeploymentRequestBody, InitializeDeploymentResponseBody } from "@trigger.dev/core/v3";
+import {
+  InitializeDeploymentRequestBody,
+  InitializeDeploymentResponseBody,
+} from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
 import { authenticateApiRequest } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
@@ -37,13 +40,11 @@ export async function action({ request, params }: ActionFunctionArgs) {
     contentHash: deployment.contentHash,
     shortCode: deployment.shortCode,
     version: deployment.version,
-    externalBuildData: deployment.externalBuildData as InitializeDeploymentResponseBody["externalBuildData"],
+    externalBuildData:
+      deployment.externalBuildData as InitializeDeploymentResponseBody["externalBuildData"],
     imageTag,
-    registryHost: env.DEPLOY_REGISTRY_HOST
-  }
+    registryHost: env.DEPLOY_REGISTRY_HOST,
+  };
 
-  return json(
-    responseBody,
-    { status: 200 }
-  );
+  return json(responseBody, { status: 200 });
 }

apps/webapp/app/routes/api.v1.projects.$projectRef.envvars.$slug.$name.ts

@@ -123,7 +123,7 @@ export async function loader({ params, request }: LoaderFunctionArgs) {
 
   const repository = new EnvironmentVariablesRepository();
 
-  const variables = await repository.getEnvironment(environment.project.id, environment.id, true);
+  const variables = await repository.getEnvironment(environment.project.id, environment.id);
 
   const environmentVariable = variables.find((v) => v.key === parsedParams.data.name);
 

apps/webapp/app/routes/api.v1.projects.$projectRef.envvars.$slug.ts

@@ -80,7 +80,7 @@ export async function loader({ params, request }: LoaderFunctionArgs) {
 
   const repository = new EnvironmentVariablesRepository();
 
-  const variables = await repository.getEnvironment(environment.project.id, environment.id, true);
+  const variables = await repository.getEnvironment(environment.project.id, environment.id);
 
   return json(variables.map((variable) => ({ name: variable.key, value: variable.value })));
 }

apps/webapp/app/routes/api.v1.projects.$projectRef.envvars.ts

@@ -2,7 +2,7 @@ import { LoaderFunctionArgs, json } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { prisma } from "~/db.server";
 import { authenticateApiRequest } from "~/services/apiAuth.server";
-import { EnvironmentVariablesRepository } from "~/v3/environmentVariables/environmentVariablesRepository.server";
+import { resolveVariablesForEnvironment } from "~/v3/environmentVariables/environmentVariablesRepository.server";
 
 const ParamsSchema = z.object({
   projectRef: z.string(),
@@ -41,9 +41,7 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
     return json({ error: "Project not found" }, { status: 404 });
   }
 
-  const repository = new EnvironmentVariablesRepository();
-
-  const variables = await repository.getEnvironmentVariables(project.id, authenticatedEnv.id);
+  const variables = await resolveVariablesForEnvironment(authenticatedEnv);
 
   return json({
     variables: variables.reduce((acc: Record<string, string>, variable) => {

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

@@ -0,0 +1,97 @@
+import { ActionFunctionArgs } from "@remix-run/server-runtime";
+import { MachinePresetName } from "@trigger.dev/core/v3";
+import { z } from "zod";
+import { prisma } from "~/db.server";
+import { validateJWTTokenAndRenew } 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";
+
+const JWTPayloadSchema = z.object({
+  environment_id: z.string(),
+  org_id: z.string(),
+  project_id: z.string(),
+  run_id: z.string(),
+  machine_preset: z.string(),
+});
+
+const BodySchema = z.object({
+  durationMs: z.number(),
+});
+
+export async function action({ request }: ActionFunctionArgs) {
+  // Ensure this is a POST request
+  if (request.method.toUpperCase() !== "POST") {
+    return { status: 405, body: "Method Not Allowed" };
+  }
+
+  const jwtResult = await validateJWTTokenAndRenew(request, JWTPayloadSchema);
+
+  if (!jwtResult) {
+    return { status: 401, body: "Unauthorized" };
+  }
+
+  const rawJson = await request.json();
+
+  const json = BodySchema.safeParse(rawJson);
+
+  if (!json.success) {
+    logger.error("Failed to parse request body", { rawJson });
+
+    return { status: 400, body: "Bad Request" };
+  }
+
+  const preset = machinePresetFromName(jwtResult.payload.machine_preset as MachinePresetName);
+
+  logger.debug("[/api/v1/usage/ingest] Reporting usage", { jwtResult, json: json.data, preset });
+
+  if (json.data.durationMs > 0) {
+    const costInCents = json.data.durationMs * preset.centsPerMs;
+
+    await prisma.taskRun.update({
+      where: {
+        id: jwtResult.payload.run_id,
+      },
+      data: {
+        usageDurationMs: {
+          increment: json.data.durationMs,
+        },
+        costInCents: {
+          increment: json.data.durationMs * preset.centsPerMs,
+        },
+      },
+    });
+
+    try {
+      await reportUsageEvent({
+        source: "webapp",
+        type: "usage",
+        subject: jwtResult.payload.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: jwtResult.payload.org_id,
+        data: {
+          costInCents: String(costInCents),
+        },
+        additionalData: {
+          durationMs: json.data.durationMs,
+        },
+      });
+    }
+  }
+
+  return new Response(null, {
+    status: 200,
+    headers: {
+      "x-trigger-jwt": jwtResult.jwt,
+    },
+  });
+}