Add PID file tracking and unsloth studio stop command (#4598)

* Add PID file tracking and `unsloth studio stop` command

On macOS the .app shortcut launches Studio via osascript into a
Terminal window, then the launcher script exits. The server process
runs outside of the launcher's context with no PID file, so there
is no straightforward way to find or stop it.

This adds:
- PID file at ~/.unsloth/studio/studio.pid, written after the
  server starts and removed on graceful shutdown or via atexit
- `unsloth studio stop` command that reads the PID file and sends
  SIGTERM (or taskkill on Windows) to shut down the server

The PID file is only removed if it still contains the current
process ID, avoiding races when a new server instance replaces
a crashed one.

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

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

* Move atexit PID cleanup into run_server()

The atexit registration was only in the __main__ block, so it
did not cover the `unsloth studio` CLI path that calls
run_server() directly via studio_default(). Moving it into
run_server() ensures the PID file is cleaned up on unexpected
exit regardless of entry point.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-03-25 08:27:27 -07:00 committed by GitHub
commit 6d6008a1ef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 97 additions and 0 deletions

View file

@ -158,6 +158,29 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int:
)
_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid"
def _write_pid_file():
"""Write the current process PID to the studio PID file."""
try:
_PID_FILE.parent.mkdir(parents = True, exist_ok = True)
_PID_FILE.write_text(str(os.getpid()))
except OSError:
pass
def _remove_pid_file():
"""Remove the PID file if it belongs to this process."""
try:
if _PID_FILE.is_file():
stored = _PID_FILE.read_text().strip()
if stored == str(os.getpid()):
_PID_FILE.unlink(missing_ok = True)
except OSError:
pass
def _graceful_shutdown(server = None):
"""Explicitly shut down all subprocess backends and the uvicorn server.
@ -165,6 +188,7 @@ def _graceful_shutdown(server = None):
before the parent exits. This is critical on Windows where atexit
handlers are unreliable after Ctrl+C.
"""
_remove_pid_file()
logger.info("Graceful shutdown initiated — cleaning up subprocesses...")
# 1. Shut down uvicorn server (releases the listening socket)
@ -307,6 +331,11 @@ def run_server(
thread.start()
time.sleep(3)
_write_pid_file()
import atexit
atexit.register(_remove_pid_file)
if not silent:
display_host = _resolve_external_ip() if host == "0.0.0.0" else host

View file

@ -166,6 +166,74 @@ def studio_default(
typer.echo("\nShutting down...")
# ── unsloth studio stop ───────────────────────────────────────────────
_PID_FILE = STUDIO_HOME / "studio.pid"
@studio_app.command()
def stop():
"""Stop a running Unsloth Studio server.
Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM
(or TerminateProcess on Windows) to shut it down gracefully.
"""
import signal as _signal
if not _PID_FILE.is_file():
typer.echo("No running Studio server found (no PID file).")
raise typer.Exit(0)
pid_text = _PID_FILE.read_text().strip()
if not pid_text.isdigit():
typer.echo(f"Invalid PID file contents: {pid_text}")
_PID_FILE.unlink(missing_ok = True)
raise typer.Exit(1)
pid = int(pid_text)
# Check if the process is still alive
try:
os.kill(pid, 0)
except ProcessLookupError:
typer.echo(
f"Studio server (PID {pid}) is not running. Cleaning up stale PID file."
)
_PID_FILE.unlink(missing_ok = True)
raise typer.Exit(0)
except PermissionError:
pass # process exists but we may not own it; try to signal anyway
# Send SIGTERM (graceful shutdown) or TerminateProcess on Windows
try:
if sys.platform == "win32":
subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True)
else:
os.kill(pid, _signal.SIGTERM)
typer.echo(f"Sent shutdown signal to Studio server (PID {pid}).")
except ProcessLookupError:
typer.echo(f"Studio server (PID {pid}) already exited.")
_PID_FILE.unlink(missing_ok = True)
raise typer.Exit(0)
except Exception as e:
typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True)
raise typer.Exit(1)
# Wait briefly for the process to exit and clean up
for _ in range(10):
time.sleep(0.5)
try:
os.kill(pid, 0)
except ProcessLookupError:
_PID_FILE.unlink(missing_ok = True)
typer.echo("Studio server stopped.")
raise typer.Exit(0)
except PermissionError:
break
typer.echo("Studio server is shutting down (may take a few seconds).")
# ── unsloth studio setup / update ─────────────────────────────────────