add pre-dequeue callback to determine max resources

Author: nicktrn

apps/webapp/app/routes/api.v1.worker-actions.dequeue.ts

@@ -1,10 +1,16 @@
 import { json, TypedResponse } from "@remix-run/server-runtime";
-import { WorkerApiDequeueResponseBody } from "@trigger.dev/worker";
-import { createLoaderWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+import { WorkerApiDequeueRequestBody, WorkerApiDequeueResponseBody } from "@trigger.dev/worker";
+import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
-export const loader = createLoaderWorkerApiRoute(
-  {},
-  async ({ authenticatedWorker }): Promise<TypedResponse<WorkerApiDequeueResponseBody>> => {
-    return json(await authenticatedWorker.dequeue());
+export const action = createActionWorkerApiRoute(
+  {
+    body: WorkerApiDequeueRequestBody,
+  },
+  async ({ authenticatedWorker, body }): Promise<TypedResponse<WorkerApiDequeueResponseBody>> => {
+    return json(
+      await authenticatedWorker.dequeue({
+        maxResources: body.maxResources,
+      })
+    );
   }
 );

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

@@ -18,6 +18,7 @@ import {
   ExecutionResult,
   MachinePreset,
   WaitForDurationResult,
+  MachineResources,
 } from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
 import { $transaction } from "~/db.server";
@@ -525,12 +526,19 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
     });
   }
 
-  async dequeue(maxRunCount = 10): Promise<DequeuedMessage[]> {
+  async dequeue({
+    maxRunCount = 10,
+    maxResources,
+  }: {
+    maxRunCount?: number;
+    maxResources?: MachineResources;
+  } = {}): Promise<DequeuedMessage[]> {
     if (this.type === WorkerInstanceGroupType.MANAGED) {
       return await this._engine.dequeueFromMasterQueue({
         consumerId: this.workerInstanceId,
         masterQueue: this.masterQueue,
         maxRunCount,
+        maxResources,
       });
     }
 

packages/worker/src/supervisor/http.ts

@@ -3,6 +3,7 @@ import { zodfetch, ApiError } from "@trigger.dev/core/v3/zodfetch";
 import {
   WorkerApiConnectRequestBody,
   WorkerApiConnectResponseBody,
+  WorkerApiDequeueRequestBody,
   WorkerApiDequeueResponseBody,
   WorkerApiHeartbeatRequestBody,
   WorkerApiHeartbeatResponseBody,
@@ -61,14 +62,17 @@ export class SupervisorHttpClient {
     );
   }
 
-  async dequeue() {
+  async dequeue(body: WorkerApiDequeueRequestBody) {
     return wrapZodFetch(
       WorkerApiDequeueResponseBody,
       `${this.apiUrl}/api/v1/worker-actions/dequeue`,
       {
+        method: "POST",
         headers: {
           ...this.defaultHeaders,
+          "Content-Type": "application/json",
         },
+        body: JSON.stringify(body),
       }
     );
   }

packages/worker/src/supervisor/queueConsumer.ts

@@ -1,14 +1,21 @@
+import { MachineResources } from "@trigger.dev/core/v3";
 import { SupervisorHttpClient } from "./http.js";
 import { WorkerApiDequeueResponseBody } from "./schemas.js";
 
 type RunQueueConsumerOptions = {
   client: SupervisorHttpClient;
   intervalMs?: number;
+  preDequeue?: () => Promise<{
+    maxResources?: MachineResources;
+  }>;
   onDequeue: (messages: WorkerApiDequeueResponseBody) => Promise<void>;
 };
 
 export class RunQueueConsumer {
   private readonly client: SupervisorHttpClient;
+  private readonly preDequeue?: () => Promise<{
+    maxResources?: MachineResources;
+  }>;
   private readonly onDequeue: (messages: WorkerApiDequeueResponseBody) => Promise<void>;
 
   private intervalMs: number;
@@ -17,6 +24,7 @@ export class RunQueueConsumer {
   constructor(opts: RunQueueConsumerOptions) {
     this.isEnabled = false;
     this.intervalMs = opts.intervalMs ?? 5_000;
+    this.preDequeue = opts.preDequeue;
     this.onDequeue = opts.onDequeue;
     this.client = opts.client;
   }
@@ -46,8 +54,18 @@ export class RunQueueConsumer {
       return;
     }
 
+    let maxResources: MachineResources | undefined;
+    if (this.preDequeue) {
+      try {
+        const preDequeueResult = await this.preDequeue();
+        maxResources = preDequeueResult.maxResources;
+      } catch (preDequeueError) {
+        console.error("[RunQueueConsumer] preDequeue error", { error: preDequeueError });
+      }
+    }
+
     try {
-      const response = await this.client.dequeue();
+      const response = await this.client.dequeue({ maxResources });
 
       if (!response.success) {
         console.error("[RunQueueConsumer] Failed to dequeue", { error: response.error });

packages/worker/src/supervisor/schemas.ts

@@ -2,6 +2,7 @@ import { z } from "zod";
 import {
   CompleteRunAttemptResult,
   DequeuedMessage,
+  MachineResources,
   RunExecutionData,
   StartRunAttemptResult,
   TaskRunExecutionResult,
@@ -40,6 +41,11 @@ export const WorkerApiConnectResponseBody = z.object({
 });
 export type WorkerApiConnectResponseBody = z.infer<typeof WorkerApiConnectResponseBody>;
 
+export const WorkerApiDequeueRequestBody = z.object({
+  maxResources: MachineResources.optional(),
+});
+export type WorkerApiDequeueRequestBody = z.infer<typeof WorkerApiDequeueRequestBody>;
+
 export const WorkerApiDequeueResponseBody = DequeuedMessage.array();
 export type WorkerApiDequeueResponseBody = z.infer<typeof WorkerApiDequeueResponseBody>;
 

packages/worker/src/supervisor/session.ts

@@ -1,4 +1,4 @@
-import { HeartbeatService } from "@trigger.dev/core/v3";
+import { HeartbeatService, MachineResources } from "@trigger.dev/core/v3";
 import { SupervisorHttpClient } from "./http.js";
 import { SupervisorClientCommonOptions } from "./types.js";
 import { WorkerApiDequeueResponseBody, WorkerApiHeartbeatRequestBody } from "./schemas.js";
@@ -13,6 +13,9 @@ import { getDefaultWorkerHeaders } from "./util.js";
 type SupervisorSessionOptions = SupervisorClientCommonOptions & {
   heartbeatIntervalSeconds?: number;
   dequeueIntervalMs?: number;
+  preDequeue?: () => Promise<{
+    maxResources?: MachineResources;
+  }>;
 };
 
 export class SupervisorSession extends EventEmitter<WorkerEvents> {
@@ -30,6 +33,7 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
     this.httpClient = new SupervisorHttpClient(opts);
     this.queueConsumer = new RunQueueConsumer({
       client: this.httpClient,
+      preDequeue: opts.preDequeue,
       onDequeue: this.onDequeue.bind(this),
       intervalMs: opts.dequeueIntervalMs,
     });