Studio: stop leaking child processes on an abnormal exit

studio/backend/cloudflare_tunnel.py

@@ -91,6 +91,24 @@ def _lifetime_kwargs() -> dict:
         return {}
 
 
+def _adopt_pid(pid: int) -> None:
+    """Record cloudflared so a force quit does not strand it (macOS has no
+    PDEATHSIG). Best-effort, like _lifetime_kwargs above."""
+    try:
+        from utils.process_lifetime import adopt_pid
+        adopt_pid(pid)
+    except Exception:
+        pass
+
+
+def _forget_pid(pid: int) -> None:
+    try:
+        from utils.process_lifetime import forget_pid
+        forget_pid(pid)
+    except Exception:
+        pass
+
+
 def _spawn_child(spawn):
     """Fork on a process-lifetime thread so the PDEATHSIG above means "die with
     the parent process", not "die when the worker thread that forked me returns"."""
@@ -447,6 +465,10 @@ def start(self) -> None:
             except Exception:
                 _set_studio_tunnel_runtime_active(self, False)
                 raise
+            # Adopted before the lock drops: a stop() that got in first would
+            # otherwise reap and forget it while nothing was tracked, and this
+            # would then record whatever inherited the pid.
+            _adopt_pid(proc.pid)
             self._proc = proc
         threading.Thread(
             target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True
@@ -522,6 +544,7 @@ def stop(self) -> bool:
         except Exception:
             pass
         if _process_exited(proc):
+            _forget_pid(proc.pid)
             _set_studio_tunnel_runtime_active(self, False)
             return True
         else:

studio/backend/core/inference/llama_cpp.py

@@ -7290,9 +7290,21 @@ def _start_diffusion_server(
             encoding = "utf-8",
             errors = "replace",
             env = utf8_child_env(env),
+            # Deliberately NOT start_new_session, as with the component
+            # installer: the desktop stops this backend by signalling its
+            # process group and force-kills it after five seconds, so a session
+            # of its own would leave the shim and the visual server holding the
+            # GPU until the next launch sweeps them.
             **_windows_hidden_subprocess_kwargs(),
             **_child_popen_kwargs(),
         )
+        # macOS has no parent-death signal, so the kwargs above are empty there and
+        # only this record lets the next startup reap a runner holding the GPU.
+        try:
+            from utils.process_lifetime import adopt_pid
+            adopt_pid(self._process.pid)
+        except Exception as e:
+            logger.debug(f"Could not track diffusion runner for lifetime sweep: {e}")
         self._stdout_thread = threading.Thread(
             target = self._drain_stdout, daemon = True, name = "diffusion-stdout"
         )
@@ -13284,6 +13296,47 @@ def unload_model(self) -> bool:
                     torch.cuda.empty_cache()
             return True
 
+    @staticmethod
+    def _leading_process_group(pid):
+        """The pid's own process group, when it leads one. None otherwise."""
+        if not pid or os.name != "posix" or not hasattr(os, "getpgid"):
+            return None
+        try:
+            pgid = os.getpgid(pid)
+        except OSError:
+            return None
+        return pgid if pgid == pid else None
+
+    @staticmethod
+    def _kill_process_group(pgid):
+        """Take down what the leader left behind, if anything is still there."""
+        if pgid is None or not hasattr(os, "killpg"):
+            return
+        try:
+            os.killpg(pgid, signal.SIGKILL)
+        except OSError:
+            pass
+
+    @staticmethod
+    def _collect_descendants(pid):
+        """The server's own children, for the kill below. Empty when unreadable."""
+        try:
+            from utils.process_lifetime import collect_descendants
+            return collect_descendants(pid)
+        except Exception:
+            return []
+
+    @staticmethod
+    def _terminate_descendants(collected):
+        """The diffusion shim's visual server, and anything else it started."""
+        if not collected:
+            return
+        try:
+            from utils.process_lifetime import terminate_descendants
+            terminate_descendants(collected, timeout = 5.0)
+        except Exception as e:
+            logger.debug(f"Could not terminate server descendants: {e}")
+
     def _kill_process(self):
         """Terminate the subprocess if running."""
         # Stop the watchdog before a deliberate kill so a planned reload/unload
@@ -13303,6 +13356,13 @@ def _kill_process(self):
         terminable = hasattr(self._process, "terminate")
         if not terminable:
             logger.debug("no terminable llama-server process to kill; clearing state")
+        # Both read before the terminate below: getpgid stops answering once the
+        # wait reaps the leader, and the shim's children are reparented the
+        # moment it exits. This is the only stop the desktop shutdown waits for,
+        # so the visual server has to be named while that link still exists.
+        _pid = getattr(self._process, "pid", None)
+        _pgid = self._leading_process_group(_pid)
+        _descendants = self._collect_descendants(_pid)
         try:
             if terminable:
                 self._process.terminate()
@@ -13323,11 +13383,25 @@ def _kill_process(self):
         except Exception as e:
             logger.warning(f"Error killing llama-server process: {e}")
         finally:
+            self._kill_process_group(_pgid)
+            self._terminate_descendants(_descendants)
             # getattr: teardown must tolerate a partially-built backend (failed
             # __init__ or a __new__-built instance), as with _llama_log_fh below.
             if getattr(self, "_stats_logger", None) is not None:
                 self._stats_logger.stop()
                 self._stats_logger = None
+            # Drop it from the lifetime record only once it is confirmed gone.
+            # A server that survived a failed kill has to stay recorded, or the
+            # next startup sweep cannot reap it. The record stores a start-time
+            # identity, so a recycled pid is never signalled either way.
+            _killed_pid = getattr(self._process, "pid", None)
+            _exited = getattr(self._process, "poll", lambda: None)() is not None
+            if _killed_pid is not None and _exited:
+                try:
+                    from utils.process_lifetime import forget_pid
+                    forget_pid(_killed_pid)
+                except Exception:
+                    pass
             self._process = None
             self._clear_server_pid()
             # Clear healthy so a /load during the replacement's warm-up can't
@@ -13371,6 +13445,14 @@ def _record_server_pid(cls, pid: int) -> None:
         since been recycled to a different process (see ``_pid_start_identity``).
         A bare ``pid`` (no identity) is still accepted on read for compatibility.
         """
+        # Track it generically too: the pidfile holds one server, while the
+        # process-lifetime record covers every child and is what the startup
+        # sweep reads where there is no parent-death signal (macOS).
+        try:
+            from utils.process_lifetime import adopt_pid
+            adopt_pid(pid)
+        except Exception as e:
+            logger.debug(f"Could not track llama-server for lifetime sweep: {e}")
         path = cls._server_pidfile_path()
         if path is None:
             return

studio/backend/core/inference/sd_cpp_engine.py

@@ -25,7 +25,7 @@
 from pathlib import Path
 from typing import Callable, Optional
 
-from utils.process_lifetime import child_popen_kwargs
+from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
 from utils.native_path_leases import child_env_without_native_path_secret
 from core.inference.sd_cpp_args import (
     SdCppGenParams,
@@ -413,6 +413,9 @@ def _run(
             # Bind the child to the parent's lifetime (PR_SET_PDEATHSIG) so a parent crash cannot orphan sd-cli holding VRAM/RAM.
             **child_popen_kwargs(),
         )
+        # The kwargs above are empty on macOS, so record it too: a crash mid-generation
+        # would otherwise leave sd-cli holding VRAM with nothing able to find it.
+        adopt_pid(proc.pid)
         # Drain stdout on a reader thread so the timeout holds even when the child hangs WITHOUT printing (a plain `for line in proc.stdout` blocks until EOF). Lines, then a None sentinel, go to a queue the main loop polls against a wall-clock deadline.
         tail: list[str] = []
         line_q: "queue.Queue[Optional[str]]" = queue.Queue()
@@ -459,6 +462,10 @@ def _drain() -> None:
         finally:
             if proc.poll() is None:
                 _terminate(proc)
+            # Only once it has actually exited: a pid still running has to stay
+            # recorded, or the next startup has no handle on it.
+            if proc.poll() is not None:
+                forget_pid(proc.pid)
 
         if ret != 0:
             raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:]))