""" Compute provider selection (Wave A — Batch 6 / HP-1 + HP-2). Resolves which ComputeProvider serves generation, per ``config.HOMEPILOT_COMPUTE_MODE``: * ``local`` → LocalComputeProvider (today's behaviour; the default) * ``ollabridge_cloud`` → OllaBridgeCloudComputeProvider * ``auto`` → local GPU when healthy, else a linked OllaBridge device, else local (offline — status says so) ``compute_status()`` backs the plain-language status UX (HP-2): normal users see "Using this PC — Private GPU" or an honest "your PC is offline" message, never endpoint/API-key configuration. """ from __future__ import annotations from typing import Any, Optional from .base import ComputeProvider, GeneratedMedia from .local import LocalComputeProvider from .ollabridge_cloud import OllaBridgeCloudComputeProvider __all__ = [ "ComputeProvider", "GeneratedMedia", "LocalComputeProvider", "OllaBridgeCloudComputeProvider", "ComputeRouter", "route_chat", "get_compute_provider", "resolve_mode", "compute_status", ] def __getattr__(name: str) -> Any: # Lazily expose the router without importing it at package import time # (keeps the import graph light for callers that only need status helpers). if name in ("ComputeRouter", "route_chat"): from .router import ComputeRouter, route_chat return {"ComputeRouter": ComputeRouter, "route_chat": route_chat}[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def _build_cloud() -> Optional[OllaBridgeCloudComputeProvider]: from app.config import ( OLLABRIDGE_CLOUD_IMAGE_MODEL, OLLABRIDGE_CLOUD_TIMEOUT, OLLABRIDGE_CLOUD_TOKEN, OLLABRIDGE_CLOUD_URL, OLLABRIDGE_CLOUD_VIDEO_MODEL, ) if not OLLABRIDGE_CLOUD_URL: return None return OllaBridgeCloudComputeProvider( OLLABRIDGE_CLOUD_URL, OLLABRIDGE_CLOUD_TOKEN, image_model=OLLABRIDGE_CLOUD_IMAGE_MODEL, video_model=OLLABRIDGE_CLOUD_VIDEO_MODEL, timeout=OLLABRIDGE_CLOUD_TIMEOUT, ) def _configured_mode() -> str: from app.config import HOMEPILOT_COMPUTE_MODE mode = (HOMEPILOT_COMPUTE_MODE or "local").strip().lower() return mode if mode in ("local", "ollabridge_cloud", "auto") else "local" def _cloud_configured() -> bool: """A cloud link is usable only when both a URL and a token are set — without a token the job API would 401.""" from app.config import OLLABRIDGE_CLOUD_TOKEN, OLLABRIDGE_CLOUD_URL return bool(OLLABRIDGE_CLOUD_URL and OLLABRIDGE_CLOUD_TOKEN) def _burst_allowed() -> bool: """Whether `auto` may burst to a cloud GPU when the local one is offline. Free for everyone unless the operator gates it to premium (MB6). The local GPU is never affected.""" from app.config import COMPUTE_BURST_REQUIRES_PREMIUM, PREMIUM_COMPUTE_ENABLED return (not COMPUTE_BURST_REQUIRES_PREMIUM) or PREMIUM_COMPUTE_ENABLED async def resolve_mode(modality: str | None = None) -> str: """Resolve the effective mode (never returns ``auto``). ``modality`` makes ``auto`` decide against the runtime the request actually needs (Ollama for chat, ComfyUI for image/video) rather than a single global "is the local GPU up" signal. """ configured = _configured_mode() if configured in ("local", "ollabridge_cloud"): return configured # auto: prefer the local runtime for this modality, then — if bursting is # allowed — a *configured* linked OllaBridge device. if await LocalComputeProvider().available(modality): return "local" if _burst_allowed(): cloud = _build_cloud() if cloud is not None and _cloud_configured() and await cloud.available(modality): return "ollabridge_cloud" return "local" # offline (or burst-gated) — compute_status() explains async def get_compute_provider( mode: str | None = None, modality: str | None = None ) -> ComputeProvider: """Return the provider for the given (or resolved) mode and modality.""" mode = mode or await resolve_mode(modality) if mode == "ollabridge_cloud": cloud = _build_cloud() if cloud is not None: return cloud return LocalComputeProvider() async def compute_status() -> dict[str, Any]: """Backs the HP-2 status UX — plain language, no infrastructure jargon.""" from app.config import ( OLLABRIDGE_CLOUD_TOKEN, OLLABRIDGE_CLOUD_URL, PREMIUM_COMPUTE_ENABLED, ) local_ok = await LocalComputeProvider().available() cloud = _build_cloud() cloud_configured = bool(OLLABRIDGE_CLOUD_URL and OLLABRIDGE_CLOUD_TOKEN) cloud_ok = await cloud.available() if cloud is not None else False mode = await resolve_mode() configured = _configured_mode() # Burst = we're on cloud because the local GPU is offline (auto fallback). burst = mode == "ollabridge_cloud" and configured == "auto" and not local_ok # Burst-gated = local offline + a reachable cloud we could use, but premium is # required and not granted. burst_gated = ( configured == "auto" and not local_ok and cloud_configured and cloud_ok and not _burst_allowed() ) if mode == "local" and local_ok: label = "Private GPU" message = "Compute: Connected — Using this PC — Mode: Private GPU" elif mode == "ollabridge_cloud" and cloud_ok: label = "OllaBridge Cloud" message = ( "Your PC is offline — running on a cloud GPU (premium)." if burst else "Compute: Connected — Using your paired GPU via OllaBridge Cloud" ) elif burst_gated: label = "Offline" message = "Your PC is offline — upgrade to premium to run on a cloud GPU, or wait for your PC." else: label = "Offline" message = ( "Your PC is offline — wait for it to come back, use the free cloud " "queue, or use a premium GPU." ) return { "mode": mode, "configured_mode": configured, "local_gpu_available": local_ok, "cloud_configured": cloud_configured, "cloud_reachable": cloud_ok, "premium": PREMIUM_COMPUTE_ENABLED, "burst": burst, "burst_gated": burst_gated, "label": label, "message": message, }