EvalHarness/build/lib/evalharness/eval/runner.py

"""The evaluation runner: Dataset x predictions -> EvalReport.

Pure orchestration, no I/O hidden inside: predictions arrive as a list
(loaded from a jsonl of model outputs, a Session store, or built inline),
results aggregate into an EvalReport that visualizers consume.

Judge wiring: pass judge=<callable(messages)->str> once a ModelAdapter
exists; llm_judge recipes work immediately after that, no recipe change.
"""

import traceback
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union

from ..data.dataset import Dataset
from ..data.sample import Sample
from .aggregator import mean as _mean_agg
from .recipe import EvalRecipe
from .record import EvalReport, SampleResult
from .scorer import ScoreContext


def evaluate(
    dataset: Union[Dataset, List[Sample]],
    predictions: Sequence[Union[str, Dict]],
    recipe: Optional[EvalRecipe] = None,
    *,
    model: str = '',
    judge: Optional[Callable] = None,
    extra_metadata: Optional[Dict] = None,
) -> EvalReport:
    """Score a dataset against raw predictions.

    dataset: a Dataset or a plain list of Samples (views/slices).
    predictions: str per sample (raw model output) or dicts with
    {'raw': str, 'group_key': ..., 'metadata': {...}} overrides.
    """
    samples: List[Sample] = list(dataset)
    spec = getattr(dataset, 'spec', None)
    ds_name = spec.name if spec is not None else samples[0].metadata.get('dataset', 'adhoc') if samples else 'adhoc'
    ds_subset = spec.subset if spec is not None else ''
    if len(predictions) != len(samples):
        raise ValueError(f'{len(predictions)} predictions for {len(samples)} samples')

    if recipe is None:
        from .recipe import get_eval

        recipe = get_eval(ds_name)
    extractor = recipe.resolve_extract()
    scorers = recipe.resolve_scorers()
    aggregators = recipe.resolve_aggregators()
    ctx = ScoreContext(judge=judge, params={})

    # If any scorer executes in docker with per-sample images, overlap pulls
    # with scoring (run sample N while N+1..N+lookahead images download).
    bp = None
    if _needs_bg_prefetch(recipe, samples):
        from ..sandbox import BackgroundPrefetcher, images_for_samples

        bp = BackgroundPrefetcher(images_for_samples(samples), workers=4, lookahead=8)
        bp.__enter__()

    results: List[SampleResult] = []

    def judge_one(sample, pred) -> SampleResult:
        """Extract + score ONE sample (thread-safe: everything here is local
        except docker/subprocess execution, which parallelizes perfectly --
        each sample gets its own container/workdir)."""
        raw = pred if isinstance(pred, str) else str(pred.get('raw', ''))
        override = {} if isinstance(pred, str) else pred
        result = SampleResult(
            sample_id=sample.id,
            dataset=ds_name,
            subset=ds_subset,
            task_type=sample.task_type,
            raw_prediction=raw,
            target=sample.target,
            group_key=str(override.get('group_key')
                          or sample.metadata.get('group_key')
                          or (sample.metadata.get('task_id') or sample.metadata.get('id') or '')),
            metadata={k: v for k, v in (sample.metadata or {}).items()
                      if k in ('category', 'subject', 'test_category', 'bin', 'difficulty')},
        )
        if isinstance(pred, dict) and pred.get('metadata'):
            result.metadata.update(pred['metadata'])
        if isinstance(pred, dict) and pred.get('trajectory'):
            result.trajectory = pred['trajectory']
        if isinstance(pred, dict) and pred.get('env_state'):
            result.env_state = pred['env_state']
        if isinstance(pred, dict) and pred.get('usage'):
            result.usage = pred['usage']
        try:
            if bp is not None and sample.sandbox and sample.sandbox.image:
                bp.ensure(sample.sandbox.image)  # wait only if this one still pulling
            value, ok, note = extractor(raw, sample)
            result.extracted_prediction = value
            result.extraction_ok = ok
            result.extraction_note = note
            if not ok:
                result.extraction_note = note or 'extractor returned not-ok'
            for metric, scorer in scorers.items():
                try:
                    sctx = ctx
                    if result.env_state and 'env_state' not in ctx.params:
                        sctx = ScoreContext(judge=ctx.judge, judge_model=ctx.judge_model,
                                            params={**ctx.params,
                                                    'env_state': result.env_state})
                    scores, details = scorer(value if ok else '', sample.target, sample, sctx)
                    result.scores.update(scores)
                    result.score_details.update(details)
                except Exception as e:  # one metric failing must not kill the run
                    result.scores[metric] = 0.0
                    result.score_details[metric] = {'error': f'{type(e).__name__}: {e}'}
        except Exception as e:
            result.error = f'{type(e).__name__}: {e}\n{traceback.format_exc(limit=2)}'
        return result

    workers = getattr(recipe, 'exec_workers', 1)
    if workers > 1 and len(samples) > 1:
        # parallel judging: docker/subprocess execution is embarrassingly
        # parallel (one container per sample); text scorers are cheap and
        # thread-safe enough. Serializes again for judge/dict-dependent runs.
        import concurrent.futures

        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
            results = list(pool.map(judge_one, samples, predictions))
    else:
        for sample, pred in zip(samples, predictions):
            results.append(judge_one(sample, pred))

    report = EvalReport(
        dataset=ds_name,
        recipe=recipe.name or dataset.spec.name,
        model=model,
        num_samples=len(results),
        num_failed_extractions=sum(1 for r in results if not r.extraction_ok),
        samples=results,
    )
    _aggregate_into(report, results, recipe, aggregators)
    if bp is not None:
        report.metric_groups['run_info'] = {
            **report.metric_groups.get('run_info', {}),
            **{f'img_{k}': v for k, v in bp.stats().items()},
        }
        bp.__exit__(None, None, None)
    if extra_metadata:
        report.metric_groups['run_info'] = {k: v for k, v in extra_metadata.items()
                                            if isinstance(v, (int, float, str))}
    return report


def _needs_bg_prefetch(recipe, samples) -> bool:
    """True when the recipe executes in docker AND samples declare images."""
    try:
        for spec in recipe.scorers.values():
            params = spec if isinstance(spec, dict) else {}
            if params.get('name') == 'execution' and params.get('sandbox') == 'docker':
                return any(s.sandbox and s.sandbox.image for s in samples[:50])
    except Exception:
        return False
    return False


def _aggregate_into(report: EvalReport, results, recipe: EvalRecipe, aggregators) -> None:
    for metric in recipe.scorers:
        agg = aggregators.get(metric)
        if agg is None:
            agg = _mean_agg
        try:
            out = agg(results, metric)
        except Exception as e:
            report.metric_groups[f'agg_error_{metric}'] = {'error': str(e)[:200]}
            continue
        if isinstance(out, dict):
            report.metric_groups[metric] = out
            # primary metric = the aggregator's same-named entry (e.g.
            # simpleqa_official returns is_correct/is_incorrect/...); the
            # old mean-of-all-values fallback invented nonsense like
            # mean(0.035, 0.945, 0.02, 0.98) for is_correct
            if metric in out and isinstance(out[metric], (int, float)):
                report.metrics[metric] = float(out[metric])
            else:
                vals = [v for v in out.values() if isinstance(v, (int, float))]
                if vals:
                    report.metrics[metric] = sum(vals) / len(vals)
        else:
            report.metrics[metric] = float(out)
    report.metrics['extraction_failure_rate'] = (
        report.num_failed_extractions / report.num_samples if report.num_samples else 0.0
    )