feat: support retrying timed out evaluations

runner/configuration/constants.ts

@@ -33,6 +33,9 @@ export const DEFAULT_MAX_BUILD_REPAIR_ATTEMPTS = 1;
  */
 export const DEFAULT_MAX_TEST_REPAIR_ATTEMPTS = 0;
 
+/** Default number of retries when a prompt evaluation timed out. */
+export const DEFAULT_PROMPT_TIMEOUT_RETRIES = 1;
+
 /** Name of the folder where we store all generated reports */
 export const REPORTS_ROOT_DIR = join(rootDir, 'reports');
 

runner/eval-cli.ts

@@ -6,6 +6,7 @@ import {
   DEFAULT_MAX_BUILD_REPAIR_ATTEMPTS,
   DEFAULT_MAX_TEST_REPAIR_ATTEMPTS,
   DEFAULT_MODEL_NAME,
+  DEFAULT_PROMPT_TIMEOUT_RETRIES,
 } from './configuration/constants.js';
 import {generateCodeAndAssess} from './orchestration/generate.js';
 import {logReportToConsole, writeReportToDisk} from './reporting/report-logging.js';
@@ -42,6 +43,7 @@ interface Options {
   skipLighthouse?: boolean;
   maxTestRepairAttempts?: number;
   maxBuildRepairAttempts?: number;
+  promptTimeoutRetries?: number;
 }
 
 function builder(argv: Argv): Argv<Options> {
@@ -168,6 +170,12 @@ function builder(argv: Argv): Argv<Options> {
         description:
           'Number of repair attempts for discovered test failures (including a11y violations and ones from testCommand)',
       })
+      .option('prompt-timeout-retries', {
+        type: 'number',
+        default: DEFAULT_PROMPT_TIMEOUT_RETRIES,
+        description:
+          'Maximum number of times to retry a prompt evaluation after it fails due to a timeout.',
+      })
       .strict()
       .version(false)
       .help()
@@ -221,6 +229,7 @@ async function handler(cliArgs: Arguments<Options>): Promise<void> {
       skipLighthouse: cliArgs.skipLighthouse,
       maxBuildRepairAttempts: cliArgs.maxBuildRepairAttempts,
       maxTestRepairAttempts: cliArgs.maxTestRepairAttempts,
+      promptTimeoutRetries: cliArgs.promptTimeoutRetries,
       abortSignal: abortCtrl.signal,
     });
 

runner/orchestration/generate.ts

@@ -21,7 +21,7 @@ import {
 } from '../shared-interfaces.js';
 import {UserFacingError} from '../utils/errors.js';
 import {executeCommand} from '../utils/exec.js';
-import {callWithTimeout} from '../utils/timeout.js';
+import {callWithTimeout, TimeoutError} from '../utils/timeout.js';
 import {LocalExecutor} from './executors/local-executor.js';
 import {startEvaluationTask} from './generate-eval-task.js';
 import {prepareSummary} from './generate-summary.js';
@@ -145,55 +145,74 @@ export async function generateCodeAndAssess(options: AssessmentConfig): Promise<
     for (const rootPromptDef of promptsToProcess) {
       allTasks.push(
         appConcurrencyQueue.add(async () => {
-          const evalID = await env.executor.initializeEval();
-          let results: AssessmentResult[] | undefined;
-
-          try {
-            results = await callWithTimeout(
-              `Evaluation of ${rootPromptDef.name}`,
-              async timeoutAbortSignal =>
-                startEvaluationTask(
-                  options,
-                  evalID,
-                  env,
-                  autoraterLlm,
-                  cujGenerationLlm,
-                  rootPromptDef,
-                  combineAbortSignals(
-                    allTasksAbortCtrl.signal,
-                    timeoutAbortSignal,
-                    options.abortSignal,
+          const evaluate = async () => {
+            const evalID = await env.executor.initializeEval();
+            let results: AssessmentResult[] | undefined;
+
+            try {
+              results = await callWithTimeout(
+                `Evaluation of ${rootPromptDef.name}`,
+                async timeoutAbortSignal =>
+                  startEvaluationTask(
+                    options,
+                    evalID,
+                    env,
+                    autoraterLlm,
+                    cujGenerationLlm,
+                    rootPromptDef,
+                    combineAbortSignals(
+                      allTasksAbortCtrl.signal,
+                      timeoutAbortSignal,
+                      options.abortSignal,
+                    ),
+                    workerConcurrencyQueue,
+                    progress,
                   ),
-                  workerConcurrencyQueue,
-                  progress,
-                ),
-              // A timeout is used to prevent from stuck evaluations.
-              env.promptTimeoutMinutes ?? 10,
-            );
-            return results;
-          } catch (e: unknown) {
-            failedPrompts.push({
-              promptName: rootPromptDef.name,
-              error: `${e}`,
-              stack: e instanceof Error ? e.stack : undefined,
-            });
-
-            let details = `Error: ${e}`;
-            if (e instanceof Error && e.stack) {
-              details += `\nStack: ${e.stack}`;
+                // A timeout is used to prevent from stuck evaluations.
+                env.promptTimeoutMinutes ?? 10,
+              );
+              return results;
+            } catch (e: unknown) {
+              failedPrompts.push({
+                promptName: rootPromptDef.name,
+                error: `${e}`,
+                stack: e instanceof Error ? e.stack : undefined,
+              });
+
+              let details = `Error: ${e}`;
+              if (e instanceof Error && e.stack) {
+                details += `\nStack: ${e.stack}`;
+              }
+
+              progress.log(rootPromptDef, 'error', 'Failed to evaluate code', details);
+              return [] satisfies AssessmentResult[];
+            } finally {
+              // Gracefully finalize the eval. Errors in finalization should not propagate.
+              try {
+                await env.executor.finalizeEval(evalID);
+              } catch (e) {
+                progress.log(rootPromptDef, 'error', 'Failed to finalize eval', `${e}`);
+              }
+              progress.evalFinished(rootPromptDef, results || []);
             }
+          };
 
-            progress.log(rootPromptDef, 'error', 'Failed to evaluate code', details);
-            return [] satisfies AssessmentResult[];
-          } finally {
-            // Gracefully finalize the eval. Errors in finalization should not propagate.
+          // Retries + initial attempt.
+          const maxAttempts = (options.promptTimeoutRetries ?? 0) + 1;
+          for (let attemptIdx = 0; attemptIdx < maxAttempts; attemptIdx++) {
             try {
-              await env.executor.finalizeEval(evalID);
-            } catch (e) {
-              progress.log(rootPromptDef, 'error', 'Failed to finalize eval', `${e}`);
+              return await evaluate();
+            } catch (e: unknown) {
+              if (e instanceof TimeoutError && attemptIdx < maxAttempts) {
+                continue;
+              }
+              throw e;
             }
-            progress.evalFinished(rootPromptDef, results || []);
           }
+          throw new Error(
+            `Unexpected code path. ` +
+              `There were ${maxAttempts} attempts for evaluating: ${rootPromptDef.name}`,
+          );
         }),
       );
     }

runner/shared-interfaces.ts

@@ -30,6 +30,7 @@ export interface AssessmentConfig {
   skipLighthouse?: boolean;
   maxTestRepairAttempts?: number;
   maxBuildRepairAttempts?: number;
+  promptTimeoutRetries?: number;
   abortSignal?: AbortSignal;
 }