diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index d2e0f6f3df..1ffd2054fe 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -148,6 +148,29 @@ def metal_text_encoder_flags() -> list[str]: return ["--clip-on-cpu"] +# Everything on the CPU backend. sd.cpp's default preference is GPU -> integrated GPU -> CPU, and +# only `--backend` changes which backend EXECUTES the graph (`--offload-to-cpu` moves parameters, +# not compute), so this is the one flag that takes ggml-metal out of the picture entirely. +CPU_BACKEND_FLAGS: tuple[str, ...] = ("--backend", "cpu") + +# The ggml signature that means "this graph cannot run on this backend at all": ggml-metal checks +# ggml_metal_device_supports_op() per node in ggml_metal_op_encode_impl() and calls GGML_ABORT when +# it returns false, because a single-backend graph has nowhere to put the node. SIGABRT (-6 on +# POSIX, 134 through a shell) takes the whole sd-server down mid-generation. +_GGML_UNSUPPORTED_OP_MARKERS = ("unsupported op", "ggml_abort") + + +def is_ggml_unsupported_op_abort(message: str) -> bool: + """True if ``message`` is a captured sd.cpp log tail carrying a ggml unsupported-op abort. + + Used to decide whether a dead sd-server is worth restarting on the CPU backend: an abort with + this signature is deterministic for the graph in question, so a plain retry on the same backend + would fail identically, while a CPU restart runs it. Any other death (OOM kill, a real bug, + a corrupt checkpoint) must NOT be silently retried.""" + text = (message or "").lower() + return all(marker in text for marker in _GGML_UNSUPPORTED_OP_MARKERS) + + def offload_flags( policy: str, *, diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index b4a97a3fba..ffb562cfd9 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -29,7 +29,7 @@ import os import re import threading import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Optional @@ -52,9 +52,11 @@ from core.inference.diffusion_memory import ( OFFLOAD_SEQUENTIAL, ) from core.inference.sd_cpp_args import ( + CPU_BACKEND_FLAGS, SdCppGenParams, SdCppModelFiles, build_img_gen_request, + is_ggml_unsupported_op_abort, offload_flags, ) from core.inference.sd_cpp_engine import ( @@ -307,6 +309,10 @@ class SdCppDiffusionBackend: # superseding load can stop it mid-startup instead of waiting out the timeout. self._pending_server: Optional[SdCppServer] = None self._gen: Optional[_SdGen] = None + # Set once this load's graph proved unrunnable on the GPU backend (a ggml unsupported-op + # abort), so the CPU restart happens once per load and every later generation goes straight + # to the backend that works. Cleared by each load. + self._cpu_backend_forced = False @property def is_loaded(self) -> bool: @@ -516,6 +522,9 @@ class SdCppDiffusionBackend: self._state = None # the old model is being torn down if old_state is not None and old_state.server is not None: old_state.server.stop() + # A new checkpoint earns a fresh attempt on the GPU backend: the previous one's + # abort says nothing about this graph. + self._cpu_backend_forced = False server: Optional[SdCppServer] = None if mode == "server": assert server_binary is not None @@ -951,12 +960,31 @@ class SdCppDiffusionBackend: distilled_guidance = flux_guidance, lora = lora_payload, ) - blobs = state.server.img_gen( - payload, - on_step = self._on_log, - cancel_event = cancel, - total_timeout = max(deadline - time.monotonic(), 1.0), - ) + try: + blobs = state.server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = max(deadline - time.monotonic(), 1.0), + ) + except RuntimeError as exc: + # A ggml unsupported-op abort killed the server: this graph cannot run on the + # GPU backend at all, so re-running it there would abort identically. Restart + # the model on the CPU backend once and retry this chunk. Any other death + # propagates unchanged. + server = self._restart_server_on_cpu_backend(state, str(exc), cancel) + if server is None: + raise + state = replace(state, server = server) + with self._lock: + if self._state is not None and self._state.server is not None: + self._state = state + blobs = server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = max(deadline - time.monotonic(), 1.0), + ) # All-or-nothing per chunk: fail rather than silently drop images from the batch. if not cancel.is_set() and len(blobs) != count: raise RuntimeError( @@ -970,6 +998,67 @@ class SdCppDiffusionBackend: shutil.rmtree(lora_stage, ignore_errors = True) return images, seeds + def _restart_server_on_cpu_backend( + self, + state: _SdState, + error_text: str, + cancel: threading.Event, + ) -> Optional[SdCppServer]: + """Relaunch this checkpoint's sd-server with ``--backend cpu``; None if that does not apply. + + ggml's Metal backend checks every node against ``ggml_metal_device_supports_op`` and calls + ``GGML_ABORT`` when one is not implemented for that device, because a single-backend graph + has nowhere else to put the node -- there is no per-op CPU fallback. The whole sd-server + dies with SIGABRT mid-generation, so the user sees "the native image renderer stopped + unexpectedly" with no way forward. Observed on macos-14 arm64 with FLUX.2-klein-4B Q2_K: + the encoder is already pinned to CPU (see metal_text_encoder_flags), and the abort then + moves into the denoise loop: + + ggml_metal_op_encode_impl: error: unsupported op 'MUL_MAT' -> ggml_abort + StableDiffusionGGML::sample -> sample_k_diffusion + + ``--backend cpu`` is the only flag that changes which backend EXECUTES the graph + (``--offload-to-cpu`` moves parameters, not compute), so the restart runs the same + checkpoint slower rather than not at all. Done once per load: the second abort, or any + other cause of death, is surfaced to the caller.""" + if not is_ggml_unsupported_op_abort(error_text): + return None + if self._cpu_backend_forced or state.device == "cpu": + return None # already on CPU: the abort is not a backend-placement problem + if state.server is None or cancel.is_set(): + return None + server_binary = find_sd_server_binary() + if not server_binary: + return None + logger.warning( + "sd-server aborted on an op the '%s' backend cannot run; restarting on the CPU " + "backend (slower, but it completes). Details: %s", + state.device, + error_text[:300], + ) + self._cpu_backend_forced = True + state.server.stop() + server = SdCppServer(server_binary) + with self._lock: + self._pending_server = server + try: + server.start( + state.files, + vae_format = state.vae_format, + offload = list(state.offload_flags), + native_speed = state.native_speed, + threads = state.threads, + extra_args = list(CPU_BACKEND_FLAGS), + ) + except Exception: # noqa: BLE001 -- the original abort is the more useful error + server.stop() + return None + finally: + with self._lock: + if self._pending_server is server: + self._pending_server = None + return server + def _generate_oneshot( self, state: _SdState, diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index c9148eae4e..25a1c4f5cd 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -175,12 +175,14 @@ class SdCppServer: threads: Optional[int] = None, env: Optional[dict[str, str]] = None, startup_timeout: float = 600.0, + extra_args: Optional[list[str]] = None, ) -> None: """Spawn the server (which loads the model) and block until it is ready. Raises ``RuntimeError`` (with the captured log tail) if the process exits during startup or never answers within ``startup_timeout``. Holds the lifecycle lock so - a concurrent start/stop can't interleave. + a concurrent start/stop can't interleave. ``extra_args`` is appended last (last + wins), which is how the CPU-backend restart pins the graph off the GPU. """ with self._lifecycle_lock: # A stop()/unload that raced in before start() took the lock already set _abort and closed the @@ -202,6 +204,7 @@ class SdCppServer: threads = threads, scratch_dir = self._scratch_dir, verbose = True, # sd-server prints the per-step sampling lines we parse + extra_args = list(extra_args or []), ) run_env = runtime_env(self.binary, child_env_without_native_path_secret()) if env: diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 7f168ddcd9..f93051edc9 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -18,6 +18,7 @@ from core.inference.diffusion_memory import ( OFFLOAD_SEQUENTIAL, ) from core.inference.sd_cpp_args import ( + CPU_BACKEND_FLAGS, SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams, @@ -25,6 +26,7 @@ from core.inference.sd_cpp_args import ( build_sd_cpp_command, build_sd_cpp_server_command, build_sd_cpp_upscale_command, + is_ggml_unsupported_op_abort, metal_text_encoder_flags, native_speed_flags, offload_flags, @@ -526,3 +528,39 @@ def test_img_gen_request_flux_uses_distilled_guidance(): def test_img_gen_request_requires_prompt(): with pytest.raises(ValueError): build_img_gen_request(prompt = " ", steps = 4) + + +def test_ggml_unsupported_op_abort_is_recognised_only_with_both_markers(): + # The CPU-backend rescue must fire for the deterministic "this backend cannot run this graph" + # abort and for nothing else: an OOM kill or a plain crash has to surface as itself. + abort = ( + "sd-server connection lost during img_gen poll (process exited, code -6)\n" + "Last output:\n" + "[ERROR] ggml_extend.hpp:70 - ggml_metal_op_encode_impl: error: unsupported op 'MUL_MAT'\n" + "1 sd-server 0x00000001044f8df4 ggml_abort + 156" + ) + assert is_ggml_unsupported_op_abort(abort) is True + # The RMS_NORM shape of the same abort (text encoder) counts too. + assert is_ggml_unsupported_op_abort( + "error: unsupported op 'RMS_NORM'\nggml_abort + 156" + ) is True + # Neither marker alone is enough, and an unrelated death is never a match. + assert is_ggml_unsupported_op_abort("unsupported op 'MUL_MAT'") is False + assert is_ggml_unsupported_op_abort("ggml_abort + 156") is False + assert is_ggml_unsupported_op_abort("process exited, code -9") is False + assert is_ggml_unsupported_op_abort("") is False + + +def test_server_command_appends_extra_args_last(): + # --backend cpu is passed as extra_args by the abort rescue, and sd.cpp is last-wins, so it has + # to land after every flag the normal build emits. + cmd = build_sd_cpp_server_command( + "/x/sd-server", + SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft"), + host = "127.0.0.1", + port = 1234, + native_speed = "default", + extra_args = list(CPU_BACKEND_FLAGS), + ) + assert cmd[-2:] == ["--backend", "cpu"] + assert cmd[0] == "/x/sd-server" diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 2ffba5bb6e..5fa14c1e0d 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -133,6 +133,8 @@ class _FakeServer: self.timeouts = [] self.alive = True self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir + # Raised by the next img_gen (one shot) so a test can stage a mid-generation server death. + self.img_gen_error = None def is_alive(self): return self.alive and not self.stopped @@ -145,6 +147,7 @@ class _FakeServer: offload = None, native_speed = None, threads = None, + extra_args = None, ): self.started = dict( files = files, @@ -152,6 +155,7 @@ class _FakeServer: offload = offload, native_speed = native_speed, threads = threads, + extra_args = list(extra_args or []), ) def img_gen( @@ -166,6 +170,10 @@ class _FakeServer: self.payloads.append(payload) self.timeouts.append(total_timeout) + if self.img_gen_error is not None: + err, self.img_gen_error = self.img_gen_error, None + self.alive = False # a ggml abort takes the process down + raise err if on_step is not None: steps = payload.get("sample_params", {}).get("sample_steps", 0) on_step(f" {steps}/{steps}") @@ -509,6 +517,7 @@ def _run_server_load( b, servers, fam_name = "z-image", + device = "cpu", ): fam = detect_family(fam_name) monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") @@ -529,7 +538,7 @@ def _run_server_load( lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, ) monkeypatch.setattr( - bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") + bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = device) ) b._load_token = 1 b._run_load( @@ -567,6 +576,79 @@ def test_server_generate_uses_one_request_for_whole_batch(monkeypatch): assert b._gen is None # cleared after generate +_GGML_ABORT = ( + "sd-server connection lost during img_gen poll (process exited, code -6)\n" + "Last output:\n" + "[ERROR] ggml_extend.hpp:70 - ggml_metal_op_encode_impl: error: unsupported op 'MUL_MAT'\n" + "1 sd-server 0x00000001044f8df4 ggml_abort + 156\n" + "10 sd-server 0x00000001043f6ce8 StableDiffusionGGML::sample" +) + + +def test_server_generation_restarts_on_the_cpu_backend_after_a_ggml_abort(monkeypatch): + # ggml checks every node against the device's supports_op and calls GGML_ABORT when one is not + # implemented, killing sd-server mid-generation with no per-op CPU fallback. Retrying on the + # same backend would abort identically, so the load is restarted with --backend cpu and the + # generation completes (slower) instead of failing. + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers, device = "mps") + servers[0].img_gen_error = RuntimeError(_GGML_ABORT) + + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 4, seed = 3) + + assert len(out["images"]) == 1 # the retry produced the image + assert len(servers) == 2 and servers[0].stopped is True + assert servers[1].started["extra_args"] == ["--backend", "cpu"] + # The same checkpoint and run settings are reused; only the backend placement changed. + assert servers[1].started["files"] is servers[0].started["files"] + assert servers[1].started["native_speed"] == servers[0].started["native_speed"] + # The live state points at the replacement, so the next generation does not touch the dead one. + assert b._state is not None and b._state.server is servers[1] + + +def test_cpu_backend_restart_happens_once_per_load(monkeypatch): + # The restart is a one-shot rescue: if the CPU backend aborts too, the error surfaces rather + # than spawning servers forever. + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers, device = "mps") + servers[0].img_gen_error = RuntimeError(_GGML_ABORT) + b.generate(prompt = "x", steps = 4, seed = 1) + assert len(servers) == 2 + + servers[1].img_gen_error = RuntimeError(_GGML_ABORT) + with pytest.raises(RuntimeError, match = "unsupported op"): + b.generate(prompt = "x", steps = 4, seed = 1) + assert len(servers) == 2 # no third spawn + + +def test_server_death_without_the_abort_signature_is_not_retried(monkeypatch): + # An OOM kill, a corrupt checkpoint or a genuine bug must not be silently retried on another + # backend: only the deterministic unsupported-op abort earns the CPU restart. + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers, device = "mps") + servers[0].img_gen_error = RuntimeError( + "sd-server connection lost during img_gen poll (process exited, code -9)" + ) + with pytest.raises(RuntimeError, match = "code -9"): + b.generate(prompt = "x", steps = 4, seed = 1) + assert len(servers) == 1 + + +def test_cpu_device_does_not_restart_on_an_abort(monkeypatch): + # Already on CPU: the abort is not a backend-placement problem, so restarting would just repeat + # it. Surface the error instead. + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers, device = "cpu") + servers[0].img_gen_error = RuntimeError(_GGML_ABORT) + with pytest.raises(RuntimeError, match = "unsupported op"): + b.generate(prompt = "x", steps = 4, seed = 1) + assert len(servers) == 1 + + def test_server_generate_splits_batches_above_server_limit(monkeypatch): # A batch above the server's per-job limit is chunked (the one-shot path did these image-by-image); # each chunk gets a timeout proportional to its image count.