workers.py

# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Module for isolated workers that run promptfoo tests."""

import collections.abc
import copy
import contextlib
import dataclasses
import json
import logging
import pathlib
import queue
import shutil
import subprocess
import sys
import tempfile
import threading
import time

import checkout_helpers
import constants
import eval_config
import promptfoo_installation
import results

sys.path.append(str(constants.CHROMIUM_SRC))
from agents.common import tempfile_ext

_ALL_QUEUED_TESTS_RUN_POLLING_SLEEP_DURATION = 0.5
_AVAILABLE_TEST_POLLING_SLEEP_DURATION = 0.1


class WorkDir(contextlib.AbstractContextManager):
    """A WorkDir used for testing destructive changes by an agent.

    Each workdir acts like a local shallow clone and has its own isolated
    checkout state (staging, untracked files, `//.gemini/extensions/`).
    """

    def __init__(
        self,
        name: str,
        src_root_dir: pathlib.Path,
        clean: bool,
        verbose: bool,
        force: bool,
    ):
        self.path = src_root_dir.parent.joinpath(name)
        self.src_root_dir = src_root_dir
        self.clean = clean
        self.verbose = verbose
        self.force = force
        self.btrfs = checkout_helpers.check_btrfs(src_root_dir)

    def __enter__(self) -> 'WorkDir':
        if self.path.exists():
            self._clean()

        logging.info('Creating new workdir: %s', self.path)
        start_time = time.time()
        if self.btrfs:
            # gclient-new-workdir does the same thing but reflinks the .git dirs
            # which we don't need to waste time on
            subprocess.check_call(
                ['btrfs', 'subvol', 'snapshot', self.src_root_dir, self.path],
                stdout=None if self.verbose else subprocess.DEVNULL,
                stderr=subprocess.STDOUT,
            )
        else:
            subprocess.check_call(
                ['gclient-new-workdir.py', self.src_root_dir, self.path],
                stdout=None if self.verbose else subprocess.DEVNULL,
                stderr=subprocess.STDOUT,
            )
        logging.debug('Took %s seconds', time.time() - start_time)
        return self

    def __exit__(self, *_exc_info) -> None:
        if self.clean:
            self._clean()

    def _clean(self) -> None:
        logging.info('Removing existing workdir: %s', self.path)
        cmd = ['sudo', 'btrfs', 'subvolume', 'delete', self.path]
        start_time = time.time()
        if self.btrfs:
            if self.force:
                cmd.insert(1, '-n')
            result = subprocess.call(
                cmd,
                stdout=None if self.verbose else subprocess.DEVNULL,
                stderr=subprocess.STDOUT,
            )
            if result != 0:
                logging.debug('Failed to remove with subvolume delete.')
        if not self.btrfs or result != 0:
            logging.debug('Removing with shutil')
            shutil.rmtree(self.path)
        logging.debug('Took %s seconds', time.time() - start_time)


@dataclasses.dataclass
class WorkerOptions:
    """Options for configuring a single worker."""

    # Whether to clean the workdir after a test.
    clean: bool
    # Whether to log verbosely.
    verbose: bool
    # Whether to force cleaning.
    force: bool
    # Whether to run tests in a sandbox.
    sandbox: bool
    # An optional path to a gemini-cli binary to use.
    gemini_cli_bin: pathlib.Path | None = None
    # An optional path to a nodejs binary to use.
    node_bin: pathlib.Path | None = None


class WorkerPool:
    """Abstracts away one or more WorkerThreads and a ResultThread."""

    def __init__(
        self,
        num_workers: int,
        promptfoo: promptfoo_installation.PromptfooInstallation,
        worker_options: WorkerOptions,
        result_options: results.ResultOptions,
    ):
        """
        Args:
            num_workers: The number of workers to use to run tests.
            promptfoo: A PromptfooInstallation to use when running tests.
            worker_options: A WorkerOptions instance whose attributes will be
                used when setting up workers.
            result_options: A ResultOptions instance whose attributes will be
                used when handling test results.
        """
        assert num_workers > 0
        # Create a copy so that options cannot be externally modified.
        # This is not done for result_options because its result_handlers are
        # liable to contain callables that use locks under the hood for thread
        # safety, which causes errors with deepcopy due to them being
        # un-picklable.
        worker_options = copy.deepcopy(worker_options)

        self._result_thread = results.ResultThread(
            result_options=result_options
        )
        self._result_thread.start()

        self._total_tests_queued = 0
        self._test_input_queue = queue.Queue()
        self._workers = []
        for i in range(num_workers):
            worker_thread = WorkerThread(
                i,
                promptfoo,
                worker_options,
                test_input_queue=self._test_input_queue,
                test_result_queue=self._result_thread.result_input_queue,
            )
            worker_thread.start()
            self._workers.append(worker_thread)

    def __del__(self):
        self.shutdown_blocking(2)

    def queue_tests(
        self, tests: collections.abc.Collection[eval_config.TestConfig]
    ) -> None:
        """Queues the provided tests to be run.

        Args:
            tests: A Collection of paths to promptfoo test configs to run.
        """
        self._total_tests_queued += len(tests)
        for t in tests:
            self._test_input_queue.put(t)

    def wait_for_all_queued_tests(self) -> list[results.TestResult]:
        """Waits for all tests that have been queued to finish.

        Returns:
            A list of failed TestResults that were produced since the last time
            this method was called.
        """
        while (
            self._result_thread.total_results_reported.get()
            != self._total_tests_queued
        ):
            self._result_thread.maybe_reraise_fatal_exception()
            for w in self._workers:
                w.maybe_reraise_fatal_exception()
            time.sleep(_ALL_QUEUED_TESTS_RUN_POLLING_SLEEP_DURATION)

        failed_tests = []
        failed_test_queue = self._result_thread.failed_result_output_queue
        while not failed_test_queue.empty():
            failed_tests.append(failed_test_queue.get())
        return failed_tests

    def shutdown_blocking(self, timeout: float | None = None) -> None:
        """Gracefully shuts down stored threads and waits for them to finish.

        Args:
            timeout: An optional timeout to use when joining underlying
                threads.
        """
        threads_to_shutdown = self._workers + [self._result_thread]
        for t in threads_to_shutdown:
            t.shutdown()
        for t in threads_to_shutdown:
            t.join(timeout)
            if t.is_alive():
                logging.error(
                    'Failed to gracefully shut down thread %s in a WorkerPool',
                    t.native_id,
                )


def _parse_test_log_results(results_json) -> str:
    """Extracts a summary of the test run for displaying

    Args:
        results_json: The decoded JSON from a promptfoo results file.

    Returns:
        A string summarizing the test/eval.
    """
    if results_json is None:
        return ''

    # Display the assertion failures
    results_list = results_json.get('results', {}).get('results')
    if not results_list:
        return 'No results found in promptfoo output.'
    run_result = results_list[0]
    assert_results = []
    grading_result = run_result.get('gradingResult')
    if grading_result:
        for componentResult in grading_result.get('componentResults', []):
            assert_results.append(
                f"pass: {componentResult['pass']}\n"
                f"reason: {componentResult['reason']}\n"
                f"score: {componentResult['score']}\n\n"
            )
    response = run_result.get('response', {})
    return (
        f"Input prompt: {response.get('metrics', {}).get('user_prompt')}\n"
        f"Response: {response.get('metrics', {}).get('full_output')}\n"
        "Assertion results:\n" + ''.join(assert_results)
    )


def _extract_metrics_from_promptfoo_results(
    results_json: dict,
) -> dict[str, dict | float]:
    """Extracts relevant metrics from promptfoo results.

    Args:
        results_json: The decoded JSON from a promptfoo results file.

    Returns:
        A potentially empty dict of extracted metrics.
    """
    if not results_json:
        return {}

    extracted_metrics = {
        'token_usage': _extract_token_usage_from_promptfoo_results(
            results_json
        ),
    }
    score = _extract_score_from_promptfoo_results(results_json)
    if score is not None:
        extracted_metrics['score'] = score
    return extracted_metrics


def _parse_input_prompt(results_json: dict[str, any]) -> str | None:
    prompts_json = results_json.get('results', {}).get('prompts')
    if not prompts_json:
        logging.error('Did not find the prompts in the test result.')
        return None
    prompt = prompts_json[0].get('raw')
    if not prompt:
        logging.error('Did not find prompt in the test result.')
    return prompt


def _parse_output(results_json: dict[str, any]) -> str | None:
    results_json = results_json.get('results', {}).get('results')
    if not results_json:
        logging.error('Did not find the response in the test result.')
        return None
    output = results_json[0].get('response', {}).get('output')
    if not output:
        logging.error('Did not find output response in the test result.')
    return output


def _load_promptfoo_results(results_file: pathlib.Path) -> dict[str, any]:
    """Loads a promptfoo results file into memory.

    Args:
        results_file: A path to a file containing promptfoo results for a test.

    Returns:
        The decoded JSON content of |results_file|. If loading fails for any
        reason, this will be an empty dict.
    """
    with open(results_file, encoding='utf-8') as infile:
        try:
            return json.load(infile)
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            logging.error(
                'Error when parsing promptfoo results. This is expected if '
                'promptfoo failed catastrophically: %s',
                e,
            )
            return {}


def _extract_token_usage_from_promptfoo_results(
    results_json: dict[str, any],
) -> dict[str, int]:
    """Extracts gemini-cli token usage from promptfoo JSON results.

    Args:
        results_json: The decoded JSON from a promptfoo results file.

    Returns:
        A dict mapping gemini-cli token type to tokens used at the end of the
        test.
    """
    token_usage = {}
    results_list = results_json.get('results', {}).get('results', [])
    if not results_list:
        logging.error(
            'Did not find promptfoo result information, cannot extract token '
            'usage. This is not expected to ever happen.'
        )
        return token_usage

    if len(results_list) > 1:
        logging.warning(
            'Unexpectedly got %d results from promptfoo when 1 is expected. '
            'Only using the first result for token usage.',
            len(results_list),
        )
    r = results_list[0]
    token_usage = (
        r.get('response', {})
        .get('metrics', {})
        .get(constants.GEMINI_CLI_TOKEN_USAGE, {})
    )
    if not token_usage:
        logging.warning(
            'Did not find gemini-cli token usage. This is not expected to '
            'ever happen'
        )
    return token_usage


def _extract_score_from_promptfoo_results(
    results_json: dict[str, any],
) -> float | None:
    """Extracts the test score from promptfoo JSON results.

    Args:
        results_json: The decoded JSON from a promptfoo results file.

    Returns:
        None if the score cannot be extracted. Otherwise, a float representing
        the test score reported by promptfoo. This value is expected to be
        between 0 and 1, inclusive.
    """
    results_list = results_json.get('results', {}).get('results', [])
    if not results_list:
        logging.error(
            'Did not find promptfoo result information, cannot extract score. '
            'This is not expected to ever happen.'
        )
        return None

    if len(results_list) > 1:
        logging.warning(
            'Unexpectedly got %d results from promptfoo when 1 is expected. '
            'Only using the first result for score.',
            len(results_list),
        )
    r = results_list[0]
    score = r.get('score')
    if score is None:
        logging.warning(
            'Did not find reported score. This is not expected to ever happen'
        )
    return score


class WorkerThread(threading.Thread):
    """Class for running tests from a queue in an isolated environment."""

    def __init__(
        self,
        worker_index: int,
        promptfoo: promptfoo_installation.PromptfooInstallation,
        worker_options: WorkerOptions,
        test_input_queue: queue.Queue[eval_config.TestConfig],
        test_result_queue: queue.Queue[results.TestResult],
        **kwargs,
    ):
        """
        Args:
            worker_index: The unique index of this thread.
            promptfoo: The PromptfooInstallation to use when running tests.
            worker_options: A WorkerOptions instance whose attributes will be
                used when configuring this object.
            test_input_queue: A Queue that will be read from to get tests to
                run.
            test_result_queue: A Queue that will receive TestResults for
                completed tests.
        """
        super().__init__(daemon=True, **kwargs)
        self._worker_index = worker_index
        self._promptfoo = promptfoo
        self._worker_options = worker_options
        self._console_width = shutil.get_terminal_size().columns

        self._test_input_queue = test_input_queue
        self._test_result_queue = test_result_queue
        self._shutdown_event = threading.Event()
        self._fatal_exception = None

    def run(self) -> None:
        try:
            self._run_incoming_tests_until_shutdown()
        except Exception as e:
            self._fatal_exception = e

    def _run_incoming_tests_until_shutdown(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                config = self._test_input_queue.get(
                    timeout=_AVAILABLE_TEST_POLLING_SLEEP_DURATION
                )
            except queue.Empty:
                continue
            self._run_one_config(config)

    def _run_one_config(self, config: eval_config.TestConfig) -> None:
        """Runs a single test config and queues a TestResult.

        Args:
            config: The TestConfig object for the test to run.
        """
        successful_runs = 0
        all_iteration_results = []

        for i in range(config.runs_per_test):
            logging.info(
                'Running test %s (iteration %d of up to %d)',
                config.test_file,
                i + 1,
                config.runs_per_test,
            )
            iteration_result = self._run_single_iteration(config)
            all_iteration_results.append(iteration_result)

            if iteration_result.success:
                successful_runs += 1

            # Exit early if the test has already passed.
            if successful_runs >= config.pass_k_threshold:
                break

            # Exit early if the test can no longer pass.
            num_failures = (i + 1) - successful_runs
            max_failures_allowed = (
                config.runs_per_test - config.pass_k_threshold
            )
            if num_failures > max_failures_allowed:
                break

        success = successful_runs >= config.pass_k_threshold
        r = results.TestResult(
            config=config,
            success=success,
            iteration_results=all_iteration_results,
        )
        self._test_result_queue.put(r)

    def _run_single_iteration(
        self, config: eval_config.TestConfig
    ) -> results.IterationResult:
        """Runs a single iteration of a test and returns an IterationResult.

        Args:
            config: The TestConfig object for the test to run.
        """
        with (
            WorkDir(
                f'workdir-{self._worker_index}',
                checkout_helpers.get_gclient_root(),
                self._worker_options.clean,
                self._worker_options.verbose,
                self._worker_options.force,
            ) as workdir,
            tempfile.TemporaryDirectory() as home_dir,
            tempfile_ext.mkstemp_closed(suffix='.json') as promptfoo_output,
        ):
            command = [
                'eval',
                '-j',
                '1',
                '--no-cache',
                # Not useful since we're running one test per eval and the
                # tables don't render properly in captured logs.
                '--no-table',
                '-c',
                str(config.test_file),
                '--var',
                f'console_width={self._console_width}',
                '--var',
                f'home_dir={home_dir}',
                '--output',
                str(promptfoo_output),
            ]
            if self._worker_options.sandbox:
                command.extend(['--var', 'sandbox=True'])
            if self._worker_options.verbose:
                command.extend(['--var', 'verbose=True'])
            if self._worker_options.gemini_cli_bin:
                command.extend(
                    [
                        '--var',
                        f'gemini_cli_bin={self._worker_options.gemini_cli_bin}',
                    ]
                )
            if self._worker_options.node_bin:
                command.extend(
                    ['--var', f'node_bin={self._worker_options.node_bin}']
                )

            start_time = time.time()
            proc = self._promptfoo.run(command, cwd=workdir.path / 'src')
            duration = time.time() - start_time

            results_json = _load_promptfoo_results(promptfoo_output)
            metrics = _extract_metrics_from_promptfoo_results(results_json)
            test_log = _parse_test_log_results(results_json)
            prompt = _parse_input_prompt(results_json)
            response = _parse_output(results_json)
            if proc.returncode != 0 and proc.stdout:
                test_log += f'\npromptfoo stdout:\n{proc.stdout}'
            return results.IterationResult(
                success=not proc.returncode,
                duration=duration,
                test_log=test_log,
                metrics=metrics,
                prompt=prompt,
                response=response,
            )

    def shutdown(self) -> None:
        """Tells the thread to shut down gracefully."""
        self._shutdown_event.set()

    def maybe_reraise_fatal_exception(self) -> None:
        """Reraises the fatal exception that caused the thread to die.

        No-op if no exception is stored.
        """
        if self._fatal_exception:
            raise self._fatal_exception