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:
parent
561f0f39be
commit
6d6008a1ef
2 changed files with 97 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ─────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue