Studio: fix Ctrl+C shutdown ordering (installer shell + uvicorn thread wait) (#6566)

* Installer: respect a declined Studio auto-start and keep Ctrl+C shutdown logs ordered

The `curl | sh` Studio auto-start prompt had two issues on Linux/macOS/WSL
(install.sh). install.ps1 already gates on input redirection, so Windows is
unaffected.

1. Typing n, or any closed/EOF /dev/tty, still launched Studio. The read
   fallbacks defaulted to "y" (read failure, and the no-tty branch), so any
   answer other than a cleanly delivered y/n line auto-started a blocking
   foreground server. Default those to "n"; a real Enter still counts as yes
   via ${_reply:-y}.

2. On Ctrl+C the shell prompt printed in the middle of Studio's shutdown logs.
   The non-interactive installer shell took the default SIGINT action and died
   before the child finished its graceful shutdown, so the prompt raced ahead
   of "All subprocesses cleaned up". trap '' INT in the installer shell so it
   waits for Studio's own graceful shutdown.

* Studio: wait for the uvicorn thread before the terminal returns on Ctrl+C

Builds on #6565 by @Imagineer99. The studio server runs uvicorn in a daemon
thread, so on Ctrl+C the process could return to the shell while that thread
was still writing its shutdown logs, interleaving them with the prompt.

Retain the uvicorn thread and join it (flushing stdout/stderr) before terminal
entrypoints return, from run.py's main shutdown path and the CLI shutdown paths.

Refinements over #6565:
- Bound the join at 5s (_SERVER_SHUTDOWN_JOIN_TIMEOUT, matching the existing
  _graceful_shutdown subprocess timeouts) so a stalled uvicorn shutdown cannot
  hang the terminal; the timeout warning branch is now reachable.
- Restore SIG_DFL for SIGINT/SIGTERM at the start of the signal handler so a
  second Ctrl+C force-quits, and drop the redundant in-handler wait (the
  post-loop wait already covers the signal path).

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: keep child Ctrl+C working and restore SIGBREAK

- install.sh: run studio in a subshell that resets INT to default
  (trap - INT; exec ...) so the foreground child does not inherit the
  installer shell's ignored SIGINT, which would otherwise swallow the
  studio process's own Ctrl+C and graceful shutdown.
- run.py: also restore SIGBREAK to SIG_DFL in the signal handler so a
  second Ctrl+Break force-quits on Windows, matching SIGINT/SIGTERM.

* install.sh: capture studio exit with || under set -e so the migration hint still prints

* Trim shutdown-fix comments to be terser (comments only, no code change)

* Dedup CLI shutdown-wait into finally blocks (review follow-up)

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-22 07:41:10 -07:00 committed by GitHub
commit dbc13f02c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 183 additions and 5 deletions

View file

@ -3126,10 +3126,11 @@ echo ""
if [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
if [ -r /dev/tty ]; then
read -r _reply </dev/tty || _reply="y"
read -r _reply </dev/tty || _reply="n"
else
_reply="y"
_reply="n"
fi
case "${_reply:-y}" in
[Yy]*|"")
@ -3137,8 +3138,12 @@ if [ -t 1 ]; then
# Detach stdin from the `curl | sh` pipe: as a foreground server the
# studio would otherwise drain the rest of this piped script, leaving
# the shell to die parsing the now-truncated tail (`unexpected fi`).
"$VENV_DIR/bin/unsloth" studio -p 8888 </dev/null
_LAUNCH_EXIT=$?
# trap '' INT: wait for studio's shutdown instead of racing the prompt.
# Subshell resets INT so the child still gets Ctrl+C (no inherited ignore).
trap '' INT
# `|| ...`: capture the exit code without set -e aborting first.
_LAUNCH_EXIT=0
(trap - INT; exec "$VENV_DIR/bin/unsloth" studio -p 8888 </dev/null) || _LAUNCH_EXIT=$?
if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
echo ""
echo "⚠️ Unsloth Studio failed to start after migration."

View file

@ -677,9 +677,37 @@ def _graceful_shutdown(server = None):
logger.info("All subprocesses cleaned up")
# Bound the join so a stuck uvicorn shutdown cannot hang the terminal.
_SERVER_SHUTDOWN_JOIN_TIMEOUT = 5.0
def _flush_standard_streams() -> None:
for stream in (sys.stdout, sys.stderr):
try:
stream.flush()
except Exception:
pass
def _wait_for_server_shutdown(timeout: Optional[float] = _SERVER_SHUTDOWN_JOIN_TIMEOUT) -> None:
"""Join the uvicorn thread so the prompt returns only after its shutdown logs
flush. Skip the self-join when called from the server thread."""
import threading
thread = _server_thread
if thread is None or thread is threading.current_thread():
_flush_standard_streams()
return
thread.join(timeout = timeout)
if thread.is_alive():
logger.warning("Timed out waiting for uvicorn server thread to stop")
_flush_standard_streams()
# The uvicorn server instance -- set by run_server(), used by callers
# that tell the server to exit (e.g. signal handlers).
_server = None
_server_thread = None
# Shutdown event -- wakes the main loop on signal.
_shutdown_event = None
@ -909,7 +937,7 @@ def run_server(
Signal handlers are NOT registered here so embedders (e.g. Colab) keep
their own interrupt semantics; standalone callers register them after.
"""
global _server, _shutdown_event
global _server, _server_thread, _shutdown_event
# Reap every child if the parent dies abnormally (terminal close, Task
# Manager kill, SIGKILL); must run before any child can spawn.
@ -1102,6 +1130,7 @@ def run_server(
startup_failed.set()
thread = Thread(target = _run, daemon = True)
_server_thread = thread
thread.start()
# Wait until uvicorn finishes lifespan startup and binds sockets, or until it
@ -1310,6 +1339,11 @@ if __name__ == "__main__":
# Signal handler -- ensures subprocess cleanup on Ctrl+C.
def _signal_handler(signum, frame):
# Restore defaults so a second signal force-quits if shutdown stalls.
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.signal(signal.SIGTERM, signal.SIG_DFL)
if hasattr(signal, "SIGBREAK"):
signal.signal(signal.SIGBREAK, signal.SIG_DFL)
_graceful_shutdown(_server)
_shutdown_event.set()
@ -1325,3 +1359,4 @@ if __name__ == "__main__":
# lets the interpreter process pending signals.
while not _shutdown_event.is_set():
_shutdown_event.wait(timeout = 1)
_wait_for_server_shutdown()

View file

@ -0,0 +1,133 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Source-level regression tests for terminal shutdown ordering (no backend import)."""
import ast
from pathlib import Path
_ROOT = Path(__file__).resolve().parent.parent
_RUN_PY = _ROOT / "studio" / "backend" / "run.py"
_STUDIO_CLI_PY = _ROOT / "unsloth_cli" / "commands" / "studio.py"
def _parse(path: Path) -> ast.Module:
return ast.parse(path.read_text(encoding = "utf-8"))
def _function(tree: ast.AST, name: str) -> ast.FunctionDef:
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == name:
return node
raise AssertionError(f"missing function {name}")
def _calls_name(tree: ast.AST, name: str) -> int:
return sum(
1
for node in ast.walk(tree)
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name
)
def _calls_shutdown_wait_getattr(tree: ast.AST) -> int:
count = 0
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if not isinstance(func, ast.Call):
continue
if not (isinstance(func.func, ast.Name) and func.func.id == "getattr"):
continue
if len(func.args) < 2:
continue
target, attr = func.args[:2]
if (
isinstance(target, ast.Name)
and target.id == "run_mod"
and isinstance(attr, ast.Constant)
and attr.value == "_wait_for_server_shutdown"
):
count += 1
return count
def test_run_server_records_uvicorn_thread_for_terminal_shutdown_wait():
tree = _parse(_RUN_PY)
run_server = _function(tree, "run_server")
assigns_thread_global = any(
isinstance(node, ast.Assign)
and any(
isinstance(target, ast.Name) and target.id == "_server_thread"
for target in node.targets
)
for node in ast.walk(run_server)
)
assert (
assigns_thread_global
), "run_server must retain the uvicorn thread so terminal shutdown can join it"
def test_wait_for_server_shutdown_joins_uvicorn_thread():
tree = _parse(_RUN_PY)
wait_func = _function(tree, "_wait_for_server_shutdown")
joins_thread = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "join"
for node in ast.walk(wait_func)
)
assert (
joins_thread
), "_wait_for_server_shutdown must join the uvicorn thread before process exit"
def test_wait_for_server_shutdown_join_is_bounded():
tree = _parse(_RUN_PY)
wait_func = _function(tree, "_wait_for_server_shutdown")
bounded_join = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "join"
and any(kw.arg == "timeout" for kw in node.keywords)
for node in ast.walk(wait_func)
)
assert (
bounded_join
), "the join must pass a timeout so a stalled uvicorn shutdown cannot hang the terminal"
def test_direct_backend_entrypoint_waits_before_returning_to_shell():
tree = _parse(_RUN_PY)
assert (
_calls_name(tree, "_wait_for_server_shutdown") >= 1
), "run.py must wait after the main shutdown event loop before returning to the shell"
def test_signal_handler_restores_default_handlers_for_force_quit():
tree = _parse(_RUN_PY)
handler = _function(tree, "_signal_handler")
restores_default = any(
isinstance(node, ast.Attribute) and node.attr == "SIG_DFL" for node in ast.walk(handler)
)
assert (
restores_default
), "the signal handler must restore SIG_DFL so a second Ctrl+C can force-quit"
def test_cli_entrypoints_wait_before_returning_to_shell():
tree = _parse(_STUDIO_CLI_PY)
assert (
_calls_shutdown_wait_getattr(tree) >= 3
), "Studio CLI terminal paths must wait for the backend thread after requesting shutdown"

View file

@ -895,6 +895,8 @@ def studio_default(
except KeyboardInterrupt:
run_mod._graceful_shutdown(run_mod._server)
typer.echo("\nShutting down...")
finally:
getattr(run_mod, "_wait_for_server_shutdown", lambda: None)()
# ── unsloth studio run ───────────────────────────────────────────────
@ -1304,6 +1306,7 @@ def run(
raise typer.Exit(1)
except BaseException:
_graceful_shutdown(_server)
getattr(run_mod, "_wait_for_server_shutdown", lambda: None)()
raise
loaded_model = result.get("model", model)
@ -1410,6 +1413,8 @@ def run(
except KeyboardInterrupt:
run_mod._graceful_shutdown(run_mod._server)
typer.echo("\nShutting down...")
finally:
getattr(run_mod, "_wait_for_server_shutdown", lambda: None)()
# ── unsloth studio stop ───────────────────────────────────────────────