diff --git a/studio/backend/run.py b/studio/backend/run.py index e32b912c37..b892037565 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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 diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index c6d398eebd..a2f0873e22 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -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 ─────────────────────────────────────