ReMe/deploy/benchmark_docker/pool.py

"""Bounded Docker workers for independent benchmark cases.

Each case completes construction and its queries before its worker is assigned
another case. A dependent case may reuse the workspace captured immediately
after its parent's construction, without sharing mutable container state.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any

from ._cleanup import finish_cleanup, format_error
from .docker_sandbox import DockerReMeSandbox, DockerReMeSandboxFactory
from .models import EvaluationQuery, JobRequest


@dataclass(frozen=True, slots=True)
class CasePlan:
    """One case, optionally seeded by a host workspace or another case."""

    case_id: str
    build_jobs: tuple[JobRequest, ...] = ()
    queries: tuple[EvaluationQuery, ...] = ()
    workspace_source: str | Path | None = None
    reuse_workspace_from: str | None = None
    export_workspace_to: str | Path | None = None
    export_build_log_to: str | Path | None = None
    export_evaluation_to: str | Path | None = None


@dataclass(slots=True)
class CaseResult:
    """Structured build and answer results; judging stays on the host."""

    case_id: str
    build: dict[str, Any] | None = None
    queries: dict[str, Any] | None = None
    workspace_snapshot: Path | None = None
    build_log: str | None = None
    error: str | None = None

    @property
    def success(self) -> bool:
        """Whether all requested phases and owned-worker cleanup succeeded."""
        return (
            self.error is None
            and (self.build is None or self.build.get("success") is True)
            and (self.queries is None or self.queries.get("success") is True)
        )


@dataclass(slots=True)
class _WorkerSlots:
    """Mutable worker availability owned by one pool run."""

    workers: list[DockerReMeSandbox]
    available: asyncio.Queue[DockerReMeSandbox | None]
    unclaimed: set[str]
    live_slots: int
    next_worker_index: int
    exhaustion_error: str = "no available workers"


class BenchmarkDockerPool:
    """Reuse at most ``concurrency`` containers across case plans.

    ``run`` owns and closes its containers. It never stages all construction
    phases ahead of queries: each available worker runs build then queries for
    one case. Dependencies await a completed parent without holding a worker.
    """

    def __init__(self, factory: DockerReMeSandboxFactory, *, concurrency: int = 4) -> None:
        if concurrency < 1:
            raise ValueError("concurrency must be at least 1")
        self.factory = factory
        self.concurrency = concurrency

    async def run(self, plans: list[CasePlan]) -> list[CaseResult]:
        """Return results in input order, closing every container on exit."""
        if not plans:
            return []
        self._validate_plans(plans)
        reused_case_ids = {plan.reuse_workspace_from for plan in plans if plan.reuse_workspace_from is not None}

        workers = await self.factory.create_cases(
            [f"pool-{index}" for index in range(min(self.concurrency, len(plans)))],
            concurrency=self.concurrency,
        )
        available: asyncio.Queue[DockerReMeSandbox | None] = asyncio.Queue()
        for worker in workers:
            available.put_nowait(worker)
        slots = _WorkerSlots(workers, available, {plan.case_id for plan in plans}, len(workers), len(workers))
        worker_results: dict[int, list[CaseResult]] = {}
        results: list[CaseResult] = []
        run_error: BaseException | None = None

        try:
            with TemporaryDirectory(prefix="reme-benchmark-pool-") as temporary:
                snapshot_root = Path(temporary)
                # Tasks are started together; awaiting a parent never consumes
                # a container. The graph check above rules out deadlocks.
                tasks: dict[str, asyncio.Task[CaseResult]] = {}

                async def execute(plan: CasePlan) -> CaseResult:
                    result = CaseResult(plan.case_id)
                    source = plan.workspace_source
                    if plan.reuse_workspace_from is not None:
                        parent = await tasks[plan.reuse_workspace_from]
                        if parent.workspace_snapshot is None:
                            slots.unclaimed.discard(plan.case_id)
                            return CaseResult(plan.case_id, error=f"source case {parent.case_id!r} failed")
                        source = parent.workspace_snapshot

                    worker = await available.get()
                    slots.unclaimed.discard(plan.case_id)
                    if worker is None:
                        return CaseResult(plan.case_id, error=slots.exhaustion_error)
                    worker_results.setdefault(id(worker), []).append(result)

                    try:
                        await worker.reset_case(plan.case_id)
                        if source is not None:
                            await worker.upload_workspace(source)
                        if plan.build_jobs:
                            try:
                                result.build = await worker.run_build(list(plan.build_jobs))
                            except Exception as exc:
                                result.error = format_error(exc)
                                await self._retain_failed_build_log(worker, plan, result)
                                return result
                            if not result.build.get("success"):
                                result.error = "build failed"
                                await self._retain_failed_build_log(worker, plan, result)
                                return result
                            if plan.export_build_log_to is not None:
                                await worker.export_build_log(plan.export_build_log_to)
                        # Capture before queries, since answering may change the
                        # workspace. Dependents receive construction state.
                        if plan.export_workspace_to is not None or plan.case_id in reused_case_ids:
                            destination = plan.export_workspace_to or snapshot_root / f"{plan.case_id}.tar.gz"
                            result.workspace_snapshot = await worker.export_workspace(destination)
                        if plan.queries:
                            result.queries = await worker.run_queries(list(plan.queries))
                            if not result.queries.get("success"):
                                result.error = "query failed"
                        if plan.export_evaluation_to is not None:
                            await worker.export_evaluation(plan.export_evaluation_to)
                        return result
                    except Exception as exc:  # Keep other cases running and release this worker.
                        result.error = format_error(exc)
                        return result
                    finally:
                        await self._release_worker(worker, result, slots)

                for plan in plans:
                    tasks[plan.case_id] = asyncio.create_task(execute(plan))
                try:
                    results = await asyncio.gather(*(tasks[plan.case_id] for plan in plans))
                finally:
                    for task in tasks.values():
                        if not task.done() and not task.cancelling():
                            task.cancel()
                    await asyncio.gather(*tasks.values(), return_exceptions=True)
                for plan, result in zip(plans, results):
                    if plan.export_workspace_to is None:
                        result.workspace_snapshot = None  # The internal temporary archive is about to be deleted.
                return results
        except BaseException as exc:
            run_error = exc
            raise
        finally:
            await self._close_workers(workers, worker_results, results, run_error)

    async def _release_worker(self, worker: DockerReMeSandbox, result: CaseResult, slots: _WorkerSlots) -> None:
        """Return a reusable worker or retire its slot after confirmed cleanup."""
        if worker.reusable:
            slots.available.put_nowait(worker)
            return

        def record_lifecycle_error(phase: str, exc: Exception) -> None:
            detail = f"{phase} failed: {format_error(exc)}"
            result.error = f"{result.error}; {detail}" if result.error else detail
            slots.exhaustion_error = f"no available workers: {detail}"

        # Do not reuse the slot while its old container may still be running.
        try:
            await worker.close()
        except Exception as exc:
            record_lifecycle_error("worker close", exc)
        else:
            if slots.unclaimed and not asyncio.current_task().cancelling():
                worker_index = slots.next_worker_index
                slots.next_worker_index += 1
                try:
                    replacement = await self.factory.create_case(f"pool-{worker_index}")
                except Exception as exc:
                    record_lifecycle_error("replacement", exc)
                else:
                    slots.workers.append(replacement)
                    slots.available.put_nowait(replacement)
                    return

        slots.live_slots -= 1
        if slots.live_slots == 0:
            # Also wake cases that are still waiting for a parent to complete.
            for _ in slots.unclaimed:
                slots.available.put_nowait(None)

    @staticmethod
    async def _retain_failed_build_log(worker: DockerReMeSandbox, plan: CasePlan, result: CaseResult) -> None:
        """Keep diagnostics without replacing the original build failure."""
        try:
            result.build_log = await worker.read_build_log()
        except Exception as exc:
            result.error += f"; build log unavailable: {type(exc).__name__}: {exc}"
        if plan.export_build_log_to is not None:
            try:
                await worker.export_build_log(plan.export_build_log_to)
            except Exception as exc:
                result.error += f"; build log export failed: {type(exc).__name__}: {exc}"

    @staticmethod
    async def _close_workers(
        workers: list[DockerReMeSandbox],
        worker_results: dict[int, list[CaseResult]],
        results: list[CaseResult],
        run_error: BaseException | None,
    ) -> None:
        """Finish all teardown and attach failures to results or the run error."""
        outcomes, cleanup_cancel = await finish_cleanup(worker.close() for worker in workers)
        for worker, outcome in zip(workers, outcomes):
            if isinstance(outcome, BaseException):
                detail = f"worker {worker.container_id!r} final close failed: {format_error(outcome)}"
                error = run_error if run_error is not None else cleanup_cancel
                if error is not None:
                    error.add_note(detail)
                # An unused worker is still owned by this run; report its
                # failure on all results rather than silently dropping it.
                for result in worker_results.get(id(worker)) or results:
                    result.error = f"{result.error}; {detail}" if result.error else detail
        if cleanup_cancel is not None and run_error is None:
            raise cleanup_cancel

    @staticmethod
    def _validate_plans(plans: list[CasePlan]) -> None:
        """Reject invalid plans and dependencies before creating containers."""
        by_id = {plan.case_id: plan for plan in plans}
        if len(by_id) != len(plans):
            raise ValueError("case IDs must be unique")
        for plan in plans:
            if plan.workspace_source is not None and plan.reuse_workspace_from is not None:
                raise ValueError(f"case {plan.case_id!r} has two workspace sources")
            if not plan.build_jobs and plan.workspace_source is None and plan.reuse_workspace_from is None:
                raise ValueError(f"case {plan.case_id!r} needs build jobs or a workspace source")
            if plan.reuse_workspace_from is not None and plan.reuse_workspace_from not in by_id:
                raise ValueError(f"unknown workspace source: {plan.reuse_workspace_from!r}")
            if plan.export_build_log_to is not None and not plan.build_jobs:
                raise ValueError("build log export requires build jobs")
            if plan.export_evaluation_to is not None and not plan.queries:
                raise ValueError("evaluation export requires queries")
        BenchmarkDockerPool._check_acyclic(by_id)

    @staticmethod
    def _check_acyclic(plans: dict[str, CasePlan]) -> None:
        visiting: set[str] = set()
        visited: set[str] = set()

        def visit(case_id: str) -> None:
            if case_id in visiting:
                raise ValueError(f"workspace reuse cycle at {case_id!r}")
            if case_id in visited:
                return
            visiting.add(case_id)
            parent = plans[case_id].reuse_workspace_from
            if parent is not None:
                visit(parent)
            visiting.remove(case_id)
            visited.add(case_id)

        for case_id in plans:
            visit(case_id)