utils.py

"""Utility for Terminal-Bench server."""

import logging
import os
import re
import threading
from contextlib import contextmanager

import docker
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

logger = logging.getLogger(__name__)


class ImageLRUCache:
    """Simple thread-safe tracker for Docker image usage (for compatibility)."""

    def __init__(self, max_size: int | None = None) -> None:
        # max_size parameter kept for compatibility but not used
        self.lock = threading.Lock()

    @contextmanager
    def use(self, image_key: str):
        """Context manager for image usage tracking."""
        with self.lock:
            pass  # Just acquire lock for thread safety
        try:
            yield
        finally:
            pass  # No cleanup needed


@contextmanager
def get_docker_client():
    docker_host = os.getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
    client = docker.DockerClient(base_url=docker_host)
    try:
        yield client
    finally:
        try:
            client.close()
        except Exception as e:
            logger.warning(f"Failed to close Docker client: {e}")


def get_container_network_mode() -> str:
    """Return the Docker network mode used for task and evaluator containers."""
    network_mode = os.getenv("CONTAINER_NETWORK_MODE", "bridge").strip()
    return network_mode or "bridge"


def get_image_candidates(task_name: str, docker_image_from_toml: str | None = None) -> list[str]:
    """Return candidate image names ordered by preference."""
    custom_registry = os.getenv("DOCKER_REGISTRY", "").strip()
    candidates: list[str] = []

    if custom_registry:
        if docker_image_from_toml:
            match = re.match(r"^[^/]+/([^:]+):(.+)$", docker_image_from_toml)
            if match:
                image_task_name = match.group(1)
                version = match.group(2)
                candidates.append(f"{custom_registry}:{image_task_name}-{version}")
            else:
                logger.warning(f"Could not parse docker_image format: {docker_image_from_toml}, using task_name")
                candidates.append(f"{custom_registry}:{task_name}")
        else:
            candidates.append(f"{custom_registry}:{task_name}")

    if docker_image_from_toml:
        candidates.append(docker_image_from_toml)
    else:
        default_registry = "registry.h.pjlab.org.cn/ailab-opencompass-opencompass_proxy/terminal_bench_2"
        candidates.append(f"{default_registry}:{task_name}")

    deduped: list[str] = []
    seen: set[str] = set()
    for candidate in candidates:
        if candidate not in seen:
            deduped.append(candidate)
            seen.add(candidate)
    return deduped


def get_image_name(task_name: str, docker_image_from_toml: str | None = None) -> str:
    """
    Get the preferred Docker image name based on configuration.

    This keeps the existing API for callers that only need the first-choice name.
    For robust execution, use resolve_image_name() instead.
    """
    return get_image_candidates(task_name, docker_image_from_toml)[0]


def _ensure_single_image_exists(client: docker.DockerClient, image_name: str) -> None:
    """Ensure one specific image exists locally, pulling if necessary."""
    pull_if_not_exists = os.getenv("PULL_IF_NOT_EXISTS", "true").lower() == "true"

    image_exists = False
    try:
        client.images.get(image_name)
        image_exists = True
        logger.info(f"Docker image found locally: {image_name}")
    except docker.errors.ImageNotFound:
        logger.info(f"Docker image not found locally: {image_name}")

    # If image doesn't exist locally, try to pull
    if not image_exists:
        if pull_if_not_exists:
            logger.info(f"Pulling Docker image: {image_name}")
            try:
                client.images.pull(image_name)
                logger.info(f"Successfully pulled image: {image_name}")
            except docker.errors.APIError as e:
                raise RuntimeError(
                    f"Failed to pull Docker image: {image_name}. Error: {e}\n"
                    "Please either:\n"
                    "1. Run build_images.sh to build images locally, or\n"
                    "2. Ensure the image is available in the registry"
                )
        else:
            raise RuntimeError(
                f"Docker image not found locally: {image_name}\n"
                "Please either:\n"
                "1. Run build_images.sh to build images locally, or\n"
                "2. Set PULL_IF_NOT_EXISTS=true to allow pulling from registry"
            )


def ensure_image_exists(client: docker.DockerClient, image_name: str) -> None:
    """Backward-compatible wrapper for callers that pass a single image name."""
    _ensure_single_image_exists(client, image_name)


def resolve_image_name(
    client: docker.DockerClient,
    task_name: str,
    docker_image_from_toml: str | None = None,
) -> str:
    """
    Resolve a usable Docker image, trying registry-mapped and original names in order.

    Returns the image name that is available locally or was pulled successfully.
    Raises the last error if all candidates fail.
    """
    candidates = get_image_candidates(task_name, docker_image_from_toml)
    errors: list[str] = []

    for idx, image_name in enumerate(candidates):
        try:
            if idx > 0:
                logger.warning("Falling back to alternate Docker image candidate: %s", image_name)
            _ensure_single_image_exists(client, image_name)
            return image_name
        except RuntimeError as exc:
            errors.append(str(exc))
            logger.warning("Docker image candidate failed: %s", image_name)

    raise RuntimeError(
        "Failed to resolve any usable Docker image. Tried:\n"
        + "\n".join(f"- {candidate}" for candidate in candidates)
        + "\nErrors:\n"
        + "\n---\n".join(errors)
    )