From e54aa852d4b14b6d0c658eb884d3e110d55e433d Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Mon, 28 Oct 2024 15:47:57 +0000
Subject: [PATCH] enable task monitor to update fatal errors

---
 apps/kubernetes-provider/src/taskMonitor.ts        | 10 ++++++----
 apps/webapp/app/v3/handleSocketIo.server.ts        | 13 ++++++++-----
 apps/webapp/app/v3/services/crashTaskRun.server.ts |  5 +++++
 packages/core/src/v3/schemas/messages.ts           |  1 +
 4 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/apps/kubernetes-provider/src/taskMonitor.ts b/apps/kubernetes-provider/src/taskMonitor.ts
index 9f1a8c092..aadcef18d 100644
--- a/apps/kubernetes-provider/src/taskMonitor.ts
+++ b/apps/kubernetes-provider/src/taskMonitor.ts
@@ -140,7 +140,7 @@ export class TaskMonitor {
     const containerState = this.#getContainerStateSummary(containerStatus.state);
     const exitCode = containerState.exitCode ?? -1;
 
-    if (exitCode === EXIT_CODE_ALREADY_HANDLED || exitCode === EXIT_CODE_CHILD_NONZERO) {
+    if (exitCode === EXIT_CODE_ALREADY_HANDLED) {
       this.#logger.debug("Ignoring pod failure, already handled by worker", {
         podName,
       });
@@ -160,7 +160,10 @@ export class TaskMonitor {
 
     let reason = rawReason || "Unknown error";
     let logs = rawLogs || "";
-    let overrideCompletion = false;
+
+    /** This will only override existing task errors. It will not crash the run.  */
+    let onlyOverrideExistingError = exitCode === EXIT_CODE_CHILD_NONZERO;
+
     let errorCode: TaskRunInternalError["code"] = TaskRunErrorCodes.POD_UNKNOWN_ERROR;
 
     switch (rawReason) {
@@ -185,7 +188,6 @@ export class TaskMonitor {
         }
         break;
       case "OOMKilled":
-        overrideCompletion = true;
         reason =
           "[TaskMonitor] Your task ran out of memory. Try increasing the machine specs. If this doesn't fix it there might be a memory leak.";
         errorCode = TaskRunErrorCodes.TASK_PROCESS_OOM_KILLED;
@@ -198,7 +200,7 @@ export class TaskMonitor {
       exitCode,
       reason,
       logs,
-      overrideCompletion,
+      overrideCompletion: onlyOverrideExistingError,
       errorCode,
     } satisfies FailureDetails;
 
diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts
index 29f753d6d..901ab178a 100644
--- a/apps/webapp/app/v3/handleSocketIo.server.ts
+++ b/apps/webapp/app/v3/handleSocketIo.server.ts
@@ -24,6 +24,7 @@ import { Redis } from "ioredis";
 import { createAdapter } from "@socket.io/redis-adapter";
 import { CrashTaskRunService } from "./services/crashTaskRun.server";
 import { CreateTaskRunAttemptService } from "./services/createTaskRunAttempt.server";
+import { UpdateFatalRunErrorService } from "./services/updateFatalRunError.server";
 
 export const socketIo = singleton("socketIo", initalizeIoServer);
 
@@ -302,11 +303,13 @@ function createProviderNamespace(io: Server) {
     handlers: {
       WORKER_CRASHED: async (message) => {
         try {
-          const service = new CrashTaskRunService();
-
-          await service.call(message.runId, {
-            ...message,
-          });
+          if (message.overrideCompletion) {
+            const updateErrorService = new UpdateFatalRunErrorService();
+            await updateErrorService.call(message.runId, { ...message });
+          } else {
+            const crashRunService = new CrashTaskRunService();
+            await crashRunService.call(message.runId, { ...message });
+          }
         } catch (error) {
           logger.error("Error while handling crashed worker", { error });
         }
diff --git a/apps/webapp/app/v3/services/crashTaskRun.server.ts b/apps/webapp/app/v3/services/crashTaskRun.server.ts
index 4c94e841c..8f9aaa190 100644
--- a/apps/webapp/app/v3/services/crashTaskRun.server.ts
+++ b/apps/webapp/app/v3/services/crashTaskRun.server.ts
@@ -29,6 +29,11 @@ export class CrashTaskRunService extends BaseService {
 
     logger.debug("CrashTaskRunService.call", { runId, opts });
 
+    if (options?.overrideCompletion) {
+      logger.error("CrashTaskRunService.call: overrideCompletion is deprecated", { runId });
+      return;
+    }
+
     const taskRun = await this._prisma.taskRun.findFirst({
       where: {
         id: runId,
diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts
index 330f13e63..31d6dc3e4 100644
--- a/packages/core/src/v3/schemas/messages.ts
+++ b/packages/core/src/v3/schemas/messages.ts
@@ -252,6 +252,7 @@ export const ProviderToPlatformMessages = {
       exitCode: z.number().optional(),
       message: z.string().optional(),
       logs: z.string().optional(),
+      /** This means we should only update the error if one exists */
       overrideCompletion: z.boolean().optional(),
       errorCode: TaskRunInternalError.shape.code.optional(),
     }),