File size: 6,490 Bytes
afff449 a42bc3b afff449 a42bc3b afff449 a42bc3b afff449 a42bc3b afff449 a42bc3b afff449 a42bc3b afff449 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """
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,
}
|