Fix OOM detection, again

Author: nicktrn

apps/kubernetes-provider/src/index.ts

@@ -529,27 +529,27 @@ provider.listen();
 
 const taskMonitor = new TaskMonitor({
   runtimeEnv: RUNTIME_ENV,
-  onIndexFailure: async (deploymentId, failureInfo) => {
-    logger.log("Indexing failed", { deploymentId, failureInfo });
+  onIndexFailure: async (deploymentId, details) => {
+    logger.log("Indexing failed", { deploymentId, details });
 
     try {
       provider.platformSocket.send("INDEXING_FAILED", {
         deploymentId,
         error: {
-          name: `Crashed with exit code ${failureInfo.exitCode}`,
-          message: failureInfo.reason,
-          stack: failureInfo.logs,
+          name: `Crashed with exit code ${details.exitCode}`,
+          message: details.reason,
+          stack: details.logs,
         },
       });
     } catch (error) {
       logger.error(error);
     }
   },
-  onRunFailure: async (runId, failureInfo) => {
-    logger.log("Run failed:", { runId, failureInfo });
+  onRunFailure: async (runId, details) => {
+    logger.log("Run failed:", { runId, details });
 
     try {
-      provider.platformSocket.send("WORKER_CRASHED", { runId, ...failureInfo });
+      provider.platformSocket.send("WORKER_CRASHED", { runId, ...details });
     } catch (error) {
       logger.error(error);
     }

apps/kubernetes-provider/src/taskMonitor.ts

@@ -1,25 +1,20 @@
 import * as k8s from "@kubernetes/client-node";
 import { SimpleLogger } from "@trigger.dev/core-apps";
+import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process";
 import { setTimeout } from "timers/promises";
 import PQueue from "p-queue";
+import type { Prettify } from "@trigger.dev/core/v3";
 
-type IndexFailureHandler = (
-  deploymentId: string,
-  failureInfo: {
-    exitCode: number;
-    reason: string;
-    logs: string;
-  }
-) => Promise<any>;
-
-type RunFailureHandler = (
-  runId: string,
-  failureInfo: {
-    exitCode: number;
-    reason: string;
-    logs: string;
-  }
-) => Promise<any>;
+type FailureDetails = Prettify<{
+  exitCode: number;
+  reason: string;
+  logs: string;
+  overrideCompletion: boolean;
+}>;
+
+type IndexFailureHandler = (deploymentId: string, details: FailureDetails) => Promise<any>;
+
+type RunFailureHandler = (runId: string, details: FailureDetails) => Promise<any>;
 
 type TaskMonitorOptions = {
   runtimeEnv: "local" | "kubernetes";
@@ -144,8 +139,7 @@ export class TaskMonitor {
     const containerState = this.#getContainerStateSummary(containerStatus.state);
     const exitCode = containerState.exitCode ?? -1;
 
-    // We use this special exit code to signal any errors were already handled elsewhere
-    if (exitCode === 111) {
+    if (exitCode === EXIT_CODE_ALREADY_HANDLED) {
       return;
     }
 
@@ -162,6 +156,7 @@ export class TaskMonitor {
 
     let reason = rawReason || "Unknown error";
     let logs = rawLogs || "";
+    let overrideCompletion = false;
 
     switch (rawReason) {
       case "Error":
@@ -181,8 +176,10 @@ export class TaskMonitor {
         }
         break;
       case "OOMKilled":
-        reason =
-          "Process ran out of memory! Try choosing a machine preset with more memory for this task.";
+        overrideCompletion = true;
+        reason = `${
+          exitCode === EXIT_CODE_CHILD_NONZERO ? "Child process" : "Parent process"
+        } ran out of memory! Try choosing a machine preset with more memory for this task.`;
         break;
       default:
         break;
@@ -192,7 +189,8 @@ export class TaskMonitor {
       exitCode,
       reason,
       logs,
-    };
+      overrideCompletion,
+    } satisfies FailureDetails;
 
     const app = pod.metadata?.labels?.app;
 

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

@@ -13,6 +13,7 @@ export type CrashTaskRunServiceOptions = {
   logs?: string;
   crashAttempts?: boolean;
   crashedAt?: Date;
+  overrideCompletion?: boolean;
 };
 
 export class CrashTaskRunService extends BaseService {
@@ -36,7 +37,7 @@ export class CrashTaskRunService extends BaseService {
     }
 
     // Make sure the task run is in a crashable state
-    if (!isCrashableRunStatus(taskRun.status)) {
+    if (!opts.overrideCompletion && !isCrashableRunStatus(taskRun.status)) {
       logger.error("Task run is not in a crashable state", { runId, status: taskRun.status });
       return;
     }

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

@@ -18,7 +18,8 @@ export class DeploymentIndexFailed extends BaseService {
       message: string;
       stack?: string;
       stderr?: string;
-    }
+    },
+    overrideCompletion = false
   ) {
     const isFriendlyId = maybeFriendlyId.startsWith("deployment_");
 
@@ -38,6 +39,15 @@ export class DeploymentIndexFailed extends BaseService {
     }
 
     if (FINAL_DEPLOYMENT_STATUSES.includes(deployment.status)) {
+      if (overrideCompletion) {
+        logger.error("No support for overriding final deployment statuses just yet", {
+          id: deployment.id,
+          status: deployment.status,
+          previousError: deployment.errorData,
+          incomingError: error,
+        });
+      }
+
       logger.error("Worker deployment already in final state", {
         id: deployment.id,
         status: deployment.status,

packages/cli-v3/src/workers/prod/entry-point.ts

@@ -5,12 +5,14 @@ import {
   PreStopCauses,
   ProdWorkerToCoordinatorMessages,
   TaskResource,
+  TaskRunErrorCodes,
   TaskRunFailedExecutionResult,
   WaitReason,
 } from "@trigger.dev/core/v3";
 import { ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket";
 import { HttpReply, getRandomPortNumber } from "@trigger.dev/core-apps/http";
 import { SimpleLogger } from "@trigger.dev/core-apps/logger";
+import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process";
 import { readFile } from "node:fs/promises";
 import { createServer } from "node:http";
 import { ProdBackgroundWorker } from "./backgroundWorker";
@@ -97,12 +99,12 @@ class ProdWorker {
     logger.log("Unhandled signal", { signal });
   }
 
-  async #exitGracefully(gracefulExitTimeoutElapsed = false) {
+  async #exitGracefully(gracefulExitTimeoutElapsed = false, exitCode = 0) {
     await this.#backgroundWorker.close(gracefulExitTimeoutElapsed);
 
     if (!gracefulExitTimeoutElapsed) {
       // TODO: Maybe add a sensible timeout instead of a conditional to avoid zombies
-      process.exit(0);
+      process.exit(exitCode);
     }
   }
 
@@ -315,7 +317,11 @@ class ProdWorker {
     }
   }
 
-  async #prepareForRetry(willCheckpointAndRestore: boolean, shouldExit: boolean) {
+  async #prepareForRetry(
+    willCheckpointAndRestore: boolean,
+    shouldExit: boolean,
+    exitCode?: number
+  ) {
     logger.log("prepare for retry", { willCheckpointAndRestore, shouldExit });
 
     // Graceful shutdown on final attempt
@@ -324,7 +330,7 @@ class ProdWorker {
         logger.log("WARNING: Will checkpoint but also requested exit. This won't end well.");
       }
 
-      await this.#exitGracefully();
+      await this.#exitGracefully(false, exitCode);
       return;
     }
 
@@ -516,7 +522,14 @@ class ProdWorker {
 
             logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit });
 
-            this.#prepareForRetry(willCheckpointAndRestore, shouldExit);
+            const exitCode =
+              !completion.ok &&
+              completion.error.type === "INTERNAL_ERROR" &&
+              completion.error.code === TaskRunErrorCodes.TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE
+                ? EXIT_CODE_CHILD_NONZERO
+                : 0;
+
+            this.#prepareForRetry(willCheckpointAndRestore, shouldExit, exitCode);
           } catch (error) {
             const completion: TaskRunFailedExecutionResult = {
               ok: false,
@@ -709,8 +722,8 @@ class ProdWorker {
             }
 
             await setTimeout(200);
-            // Use exit code 111 so we can ignore those failures in the task monitor
-            process.exit(111);
+
+            process.exit(EXIT_CODE_ALREADY_HANDLED);
           }
         }
 

packages/core-apps/src/process.ts

@@ -0,0 +1,2 @@
+export const EXIT_CODE_ALREADY_HANDLED = 111;
+export const EXIT_CODE_CHILD_NONZERO = 112;

packages/core-apps/src/provider.ts

@@ -14,6 +14,7 @@ import { getRandomPortNumber, HttpReply, getTextBody } from "./http";
 import { SimpleLogger } from "./logger";
 import { isExecaChildProcess } from "./checkpoints";
 import { setTimeout } from "node:timers/promises";
+import { EXIT_CODE_ALREADY_HANDLED } from "./process";
 
 const HTTP_SERVER_PORT = Number(process.env.HTTP_SERVER_PORT || getRandomPortNumber());
 const MACHINE_NAME = process.env.MACHINE_NAME || "local";
@@ -198,7 +199,7 @@ export class ProviderShell implements Provider {
                 stderr: error.stderr,
               });
 
-              if (error.exitCode === 111) {
+              if (error.exitCode === EXIT_CODE_ALREADY_HANDLED) {
                 logger.error("Index failure already reported by the worker", {
                   socketMessage: message,
                 });

packages/core/src/v3/schemas/messages.ts

@@ -334,6 +334,7 @@ export const ProviderToPlatformMessages = {
       exitCode: z.number().optional(),
       message: z.string().optional(),
       logs: z.string().optional(),
+      overrideCompletion: z.boolean().optional(),
     }),
   },
   INDEXING_FAILED: {
@@ -346,6 +347,7 @@ export const ProviderToPlatformMessages = {
         stack: z.string().optional(),
         stderr: z.string().optional(),
       }),
+      overrideCompletion: z.string().optional(),
     }),
   },
 };