diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index b3b921419b..16b3207256 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -27,6 +27,7 @@ import os import shutil import subprocess import sys +import threading import time from pathlib import Path from typing import Callable, Optional @@ -232,8 +233,14 @@ class SdCppEngine: errors = "replace", env = run_env, ) + # Drain stdout on a background thread and wait on the PROCESS, not the stream: + # iterating proc.stdout directly blocks until the stream closes, so a sd-cli that + # hangs without producing output (or closing stdout) would never reach proc.wait + # and the timeout would be silently bypassed. With the reader on its own thread the + # main thread always enforces the wall-clock timeout and kills a hung process. tail: list[str] = [] - try: + + def _drain() -> None: assert proc.stdout is not None for line in proc.stdout: line = line.rstrip("\n") @@ -242,13 +249,20 @@ class SdCppEngine: tail.pop(0) if on_log is not None: on_log(line) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + try: ret = proc.wait(timeout = timeout) except subprocess.TimeoutExpired: proc.kill() + reader.join(timeout = 5.0) raise RuntimeError(f"sd-cli timed out after {timeout}s") finally: if proc.poll() is None: proc.kill() + # The process has exited; let the reader finish draining the buffered output. + reader.join(timeout = 5.0) if ret != 0: raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 25dd19cec6..ac63596a97 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -77,7 +77,10 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): # ── availability / version ────────────────────────────────────────────────── -def test_engine_unavailable_when_no_binary(): +def test_engine_unavailable_when_no_binary(monkeypatch): + # Hermetic: force discovery to find nothing so a real sd-cli installed on the host + # (e.g. ~/.unsloth) can't leak in and make binary=None resolve to a real binary. + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) assert e.is_available() is False assert e.version() is None @@ -242,6 +245,62 @@ def test_generate_raises_when_binary_missing(): ) +def test_generate_times_out_even_when_stdout_blocks(tmp_path, monkeypatch): + # A sd-cli that hangs WITHOUT closing stdout must still hit the wall-clock timeout: + # the reader drains on a thread while the main thread waits on the PROCESS, so the + # timeout can no longer be bypassed by an unending stdout stream. + import subprocess as _sp + import threading as _threading + + released = _threading.Event() + + class _Block: + def __iter__(self): + return self + + def __next__(self): + # Models a hung stream that only ends once the process is killed. + if not released.wait(5.0): + raise AssertionError("stdout was never released by kill()") + raise StopIteration + + class _HangingPopen: + def __init__(self, cmd, **kw): + self.killed = False + + @property + def stdout(self): + return _Block() + + def wait(self, timeout = None): + raise _sp.TimeoutExpired(cmd = "sd-cli", timeout = timeout) + + def poll(self): + return 0 if self.killed else None + + def kill(self): + self.killed = True + released.set() # killing closes the pipe, so the reader unblocks + + holder: dict = {} + + def _factory(cmd, **kw): + holder["proc"] = _HangingPopen(cmd, **kw) + return holder["proc"] + + monkeypatch.setattr(eng.subprocess, "Popen", _factory) + + e = _engine(tmp_path) + with pytest.raises(RuntimeError, match = "timed out"): + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(tmp_path / "o.png"), + timeout = 0.01, + ) + assert holder["proc"].killed is True + + # ── engine routing ──────────────────────────────────────────────────────────