From 83dff56a8d365d3ce5efa7e9dd425f23b7fffc9f Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Wed, 12 Feb 2025 17:40:13 +0000
Subject: [PATCH] handle exit for case where we already retried after oom

---
 .../app/v3/services/completeAttempt.server.ts | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts
index a1105717f..2a595000a 100644
--- a/apps/webapp/app/v3/services/completeAttempt.server.ts
+++ b/apps/webapp/app/v3/services/completeAttempt.server.ts
@@ -254,9 +254,11 @@ export class CompleteAttemptService extends BaseService {
 
     let retriableError = shouldRetryError(taskRunErrorEnhancer(completion.error));
     let isOOMRetry = false;
+    let isOOMAttempt = isOOMError(completion.error);
+    let isOnMaxOOMMachine = false;
 
-    //OOM errors should retry (if an OOM machine is specified)
-    if (isOOMError(completion.error)) {
+    //OOM errors should retry (if an OOM machine is specified, and we're not already on it)
+    if (isOOMAttempt) {
       const retryConfig = FailedTaskRunRetryHelper.getRetryConfig({
         run: {
           ...taskRunAttempt.taskRun,
@@ -266,10 +268,10 @@ export class CompleteAttemptService extends BaseService {
         execution,
       });
 
-      if (
-        retryConfig?.outOfMemory?.machine &&
-        retryConfig.outOfMemory.machine !== taskRunAttempt.taskRun.machinePreset
-      ) {
+      isOnMaxOOMMachine =
+        retryConfig?.outOfMemory?.machine === taskRunAttempt.taskRun.machinePreset;
+
+      if (retryConfig?.outOfMemory?.machine && !isOnMaxOOMMachine) {
         //we will retry
         isOOMRetry = true;
         retriableError = true;
@@ -312,6 +314,11 @@ export class CompleteAttemptService extends BaseService {
 
     // The attempt has failed and we won't retry
 
+    if (isOOMAttempt && isOnMaxOOMMachine) {
+      // The attempt failed due to an OOM error but we're already on the machine we should retry on
+      exitRun(taskRunAttempt.taskRunId);
+    }
+
     // Now we need to "complete" the task run event/span
     await eventRepository.completeEvent(
       getTaskEventStoreTableForRun(taskRunAttempt.taskRun),
@@ -508,10 +515,7 @@ export class CompleteAttemptService extends BaseService {
 
       // The run won't know it should shut down as we make the decision to force requeue here
       // This also ensures that this change is backwards compatible with older workers
-      socketIo.coordinatorNamespace.emit("REQUEST_RUN_CANCELLATION", {
-        version: "v1",
-        runId: run.id,
-      });
+      exitRun(run.id);
 
       await retryViaQueue();
       return;
@@ -759,3 +763,10 @@ function isOOMError(error: TaskRunError) {
 
   return false;
 }
+
+function exitRun(runId: string) {
+  socketIo.coordinatorNamespace.emit("REQUEST_RUN_CANCELLATION", {
+    version: "v1",
+    runId,
+  });
+}