298 lines
11 KiB
Python
298 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Lease-aware FLUX image generation for the shared titan-24 GPU."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import gc
|
|
import io
|
|
import json
|
|
import os
|
|
import ssl
|
|
import threading
|
|
from datetime import UTC, datetime
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.error import HTTPError
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
HOST = os.environ.get("HERMES_LOCAL_IMAGE_HOST", "0.0.0.0")
|
|
PORT = int(os.environ.get("HERMES_LOCAL_IMAGE_PORT", "9004"))
|
|
MODEL_ID = os.environ.get(
|
|
"HERMES_LOCAL_IMAGE_MODEL", "black-forest-labs/FLUX.2-klein-4B"
|
|
)
|
|
MODEL_REVISION = os.environ.get(
|
|
"HERMES_LOCAL_IMAGE_REVISION", "e7b7dc27f91deacad38e78976d1f2b499d76a294"
|
|
)
|
|
MODEL_CACHE = os.environ.get("HF_HOME", "/models/huggingface")
|
|
LEASE_NAMESPACE = os.environ.get("LEASE_NAMESPACE", "hermes")
|
|
LEASE_NAME = os.environ.get("LEASE_NAME", "titan-24-gpu-owner")
|
|
LEASE_IDLE_OWNER = os.environ.get("LEASE_IDLE_OWNER", "hermes")
|
|
LEASE_IMAGE_OWNER = os.environ.get("LEASE_IMAGE_OWNER", "hermes-image")
|
|
MAX_BODY_BYTES = int(os.environ.get("HERMES_LOCAL_IMAGE_MAX_BODY", str(96 << 20)))
|
|
QUEUE_TIMEOUT_SECONDS = float(os.environ.get("HERMES_LOCAL_IMAGE_QUEUE_TIMEOUT", "1200"))
|
|
API_HOST = os.environ.get("KUBERNETES_SERVICE_HOST", "kubernetes.default.svc")
|
|
API_PORT = os.environ.get("KUBERNETES_SERVICE_PORT_HTTPS", "443")
|
|
TOKEN_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/token")
|
|
CA_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/ca.crt")
|
|
LEASE_URL = (
|
|
f"https://{API_HOST}:{API_PORT}/apis/coordination.k8s.io/v1/"
|
|
f"namespaces/{LEASE_NAMESPACE}/leases/{LEASE_NAME}"
|
|
)
|
|
ASPECTS = {
|
|
"square": (1024, 1024),
|
|
"landscape": (768, 1280),
|
|
"portrait": (1280, 768),
|
|
}
|
|
|
|
_generation_lock = threading.Lock()
|
|
_state_lock = threading.Lock()
|
|
_state: dict[str, Any] = {
|
|
"busy": False,
|
|
"phase": "idle",
|
|
"last_error": None,
|
|
"last_model": None,
|
|
}
|
|
|
|
|
|
def _set_state(**values: Any) -> None:
|
|
"""Update health state without exposing prompts or image content."""
|
|
with _state_lock:
|
|
_state.update(values)
|
|
|
|
|
|
def _state_snapshot() -> dict[str, Any]:
|
|
"""Return a stable copy for health responses."""
|
|
with _state_lock:
|
|
return dict(_state)
|
|
|
|
|
|
def _kube_request(method: str, body: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
"""Call the Lease API using the pod service account."""
|
|
token = TOKEN_PATH.read_text(encoding="utf-8").strip()
|
|
raw = None if body is None else json.dumps(body).encode("utf-8")
|
|
headers = {"Authorization": f"Bearer {token}"}
|
|
if raw is not None:
|
|
headers["Content-Type"] = "application/merge-patch+json"
|
|
request = Request(LEASE_URL, data=raw, headers=headers, method=method)
|
|
context = ssl.create_default_context(cafile=str(CA_PATH))
|
|
with urlopen(request, timeout=5, context=context) as response:
|
|
return json.load(response)
|
|
|
|
|
|
def _lease_owner() -> str:
|
|
"""Read the current GPU owner, failing closed on API errors."""
|
|
payload = _kube_request("GET")
|
|
return str((payload.get("spec") or {}).get("holderIdentity") or "unavailable")
|
|
|
|
|
|
def _change_owner(expected: str, desired: str) -> bool:
|
|
"""Change owner only if the resource version and holder still match."""
|
|
current = _kube_request("GET")
|
|
if str((current.get("spec") or {}).get("holderIdentity")) != expected:
|
|
return False
|
|
resource_version = str((current.get("metadata") or {}).get("resourceVersion") or "")
|
|
if not resource_version:
|
|
return False
|
|
patch = {
|
|
"metadata": {"resourceVersion": resource_version},
|
|
"spec": {
|
|
"holderIdentity": desired,
|
|
"renewTime": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
|
|
},
|
|
}
|
|
try:
|
|
_kube_request("PATCH", patch)
|
|
except HTTPError as exc:
|
|
if exc.code == 409:
|
|
return False
|
|
raise
|
|
return True
|
|
|
|
|
|
def _recover_stale_image_lease() -> None:
|
|
"""Release this deployment's lease after a prior renderer crash."""
|
|
try:
|
|
if _lease_owner() == LEASE_IMAGE_OWNER:
|
|
_change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER)
|
|
except Exception as exc:
|
|
_set_state(last_error=f"lease recovery failed: {type(exc).__name__}: {exc}")
|
|
|
|
|
|
def _decode_image(value: str):
|
|
"""Decode a bounded image data URL for FLUX editing."""
|
|
from PIL import Image
|
|
|
|
if not value.startswith("data:image/") or ";base64," not in value:
|
|
raise ValueError("local reference images must be data URLs")
|
|
raw = base64.b64decode(value.split(",", 1)[1], validate=True)
|
|
if not raw or len(raw) > 25 << 20:
|
|
raise ValueError("reference image must be between 1 byte and 25 MiB")
|
|
return Image.open(io.BytesIO(raw)).convert("RGB")
|
|
|
|
|
|
def _render(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Load FLUX for one request, render, then fully release its GPU memory."""
|
|
import torch
|
|
from diffusers import Flux2KleinPipeline
|
|
|
|
prompt = str(payload.get("prompt") or "").strip()
|
|
if not prompt:
|
|
raise ValueError("prompt is required")
|
|
aspect = str(payload.get("aspect_ratio") or "landscape").lower()
|
|
if aspect not in ASPECTS:
|
|
raise ValueError("aspect_ratio must be landscape, square, or portrait")
|
|
height, width = ASPECTS[aspect]
|
|
inputs = []
|
|
if payload.get("image_url"):
|
|
inputs.append(_decode_image(str(payload["image_url"])))
|
|
for item in payload.get("reference_image_urls") or []:
|
|
inputs.append(_decode_image(str(item)))
|
|
inputs = inputs[:4]
|
|
|
|
_set_state(phase="loading-flux", last_model=MODEL_ID)
|
|
pipe = None
|
|
try:
|
|
pipe = Flux2KleinPipeline.from_pretrained(
|
|
MODEL_ID,
|
|
revision=MODEL_REVISION,
|
|
torch_dtype=torch.bfloat16,
|
|
cache_dir=MODEL_CACHE,
|
|
)
|
|
pipe.enable_model_cpu_offload()
|
|
_set_state(phase="rendering")
|
|
arguments: dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"height": height,
|
|
"width": width,
|
|
"guidance_scale": 1.0,
|
|
"num_inference_steps": 4,
|
|
}
|
|
if inputs:
|
|
arguments["image"] = inputs
|
|
image = pipe(**arguments).images[0]
|
|
output = io.BytesIO()
|
|
image.save(output, format="PNG", optimize=True)
|
|
return {
|
|
"success": True,
|
|
"image_b64": base64.b64encode(output.getvalue()).decode("ascii"),
|
|
"mime_type": "image/png",
|
|
"model": "flux-2-klein-4b-local",
|
|
"quality": "local-high",
|
|
"size": f"{width}x{height}",
|
|
"input_image_count": len(inputs),
|
|
"route": "local",
|
|
}
|
|
finally:
|
|
del pipe
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
|
|
|
|
def _generate(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Serialize renders and perform an exclusive lease-backed GPU handoff."""
|
|
acquired = _generation_lock.acquire(timeout=QUEUE_TIMEOUT_SECONDS)
|
|
if not acquired:
|
|
raise TimeoutError("local image queue is busy; retry shortly")
|
|
claimed = False
|
|
_set_state(busy=True, phase="claiming-gpu", last_error=None)
|
|
try:
|
|
claimed = _change_owner(LEASE_IDLE_OWNER, LEASE_IMAGE_OWNER)
|
|
if not claimed:
|
|
owner = _lease_owner()
|
|
raise RuntimeError(f"local image GPU unavailable while titan-24 owner is {owner}")
|
|
result = _render(payload)
|
|
_set_state(phase="releasing-flux")
|
|
if _lease_owner() == LEASE_IMAGE_OWNER:
|
|
if not _change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER):
|
|
raise RuntimeError("GPU ownership changed while releasing the image lane")
|
|
return result
|
|
except Exception as exc:
|
|
_set_state(last_error=f"{type(exc).__name__}: {exc}")
|
|
raise
|
|
finally:
|
|
# If Wolf has moved the lease to wolf-draining, never overwrite that
|
|
# transition. Ariadne waits for busy=false before using the GPU.
|
|
if claimed:
|
|
try:
|
|
if _lease_owner() == LEASE_IMAGE_OWNER:
|
|
_change_owner(LEASE_IMAGE_OWNER, LEASE_IDLE_OWNER)
|
|
except Exception as exc:
|
|
_set_state(last_error=f"restore failed: {type(exc).__name__}: {exc}")
|
|
_set_state(busy=False, phase="idle")
|
|
_generation_lock.release()
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
"""Internal JSON interface used only by the owner image broker."""
|
|
|
|
server_version = "HermesLocalImage/1"
|
|
|
|
def _json(self, status: int, value: dict[str, Any]) -> None:
|
|
body = json.dumps(value, separators=(",", ":")).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.send_header("Cache-Control", "no-store")
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def do_GET(self) -> None: # noqa: N802
|
|
if self.path != "/health":
|
|
self._json(404, {"success": False, "error": "not found"})
|
|
return
|
|
state = _state_snapshot()
|
|
try:
|
|
owner = _lease_owner()
|
|
except Exception:
|
|
owner = "unavailable"
|
|
self._json(
|
|
200,
|
|
{
|
|
"success": True,
|
|
"available": owner == LEASE_IDLE_OWNER and not state["busy"],
|
|
"gpu_owner": owner,
|
|
"model": "flux-2-klein-4b-local",
|
|
**state,
|
|
},
|
|
)
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
if self.path != "/v1/images/generations":
|
|
self._json(404, {"success": False, "error": "not found"})
|
|
return
|
|
try:
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
except ValueError:
|
|
length = 0
|
|
if length <= 0 or length > MAX_BODY_BYTES:
|
|
self._json(413, {"success": False, "error": "invalid request size"})
|
|
return
|
|
try:
|
|
payload = json.loads(self.rfile.read(length))
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("JSON object required")
|
|
result = _generate(payload)
|
|
self._json(200, result)
|
|
except ValueError as exc:
|
|
self._json(400, {"success": False, "error": str(exc)})
|
|
except TimeoutError as exc:
|
|
self._json(503, {"success": False, "error": str(exc)})
|
|
except Exception as exc:
|
|
self._json(
|
|
503,
|
|
{"success": False, "error": f"local image generation failed: {type(exc).__name__}: {exc}"},
|
|
)
|
|
|
|
def log_message(self, format_string: str, *args: Any) -> None:
|
|
print(f"local-image {self.address_string()} {format_string % args}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_recover_stale_image_lease()
|
|
ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()
|