Blame for src/elbench/execution/runner.py

from __future__ import annotations

import asyncio
import json
import logging
import re
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Iterable

from elbench.config import load_project_config
from elbench.judges import JudgeRouter
from elbench.loaders import LoaderFactory
from elbench.persistence import CheckpointStore, JsonlWriter, OutputPaths, configure_run_logger
from elbench.providers import ProviderFactory
from elbench.registry import FileRegistry
from elbench.schemas.config import ModelConfig, ProjectConfig, RateLimitConfig, RetryConfig
from elbench.schemas.evaluation import EvalResult, FailureRecord, GenerationRequest, ModelResponse, Sample
from elbench.summary import build_summary

from .basic_education import BasicEducationExecutor, DEFAULT_BASIC_MODULE_NAME
from .selection import resolve_module_selection
from .rate_limit import SlidingWindowRateLimiter
from .response_format import (
    ResponseFormatError,
    attach_format_metadata,
    augment_prompt_for_format,
    build_concise_retry_prompt,
    build_retry_followup,
    get_response_format_spec,
)
from .retry import is_retryable_exception, sleep_before_retry


ANSWER_COMPLETION_TASKS = {"math_500", "aime24", "aime25", "aime26"}


@dataclass(slots=True)
class RunOptions:
    model_id: str
    run_id: str
    modules: set[str] | None = None
    subsets: set[str] | None = None
    source_files: set[str] | None = None
    dimensions: set[str] | None = None
    max_samples: int | None = None
    # Cap per source file (per scenario for basic education), applied before max_samples.
    max_samples_per_file: int | None = None
    max_concurrency: int | None = None
    resume: bool = True
    judge_enabled: bool = True
    progress_enabled: bool = False


class BenchmarkRunner:
    def __init__(self, config: ProjectConfig | None = None) -> None:
        self.config = config or load_project_config()

    async def run(self, options: RunOptions) -> dict[str, object]:
        model_config = self._get_model_config(options.model_id)
        output_paths = OutputPaths.build(self.config.app.output_root, options.run_id, options.model_id)
        logger = configure_run_logger(output_paths.log_path)
        flush_interval = self.config.app.default_run.checkpoint_interval
        checkpoint = CheckpointStore(output_paths.checkpoint_path, flush_interval=flush_interval)
        if options.resume:
            checkpoint.load()

        raw_writer = JsonlWriter(output_paths.raw_path, flush_interval=flush_interval)
        judged_writer = JsonlWriter(output_paths.judged_path, flush_interval=flush_interval)
        failure_writer = JsonlWriter(output_paths.failures_path, flush_interval=flush_interval)
        retry_writer = JsonlWriter(output_paths.retries_path, flush_interval=flush_interval)

        module_selection = resolve_module_selection(self.config, options.modules)
        include_basic_education = module_selection.include_basic_education
        standard_modules = set(module_selection.standard_modules)
        if include_basic_education and standard_modules and options.max_samples is not None:
            logger.warning(
                "Mixed module run with max_samples=%s: standard modules consume the budget first; "
                "basic education receives the remaining budget.",
                options.max_samples,
            )

        standard_loaded = 0
        standard_failures = 0
        if standard_modules:
            standard_stats = await self._run_standard_samples(
                logger=logger,
                checkpoint=checkpoint,
                raw_writer=raw_writer,
                judged_writer=judged_writer,
                failure_writer=failure_writer,
                retry_writer=retry_writer,
                model_config=model_config,
                options=options,
                modules=standard_modules,
            )
            standard_loaded = standard_stats["loaded_samples"]
            standard_failures = standard_stats["failed_samples"]
            standard_quota_exhausted = bool(standard_stats.get("quota_exhausted", 0))
        else:
            standard_quota_exhausted = False

        basic_stats = {
            "scenario_count": 0,
            "loaded_samples": 0,
            "completed_samples": 0,
            "failed_samples": 0,
        }
        remaining_for_basic = options.max_samples
        if remaining_for_basic is not None:
            remaining_for_basic = max(0, remaining_for_basic - standard_loaded)
        if include_basic_education:
            if standard_quota_exhausted:
                logger.info("Skip basic education branch because the target API quota is exhausted.")
            elif remaining_for_basic == 0:
                logger.info("Skip basic education branch because max_samples budget is exhausted.")
            else:
                basic_runner = BasicEducationExecutor(
                    project_config=self.config,
                    model_config=model_config,
                    logger=logger,
                )
                basic_result = await basic_runner.run(
                    run_id=options.run_id,
                    checkpoint=checkpoint,
                    raw_writer=raw_writer,
                    judged_writer=judged_writer,
                    failure_writer=failure_writer,
                    judge_enabled=options.judge_enabled,
                    subsets=options.subsets,
                    source_files=options.source_files,
                    dimensions=options.dimensions,
                    max_samples=remaining_for_basic,
                    max_samples_per_file=options.max_samples_per_file,
                )
                basic_stats = {
                    "scenario_count": basic_result.scenario_count,
                    "loaded_samples": basic_result.loaded_samples,
                    "completed_samples": basic_result.completed_samples,
                    "failed_samples": basic_result.failed_samples,
                }

        await raw_writer.flush()
        await judged_writer.flush()
        await failure_writer.flush()
        await retry_writer.flush()
        await checkpoint.flush()

        summary = build_summary(output_paths.judged_path, output_paths.failures_path)
        output_paths.summary_path.write_text(
            json.dumps(summary, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("Run finished. summary_path=%s", output_paths.summary_path)
        return {
            "run_id": options.run_id,
            "summary_path": str(output_paths.summary_path),
            "raw_path": str(output_paths.raw_path),
            "judged_path": str(output_paths.judged_path),
            "failures_path": str(output_paths.failures_path),
            "standard_loaded_samples": standard_loaded,
            "standard_failed_samples": standard_failures,
            "quota_exhausted": standard_quota_exhausted,
            "basic_education": basic_stats,
        }

    async def _run_standard_samples(
        self,
        *,
        logger: logging.Logger,
        checkpoint: CheckpointStore,
        raw_writer: JsonlWriter,
        judged_writer: JsonlWriter,
        failure_writer: JsonlWriter,
        retry_writer: JsonlWriter,
        model_config: ModelConfig,
        options: RunOptions,
        modules: set[str] | None,
    ) -> dict[str, int]:
        provider_config = self.config.providers[model_config.provider_name]
        registry = FileRegistry(self.config)
        resolved_items = registry.resolve(
            modules=modules,
            subsets=options.subsets,
            source_files=options.source_files,
        )
        samples = list(
            self._load_samples(
                resolved_items, options.dimensions, options.max_samples, options.max_samples_per_file
            )
        )
        logger.info("Resolved %s samples from %s files.", len(samples), len(resolved_items))
        if not samples:
            return {"loaded_samples": 0, "failed_samples": 0}

        judge_router = JudgeRouter(self.config)
        client = ProviderFactory.create(provider_config, model_config)
        merged_rate_limits = self._merge_rate_limits(
            self.config.app.default_run.provider_rate_limits.get(provider_config.provider_name)
            or self.config.app.default_run.provider_rate_limits.get("default")
            or RateLimitConfig(),
            provider_config.rate_limits,
            model_config.rate_limits,
        )
        if options.max_concurrency is not None:
            merged_rate_limits.max_concurrency = options.max_concurrency

        limiter = SlidingWindowRateLimiter(merged_rate_limits)
        queue: asyncio.Queue[Sample | None] = asyncio.Queue()
        for sample in samples:
            await queue.put(sample)

        worker_count = options.max_concurrency or merged_rate_limits.max_concurrency or 100
        worker_count = min(worker_count, max(1, len(samples)))
        for _ in range(worker_count):
            await queue.put(None)

        progress = {"done": 0, "failed": 0, "skipped": 0, "total": len(samples), "quota_exhausted": 0}
        progress_reporter = _ProgressReporter(
            enabled=options.progress_enabled,
            label=model_config.model_id,
            total=len(samples),
        )
        tasks = [
            asyncio.create_task(
                self._worker(
                    queue=queue,
                    logger=logger,
                    checkpoint=checkpoint,
                    raw_writer=raw_writer,
                    judged_writer=judged_writer,
                    failure_writer=failure_writer,
                    retry_writer=retry_writer,
                    client=client,
                    model_config=model_config,
                    limiter=limiter,
                    judge_router=judge_router,
                    judge_enabled=options.judge_enabled,
                    progress=progress,
                    progress_reporter=progress_reporter,
                )
            )
            for _ in range(worker_count)
        ]
        await asyncio.gather(*tasks)
        progress_reporter.close()
        await client.aclose()
        await judge_router.aclose()
        return {
            "loaded_samples": len(samples),
            "failed_samples": progress["failed"],
            "quota_exhausted": progress["quota_exhausted"],
        }

    def _load_samples(
        self,
        resolved_items: Iterable,
        dimensions: set[str] | None,
        max_samples: int | None,
        max_samples_per_file: int | None = None,
    ) -> Iterable[Sample]:
        loaded = 0
        for item in resolved_items:
            loader = LoaderFactory.create(item.entry.loader_name)
            from_file = 0
            for sample in loader.iter_samples(item):
                if dimensions and sample.dimension not in dimensions:
                    continue
                if max_samples_per_file is not None and from_file >= max_samples_per_file:
                    break
                yield sample
                loaded += 1
                from_file += 1
                if max_samples is not None and loaded >= max_samples:
                    return

    async def _worker(
        self,
        queue: asyncio.Queue[Sample | None],
        logger: logging.Logger,
        checkpoint: CheckpointStore,
        raw_writer: JsonlWriter,
        judged_writer: JsonlWriter,
        failure_writer: JsonlWriter,
        retry_writer: JsonlWriter,
        client,
        model_config: ModelConfig,
        limiter: SlidingWindowRateLimiter,
        judge_router: JudgeRouter,
        judge_enabled: bool,
        progress: dict[str, int],
        progress_reporter: "_ProgressReporter",
    ) -> None:
        while True:
            sample = await queue.get()
            if sample is None:
                queue.task_done()
                return
            sample_key = self._sample_key(sample)
            if progress.get("quota_exhausted"):
                progress["skipped"] += 1
                progress_reporter.update(progress)
                queue.task_done()
                continue
            if checkpoint.is_completed(sample_key):
                progress["skipped"] += 1
                progress_reporter.update(progress)
                queue.task_done()
                continue
            try:
                result = await self._process_sample(
                    sample=sample,
                    client=client,
                    model_config=model_config,
                    limiter=limiter,
                    judge_router=judge_router,
                    judge_enabled=judge_enabled,
                    retry_writer=retry_writer,
                    checkpoint=checkpoint,
                )
                await raw_writer.write(
                    {
                        "sample_id": result.sample_id,
                        "source_file": result.source_file,
                        "module": result.module,
                        "subset": result.subset,
                        "task": result.task,
                        "dimension": result.dimension,
                        "reference": result.reference,
                        "provider_name": result.provider_name,
                        "model_id": result.model_id,
                        "model_name": result.model_name,
                        "model_response": result.model_response,
                        "latency_ms": result.latency_ms,
                        "retry_count": result.retry_count,
                        "timestamp": result.timestamp.isoformat(),
                        "judge_metadata": result.judge_metadata,
                        "raw_response": result.raw_response,
                        "metadata": result.metadata,
                    }
                )
                await judged_writer.write(result.model_dump(mode="json"))
                await checkpoint.mark_completed(sample_key)
                progress["done"] += 1
                progress_reporter.update(progress)
                if progress["done"] % self.config.app.default_run.log_every_n == 0:
                    logger.info("Progress: completed=%s failed=%s", progress["done"], progress["failed"])
            except Exception as exc:  # noqa: BLE001
                if _is_quota_exhausted_error(exc):
                    progress["quota_exhausted"] = 1
                    logger.error("Quota exhausted for model %s: %s", model_config.model_id, exc)
                    progress_reporter.update(progress)
                    continue
                progress["failed"] += 1
                failure = FailureRecord(
                    sample_id=sample.sample_id,
                    source_file=sample.source_file,
                    module=sample.module,
                    subset=sample.subset,
                    provider_name=model_config.provider_name,
                    model_id=model_config.model_id,
                    retry_count=checkpoint.retry_counts.get(sample_key, 0),
                    error_type=type(exc).__name__,
                    error_message=str(exc),
                    timestamp=datetime.now(timezone.utc),
                    metadata={"dimension": sample.dimension},
                )
                await failure_writer.write(failure.model_dump(mode="json"))
                await checkpoint.mark_failed(sample_key)
                progress_reporter.update(progress)
                logger.exception("Sample failed: %s", sample_key)
            finally:
                queue.task_done()

    async def _process_sample(
        self,
        sample: Sample,
        client,
        model_config: ModelConfig,
        limiter: SlidingWindowRateLimiter,
        judge_router: JudgeRouter,
        judge_enabled: bool,
        retry_writer: JsonlWriter,
        checkpoint: CheckpointStore,
    ) -> EvalResult:
        prompt = augment_prompt_for_format(sample, sample.prompt)
        request = GenerationRequest(
            prompt=prompt,
            temperature=model_config.temperature,
            max_tokens=model_config.max_tokens,
            stream=bool(model_config.supports_stream),
            messages=[{"role": "user", "content": prompt}],
            provider_kwargs={},
        )
        response, retry_count = await self._generate_with_retry(
            sample=sample,
            client=client,
            request=request,
            retry_policy=model_config.retry or self.config.app.default_run.retry,
            limiter=limiter,
            retry_writer=retry_writer,
            checkpoint=checkpoint,
        )
        judge = judge_router.get_judge(sample)
        judged = await judge.judge(sample, response) if judge_enabled else None
        metadata = dict(sample.metadata)
        if response.format_valid is False:
            metadata.update(response.format_metadata)
        return EvalResult(
            sample_id=sample.sample_id,
            source_file=sample.source_file,
            source_path=sample.source_path,
            module=sample.module,
            subset=sample.subset,
            task=sample.task,
            dimension=sample.dimension,
            prompt=sample.prompt,
            reference=sample.reference,
            provider_name=model_config.provider_name,
            model_id=model_config.model_id,
            model_name=model_config.model_name,
            model_response=response.text,
            judge_result=judged.judge_result if judged else None,
            score=judged.score if judged else None,
            judge_reason=judged.judge_reason if judged else None,
            latency_ms=response.latency_ms,
            retry_count=retry_count,
            timestamp=datetime.now(timezone.utc),
            metadata=metadata,
            judge_metadata=judged.judge_metadata if judged else {},
            raw_response=response.raw_payload,
        )

    async def _generate_with_retry(
        self,
        sample: Sample,
        client,
        request: GenerationRequest,
        retry_policy: RetryConfig,
        limiter: SlidingWindowRateLimiter,
        retry_writer: JsonlWriter,
        checkpoint: CheckpointStore,
    ) -> tuple[ModelResponse, int]:
        attempts = retry_policy.max_attempts
        retryable_codes = set(retry_policy.retryable_status_codes)
        sample_key = self._sample_key(sample)
        format_spec = get_response_format_spec(sample)
        base_prompt = request.prompt
        request_messages = list(request.messages) if request.messages else [{"role": "user", "content": base_prompt}]
        base_max_tokens = request.max_tokens
        if sample_key in checkpoint.failed_ids and format_spec is not None:
            concise_prompt = build_concise_retry_prompt(base_prompt, format_spec)
            request_messages = [{"role": "user", "content": concise_prompt}]
            request.prompt = concise_prompt
            request.max_tokens = 64
        elif sample.task in ANSWER_COMPLETION_TASKS and sample_key in checkpoint.failed_ids:
            concise_prompt = _build_concise_answer_retry_prompt(base_prompt)
            request_messages = [{"role": "user", "content": concise_prompt}]
            request.prompt = concise_prompt
            request.max_tokens = _answer_retry_max_tokens(base_max_tokens)
        for attempt in range(1, attempts + 1):
            estimated_tokens = self._estimate_tokens(request.prompt, request.max_tokens)
            await limiter.acquire(estimated_tokens=estimated_tokens)
            try:
                request.messages = request_messages
                response = await client.generate(sample=sample, request=request)
                if _should_retry_truncated_answer(sample, response):
                    await checkpoint.set_retry_count(sample_key, attempt)
                    if attempt >= attempts:
                        response.retry_count = attempt - 1
                        return response, attempt - 1
                    delay = await sleep_before_retry(retry_policy, attempt, None)
                    concise_prompt = _build_concise_answer_retry_prompt(base_prompt)
                    request_messages = [{"role": "user", "content": concise_prompt}]
                    request.prompt = concise_prompt
                    request.max_tokens = _answer_retry_max_tokens(base_max_tokens)
                    await retry_writer.write(
                        {
                            "sample_id": sample.sample_id,
                            "source_file": sample.source_file,
                            "attempt": attempt,
                            "next_delay_seconds": delay,
                            "error_type": "LengthTruncatedResponse",
                            "error_message": "Model response ended because max_tokens was reached before a reliable final answer.",
                            "retry_reason": "answer_length_truncated",
                        }
                    )
                    continue
                response = attach_format_metadata(response, format_spec)
                response.retry_count = attempt - 1
                await checkpoint.set_retry_count(sample_key, attempt - 1)
                return response, attempt - 1
            except ResponseFormatError as exc:
                await checkpoint.set_retry_count(sample_key, attempt)
                if attempt >= attempts:
                    response = locals().get("response")
                    if isinstance(response, ModelResponse):
                        response.format_valid = False
                        response.format_metadata = {
                            **exc.metadata,
                            "failure_category": "model_instruction_following",
                            "failure_type": "response_format",
                            "error_message": str(exc),
                        }
                        response.retry_count = attempt - 1
                        await checkpoint.set_retry_count(sample_key, attempt - 1)
                        return response, attempt - 1
                    raise
                delay = await sleep_before_retry(retry_policy, attempt, exc)
                previous_response_text = ""
                response = locals().get("response")
                if "response" in locals():
                    maybe_response = locals()["response"]
                    if isinstance(maybe_response, ModelResponse):
                        previous_response_text = maybe_response.text or ""
                followup = build_retry_followup(
                    format_spec,
                    response if isinstance(response, ModelResponse) else None,
                )
                if followup.get("retry_reason") == "format_empty_after_reasoning":
                    request_messages = [
                        {
                            "role": "user",
                            "content": build_concise_retry_prompt(base_prompt, format_spec),
                        }
                    ]
                else:
                    request_messages = [{"role": "user", "content": base_prompt}]
                    assistant_content = followup.get("assistant_content")
                    if assistant_content:
                        request_messages.append({"role": "assistant", "content": assistant_content})
                    user_reminder = str(followup.get("user_reminder", "") or "").strip()
                    if user_reminder:
                        request_messages.append({"role": "user", "content": user_reminder})
                request.max_tokens = followup.get("max_tokens") or base_max_tokens
                await retry_writer.write(
                    {
                        "sample_id": sample.sample_id,
                        "source_file": sample.source_file,
                        "attempt": attempt,
                        "next_delay_seconds": delay,
                        "error_type": type(exc).__name__,
                        "error_message": str(exc),
                        "retry_reason": followup.get("retry_reason", "format_invalid"),
                        "format_metadata": exc.metadata,
                    }
                )
            except Exception as exc:  # noqa: BLE001
                await checkpoint.set_retry_count(sample_key, attempt)
                if attempt >= attempts or not is_retryable_exception(exc, retryable_codes):
                    raise
                delay = await sleep_before_retry(retry_policy, attempt, exc)
                retry_reason = "http_retry"
                if sample.task in ANSWER_COMPLETION_TASKS:
                    concise_prompt = _build_concise_answer_retry_prompt(base_prompt)
                    request_messages = [{"role": "user", "content": concise_prompt}]
                    request.prompt = concise_prompt
                    request.max_tokens = _answer_retry_max_tokens(base_max_tokens)
                    retry_reason = "answer_retry_after_error"
                await retry_writer.write(
                    {
                        "sample_id": sample.sample_id,
                        "source_file": sample.source_file,
                        "attempt": attempt,
                        "next_delay_seconds": delay,
                        "error_type": type(exc).__name__,
                        "error_message": str(exc),
                        "retry_reason": retry_reason,
                    }
                )
            finally:
                limiter.release()
        raise RuntimeError("Retry loop exited unexpectedly")

    def _get_model_config(self, model_id: str) -> ModelConfig:
        if model_id not in self.config.models:
            available = ", ".join(sorted(self.config.models))
            raise KeyError(f"Unknown model_id={model_id!r}. Available: {available}")
        return self.config.models[model_id]

    def _merge_rate_limits(self, *configs: RateLimitConfig) -> RateLimitConfig:
        merged = RateLimitConfig()
        for config in configs:
            if config.rpm is not None:
                merged.rpm = config.rpm
            if config.tpm is not None:
                merged.tpm = config.tpm
            if config.qps is not None:
                merged.qps = config.qps
            if config.max_concurrency is not None:
                merged.max_concurrency = config.max_concurrency
        return merged

    def _sample_key(self, sample: Sample) -> str:
        return f"{sample.source_file}::{sample.sample_id}"

    def _estimate_tokens(self, prompt: str, max_tokens: int | None) -> int:
        prompt_tokens = max(1, len(prompt) // 4)
        return prompt_tokens + (max_tokens or 0)


class _ProgressReporter:
    def __init__(self, *, enabled: bool, label: str, total: int) -> None:
        self.enabled = enabled
        self.label = label
        self.total = max(1, total)
        self._last_render = ""

    def update(self, progress: dict[str, int]) -> None:
        if not self.enabled:
            return
        completed = progress.get("done", 0) + progress.get("failed", 0) + progress.get("skipped", 0)
        width = 28
        filled = min(width, int(width * completed / self.total))
        bar = "#" * filled + "-" * (width - filled)
        text = (
            f"\r{self.label} [{bar}] {completed}/{self.total} "
            f"ok={progress.get('done', 0)} failed={progress.get('failed', 0)} "
            f"skipped={progress.get('skipped', 0)}"
        )
        if progress.get("quota_exhausted"):
            text += " quota_exhausted"
        if text != self._last_render:
            print(text, end="", file=sys.stderr, flush=True)
            self._last_render = text

    def close(self) -> None:
        if self.enabled and self._last_render:
            print(file=sys.stderr)


def _is_quota_exhausted_error(exc: BaseException) -> bool:
    text = str(exc).lower()
    markers = (
        "insufficient quota",
        "quota exceeded",
        "quota_exceeded",
        "balance",
        "insufficient balance",
        "\u4f59\u989d",
        "\u989d\u5ea6",
        "credit",
        "payment required",
    )
    return any(marker in text for marker in markers)


def _should_retry_truncated_answer(sample: Sample, response: ModelResponse) -> bool:
    if sample.task not in ANSWER_COMPLETION_TASKS:
        return False
    payload = response.raw_payload if isinstance(response.raw_payload, dict) else {}
    choices = payload.get("choices")
    if not isinstance(choices, list) or not choices:
        return False
    first = choices[0] if isinstance(choices[0], dict) else {}
    finish_reason = str(first.get("finish_reason") or "").strip().lower()
    return finish_reason == "length"


def _answer_retry_max_tokens(base_max_tokens: int | None) -> int:
    return base_max_tokens or 512


def _build_concise_answer_retry_prompt(base_prompt: str) -> str:
    lines = []
    for line in base_prompt.splitlines():
        cleaned = re.sub(
            r"\b(?:please\s+)?(?:reason|think)\s+step\s+by\s+step\b(?:\s+before\s+answering)?[,.]?\s*(?:and\s+)?",
            "",
            line,
            flags=re.I,
        ).strip()
        if cleaned:
            lines.append(cleaned)
    prompt = "\n".join(lines).rstrip()
    return (
        f"{prompt}\n\n"
        "Do not include any reasoning, explanation, or thinking. "
        "Output only the final answer. Put the final answer within \\boxed{}."
    ).strip()