treat immediate retries as warm starts

Author: nicktrn

packages/cli-v3/src/entryPoints/managed-run-controller.ts

@@ -579,7 +579,7 @@ class ManagedRunController {
             runId: run.friendlyId,
             message: "Run is finished, will wait for next run",
           });
-          this.waitForNextRun();
+            this.waitForNextRun();
           return;
         }
         case "QUEUED_EXECUTING":
@@ -983,7 +983,14 @@ class ManagedRunController {
       };
 
       try {
-        return await this.executeRun({ run, snapshot, envVars: taskRunEnv, execution, metrics });
+        return await this.executeRun({
+          run,
+          snapshot,
+          envVars: taskRunEnv,
+          execution,
+          metrics,
+          isWarmStart,
+        });
       } catch (error) {
         if (error instanceof SuspendedProcessError) {
           this.sendDebugLog({
@@ -1344,8 +1351,10 @@ class ManagedRunController {
     envVars,
     execution,
     metrics,
+    isWarmStart,
   }: WorkloadRunAttemptStartResponseBody & {
     metrics?: TaskRunExecutionMetrics;
+    isWarmStart?: boolean;
   }) {
     this.snapshotPoller.start();
 
@@ -1360,6 +1369,7 @@ class ManagedRunController {
           engine: "V2",
         },
         machine: execution.machine,
+        isWarmStart,
       }).initialize();
     }
 
@@ -1372,15 +1382,18 @@ class ManagedRunController {
       },
     });
 
-    const completion = await this.taskRunProcess.execute({
-      payload: {
-        execution,
-        traceContext: execution.run.traceContext ?? {},
-        metrics,
+    const completion = await this.taskRunProcess.execute(
+      {
+        payload: {
+          execution,
+          traceContext: execution.run.traceContext ?? {},
+          metrics,
+        },
+        messageId: run.friendlyId,
+        env: envVars,
       },
-      messageId: run.friendlyId,
-      env: envVars,
-    });
+      isWarmStart
+    );
 
     this.sendDebugLog({
       runId: this.runFriendlyId,
@@ -1543,6 +1556,7 @@ class ManagedRunController {
         runFriendlyId: run.friendlyId,
         snapshotFriendlyId: this.snapshotFriendlyId,
         skipLockCheckForImmediateRetry: true,
+        isWarmStart: true,
       }).finally(() => {});
       return;
     }

packages/cli-v3/src/executions/taskRunProcess.ts

@@ -215,7 +215,10 @@ export class TaskRunProcess {
     await this._ipc?.sendWithAck("FLUSH", { timeoutInMs }, timeoutInMs + 1_000);
   }
 
-  async execute(params: TaskRunProcessExecuteParams): Promise<TaskRunExecutionResult> {
+  async execute(
+    params: TaskRunProcessExecuteParams,
+    isWarmStart?: boolean
+  ): Promise<TaskRunExecutionResult> {
     this._isPreparedForNextRun = false;
 
     let resolver: (value: TaskRunExecutionResult) => void;
@@ -251,7 +254,7 @@ export class TaskRunProcess {
         metadata: this.options.serverWorker,
         metrics,
         env: params.env,
-        isWarmStart: this.options.isWarmStart,
+        isWarmStart: isWarmStart ?? this.options.isWarmStart,
       });
     }
 

packages/core/src/v3/workers/taskExecutor.ts

@@ -89,7 +89,8 @@ export class TaskExecutor {
     execution: TaskRunExecution,
     worker: ServerBackgroundWorker,
     traceContext: Record<string, unknown>,
-    signal?: AbortSignal
+    signal?: AbortSignal,
+    isWarmStart?: boolean
   ): Promise<{ result: TaskRunExecutionResult }> {
     const ctx = TaskRunContext.parse(execution);
     const attemptMessage = `Attempt ${execution.attempt.number}`;
@@ -102,7 +103,7 @@ export class TaskExecutor {
     taskContext.setGlobalTaskContext({
       ctx,
       worker,
-      isWarmStart: this._isWarmStart,
+      isWarmStart: isWarmStart ?? this._isWarmStart,
     });
 
     if (execution.run.metadata) {