src/content_engine/services/router.py

"""Job routing service — decides whether to run generation locally or in the cloud.

For v1 (single GPU, no cloud accounts), everything routes locally.
The router checks ComfyUI's queue depth and VRAM before dispatching.
"""

from __future__ import annotations

import logging
from enum import Enum

from content_engine.config import settings
from content_engine.services.comfyui_client import ComfyUIClient

logger = logging.getLogger(__name__)


class Backend(str, Enum):
    LOCAL = "local"
    CLOUD = "cloud"


class RouterService:
    """Routes generation jobs to local GPU or cloud APIs."""

    def __init__(self, comfyui_client: ComfyUIClient):
        self.comfyui = comfyui_client
        self.max_queue_depth = settings.comfyui.max_local_queue_depth
        self.min_vram_gb = settings.comfyui.min_vram_gb

    async def route(self, estimated_vram_gb: float = 4.0) -> Backend:
        """Decide where to run a generation job.

        Returns Backend.LOCAL if the local GPU is available,
        Backend.CLOUD if it should be offloaded.
        """
        # v1: No cloud providers configured — always local
        if not settings.cloud_providers:
            return Backend.LOCAL

        # Check if ComfyUI is reachable
        if not await self.comfyui.is_available():
            logger.warning("ComfyUI unavailable, routing to cloud")
            return Backend.CLOUD

        # Check queue depth
        queue_depth = await self.comfyui.get_queue_depth()
        if queue_depth >= self.max_queue_depth:
            logger.info(
                "Local queue full (%d/%d), routing to cloud",
                queue_depth,
                self.max_queue_depth,
            )
            return Backend.CLOUD

        # Check VRAM
        vram_free = await self.comfyui.get_vram_free_gb()
        if vram_free is not None and vram_free < self.min_vram_gb:
            logger.info(
                "Insufficient VRAM (%.1f GB free, need %.1f GB), routing to cloud",
                vram_free,
                self.min_vram_gb,
            )
            return Backend.CLOUD

        return Backend.LOCAL