Add dual-GPU FLUX 9B image pipeline
This commit is contained in:
@@ -25,6 +25,8 @@ ALLOWED = tuple(x.strip() for x in os.environ.get(
|
||||
LABEL_KEY = "com.mike-ai.llama-profile"
|
||||
IMAGE_LABEL_KEY = "com.mike-ai.image-worker"
|
||||
IMAGE_WORKER = os.environ.get("IMAGE_WORKER", "image")
|
||||
TTS_LABEL_KEY = "com.mike-ai.tts-worker"
|
||||
TTS_WORKER = os.environ.get("TTS_WORKER", "qwen3")
|
||||
LOCK = threading.Lock()
|
||||
log = logging.getLogger("profile-controller")
|
||||
|
||||
@@ -77,6 +79,15 @@ def image_container() -> dict:
|
||||
return matches[0]
|
||||
|
||||
|
||||
def tts_container() -> dict:
|
||||
matches = [item for item in labelled_containers(TTS_LABEL_KEY)
|
||||
if item.get("Labels", {}).get(TTS_LABEL_KEY) == TTS_WORKER]
|
||||
if len(matches) != 1:
|
||||
raise RuntimeError(
|
||||
f"expected exactly one TTS worker {TTS_WORKER!r}, found {len(matches)}")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def stop_container(item: dict, timeout: int = 120) -> None:
|
||||
if item.get("State") != "running":
|
||||
return
|
||||
@@ -85,6 +96,14 @@ def stop_container(item: dict, timeout: int = 120) -> None:
|
||||
raise RuntimeError(f"failed to stop container: HTTP {status}")
|
||||
|
||||
|
||||
def start_container(item: dict) -> None:
|
||||
if item.get("State") == "running":
|
||||
return
|
||||
status, _ = docker_request("POST", f"/containers/{item['Id']}/start")
|
||||
if status not in (204, 304):
|
||||
raise RuntimeError(f"failed to start container: HTTP {status}")
|
||||
|
||||
|
||||
def stop_inference() -> dict:
|
||||
with LOCK:
|
||||
items = containers()
|
||||
@@ -101,14 +120,17 @@ def set_image_worker(running: bool) -> dict:
|
||||
# The image worker may never overlap a llama profile on the 5080.
|
||||
for profile_item in containers().values():
|
||||
stop_container(profile_item)
|
||||
if item.get("State") != "running":
|
||||
status, _ = docker_request("POST", f"/containers/{item['Id']}/start")
|
||||
if status not in (204, 304):
|
||||
raise RuntimeError(f"failed to start image worker: HTTP {status}")
|
||||
# The 9B beta text encoder temporarily borrows the RTX 3060 from
|
||||
# Qwen3-TTS. The gateway retains Piper as a fallback meanwhile.
|
||||
stop_container(tts_container(), timeout=30)
|
||||
start_container(item)
|
||||
else:
|
||||
# CUDA/PyTorch may not react promptly to SIGTERM after an OOM.
|
||||
# Bound recovery time and let Docker issue SIGKILL afterwards.
|
||||
stop_container(item, timeout=20)
|
||||
# TTS is restored by the following profile activation. Keeping it
|
||||
# stopped here lets the router verify that both GPUs really
|
||||
# released the image model before Qwen and TTS are reloaded.
|
||||
return {"image_worker": "running" if running else "stopped"}
|
||||
|
||||
|
||||
@@ -126,6 +148,7 @@ def activate(profile: str) -> dict:
|
||||
with LOCK:
|
||||
# Defensive mutual exclusion even if a caller bypasses the router.
|
||||
stop_container(image_container())
|
||||
start_container(tts_container())
|
||||
items = containers()
|
||||
missing = [name for name in ALLOWED if name not in items]
|
||||
if missing:
|
||||
|
||||
Reference in New Issue
Block a user