Add support for JSONL format, simplify task output directory paths, update run-eval-loop.sh (GoogleCloudPlatform#614)

Author: noahlwest

dev/ci/periodics/run-eval-loop.sh

@@ -96,7 +96,9 @@ echo "Task Pattern: ${TASK_PATTERN:-"All Tasks"}"
 # Loop from 1 to the specified number of iterations
 for i in $(seq 1 $ITERATIONS)
 do
-  OUTPUT_DIR="${REPO_ROOT}/.build/k8s-ai-bench-${MODEL}-${i}"
+  # Create a sanitized version of model name: replace all '/' with '-'
+  SAFE_MODEL="${MODEL//\//-}"
+  OUTPUT_DIR="${REPO_ROOT}/.build/k8s-ai-bench-${SAFE_MODEL}-${i}"
 
   echo "Running iteration $i of $ITERATIONS..."
 
@@ -124,6 +126,7 @@ do
   # Paths for analysis files
   MARKDOWN_FILE="${OUTPUT_DIR}/k8s-ai-bench.md"
   JSON_FILE="${OUTPUT_DIR}/k8s-ai-bench.json"
+  JSONL_FILE="${OUTPUT_DIR}/k8s-ai-bench.jsonl"
 
   # Run for markdown format
   "${K8S_AI_BENCH_BIN}" analyze --input-dir="${OUTPUT_DIR}" --results-filepath="${MARKDOWN_FILE}" --output-format=markdown --show-failures
@@ -139,6 +142,13 @@ do
     exit 1
   fi
 
+  # Run for jsonl format
+  "${K8S_AI_BENCH_BIN}" analyze --input-dir="${OUTPUT_DIR}" --results-filepath="${JSONL_FILE}" --output-format=jsonl --show-failures
+  if [ $? -ne 0 ]; then
+    echo "Error on iteration $i during JSONL analysis. Aborting loop."
+    exit 1
+  fi
+
   # Extract the time value and append it to the markdown file
   if [ -n "$run_time_line" ]; then
     time_value=$(echo $run_time_line | awk '{print $4}')

k8s-ai-bench/eval.go

@@ -130,7 +130,7 @@ func runEvaluation(ctx context.Context, config EvalConfig) error {
 				for _, llmConfig := range config.LLMConfigs {
 					taskOutputDir := ""
 					if config.OutputDir != "" {
-						taskOutputDir = filepath.Join(config.OutputDir, job.taskID, llmConfig.ID)
+						taskOutputDir = filepath.Join(config.OutputDir, job.taskID)
 						if err := os.MkdirAll(taskOutputDir, 0755); err != nil {
 							errorsCh <- fmt.Errorf("creating directory %q: %w", taskOutputDir, err)
 							return
@@ -288,7 +288,7 @@ func evaluateTask(ctx context.Context, config EvalConfig, taskID string, task Ta
 	taskCtx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
-	taskOutputDir := filepath.Join(config.OutputDir, taskID, llmConfig.ID)
+	taskOutputDir := filepath.Join(config.OutputDir, taskID)
 
 	var logBuffer bytes.Buffer
 	multiWriter := io.MultiWriter(&logBuffer)
@@ -343,15 +343,7 @@ func evaluateTask(ctx context.Context, config EvalConfig, taskID string, task Ta
 		logString := logBuffer.String()
 		logTail, truncated := getLastNLines(logString, maxErrLogLines)
 		// build log file path
-		shimSegment := "shim_disabled"
-		if x.llmConfig.EnableToolUseShim {
-			shimSegment = "shim_enabled"
-		}
-		logPath := filepath.Join(
-			config.OutputDir,
-			taskID,
-			shimSegment+"-"+x.llmConfig.ProviderID+"-"+x.llmConfig.ModelID,
-		)
+		logPath := taskOutputDir
 		errorMessage := fmt.Sprintf("agent encountered error: %v\n---LOG---\n%s", err, logTail)
 		if truncated {
 			errorMessage += fmt.Sprintf("\n... (log truncated, full log at %s)", logPath)
@@ -416,15 +408,7 @@ func evaluateTask(ctx context.Context, config EvalConfig, taskID string, task Ta
 			logString := logBuffer.String()
 			logTail, truncated := getLastNLines(logString, maxLogLines)
 			// build log file path
-			shimSegment := "shim_disabled"
-			if x.llmConfig.EnableToolUseShim {
-				shimSegment = "shim_enabled"
-			}
-			logPath := filepath.Join(
-				config.OutputDir,
-				taskID,
-				shimSegment+"-"+x.llmConfig.ProviderID+"-"+x.llmConfig.ModelID,
-			)
+			logPath := taskOutputDir
 			failureMessage := fmt.Sprintf("verifier script failed: %v\n---LOG---\n%s", err, logTail)
 			if truncated {
 				failureMessage += fmt.Sprintf("\n... (log truncated, full log at %s)", logPath)

k8s-ai-bench/main.go

@@ -360,8 +360,8 @@ func runAnalyze() error {
 	}
 
 	// Check if output format is valid
-	if config.OutputFormat != "markdown" && config.OutputFormat != "json" {
-		return fmt.Errorf("invalid output format: %s, valid options are 'markdown' or 'json'", config.OutputFormat)
+	if config.OutputFormat != "markdown" && config.OutputFormat != "json" && config.OutputFormat != "jsonl" {
+		return fmt.Errorf("invalid output format: %s, valid options are 'markdown', 'json' or 'jsonl'", config.OutputFormat)
 	}
 
 	// Check if input directory exists
@@ -375,14 +375,19 @@ func runAnalyze() error {
 	}
 
 	// Format and output results
-	if config.OutputFormat == "markdown" {
+	switch config.OutputFormat {
+	case "markdown":
 		if err := printMarkdownResults(config, allResults, resultsFilePath); err != nil {
 			return fmt.Errorf("printing markdown results: %w", err)
 		}
-	} else {
+	case "json":
 		if err := printJSONResults(allResults, resultsFilePath); err != nil {
 			return fmt.Errorf("printing JSON results: %w", err)
 		}
+	case "jsonl":
+		if err := printJSONLResults(allResults, resultsFilePath); err != nil {
+			return fmt.Errorf("printing JSONL results: %w", err)
+		}
 	}
 
 	return nil
@@ -756,3 +761,28 @@ func printJSONResults(results []model.TaskResult, resultsFilePath string) error
 
 	return nil
 }
+
+func printJSONLResults(results []model.TaskResult, resultsFilePath string) error {
+	var buffer strings.Builder
+	for _, result := range results {
+		jsonData, err := json.Marshal(result)
+		if err != nil {
+			return fmt.Errorf("marshaling result to JSON: %w", err)
+		}
+		buffer.Write(jsonData)
+		buffer.WriteString("\n")
+	}
+
+	// Write to file if path is provided, otherwise print to stdout
+	if resultsFilePath != "" {
+		if err := os.WriteFile(resultsFilePath, []byte(buffer.String()), 0644); err != nil {
+			return fmt.Errorf("writing to file %q: %w", resultsFilePath, err)
+		}
+		fmt.Printf("Results written to %s\n", resultsFilePath)
+	} else {
+		// Print to stdout only if no file path is specified
+		fmt.Print(buffer.String())
+	}
+
+	return nil
+}