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:
parent
86d65f3d4a
commit
dbc13f02c9
4 changed files with 183 additions and 5 deletions
13
install.sh
13
install.sh
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
133
tests/test_studio_shutdown_thread_wait.py
Normal file
133
tests/test_studio_shutdown_thread_wait.py
Normal 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"
|
||||
|
|
@ -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 ───────────────────────────────────────────────
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue