Spaces:
Running
Running
| """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 | |