490 lines
19 KiB
Python
490 lines
19 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 subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from contextlib import redirect_stdout
|
|
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")
|
|
# Kubernetes reserves <SERVICE>_PORT for service-link discovery. Keep the
|
|
# listener setting out of that namespace so it is always numeric.
|
|
PORT = int(os.environ.get("HERMES_LOCAL_IMAGE_LISTEN_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")
|
|
OFFLOAD_MODE = os.environ.get("HERMES_LOCAL_IMAGE_OFFLOAD_MODE", "sequential")
|
|
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"))
|
|
RENDER_TIMEOUT_SECONDS = float(
|
|
os.environ.get("HERMES_LOCAL_IMAGE_RENDER_TIMEOUT", "1200")
|
|
)
|
|
GPU_ACTIVITY_URL = os.environ.get(
|
|
"HERMES_LOCAL_IMAGE_GPU_ACTIVITY_URL",
|
|
"http://nvidia-process-exporter.monitoring.svc.cluster.local:9401/metrics",
|
|
)
|
|
GPU_ACTIVITY_NODE = os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_NODE", "titan-24")
|
|
GPU_ACTIVE_SM_PERCENT = float(
|
|
os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVE_SM_PERCENT", "8")
|
|
)
|
|
GPU_MAX_EXTERNAL_MEMORY_BYTES = int(
|
|
os.environ.get("HERMES_LOCAL_IMAGE_GPU_MAX_EXTERNAL_MEMORY_BYTES", str(3 << 30))
|
|
)
|
|
GPU_ACTIVITY_SAMPLES = max(
|
|
1, int(os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_SAMPLES", "3"))
|
|
)
|
|
GPU_ACTIVITY_SAMPLE_INTERVAL = max(
|
|
0.0, float(os.environ.get("HERMES_LOCAL_IMAGE_GPU_ACTIVITY_SAMPLE_INTERVAL", "1"))
|
|
)
|
|
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,
|
|
"interactive_active": None,
|
|
"external_gpu_memory_bytes": None,
|
|
"external_gpu_sm_percent": None,
|
|
"gpu_guard_reason": None,
|
|
}
|
|
|
|
|
|
def _prometheus_labels(raw: str) -> dict[str, str]:
|
|
"""Parse the simple quoted labels emitted by nvidia-process-exporter."""
|
|
labels: dict[str, str] = {}
|
|
for item in raw.split(","):
|
|
key, separator, value = item.partition("=")
|
|
if not separator:
|
|
continue
|
|
labels[key.strip()] = value.strip().strip('"')
|
|
return labels
|
|
|
|
|
|
def _parse_gpu_activity(metrics: str) -> dict[str, Any]:
|
|
"""Summarize non-Hermes GPU processes on the configured image node."""
|
|
memory_bytes = 0
|
|
max_sm_percent = 0.0
|
|
processes: set[str] = set()
|
|
for line in metrics.splitlines():
|
|
if not line or line.startswith("#") or "{" not in line or "}" not in line:
|
|
continue
|
|
metric, remainder = line.split("{", 1)
|
|
raw_labels, separator, raw_value = remainder.partition("}")
|
|
if not separator or metric not in {
|
|
"nvidia_process_gpu_memory_used_bytes",
|
|
"nvidia_process_gpu_sm_util_percent",
|
|
}:
|
|
continue
|
|
labels = _prometheus_labels(raw_labels)
|
|
if labels.get("node") != GPU_ACTIVITY_NODE or labels.get("namespace") == "hermes":
|
|
continue
|
|
try:
|
|
value = float(raw_value.strip().split()[0])
|
|
except (ValueError, IndexError):
|
|
continue
|
|
process = labels.get("process") or "unknown"
|
|
processes.add(f"{labels.get('namespace', 'unknown')}/{process}")
|
|
if metric == "nvidia_process_gpu_memory_used_bytes":
|
|
memory_bytes += max(0, int(value))
|
|
else:
|
|
max_sm_percent = max(max_sm_percent, value)
|
|
active = (
|
|
max_sm_percent >= GPU_ACTIVE_SM_PERCENT
|
|
or memory_bytes >= GPU_MAX_EXTERNAL_MEMORY_BYTES
|
|
)
|
|
reasons = []
|
|
if max_sm_percent >= GPU_ACTIVE_SM_PERCENT:
|
|
reasons.append(f"external SM utilization is {max_sm_percent:g}%")
|
|
if memory_bytes >= GPU_MAX_EXTERNAL_MEMORY_BYTES:
|
|
reasons.append(f"external GPU memory is {memory_bytes} bytes")
|
|
return {
|
|
"interactive_active": active,
|
|
"external_gpu_memory_bytes": memory_bytes,
|
|
"external_gpu_sm_percent": max_sm_percent,
|
|
"external_gpu_processes": sorted(processes),
|
|
"gpu_guard_reason": "; ".join(reasons) or None,
|
|
}
|
|
|
|
|
|
def _gpu_activity_snapshot() -> dict[str, Any]:
|
|
"""Fetch live per-process GPU attribution, failing closed if unavailable."""
|
|
request = Request(GPU_ACTIVITY_URL, headers={"Accept": "text/plain"})
|
|
with urlopen(request, timeout=5) as response:
|
|
metrics = response.read(2 << 20).decode("utf-8", errors="replace")
|
|
return _parse_gpu_activity(metrics)
|
|
|
|
|
|
def _ensure_interactive_lane_idle() -> dict[str, Any]:
|
|
"""Reject FLUX while desktop or Wolf activity indicates interactive use."""
|
|
worst: dict[str, Any] | None = None
|
|
for sample in range(GPU_ACTIVITY_SAMPLES):
|
|
snapshot = _gpu_activity_snapshot()
|
|
if worst is None or (
|
|
snapshot["external_gpu_sm_percent"], snapshot["external_gpu_memory_bytes"]
|
|
) > (worst["external_gpu_sm_percent"], worst["external_gpu_memory_bytes"]):
|
|
worst = snapshot
|
|
if snapshot["interactive_active"]:
|
|
break
|
|
if sample + 1 < GPU_ACTIVITY_SAMPLES:
|
|
time.sleep(GPU_ACTIVITY_SAMPLE_INTERVAL)
|
|
assert worst is not None
|
|
_set_state(**{key: value for key, value in worst.items() if key != "external_gpu_processes"})
|
|
if worst["interactive_active"]:
|
|
raise RuntimeError(
|
|
"titan-24 desktop/Wolf lane is active; reserve or release it from the Atlas GPU checkout"
|
|
f" ({worst['gpu_guard_reason']})"
|
|
)
|
|
return worst
|
|
|
|
|
|
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_in_process(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Load and run FLUX inside the disposable CUDA worker process."""
|
|
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,
|
|
)
|
|
if OFFLOAD_MODE == "sequential":
|
|
# The shared 10 GiB card keeps the desktop and idle Wolf daemon
|
|
# resident. Layer-level offload preserves 1024px output quality
|
|
# without requiring those low-memory services to be torn down.
|
|
pipe.enable_sequential_cpu_offload()
|
|
elif OFFLOAD_MODE == "model":
|
|
pipe.enable_model_cpu_offload()
|
|
else:
|
|
raise ValueError("HERMES_LOCAL_IMAGE_OFFLOAD_MODE must be sequential or model")
|
|
_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 _render(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Render in a child process so process exit releases every CUDA allocation."""
|
|
_set_state(phase="render-worker", last_model=MODEL_ID)
|
|
completed = subprocess.run(
|
|
[sys.executable, str(Path(__file__).resolve()), "--render-worker"],
|
|
input=json.dumps(payload, separators=(",", ":")).encode("utf-8"),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
timeout=RENDER_TIMEOUT_SECONDS,
|
|
check=False,
|
|
)
|
|
if completed.returncode != 0:
|
|
detail = completed.stderr.decode("utf-8", errors="replace").strip()
|
|
if len(detail) > 4000:
|
|
detail = detail[-4000:]
|
|
raise RuntimeError(detail or f"FLUX worker exited {completed.returncode}")
|
|
try:
|
|
result = json.loads(completed.stdout)
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError("FLUX worker returned an invalid response") from exc
|
|
if not isinstance(result, dict) or not result.get("success"):
|
|
raise RuntimeError("FLUX worker returned an unsuccessful response")
|
|
return result
|
|
|
|
|
|
def _run_render_worker() -> int:
|
|
"""Read one bounded request and emit one JSON response for the parent."""
|
|
raw = sys.stdin.buffer.read(MAX_BODY_BYTES + 1)
|
|
if not raw or len(raw) > MAX_BODY_BYTES:
|
|
print("invalid render worker request size", file=sys.stderr)
|
|
return 2
|
|
try:
|
|
payload = json.loads(raw)
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("JSON object required")
|
|
# Keep library progress output off the machine-readable stdout channel.
|
|
with redirect_stdout(sys.stderr):
|
|
result = _render_in_process(payload)
|
|
sys.stdout.write(json.dumps(result, separators=(",", ":")))
|
|
sys.stdout.flush()
|
|
return 0
|
|
except Exception as exc:
|
|
print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
|
|
return 1
|
|
|
|
|
|
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:
|
|
if _lease_owner() != LEASE_IDLE_OWNER:
|
|
owner = _lease_owner()
|
|
raise RuntimeError(f"local image GPU unavailable while titan-24 owner is {owner}")
|
|
_set_state(phase="checking-interactive-lane")
|
|
_ensure_interactive_lane_idle()
|
|
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"
|
|
try:
|
|
activity = _gpu_activity_snapshot()
|
|
_set_state(**{key: value for key, value in activity.items() if key != "external_gpu_processes"})
|
|
state = _state_snapshot()
|
|
except Exception as exc:
|
|
activity = {"interactive_active": True, "gpu_guard_reason": f"GPU guard unavailable: {exc}"}
|
|
_set_state(**activity)
|
|
state = _state_snapshot()
|
|
self._json(
|
|
200,
|
|
{
|
|
"success": True,
|
|
"available": (
|
|
owner == LEASE_IDLE_OWNER
|
|
and not state["busy"]
|
|
and not activity["interactive_active"]
|
|
),
|
|
"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__":
|
|
if sys.argv[1:] == ["--render-worker"]:
|
|
raise SystemExit(_run_render_worker())
|
|
_recover_stale_image_lease()
|
|
ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()
|