feat(js/ai/evaluators): Support batching for JS evaluator & CLI

genkit-tools/cli/src/commands/eval-flow.ts

@@ -45,6 +45,7 @@ interface EvalFlowRunCliOptions {
   context?: string;
   evaluators?: string;
   force?: boolean;
+  batchSize?: number;
   outputFormat: string;
 }
 
@@ -81,6 +82,11 @@ export const evalFlow = new Command('eval:flow')
     '-e, --evaluators <evaluators>',
     'comma separated list of evaluators to use (by default uses all)'
   )
+  .option(
+    '--batchSize <batchSize>',
+    'batch size to use for parallel evals (default to 1, no parallelization)',
+    parseInt
+  )
   .option('-f, --force', 'Automatically accept all interactive prompts')
   .action(
     async (flowName: string, data: string, options: EvalFlowRunCliOptions) => {
@@ -153,6 +159,7 @@ export const evalFlow = new Command('eval:flow')
           manager,
           evaluatorActions,
           evalDataset,
+          batchSize: options.batchSize,
           augments: {
             actionRef: `/flow/${flowName}`,
             datasetId:

genkit-tools/cli/src/commands/eval-run.ts

@@ -35,6 +35,7 @@ interface EvalRunCliOptions {
   output?: string;
   evaluators?: string;
   force?: boolean;
+  batchSize?: number;
   outputFormat: string;
 }
 
@@ -58,6 +59,11 @@ export const evalRun = new Command('eval:run')
     '--evaluators <evaluators>',
     'comma separated list of evaluators to use (by default uses all)'
   )
+  .option(
+    '--batchSize <batchSize>',
+    'batch size to use for parallel evals (default to 1, no parallelization)',
+    parseInt
+  )
   .option('--force', 'Automatically accept all interactive prompts')
   .action(async (dataset: string, options: EvalRunCliOptions) => {
     await runWithManager(async (manager) => {
@@ -105,6 +111,7 @@ export const evalRun = new Command('eval:run')
         manager,
         evaluatorActions,
         evalDataset,
+        batchSize: options.batchSize,
       });
 
       if (options.output) {

genkit-tools/common/src/eval/evaluate.ts

@@ -124,6 +124,7 @@ export async function runNewEvaluation(
     manager,
     evaluatorActions,
     evalDataset,
+    batchSize: request.options?.batchSize,
     augments: {
       ...metadata,
       actionRef,
@@ -163,21 +164,28 @@ export async function runEvaluation(params: {
   evaluatorActions: Action[];
   evalDataset: EvalInput[];
   augments?: EvalKeyAugments;
+  batchSize?: number;
 }): Promise<EvalRun> {
-  const { manager, evaluatorActions, evalDataset, augments } = params;
+  const { manager, evaluatorActions, evalDataset, augments, batchSize } =
+    params;
   if (evalDataset.length === 0) {
     throw new Error('Cannot run evaluation, no data provided');
   }
   const evalRunId = randomUUID();
   const scores: Record<string, any> = {};
   logger.info('Running evaluation...');
+
+  const runtime = manager.getMostRecentRuntime();
+  const isNodeRuntime = runtime?.genkitVersion?.startsWith('nodejs') ?? false;
+
   for (const action of evaluatorActions) {
     const name = evaluatorName(action);
     const response = await manager.runAction({
       key: name,
       input: {
         dataset: evalDataset.filter((row) => !row.error),
         evalRunId,
+        batchSize: isNodeRuntime ? batchSize : undefined,
       },
     });
     scores[name] = response.result;

genkit-tools/common/src/types/apis.ts

@@ -161,6 +161,12 @@ export const RunNewEvaluationRequestSchema = z.object({
         .any()
         .describe('addition parameters required for inference')
         .optional(),
+      batchSize: z
+        .number()
+        .describe(
+          'Batch the dataset into smaller segments that are run in parallel'
+        )
+        .optional(),
     })
     .optional(),
 });

js/ai/src/evaluator.ts

@@ -177,6 +177,7 @@ export function defineEvaluator<
           : z.array(BaseDataPointSchema),
         options: options.configSchema ?? z.unknown(),
         evalRunId: z.string(),
+        batchSize: z.number().optional(),
       }),
       outputSchema: EvalResponsesSchema,
       metadata: {
@@ -186,17 +187,19 @@ export function defineEvaluator<
     },
     async (i) => {
       let evalResponses: EvalResponses = [];
-      for (let index = 0; index < i.dataset.length; index++) {
-        const datapoint: BaseEvalDataPoint = {
-          ...i.dataset[index],
-          testCaseId: i.dataset[index].testCaseId ?? randomUUID(),
-        };
+      // This also populates missing testCaseIds
+      const batches = getBatchedArray(i.dataset, i.batchSize);
+
+      for (let batchIndex = 0; batchIndex < batches.length; batchIndex++) {
+        const batch = batches[batchIndex];
         try {
           await runInNewSpan(
             registry,
             {
               metadata: {
-                name: `Test Case ${datapoint.testCaseId}`,
+                name: i.batchSize
+                  ? `Batch ${batchIndex}`
+                  : `Test Case ${batch[0].testCaseId}`,
                 metadata: { 'evaluator:evalRunId': i.evalRunId },
               },
               labels: {
@@ -206,38 +209,48 @@ export function defineEvaluator<
             async (metadata, otSpan) => {
               const spanId = otSpan.spanContext().spanId;
               const traceId = otSpan.spanContext().traceId;
-              try {
+              const evalRunPromises = batch.map((d, index) => {
+                const sampleIndex = i.batchSize
+                  ? i.batchSize * batchIndex + index
+                  : batchIndex;
+                const datapoint = d as BaseEvalDataPoint;
                 metadata.input = {
                   input: datapoint.input,
                   output: datapoint.output,
                   context: datapoint.context,
                 };
-                const testCaseOutput = await runner(datapoint, i.options);
-                testCaseOutput.sampleIndex = index;
-                testCaseOutput.spanId = spanId;
-                testCaseOutput.traceId = traceId;
-                metadata.output = testCaseOutput;
-                evalResponses.push(testCaseOutput);
-                return testCaseOutput;
-              } catch (e) {
-                evalResponses.push({
-                  sampleIndex: index,
-                  spanId,
-                  traceId,
-                  testCaseId: datapoint.testCaseId,
-                  evaluation: {
-                    error: `Evaluation of test case ${datapoint.testCaseId} failed: \n${(e as Error).stack}`,
-                    status: EvalStatusEnum.FAIL,
-                  },
-                });
-                // Throw to mark the span as failed.
-                throw e;
-              }
+                const evalOutputPromise = runner(datapoint, i.options)
+                  .then((result) => ({
+                    ...result,
+                    traceId,
+                    spanId,
+                    sampleIndex,
+                  }))
+                  .catch((error) => {
+                    return {
+                      sampleIndex,
+                      spanId,
+                      traceId,
+                      testCaseId: datapoint.testCaseId,
+                      evaluation: {
+                        error: `Evaluation of test case ${datapoint.testCaseId} failed: \n${error}`,
+                      },
+                    };
+                  });
+                return evalOutputPromise;
+              });
+
+              const allResults = await Promise.all(evalRunPromises);
+              metadata.output =
+                allResults.length === 1 ? allResults[0] : allResults;
+              allResults.map((result) => {
+                evalResponses.push(result);
+              });
             }
           );
         } catch (e) {
           logger.error(
-            `Evaluation of test case ${datapoint.testCaseId} failed: \n${(e as Error).stack}`
+            `Evaluation of batch ${batchIndex} failed: \n${(e as Error).stack}`
           );
           continue;
         }
@@ -317,3 +330,31 @@ export function evaluatorRef<
 ): EvaluatorReference<CustomOptionsSchema> {
   return { ...options };
 }
+
+/**
+ * Helper method to generated batched array. Also ensures each testCase has a
+ * testCaseId
+ */
+function getBatchedArray<T extends { testCaseId?: string }>(
+  arr: T[],
+  batchSize?: number
+): T[][] {
+  let size: number;
+  if (!batchSize) {
+    size = 1;
+  } else {
+    size = batchSize;
+  }
+
+  let batches: T[][] = [];
+  for (var i = 0; i < arr.length; i += size) {
+    batches.push(
+      arr.slice(i, i + size).map((d) => ({
+        ...d,
+        testCaseId: d.testCaseId ?? randomUUID(),
+      }))
+    );
+  }
+
+  return batches;
+}