diff --git a/studio/backend/run.py b/studio/backend/run.py index 8ef1ac06b8..e906e0ad42 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -744,7 +744,35 @@ def _find_free_port( from utils.paths.storage_roots import studio_root as _studio_root +# Legacy single-instance file; still read so `stop` finds an older build's server. _PID_FILE = _studio_root() / "studio.pid" +PID_FILE_GLOB = "studio-*.pid" + + +def _pid_file_for_port(port: int) -> Path: + return _studio_root() / f"studio-{port}.pid" + + +def _blocker_is_own_studio(blocker: "tuple[int, str] | None") -> bool: + """True when the process holding the port is a server we recorded.""" + return bool(blocker) and blocker[0] in _recorded_studio_pids() + + +def _recorded_studio_pids() -> "set[int]": + """PIDs recorded under this STUDIO_HOME.""" + pids: "set[int]" = set() + try: + paths = list(_studio_root().glob(PID_FILE_GLOB)) + [_PID_FILE] + except OSError: + return pids + for path in paths: + try: + text = path.read_text(encoding = "utf-8").strip() + except (OSError, UnicodeDecodeError): + continue + if text.isdigit(): + pids.add(int(text)) + return pids # Direct backend launches bypass the CLI's env re-export; do it here for # real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR @@ -770,22 +798,30 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") -def _write_pid_file(): - """Write the current process PID to the studio PID file.""" +_OWN_PID_FILE: "Path | None" = None + + +def _write_pid_file(port: int): + """Record this PID under its own port so `stop` can find every server.""" + global _OWN_PID_FILE + path = _pid_file_for_port(port) try: - _PID_FILE.parent.mkdir(parents = True, exist_ok = True) - _PID_FILE.write_text(str(os.getpid()), encoding = "utf-8") + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(str(os.getpid()), encoding = "utf-8") except OSError: - pass + return + _OWN_PID_FILE = path def _remove_pid_file(): """Remove the PID file if it belongs to this process.""" + if _OWN_PID_FILE is None: + return try: - if _PID_FILE.is_file(): - stored = _PID_FILE.read_text(encoding = "utf-8").strip() + if _OWN_PID_FILE.is_file(): + stored = _OWN_PID_FILE.read_text(encoding = "utf-8").strip() if stored == str(os.getpid()): - _PID_FILE.unlink(missing_ok = True) + _OWN_PID_FILE.unlink(missing_ok = True) except (OSError, UnicodeDecodeError): pass @@ -1533,6 +1569,16 @@ def run_server( if not _is_port_free(host, port): original_port = port blocker = _get_pid_on_port(port) + # Falling back past our own server is what creates the orphan. + if _blocker_is_own_studio(blocker): + print( + f"Error: Unsloth Studio is already running on port {port} " + f"(PID {blocker[0]}). Run `unsloth studio stop` first, or start this " + "one on a different --port.", + file = sys.stderr, + flush = True, + ) + sys.exit(1) port = _find_free_port(host, port + 1) if not silent: print("") @@ -1731,7 +1777,7 @@ def run_server( (time.perf_counter() - boot_started) * 1000, ) - _write_pid_file() + _write_pid_file(port) import atexit atexit.register(_remove_pid_file) diff --git a/studio/backend/tests/test_studio_pid_files.py b/studio/backend/tests/test_studio_pid_files.py new file mode 100644 index 0000000000..2cfbd8fc14 --- /dev/null +++ b/studio/backend/tests/test_studio_pid_files.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-port PID files, so `unsloth studio stop` can find every server. + +Imports run.py directly, so run under the Unsloth venv. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import run # noqa: E402 + + +@pytest.fixture(autouse = True) +def isolated_root(tmp_path, monkeypatch): + monkeypatch.setattr(run, "_studio_root", lambda: tmp_path) + monkeypatch.setattr(run, "_PID_FILE", tmp_path / "studio.pid") + monkeypatch.setattr(run, "_OWN_PID_FILE", None) + yield + + +def test_write_pid_file_is_per_port(tmp_path): + run._write_pid_file(8901) + + path = tmp_path / "studio-8901.pid" + assert path.read_text(encoding = "utf-8") == str(os.getpid()) + + +def test_second_port_does_not_clobber_the_first(tmp_path): + (tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8") + + run._write_pid_file(8902) + + assert (tmp_path / "studio-8901.pid").read_text(encoding = "utf-8") == "8550" + assert (tmp_path / "studio-8902.pid").read_text(encoding = "utf-8") == str(os.getpid()) + + +def test_remove_pid_file_only_removes_our_own(tmp_path): + run._write_pid_file(8901) + (tmp_path / "studio-8902.pid").write_text("8600", encoding = "utf-8") + + run._remove_pid_file() + + assert not (tmp_path / "studio-8901.pid").exists() + assert (tmp_path / "studio-8902.pid").exists() + + +def test_remove_pid_file_leaves_a_reused_entry_alone(tmp_path): + run._write_pid_file(8901) + (tmp_path / "studio-8901.pid").write_text("999999", encoding = "utf-8") + + run._remove_pid_file() + + assert (tmp_path / "studio-8901.pid").read_text(encoding = "utf-8") == "999999" + + +def test_recorded_studio_pids_reads_per_port_and_legacy_files(tmp_path): + (tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8") + (tmp_path / "studio-8902.pid").write_text("8600", encoding = "utf-8") + (tmp_path / "studio.pid").write_text("4242", encoding = "utf-8") + + assert run._recorded_studio_pids() == {8550, 8600, 4242} + + +def test_recorded_studio_pids_ignores_corrupt_files(tmp_path): + (tmp_path / "studio-8901.pid").write_text("not-a-pid", encoding = "utf-8") + + assert run._recorded_studio_pids() == set() + + +def test_own_studio_blocking_the_port_is_recognised(tmp_path): + (tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8") + + assert run._blocker_is_own_studio((8550, "python")) is True + + +def test_a_foreign_blocker_still_falls_back(tmp_path): + # jupyter-lab on 8888 must keep the fallback, not abort the launch. + (tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8") + + assert run._blocker_is_own_studio((117, "jupyter-lab")) is False + assert run._blocker_is_own_studio(None) is False diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 864941a20a..54112071c6 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -2394,6 +2394,7 @@ def run( # ── unsloth studio stop ─────────────────────────────────────────────── _PID_FILE = STUDIO_HOME / "studio.pid" +PID_FILE_GLOB = "studio-*.pid" def _pid_alive(pid: int) -> bool: @@ -2423,58 +2424,93 @@ def _pid_alive(pid: int) -> bool: return True -@studio_app.command() -def stop(): - """Stop a running Unsloth Studio server. +def _pid_file_entries() -> "list[tuple[Path, int]]": + """(path, pid) per recorded server, including the legacy studio.pid.""" + entries = [] + try: + paths = sorted(STUDIO_HOME.glob(PID_FILE_GLOB)) + [_PID_FILE] + except OSError: + paths = [_PID_FILE] + seen = set() + for path in paths: + if path in seen or not path.is_file(): + continue + seen.add(path) + try: + text = path.read_text(encoding = "utf-8").strip() + except (OSError, UnicodeDecodeError): + continue + if text.isdigit(): + entries.append((path, int(text))) + else: + typer.echo(f"Ignoring invalid PID file {path.name}: {text}") + path.unlink(missing_ok = True) + return entries - Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM - (or TerminateProcess on Windows) to shut it down gracefully. - """ + +def _signal_stop(pid: int) -> "str | None": + """SIGTERM (or taskkill) the pid. Returns an error string, or None on success.""" import signal as _signal - if not _PID_FILE.is_file(): - typer.echo("No running Unsloth server found (no PID file).") - raise typer.Exit(0) - - pid_text = _PID_FILE.read_text(encoding = "utf-8").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 still alive (os.kill(pid, 0) is invalid on Windows -- see _pid_alive). - if not _pid_alive(pid): - typer.echo(f"Unsloth server (PID {pid}) is not running. Cleaning up stale PID file.") - _PID_FILE.unlink(missing_ok = True) - raise typer.Exit(0) - - # Send SIGTERM (graceful shutdown) or TerminateProcess on Windows try: if sys.platform == "win32": # /T also stops llama-server children, which otherwise keep GPU and port. subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], check = True) else: os.kill(pid, _signal.SIGTERM) - typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).") except ProcessLookupError: - typer.echo(f"Unsloth server (PID {pid}) already exited.") - _PID_FILE.unlink(missing_ok = True) - raise typer.Exit(0) + return None except Exception as e: - typer.echo(f"Failed to stop Unsloth server (PID {pid}): {e}", err = True) - raise typer.Exit(1) + return str(e) + return None - # Wait briefly for the process to exit and clean up. - for _ in range(10): - time.sleep(0.5) + +@studio_app.command() +def stop(): + """Stop every running Unsloth Studio server for this STUDIO_HOME. + + The port fallback can leave more than one running, so stop them all. + """ + entries = _pid_file_entries() + if not entries: + typer.echo("No running Unsloth server found (no PID file).") + raise typer.Exit(0) + + signalled, failed = [], [] + for path, pid in entries: if not _pid_alive(pid): - _PID_FILE.unlink(missing_ok = True) - typer.echo("Unsloth server stopped.") - raise typer.Exit(0) + path.unlink(missing_ok = True) + continue + error = _signal_stop(pid) + if error is not None: + failed.append((pid, error)) + typer.echo(f"Failed to stop Unsloth server (PID {pid}): {error}", err = True) + continue + typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).") + signalled.append((path, pid)) - typer.echo("Unsloth server is shutting down (may take a few seconds).") + if not signalled and not failed: + typer.echo("No running Unsloth server found (cleaned up stale PID files).") + raise typer.Exit(0) + + pending = list(signalled) + for _ in range(10): + if not pending: + break + time.sleep(0.5) + for entry in list(pending): + path, pid = entry + if not _pid_alive(pid): + path.unlink(missing_ok = True) + pending.remove(entry) + + stopped = len(signalled) - len(pending) + if stopped: + typer.echo(f"Unsloth server{'s' if stopped > 1 else ''} stopped ({stopped}).") + for _path, pid in pending: + typer.echo(f"Unsloth server (PID {pid}) is shutting down (may take a few seconds).") + if failed: + raise typer.Exit(1) # ── unsloth studio setup / update ───────────────────────────────────── diff --git a/unsloth_cli/tests/test_studio_stop.py b/unsloth_cli/tests/test_studio_stop.py new file mode 100644 index 0000000000..c926b74aae --- /dev/null +++ b/unsloth_cli/tests/test_studio_stop.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth studio stop` must stop every server it started. + +With one PID file the second launch overwrote the first entry, so stop killed +the newer server, claimed success, and left the older one serving. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _studio(): + from unsloth_cli.commands import studio as _studio_mod + return _studio_mod + + +def _install(monkeypatch, tmp_path, *, alive, killed = None): + """Point the CLI at tmp_path and fake process liveness.""" + studio_mod = _studio() + monkeypatch.setattr(studio_mod, "STUDIO_HOME", tmp_path) + monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid") + monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) + + live = set(alive) + killed = killed if killed is not None else [] + + monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live) + + def fake_kill(pid, _sig): + killed.append(pid) + live.discard(pid) + + monkeypatch.setattr(studio_mod.os, "kill", fake_kill) + monkeypatch.setattr(sys, "platform", "linux") + return studio_mod, live, killed + + +def _write_pid(tmp_path, name, pid): + (tmp_path / name).write_text(str(pid), encoding = "utf-8") + + +def _run_stop(studio_mod): + import typer as _typer + + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + return CliRunner().invoke(app, ["studio", "stop"]) + + +def test_stop_kills_every_recorded_server(monkeypatch, tmp_path): + studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {8550, 8600}) + _write_pid(tmp_path, "studio-8901.pid", 8550) + _write_pid(tmp_path, "studio-8902.pid", 8600) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert sorted(killed) == [8550, 8600] + assert not list(tmp_path.glob("studio-*.pid")) + + +def test_stop_does_not_leave_the_older_instance_running(monkeypatch, tmp_path): + # The reported symptom: stop claimed success while instance A kept serving. + studio_mod, live, _killed = _install(monkeypatch, tmp_path, alive = {8550, 8600}) + _write_pid(tmp_path, "studio-8901.pid", 8550) + _write_pid(tmp_path, "studio-8902.pid", 8600) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert live == set() + + +def test_stop_reads_the_legacy_single_pid_file(monkeypatch, tmp_path): + studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {4242}) + _write_pid(tmp_path, "studio.pid", 4242) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [4242] + assert not (tmp_path / "studio.pid").exists() + + +def test_stop_reports_nothing_running_without_pid_files(monkeypatch, tmp_path): + studio_mod, _live, _killed = _install(monkeypatch, tmp_path, alive = set()) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert "no running unsloth server" in result.output.lower() + + +def test_stop_cleans_stale_pid_files_without_claiming_a_stop(monkeypatch, tmp_path): + studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = set()) + _write_pid(tmp_path, "studio-8901.pid", 8550) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [] + assert not (tmp_path / "studio-8901.pid").exists() + assert "stopped" not in result.output.lower() + + +def test_stop_does_not_claim_a_stop_while_a_server_is_still_alive(monkeypatch, tmp_path): + # SIGTERM delivered but it never exits: don't claim a stop, keep the file. + studio_mod = _studio() + monkeypatch.setattr(studio_mod, "STUDIO_HOME", tmp_path) + monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid") + monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) + monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: True) + monkeypatch.setattr(studio_mod.os, "kill", lambda pid, sig: None) + monkeypatch.setattr(sys, "platform", "linux") + _write_pid(tmp_path, "studio-8901.pid", 8550) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert "shutting down" in result.output.lower() + assert "stopped" not in result.output.lower() + assert (tmp_path / "studio-8901.pid").exists() + + +def test_stop_continues_after_one_server_fails_to_stop(monkeypatch, tmp_path): + studio_mod = _studio() + monkeypatch.setattr(studio_mod, "STUDIO_HOME", tmp_path) + monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid") + monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) + live = {8550, 8600} + monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live) + + def fake_kill(pid, _sig): + if pid == 8550: + raise PermissionError("not permitted") + live.discard(pid) + + monkeypatch.setattr(studio_mod.os, "kill", fake_kill) + monkeypatch.setattr(sys, "platform", "linux") + _write_pid(tmp_path, "studio-8901.pid", 8550) + _write_pid(tmp_path, "studio-8902.pid", 8600) + + result = _run_stop(studio_mod) + + combined = (result.output or "") + (getattr(result, "stderr", "") or "") + assert result.exit_code == 1, combined + assert 8600 not in live + assert "8550" in combined + + +def test_stop_discards_a_corrupt_pid_file(monkeypatch, tmp_path): + studio_mod, _live, _killed = _install(monkeypatch, tmp_path, alive = set()) + (tmp_path / "studio-8901.pid").write_text("not-a-pid", encoding = "utf-8") + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert not (tmp_path / "studio-8901.pid").exists()