core/executor.py

"""Serializable inference dispatch across local, GPU, and ZeroGPU runtimes."""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any, Callable, Literal

from core.errors import GatewayError
from core.loader import ModelLoader
from core.runtime import inference_context, memory_stats


logger = logging.getLogger(__name__)


def _gpu_task(**options: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Use the Space scheduler in deployment and a direct call in local tests."""
    try:
        import spaces
    except ImportError:
        return lambda function: function
    return spaces.GPU(**options)


@dataclass(frozen=True, slots=True)
class InferenceCommand:
    """Pickle-safe description of one model method invocation."""

    model_name: str
    method_name: str
    arguments: dict[str, Any]
    request_id: str
    duration_seconds: int
    gpu_size: Literal["large", "xlarge"] = "large"


@dataclass(frozen=True, slots=True)
class CommandError:
    """Serializable error returned from a forked ZeroGPU worker."""

    message: str
    status_code: int
    code: str


@dataclass(frozen=True, slots=True)
class CommandResult:
    """Serializable success or failure from a GPU allocation."""

    value: Any = None
    error: CommandError | None = None


def _duration(command: InferenceCommand) -> int:
    return command.duration_seconds


def _zero_gpu_active() -> bool:
    try:
        from spaces.config import Config
    except ImportError:
        return False
    return bool(Config.zero_gpu)


def _invoke(command: InferenceCommand) -> CommandResult:
    """Execute one command inside an allocated device process."""
    loader = ModelLoader()
    started = time.perf_counter()
    try:
        with loader.use_model(command.model_name) as model:
            method = getattr(model, command.method_name, None)
            if not callable(method):
                raise RuntimeError(
                    f"Model {command.model_name} does not implement "
                    f"{command.method_name}"
                )
            with inference_context(loader.settings, loader.device):
                value = method(**command.arguments)
    except GatewayError as exc:
        return CommandResult(
            error=CommandError(exc.message, exc.status_code, exc.code)
        )
    except Exception as exc:
        logger.exception(
            "model inference failed",
            extra={
                "model": command.model_name,
                "request_id": command.request_id,
                **memory_stats(),
            },
        )
        message = str(exc)
        if "out of memory" in message.lower():
            error = CommandError(
                f"{command.model_name} ran out of memory",
                507,
                "out_of_memory",
            )
        elif isinstance(exc, OSError):
            error = CommandError(
                f"File operation failed during {command.model_name} inference",
                500,
                "file_error",
            )
        else:
            error = CommandError(
                f"{command.model_name} inference failed: {message}",
                500,
                "inference_failed",
            )
        return CommandResult(error=error)
    finally:
        if _zero_gpu_active():
            loader.close()

    logger.info(
        "model inference completed",
        extra={
            "model": command.model_name,
            "request_id": command.request_id,
            "execution_time": round(time.perf_counter() - started, 3),
            **memory_stats(),
        },
    )
    return CommandResult(value=value)


@_gpu_task(duration=_duration, size="large")
def _invoke_large(command: InferenceCommand) -> CommandResult:
    return _invoke(command)


@_gpu_task(duration=_duration, size="xlarge")
def _invoke_xlarge(command: InferenceCommand) -> CommandResult:
    return _invoke(command)


def execute_inference(command: InferenceCommand) -> Any:
    """Acquire the requested ZeroGPU tier and restore domain errors in the parent."""
    try:
        result = (
            _invoke_xlarge(command)
            if command.gpu_size == "xlarge"
            else _invoke_large(command)
        )
    except Exception as exc:
        message = str(exc)
        normalized = message.lower()
        if "quota" in normalized or "no gpu" in normalized or "zerogpu" in normalized:
            raise GatewayError(
                "ZeroGPU is unavailable or its quota was exceeded",
                status_code=503,
                code="gpu_unavailable",
            ) from exc
        raise GatewayError(
            "ZeroGPU scheduling failed",
            status_code=503,
            code="gpu_scheduling_failed",
        ) from exc
    if result.error is not None:
        raise GatewayError(
            result.error.message,
            status_code=result.error.status_code,
            code=result.error.code,
        )
    return result.value